Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 1 | #!/usr/bin/env python |
| 2 | # |
| 3 | # Copyright 2016 - The Android Open Source Project |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | # you may not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 16 | """Common Utilities.""" |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 17 | |
| 18 | import base64 |
| 19 | import binascii |
| 20 | import errno |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 21 | import getpass |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 22 | import logging |
| 23 | import os |
| 24 | import shutil |
| 25 | import struct |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 26 | import subprocess |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 27 | import sys |
| 28 | import tarfile |
| 29 | import tempfile |
| 30 | import time |
| 31 | import uuid |
| 32 | |
| 33 | from acloud.public import errors |
| 34 | |
| 35 | logger = logging.getLogger(__name__) |
| 36 | |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 37 | SSH_KEYGEN_CMD = ["ssh-keygen", "-t", "rsa", "-b", "4096"] |
| 38 | |
| 39 | |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 40 | class TempDir(object): |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 41 | """A context manager that ceates a temporary directory. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 42 | |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 43 | Attributes: |
| 44 | path: The path of the temporary directory. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 45 | """ |
| 46 | |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 47 | def __init__(self): |
| 48 | self.path = tempfile.mkdtemp() |
| 49 | os.chmod(self.path, 0o700) |
| 50 | logger.debug("Created temporary dir %s", self.path) |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 51 | |
| 52 | def __enter__(self): |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 53 | """Enter.""" |
| 54 | return self.path |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 55 | |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 56 | def __exit__(self, exc_type, exc_value, traceback): |
| 57 | """Exit. |
| 58 | |
| 59 | Args: |
| 60 | exc_type: Exception type raised within the context manager. |
| 61 | None if no execption is raised. |
| 62 | exc_value: Exception instance raised within the context manager. |
| 63 | None if no execption is raised. |
| 64 | traceback: Traceback for exeception that is raised within |
| 65 | the context manager. |
| 66 | None if no execption is raised. |
| 67 | Raises: |
| 68 | EnvironmentError or OSError when failed to delete temp directory. |
| 69 | """ |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 70 | try: |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 71 | if self.path: |
| 72 | shutil.rmtree(self.path) |
| 73 | logger.debug("Deleted temporary dir %s", self.path) |
| 74 | except EnvironmentError as e: |
| 75 | # Ignore error if there is no exception raised |
| 76 | # within the with-clause and the EnvironementError is |
| 77 | # about problem that directory or file does not exist. |
| 78 | if not exc_type and e.errno != errno.ENOENT: |
| 79 | raise |
| 80 | except Exception as e: # pylint: disable=W0703 |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 81 | if exc_type: |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 82 | logger.error( |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 83 | "Encountered error while deleting %s: %s", |
| 84 | self.path, |
| 85 | str(e), |
| 86 | exc_info=True) |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 87 | else: |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 88 | raise |
| 89 | |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 90 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 91 | def RetryOnException(retry_checker, |
| 92 | max_retries, |
| 93 | sleep_multiplier=0, |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 94 | retry_backoff_factor=1): |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 95 | """Decorater which retries the function call if |retry_checker| returns true. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 96 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 97 | Args: |
| 98 | retry_checker: A callback function which should take an exception instance |
| 99 | and return True if functor(*args, **kwargs) should be retried |
| 100 | when such exception is raised, and return False if it should |
| 101 | not be retried. |
| 102 | max_retries: Maximum number of retries allowed. |
| 103 | sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if |
| 104 | retry_backoff_factor is 1. Will sleep |
| 105 | sleep_multiplier * ( |
| 106 | retry_backoff_factor ** (attempt_count - 1)) |
| 107 | if retry_backoff_factor != 1. |
| 108 | retry_backoff_factor: See explanation of sleep_multiplier. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 109 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 110 | Returns: |
| 111 | The function wrapper. |
| 112 | """ |
| 113 | |
| 114 | def _Wrapper(func): |
| 115 | def _FunctionWrapper(*args, **kwargs): |
| 116 | return Retry(retry_checker, max_retries, func, sleep_multiplier, |
| 117 | retry_backoff_factor, *args, **kwargs) |
| 118 | |
| 119 | return _FunctionWrapper |
| 120 | |
| 121 | return _Wrapper |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 122 | |
| 123 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 124 | def Retry(retry_checker, |
| 125 | max_retries, |
| 126 | functor, |
| 127 | sleep_multiplier, |
| 128 | retry_backoff_factor, |
| 129 | *args, |
| 130 | **kwargs): |
| 131 | """Conditionally retry a function. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 132 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 133 | Args: |
| 134 | retry_checker: A callback function which should take an exception instance |
| 135 | and return True if functor(*args, **kwargs) should be retried |
| 136 | when such exception is raised, and return False if it should |
| 137 | not be retried. |
| 138 | max_retries: Maximum number of retries allowed. |
| 139 | functor: The function to call, will call functor(*args, **kwargs). |
| 140 | sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if |
| 141 | retry_backoff_factor is 1. Will sleep |
| 142 | sleep_multiplier * ( |
| 143 | retry_backoff_factor ** (attempt_count - 1)) |
| 144 | if retry_backoff_factor != 1. |
| 145 | retry_backoff_factor: See explanation of sleep_multiplier. |
| 146 | *args: Arguments to pass to the functor. |
| 147 | **kwargs: Key-val based arguments to pass to the functor. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 148 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 149 | Returns: |
| 150 | The return value of the functor. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 151 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 152 | Raises: |
| 153 | Exception: The exception that functor(*args, **kwargs) throws. |
| 154 | """ |
| 155 | attempt_count = 0 |
| 156 | while attempt_count <= max_retries: |
| 157 | try: |
| 158 | attempt_count += 1 |
| 159 | return_value = functor(*args, **kwargs) |
| 160 | return return_value |
| 161 | except Exception as e: # pylint: disable=W0703 |
| 162 | if retry_checker(e) and attempt_count <= max_retries: |
| 163 | if retry_backoff_factor != 1: |
| 164 | sleep = sleep_multiplier * (retry_backoff_factor** |
| 165 | (attempt_count - 1)) |
| 166 | else: |
| 167 | sleep = sleep_multiplier * attempt_count |
| 168 | time.sleep(sleep) |
| 169 | else: |
| 170 | raise |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 171 | |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 172 | |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 173 | def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs): |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 174 | """Retry exception if it is one of the given types. |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 175 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 176 | Args: |
| 177 | exception_types: A tuple of exception types, e.g. (ValueError, KeyError) |
| 178 | max_retries: Max number of retries allowed. |
| 179 | functor: The function to call. Will be retried if exception is raised and |
| 180 | the exception is one of the exception_types. |
| 181 | *args: Arguments to pass to Retry function. |
| 182 | **kwargs: Key-val based arguments to pass to Retry functions. |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 183 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 184 | Returns: |
| 185 | The value returned by calling functor. |
| 186 | """ |
| 187 | return Retry(lambda e: isinstance(e, exception_types), max_retries, |
| 188 | functor, *args, **kwargs) |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 189 | |
| 190 | |
| 191 | def PollAndWait(func, expected_return, timeout_exception, timeout_secs, |
| 192 | sleep_interval_secs, *args, **kwargs): |
| 193 | """Call a function until the function returns expected value or times out. |
| 194 | |
| 195 | Args: |
| 196 | func: Function to call. |
| 197 | expected_return: The expected return value. |
| 198 | timeout_exception: Exception to raise when it hits timeout. |
| 199 | timeout_secs: Timeout seconds. |
| 200 | If 0 or less than zero, the function will run once and |
| 201 | we will not wait on it. |
| 202 | sleep_interval_secs: Time to sleep between two attemps. |
| 203 | *args: list of args to pass to func. |
| 204 | **kwargs: dictionary of keyword based args to pass to func. |
| 205 | |
| 206 | Raises: |
| 207 | timeout_exception: if the run of function times out. |
| 208 | """ |
| 209 | # TODO(fdeng): Currently this method does not kill |
| 210 | # |func|, if |func| takes longer than |timeout_secs|. |
| 211 | # We can use a more robust version from chromite. |
| 212 | start = time.time() |
| 213 | while True: |
| 214 | return_value = func(*args, **kwargs) |
| 215 | if return_value == expected_return: |
| 216 | return |
| 217 | elif time.time() - start > timeout_secs: |
| 218 | raise timeout_exception |
| 219 | else: |
| 220 | if sleep_interval_secs > 0: |
| 221 | time.sleep(sleep_interval_secs) |
| 222 | |
| 223 | |
| 224 | def GenerateUniqueName(prefix=None, suffix=None): |
| 225 | """Generate a random unque name using uuid4. |
| 226 | |
| 227 | Args: |
| 228 | prefix: String, desired prefix to prepend to the generated name. |
| 229 | suffix: String, desired suffix to append to the generated name. |
| 230 | |
| 231 | Returns: |
| 232 | String, a random name. |
| 233 | """ |
| 234 | name = uuid.uuid4().hex |
| 235 | if prefix: |
| 236 | name = "-".join([prefix, name]) |
| 237 | if suffix: |
| 238 | name = "-".join([name, suffix]) |
| 239 | return name |
| 240 | |
| 241 | |
| 242 | def MakeTarFile(src_dict, dest): |
| 243 | """Archive files in tar.gz format to a file named as |dest|. |
| 244 | |
| 245 | Args: |
| 246 | src_dict: A dictionary that maps a path to be archived |
| 247 | to the corresponding name that appears in the archive. |
| 248 | dest: String, path to output file, e.g. /tmp/myfile.tar.gz |
| 249 | """ |
| 250 | logger.info("Compressing %s into %s.", src_dict.keys(), dest) |
| 251 | with tarfile.open(dest, "w:gz") as tar: |
| 252 | for src, arcname in src_dict.iteritems(): |
| 253 | tar.add(src, arcname=arcname) |
| 254 | |
| 255 | |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 256 | def CreateSshKeyPairIfNotExist(private_key_path, public_key_path): |
| 257 | """Create the ssh key pair if they don't exist. |
| 258 | |
| 259 | Check if the public and private key pairs exist at |
| 260 | the given places. If not, create them. |
| 261 | |
| 262 | Args: |
| 263 | private_key_path: Path to the private key file. |
| 264 | e.g. ~/.ssh/acloud_rsa |
| 265 | public_key_path: Path to the public key file. |
| 266 | e.g. ~/.ssh/acloud_rsa.pub |
| 267 | Raises: |
| 268 | error.DriverError: If failed to create the key pair. |
| 269 | """ |
| 270 | public_key_path = os.path.expanduser(public_key_path) |
| 271 | private_key_path = os.path.expanduser(private_key_path) |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 272 | create_key = (not os.path.exists(public_key_path) |
| 273 | and not os.path.exists(private_key_path)) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 274 | if not create_key: |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 275 | logger.debug( |
| 276 | "The ssh private key (%s) or public key (%s) already exist," |
| 277 | "will not automatically create the key pairs.", private_key_path, |
| 278 | public_key_path) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 279 | return |
| 280 | cmd = SSH_KEYGEN_CMD + ["-C", getpass.getuser(), "-f", private_key_path] |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 281 | logger.info( |
| 282 | "The ssh private key (%s) and public key (%s) do not exist, " |
| 283 | "automatically creating key pair, calling: %s", private_key_path, |
| 284 | public_key_path, " ".join(cmd)) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 285 | try: |
| 286 | subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout) |
| 287 | except subprocess.CalledProcessError as e: |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 288 | raise errors.DriverError("Failed to create ssh key pair: %s" % str(e)) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 289 | except OSError as e: |
| 290 | raise errors.DriverError( |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 291 | "Failed to create ssh key pair, please make sure " |
| 292 | "'ssh-keygen' is installed: %s" % str(e)) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 293 | |
| 294 | # By default ssh-keygen will create a public key file |
| 295 | # by append .pub to the private key file name. Rename it |
| 296 | # to what's requested by public_key_path. |
| 297 | default_pub_key_path = "%s.pub" % private_key_path |
| 298 | try: |
| 299 | if default_pub_key_path != public_key_path: |
| 300 | os.rename(default_pub_key_path, public_key_path) |
| 301 | except OSError as e: |
| 302 | raise errors.DriverError( |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 303 | "Failed to rename %s to %s: %s" % (default_pub_key_path, |
| 304 | public_key_path, str(e))) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 305 | |
| 306 | logger.info("Created ssh private key (%s) and public key (%s)", |
| 307 | private_key_path, public_key_path) |
| 308 | |
| 309 | |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 310 | def VerifyRsaPubKey(rsa): |
| 311 | """Verify the format of rsa public key. |
| 312 | |
| 313 | Args: |
| 314 | rsa: content of rsa public key. It should follow the format of |
| 315 | ssh-rsa AAAAB3NzaC1yc2EA.... test@test.com |
| 316 | |
| 317 | Raises: |
| 318 | DriverError if the format is not correct. |
| 319 | """ |
| 320 | if not rsa or not all(ord(c) < 128 for c in rsa): |
| 321 | raise errors.DriverError( |
| 322 | "rsa key is empty or contains non-ascii character: %s" % rsa) |
| 323 | |
| 324 | elements = rsa.split() |
| 325 | if len(elements) != 3: |
| 326 | raise errors.DriverError("rsa key is invalid, wrong format: %s" % rsa) |
| 327 | |
| 328 | key_type, data, _ = elements |
| 329 | try: |
| 330 | binary_data = base64.decodestring(data) |
| 331 | # number of bytes of int type |
| 332 | int_length = 4 |
| 333 | # binary_data is like "7ssh-key..." in a binary format. |
| 334 | # The first 4 bytes should represent 7, which should be |
| 335 | # the length of the following string "ssh-key". |
| 336 | # And the next 7 bytes should be string "ssh-key". |
| 337 | # We will verify that the rsa conforms to this format. |
| 338 | # ">I" in the following line means "big-endian unsigned integer". |
| 339 | type_length = struct.unpack(">I", binary_data[:int_length])[0] |
| 340 | if binary_data[int_length:int_length + type_length] != key_type: |
| 341 | raise errors.DriverError("rsa key is invalid: %s" % rsa) |
| 342 | except (struct.error, binascii.Error) as e: |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 343 | raise errors.DriverError( |
| 344 | "rsa key is invalid: %s, error: %s" % (rsa, str(e))) |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 345 | |
| 346 | |
| 347 | class BatchHttpRequestExecutor(object): |
| 348 | """A helper class that executes requests in batch with retry. |
| 349 | |
| 350 | This executor executes http requests in a batch and retry |
| 351 | those that have failed. It iteratively updates the dictionary |
| 352 | self._final_results with latest results, which can be retrieved |
| 353 | via GetResults. |
| 354 | """ |
| 355 | |
| 356 | def __init__(self, |
| 357 | execute_once_functor, |
| 358 | requests, |
| 359 | retry_http_codes=None, |
| 360 | max_retry=None, |
| 361 | sleep=None, |
| 362 | backoff_factor=None, |
| 363 | other_retriable_errors=None): |
| 364 | """Initializes the executor. |
| 365 | |
| 366 | Args: |
| 367 | execute_once_functor: A function that execute requests in batch once. |
| 368 | It should return a dictionary like |
| 369 | {request_id: (response, exception)} |
| 370 | requests: A dictionary where key is request id picked by caller, |
| 371 | and value is a apiclient.http.HttpRequest. |
| 372 | retry_http_codes: A list of http codes to retry. |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 373 | max_retry: See utils.Retry. |
| 374 | sleep: See utils.Retry. |
| 375 | backoff_factor: See utils.Retry. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 376 | other_retriable_errors: A tuple of error types that should be retried |
| 377 | other than errors.HttpError. |
| 378 | """ |
| 379 | self._execute_once_functor = execute_once_functor |
| 380 | self._requests = requests |
| 381 | # A dictionary that maps request id to pending request. |
| 382 | self._pending_requests = {} |
| 383 | # A dictionary that maps request id to a tuple (response, exception). |
| 384 | self._final_results = {} |
| 385 | self._retry_http_codes = retry_http_codes |
| 386 | self._max_retry = max_retry |
| 387 | self._sleep = sleep |
| 388 | self._backoff_factor = backoff_factor |
| 389 | self._other_retriable_errors = other_retriable_errors |
| 390 | |
| 391 | def _ShoudRetry(self, exception): |
| 392 | """Check if an exception is retriable.""" |
| 393 | if isinstance(exception, self._other_retriable_errors): |
| 394 | return True |
| 395 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 396 | if (isinstance(exception, errors.HttpError) |
| 397 | and exception.code in self._retry_http_codes): |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 398 | return True |
| 399 | return False |
| 400 | |
| 401 | def _ExecuteOnce(self): |
| 402 | """Executes pending requests and update it with failed, retriable ones. |
| 403 | |
| 404 | Raises: |
| 405 | HasRetriableRequestsError: if some requests fail and are retriable. |
| 406 | """ |
| 407 | results = self._execute_once_functor(self._pending_requests) |
| 408 | # Update final_results with latest results. |
| 409 | self._final_results.update(results) |
| 410 | # Clear pending_requests |
| 411 | self._pending_requests.clear() |
| 412 | for request_id, result in results.iteritems(): |
| 413 | exception = result[1] |
| 414 | if exception is not None and self._ShoudRetry(exception): |
| 415 | # If this is a retriable exception, put it in pending_requests |
| 416 | self._pending_requests[request_id] = self._requests[request_id] |
| 417 | if self._pending_requests: |
| 418 | # If there is still retriable requests pending, raise an error |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 419 | # so that Retry will retry this function with pending_requests. |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 420 | raise errors.HasRetriableRequestsError( |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 421 | "Retriable errors: %s" % |
| 422 | [str(results[rid][1]) for rid in self._pending_requests]) |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 423 | |
| 424 | def Execute(self): |
| 425 | """Executes the requests and retry if necessary. |
| 426 | |
| 427 | Will populate self._final_results. |
| 428 | """ |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 429 | |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 430 | def _ShouldRetryHandler(exc): |
| 431 | """Check if |exc| is a retriable exception. |
| 432 | |
| 433 | Args: |
| 434 | exc: An exception. |
| 435 | |
| 436 | Returns: |
| 437 | True if exception is of type HasRetriableRequestsError; False otherwise. |
| 438 | """ |
| 439 | should_retry = isinstance(exc, errors.HasRetriableRequestsError) |
| 440 | if should_retry: |
| 441 | logger.info("Will retry failed requests.", exc_info=True) |
| 442 | logger.info("%s", exc) |
| 443 | return should_retry |
| 444 | |
| 445 | try: |
| 446 | self._pending_requests = self._requests.copy() |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 447 | Retry( |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame^] | 448 | _ShouldRetryHandler, |
| 449 | max_retries=self._max_retry, |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 450 | functor=self._ExecuteOnce, |
| 451 | sleep_multiplier=self._sleep, |
| 452 | retry_backoff_factor=self._backoff_factor) |
Keun Soo Yim | b293fdb | 2016-09-21 16:03:44 -0700 | [diff] [blame] | 453 | except errors.HasRetriableRequestsError: |
| 454 | logger.debug("Some requests did not succeed after retry.") |
| 455 | |
| 456 | def GetResults(self): |
| 457 | """Returns final results. |
| 458 | |
| 459 | Returns: |
| 460 | results, a dictionary in the following format |
| 461 | {request_id: (response, exception)} |
| 462 | request_ids are those from requests; response |
| 463 | is the http response for the request or None on error; |
| 464 | exception is an instance of DriverError or None if no error. |
| 465 | """ |
| 466 | return self._final_results |