Snap for 4613997 from c8f273c3621c24b31b3342a7a587cc063ad5e63e to pi-releaseandroid-wear-9.0.0_r9android-wear-9.0.0_r8android-wear-9.0.0_r7android-wear-9.0.0_r6android-wear-9.0.0_r5android-wear-9.0.0_r4android-wear-9.0.0_r34android-wear-9.0.0_r33android-wear-9.0.0_r32android-wear-9.0.0_r31android-wear-9.0.0_r30android-wear-9.0.0_r3android-wear-9.0.0_r29android-wear-9.0.0_r28android-wear-9.0.0_r27android-wear-9.0.0_r26android-wear-9.0.0_r25android-wear-9.0.0_r24android-wear-9.0.0_r23android-wear-9.0.0_r22android-wear-9.0.0_r21android-wear-9.0.0_r20android-wear-9.0.0_r2android-wear-9.0.0_r19android-wear-9.0.0_r18android-wear-9.0.0_r17android-wear-9.0.0_r16android-wear-9.0.0_r15android-wear-9.0.0_r14android-wear-9.0.0_r13android-wear-9.0.0_r12android-wear-9.0.0_r11android-wear-9.0.0_r10android-wear-9.0.0_r1android-security-9.0.0_r76android-security-9.0.0_r75android-security-9.0.0_r74android-security-9.0.0_r73android-security-9.0.0_r72android-security-9.0.0_r71android-security-9.0.0_r70android-security-9.0.0_r69android-security-9.0.0_r68android-security-9.0.0_r67android-security-9.0.0_r66android-security-9.0.0_r65android-security-9.0.0_r64android-security-9.0.0_r63android-security-9.0.0_r62android-cts-9.0_r9android-cts-9.0_r8android-cts-9.0_r7android-cts-9.0_r6android-cts-9.0_r5android-cts-9.0_r4android-cts-9.0_r3android-cts-9.0_r20android-cts-9.0_r2android-cts-9.0_r19android-cts-9.0_r18android-cts-9.0_r17android-cts-9.0_r16android-cts-9.0_r15android-cts-9.0_r14android-cts-9.0_r13android-cts-9.0_r12android-cts-9.0_r11android-cts-9.0_r10android-cts-9.0_r1android-9.0.0_r9android-9.0.0_r8android-9.0.0_r7android-9.0.0_r61android-9.0.0_r60android-9.0.0_r6android-9.0.0_r59android-9.0.0_r58android-9.0.0_r57android-9.0.0_r56android-9.0.0_r55android-9.0.0_r54android-9.0.0_r53android-9.0.0_r52android-9.0.0_r51android-9.0.0_r50android-9.0.0_r5android-9.0.0_r49android-9.0.0_r48android-9.0.0_r3android-9.0.0_r2android-9.0.0_r18android-9.0.0_r17android-9.0.0_r10android-9.0.0_r1security-pi-releasepie-security-releasepie-s2-releasepie-release-2pie-releasepie-r2-s2-releasepie-r2-s1-releasepie-r2-releasepie-platform-releasepie-gsipie-cuttlefish-testingpie-cts-release

Change-Id: Ia163b028a3749475997fa80006d4eabbc30daead

5 files changed, 168 insertions, 157 deletions

diff --git a/internal/lib/base_cloud_client.py b/internal/lib/base_cloud_client.py
index f876118a..e9dd097f 100755
--- a/internal/lib/base_cloud_client.py
+++ b/internal/lib/base_cloud_client.py
@@ -82,14 +82,15 @@ class BaseCloudApiClient(object):
             An apiclient.discovery.Resource object
         """
         http_auth = oauth2_credentials.authorize(httplib2.Http())
-        return utils.RetryException(exc_retry=cls.RETRIABLE_AUTH_ERRORS,
-                                    max_retry=cls.RETRY_COUNT,
-                                    functor=build,
-                                    sleep=cls.RETRY_SLEEP_MULTIPLIER,
-                                    backoff_factor=cls.RETRY_BACKOFF_FACTOR,
-                                    serviceName=cls.API_NAME,
-                                    version=cls.API_VERSION,
-                                    http=http_auth)
+        return utils.RetryExceptionType(
+                exception_types=cls.RETRIABLE_AUTH_ERRORS,
+                max_retries=cls.RETRY_COUNT,
+                functor=build,
+                sleep_multiplier=cls.RETRY_SLEEP_MULTIPLIER,
+                retry_backoff_factor=cls.RETRY_BACKOFF_FACTOR,
+                serviceName=cls.API_NAME,
+                version=cls.API_VERSION,
+                http=http_auth)
 
     def _ShouldRetry(self, exception, retry_http_codes,
                      other_retriable_errors):
@@ -172,9 +173,9 @@ class BaseCloudApiClient(object):
         Args:
             api: An apiclient.http.HttpRequest, representing the api to execute:
             retry_http_codes: A list of http codes to retry.
-            max_retry: See utils.GenericRetry.
-            sleep: See utils.GenericRetry.
-            backoff_factor: See utils.GenericRetry.
+            max_retry: See utils.Retry.
+            sleep: See utils.Retry.
+            backoff_factor: See utils.Retry.
             other_retriable_errors: A tuple of error types that should be retried
                                     other than errors.HttpError.
 
@@ -210,12 +211,10 @@ class BaseCloudApiClient(object):
                 return True
             return False
 
-        return utils.GenericRetry(_Handler,
-                                  max_retry=max_retry,
-                                  functor=self.ExecuteOnce,
-                                  sleep=sleep,
-                                  backoff_factor=backoff_factor,
-                                  api=api)
+        return utils.Retry(
+             _Handler, max_retries=max_retry, functor=self.ExecuteOnce,
+             sleep_multiplier=sleep, retry_backoff_factor=backoff_factor,
+             api=api)
 
     def BatchExecuteOnce(self, requests):
         """Execute requests in a batch.
@@ -259,9 +258,9 @@ class BaseCloudApiClient(object):
             requests: A dictionary where key is request id picked by caller,
                       and value is a apiclient.http.HttpRequest.
             retry_http_codes: A list of http codes to retry.
-            max_retry: See utils.GenericRetry.
-            sleep: See utils.GenericRetry.
-            backoff_factor: See utils.GenericRetry.
+            max_retry: See utils.Retry.
+            sleep: See utils.Retry.
+            backoff_factor: See utils.Retry.
             other_retriable_errors: A tuple of error types that should be retried
                                     other than errors.HttpError.
 
diff --git a/internal/lib/gstorage_client.py b/internal/lib/gstorage_client.py
index 686d8afc..ebf60871 100755
--- a/internal/lib/gstorage_client.py
+++ b/internal/lib/gstorage_client.py
@@ -157,10 +157,9 @@ class StorageClient(base_cloud_client.BaseCloudApiClient):
         Raises:
             errors.ResourceNotFoundError: when file is not found.
         """
-        item = utils.RetryException(errors.ResourceNotFoundError,
-                                    max_retry=self.GET_OBJ_MAX_RETRY,
-                                    functor=self.Get,
-                                    sleep=self.GET_OBJ_RETRY_SLEEP,
-                                    bucket_name=bucket_name,
-                                    object_name=object_name)
+        item = utils.RetryExceptionType(
+                errors.ResourceNotFoundError,
+                max_retries=self.GET_OBJ_MAX_RETRY, functor=self.Get,
+                sleep_multiplier=self.GET_OBJ_RETRY_SLEEP,
+                bucket_name=bucket_name, object_name=object_name)
         return item["selfLink"]
diff --git a/internal/lib/utils.py b/internal/lib/utils.py
index 4789d3ee..54524dac 100755
--- a/internal/lib/utils.py
+++ b/internal/lib/utils.py
@@ -18,8 +18,6 @@
 
 The following code is copied from chromite with modifications.
   - class TempDir: chromite/lib/osutils.py
-  - method GenericRetry:: chromite/lib/retry_util.py
-  - method RetryException:: chromite/lib/retry_util.py
 
 """
 
@@ -113,127 +111,95 @@ class TempDir(object):
         """Delete the object."""
         self.Cleanup()
 
+def RetryOnException(retry_checker, max_retries, sleep_multiplier=0,
+                     retry_backoff_factor=1):
+  """Decorater which retries the function call if |retry_checker| returns true.
+
+  Args:
+    retry_checker: A callback function which should take an exception instance
+                   and return True if functor(*args, **kwargs) should be retried
+                   when such exception is raised, and return False if it should
+                   not be retried.
+    max_retries: Maximum number of retries allowed.
+    sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
+                      retry_backoff_factor is 1.  Will sleep
+                      sleep_multiplier * (
+                          retry_backoff_factor ** (attempt_count -  1))
+                      if retry_backoff_factor != 1.
+    retry_backoff_factor: See explanation of sleep_multiplier.
+
+  Returns:
+    The function wrapper.
+  """
+  def _Wrapper(func):
+    def _FunctionWrapper(*args, **kwargs):
+      return Retry(retry_checker, max_retries, func, sleep_multiplier,
+                   retry_backoff_factor,
+                   *args, **kwargs)
+    return _FunctionWrapper
+  return _Wrapper
+
+
+def Retry(retry_checker, max_retries, functor, sleep_multiplier=0,
+          retry_backoff_factor=1, *args, **kwargs):
+  """Conditionally retry a function.
+
+  Args:
+    retry_checker: A callback function which should take an exception instance
+                   and return True if functor(*args, **kwargs) should be retried
+                   when such exception is raised, and return False if it should
+                   not be retried.
+    max_retries: Maximum number of retries allowed.
+    functor: The function to call, will call functor(*args, **kwargs).
+    sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
+                      retry_backoff_factor is 1.  Will sleep
+                      sleep_multiplier * (
+                          retry_backoff_factor ** (attempt_count -  1))
+                      if retry_backoff_factor != 1.
+    retry_backoff_factor: See explanation of sleep_multiplier.
+    *args: Arguments to pass to the functor.
+    **kwargs: Key-val based arguments to pass to the functor.
+
+  Returns:
+    The return value of the functor.
+
+  Raises:
+    Exception: The exception that functor(*args, **kwargs) throws.
+  """
+  attempt_count = 0
+  while attempt_count <= max_retries:
+    try:
+      attempt_count += 1
+      return_value = functor(*args, **kwargs)
+      return return_value
+    except Exception as e:  # pylint: disable=W0703
+      if retry_checker(e) and attempt_count <= max_retries:
+        if retry_backoff_factor != 1:
+          sleep = sleep_multiplier * (
+              retry_backoff_factor ** (attempt_count -  1))
+        else:
+          sleep = sleep_multiplier * attempt_count
+        time.sleep(sleep)
+      else:
+        raise
 
-def GenericRetry(handler,
-                 max_retry,
-                 functor,
-                 sleep=0,
-                 backoff_factor=1,
-                 success_functor=lambda x: None,
-                 raise_first_exception_on_failure=True,
-                 *args,
-                 **kwargs):
-    """Generic retry loop w/ optional break out depending on exceptions.
-
-    To retry based on the return value of |functor| see the timeout_util module.
-
-    Keep in mind that the total sleep time will be the triangular value of
-    max_retry multiplied by the sleep value.  e.g. max_retry=5 and sleep=10
-    will be T5 (i.e. 5+4+3+2+1) times 10, or 150 seconds total.  Rather than
-    use a large sleep value, you should lean more towards large retries and
-    lower sleep intervals, or by utilizing backoff_factor.
-
-    Args:
-        handler: A functor invoked w/ the exception instance that
-                functor(*args, **kwargs) threw.  If it returns True, then a
-                retry is attempted.  If False, the exception is re-raised.
-        max_retry: A positive integer representing how many times to retry
-                   the command before giving up.  Worst case, the command is
-                   invoked (max_retry + 1) times before failing.
-        functor: A callable to pass args and kwargs to.
-        sleep: Optional keyword.  Multiplier for how long to sleep between
-               retries; will delay (1*sleep) the first time, then (2*sleep),
-               continuing via attempt * sleep.
-        backoff_factor: Optional keyword. If supplied and > 1, subsequent sleeps
-                        will be of length (backoff_factor ^ (attempt - 1)) * sleep,
-                        rather than the default behavior of attempt * sleep.
-        success_functor: Optional functor that accepts 1 argument. Will be called
-                         after successful call to |functor|, with the argument
-                         being the number of attempts (1 = |functor| succeeded on
-                         first try).
-        raise_first_exception_on_failure: Optional boolean which determines which
-                                          exception is raised upon failure after
-                                          retries. If True, the first exception
-                                          that was encountered. If False, the
-                                          final one. Default: True.
-        *args: Positional args passed to functor.
-        **kwargs: Optional args passed to functor.
-
-    Returns:
-        Whatever functor(*args, **kwargs) returns.
-
-    Raises:
-        Exception: Whatever exceptions functor(*args, **kwargs) throws and
-                   isn't suppressed is raised.  Note that the first exception
-                   encountered is what's thrown.
-    """
-
-    if max_retry < 0:
-        raise ValueError('max_retry needs to be zero or more: %s' % max_retry)
-
-    if backoff_factor < 1:
-        raise ValueError('backoff_factor must be 1 or greater: %s' %
-                         backoff_factor)
-
-    ret, success = (None, False)
-    attempt = 0
-
-    exc_info = None
-    for attempt in xrange(max_retry + 1):
-        if attempt and sleep:
-            if backoff_factor > 1:
-                sleep_time = sleep * backoff_factor**(attempt - 1)
-            else:
-                sleep_time = sleep * attempt
-            time.sleep(sleep_time)
-        try:
-            ret = functor(*args, **kwargs)
-            success = True
-            break
-        except Exception as e:  # pylint: disable=W0703
-            # Note we're not snagging BaseException, so MemoryError/KeyboardInterrupt
-            # and friends don't enter this except block.
-            if not handler(e):
-                raise
-            # If raise_first_exception_on_failure, we intentionally ignore
-            # any failures in later attempts since we'll throw the original
-            # failure if all retries fail.
-            if exc_info is None or not raise_first_exception_on_failure:
-                exc_info = sys.exc_info()
-
-    if success:
-        success_functor(attempt + 1)
-        return ret
-
-    raise exc_info[0], exc_info[1], exc_info[2]
-
-
-def RetryException(exc_retry, max_retry, functor, *args, **kwargs):
-    """Convenience wrapper for GenericRetry based on exceptions.
-
-    Args:
-        exc_retry: A class (or tuple of classes).  If the raised exception
-                   is the given class(es), a retry will be attempted.
-                   Otherwise, the exception is raised.
-        max_retry: See GenericRetry.
-        functor: See GenericRetry.
-        *args: See GenericRetry.
-        **kwargs: See GenericRetry.
 
-    Returns:
-        Return what functor returns.
-
-    Raises:
-        TypeError, if exc_retry is of an unexpected type.
-    """
-    if not isinstance(exc_retry, (tuple, type)):
-        raise TypeError("exc_retry should be an exception (or tuple), not %r" %
-                        exc_retry)
+def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs):
+  """Retry exception if it is one of the given types.
 
-    def _Handler(exc, values=exc_retry):
-        return isinstance(exc, values)
+  Args:
+    exception_types: A tuple of exception types, e.g. (ValueError, KeyError)
+    max_retries: Max number of retries allowed.
+    functor: The function to call. Will be retried if exception is raised and
+             the exception is one of the exception_types.
+    *args: Arguments to pass to Retry function.
+    **kwargs: Key-val based arguments to pass to Retry functions.
 
-    return GenericRetry(_Handler, max_retry, functor, *args, **kwargs)
+  Returns:
+    The value returned by calling functor.
+  """
+  return Retry(lambda e: isinstance(e, exception_types), max_retries,
+               functor, *args, **kwargs)
 
 
 def PollAndWait(func, expected_return, timeout_exception, timeout_secs,
@@ -418,9 +384,9 @@ class BatchHttpRequestExecutor(object):
             requests: A dictionary where key is request id picked by caller,
                       and value is a apiclient.http.HttpRequest.
             retry_http_codes: A list of http codes to retry.
-            max_retry: See utils.GenericRetry.
-            sleep: See utils.GenericRetry.
-            backoff_factor: See utils.GenericRetry.
+            max_retry: See utils.Retry.
+            sleep: See utils.Retry.
+            backoff_factor: See utils.Retry.
             other_retriable_errors: A tuple of error types that should be retried
                                     other than errors.HttpError.
         """
@@ -464,7 +430,7 @@ class BatchHttpRequestExecutor(object):
                 self._pending_requests[request_id] = self._requests[request_id]
         if self._pending_requests:
             # If there is still retriable requests pending, raise an error
-            # so that GenericRetry will retry this function with pending_requests.
+            # so that Retry will retry this function with pending_requests.
             raise errors.HasRetriableRequestsError(
                 "Retriable errors: %s" % [str(results[rid][1])
                                           for rid in self._pending_requests])
@@ -491,11 +457,11 @@ class BatchHttpRequestExecutor(object):
 
         try:
             self._pending_requests = self._requests.copy()
-            GenericRetry(_ShouldRetryHandler,
-                         max_retry=self._max_retry,
-                         functor=self._ExecuteOnce,
-                         sleep=self._sleep,
-                         backoff_factor=self._backoff_factor)
+            Retry(
+                _ShouldRetryHandler, max_retries=self._max_retry,
+                functor=self._ExecuteOnce,
+                sleep_multiplier=self._sleep,
+                retry_backoff_factor=self._backoff_factor)
         except errors.HasRetriableRequestsError:
             logger.debug("Some requests did not succeed after retry.")
 
diff --git a/internal/lib/utils_test.py b/internal/lib/utils_test.py
index 9899d752..bb1200ed 100644
--- a/internal/lib/utils_test.py
+++ b/internal/lib/utils_test.py
@@ -19,6 +19,7 @@
 import getpass
 import os
 import subprocess
+import time
 
 import mock
 
@@ -51,6 +52,52 @@ class UtilsTest(driver_test_lib.BaseDriverTest):
         utils.SSH_KEYGEN_CMD + ["-C", getpass.getuser(), "-f", private_key],
         stdout=mock.ANY, stderr=mock.ANY)
 
+  def testRetryOnException(self):
+    def _IsValueError(exc):
+      return isinstance(exc, ValueError)
+    num_retry = 5
+
+    @utils.RetryOnException(_IsValueError, num_retry)
+    def _RaiseAndRetry(sentinel):
+      sentinel.alert()
+      raise ValueError("Fake error.")
+
+    sentinel = mock.MagicMock()
+    self.assertRaises(ValueError, _RaiseAndRetry, sentinel)
+    self.assertEqual(1 + num_retry, sentinel.alert.call_count)
+
+  def testRetryExceptionType(self):
+    """Test RetryExceptionType function."""
+    def _RaiseAndRetry(sentinel):
+      sentinel.alert()
+      raise ValueError("Fake error.")
+
+    num_retry = 5
+    sentinel = mock.MagicMock()
+    self.assertRaises(ValueError, utils.RetryExceptionType,
+                      (KeyError, ValueError), num_retry, _RaiseAndRetry,
+                      sentinel=sentinel)
+    self.assertEqual(1 + num_retry, sentinel.alert.call_count)
+
+  def testRetry(self):
+    """Test Retry."""
+    self.Patch(time, "sleep")
+    def _RaiseAndRetry(sentinel):
+      sentinel.alert()
+      raise ValueError("Fake error.")
+
+    num_retry = 5
+    sentinel = mock.MagicMock()
+    self.assertRaises(ValueError, utils.RetryExceptionType,
+                      (ValueError, KeyError), num_retry, _RaiseAndRetry,
+                      sleep_multiplier=1,
+                      retry_backoff_factor=2,
+                      sentinel=sentinel)
+
+    self.assertEqual(1 + num_retry, sentinel.alert.call_count)
+    time.sleep.assert_has_calls(
+        [mock.call(1), mock.call(2), mock.call(4), mock.call(8), mock.call(16)])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/public/acloud_kernel/kernel_swapper.py b/public/acloud_kernel/kernel_swapper.py
index ca48c266..d30eb9e6 100755
--- a/public/acloud_kernel/kernel_swapper.py
+++ b/public/acloud_kernel/kernel_swapper.py
@@ -147,8 +147,8 @@ class KernelSwapper(object):
             subprocess.CalledProcessError: For any non-zero return code of
                                            host_cmd.
         """
-        utils.GenericRetry(
-            handler=lambda e: isinstance(e, subprocess.CalledProcessError),
-            max_retry=2,
+        utils.Retry(
+            retry_checker=lambda e: isinstance(e, subprocess.CalledProcessError),
+            max_retries=2,
             functor=lambda cmd: subprocess.check_call(cmd, shell=True),
             cmd=host_cmd)