DRIVERS-3032 Add retries to download attempts (#612)

Co-authored-by: Ezra Chung <[email protected]>

Author: blink1073

.evergreen/mongodl.py

@@ -28,6 +28,7 @@
 import sys
 import tarfile
 import textwrap
+import time
 import urllib.error
 import urllib.request
 import warnings
@@ -317,6 +318,26 @@ def mdb_version_rapid(version: str) -> bool:
     return tup[1] > 0
 
 
+class DownloadRetrier:
+    """Class that handles retry logic.  It performs exponential backoff with a maximum delay of 10 minutes between retry attempts."""
+
+    def __init__(self, retries: int) -> None:
+        self.retries = retries
+        self.attempt = 0
+        assert self.retries >= 0
+
+    def retry(self) -> bool:
+        if self.attempt >= self.retries:
+            return False
+        self.attempt += 1
+        LOGGER.warning(
+            f"Download attempt failed, retrying attempt {self.attempt} of {self.retries}"
+        )
+        ten_minutes = 600
+        time.sleep(min(2 ** (self.attempt - 1), ten_minutes))
+        return True
+
+
 class CacheDB:
     """
     Abstract a mongodl cache SQLite database.
@@ -832,6 +853,7 @@ def _dl_component(
     test: bool,
     no_download: bool,
     latest_build_branch: "str|None",
+    retries: int,
 ) -> ExpandResult:
     LOGGER.info(f"Download {component} {version}-{edition} for {target}-{arch}")
     if version in ("latest-build", "latest"):
@@ -859,12 +881,23 @@ def _dl_component(
                 cache, version, target, arch, edition, component
             )
 
+    # This must go to stdout to be consumed by the calling program.
+    print(dl_url)
+    LOGGER.info("Download url: %s", dl_url)
+
     if no_download:
-        # This must go to stdout to be consumed by the calling program.
-        print(dl_url)
         return None
-    cached = cache.download_file(dl_url).path
-    return _expand_archive(cached, out_dir, pattern, strip_components, test=test)
+
+    retrier = DownloadRetrier(retries)
+    while True:
+        try:
+            cached = cache.download_file(dl_url).path
+            return _expand_archive(
+                cached, out_dir, pattern, strip_components, test=test
+            )
+        except Exception:
+            if not retrier.retry():
+                raise
 
 
 def _pathjoin(items: "Iterable[str]") -> PurePath:
@@ -1145,6 +1178,7 @@ def main(argv=None):
         'download the with "--version=latest-build"',
         metavar="BRANCH_NAME",
     )
+    dl_grp.add_argument("--retries", help="The number of times to retry", default=0)
     args = parser.parse_args(argv)
     cache = Cache.open_in(args.cache_dir)
     cache.refresh_full_json()
@@ -1184,6 +1218,7 @@ def main(argv=None):
         test=args.test,
         no_download=args.no_download,
         latest_build_branch=args.latest_build_branch,
+        retries=int(args.retries),
     )
     if result is ExpandResult.Empty and args.empty_is_error:
         sys.exit(1)

.evergreen/mongosh_dl.py

@@ -23,7 +23,13 @@
 HERE = Path(__file__).absolute().parent
 sys.path.insert(0, str(HERE))
 from mongodl import LOGGER as DL_LOGGER
-from mongodl import SSL_CONTEXT, ExpandResult, _expand_archive, infer_arch
+from mongodl import (
+    SSL_CONTEXT,
+    DownloadRetrier,
+    ExpandResult,
+    _expand_archive,
+    infer_arch,
+)
 
 
 def _get_latest_version():
@@ -75,6 +81,7 @@ def _download(
     strip_components: int,
     test: bool,
     no_download: bool,
+    retries: int,
 ) -> int:
     LOGGER.info(f"Download {version} mongosh for {target}-{arch}")
     if version == "latest":
@@ -96,24 +103,31 @@ def _download(
     dl_url = f"https://downloads.mongodb.com/compass/mongosh-{version}-{target}-{arch}{suffix}"
     # This must go to stdout to be consumed by the calling program.
     print(dl_url)
+    LOGGER.info("Download url: %s", dl_url)
 
     if no_download:
         return ExpandResult.Okay
 
     req = urllib.request.Request(dl_url)
-    resp = urllib.request.urlopen(req)
-
-    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as fp:
-        buf = resp.read(1024 * 1024 * 4)
-        while buf:
-            fp.write(buf)
-            buf = resp.read(1024 * 1024 * 4)
-        fp.close()
-        resp = _expand_archive(
-            Path(fp.name), out_dir, pattern, strip_components, test=test
-        )
-        os.remove(fp.name)
-    return resp
+    retrier = DownloadRetrier(retries)
+    while True:
+        try:
+            resp = urllib.request.urlopen(req)
+            with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as fp:
+                four_mebibytes = 1024 * 1024 * 4
+                buf = resp.read(four_mebibytes)
+                while buf:
+                    fp.write(buf)
+                    buf = resp.read(four_mebibytes)
+                fp.close()
+                resp = _expand_archive(
+                    Path(fp.name), out_dir, pattern, strip_components, test=test
+                )
+                os.remove(fp.name)
+            return resp
+        except Exception:
+            if not retrier.retry():
+                raise
 
 
 def main(argv=None):
@@ -189,6 +203,7 @@ def main(argv=None):
         help="Do not extract or place any files/directories. "
         "Only print what will be extracted without placing any files.",
     )
+    dl_grp.add_argument("--retries", help="The number of times to retry", default=0)
     args = parser.parse_args(argv)
 
     target = args.target
@@ -214,6 +229,7 @@ def main(argv=None):
         strip_components=args.strip_components,
         test=args.test,
         no_download=args.no_download,
+        retries=int(args.retries),
     )
     if result is ExpandResult.Empty:
         sys.exit(1)

.evergreen/orchestration/drivers_orchestration.py

@@ -194,7 +194,7 @@ def run(opts):
     version = opts.version
     cache_dir = DRIVERS_TOOLS / ".local/cache"
     cache_dir_str = cache_dir.as_posix()
-    default_args = f"--out {mdb_binaries_str} --cache-dir {cache_dir_str}"
+    default_args = f"--out {mdb_binaries_str} --cache-dir {cache_dir_str} --retries 5"
     if opts.quiet:
         default_args += " -q"
     elif opts.verbose:
@@ -243,7 +243,7 @@ def run(opts):
         expansion_sh.write_text(crypt_text.replace(": ", "="))
 
     # Download mongosh
-    args = f"--out {mdb_binaries_str} --strip-path-components 2"
+    args = f"--out {mdb_binaries_str} --strip-path-components 2 --retries 5"
     if opts.verbose:
         args += " -v"
     elif opts.quiet:

.evergreen/tests/test-cli.sh

@@ -24,24 +24,24 @@ else
   DOWNLOAD_DIR=$(cygpath -m $DOWNLOAD_DIR)
 fi
 
-./mongodl --edition enterprise --version 7.0 --component archive --test
-./mongodl --edition enterprise --version 7.0 --component cryptd --out ${DOWNLOAD_DIR} --strip-path-components 1
+./mongodl --edition enterprise --version 7.0 --component archive --test --retries 5
+./mongodl --edition enterprise --version 7.0 --component cryptd --out ${DOWNLOAD_DIR} --strip-path-components 1 --retries 5
 ./mongosh-dl --no-download
 ./mongosh-dl --version 2.1.1 --no-download
 
 export PATH="${DOWNLOAD_DIR}/bin:$PATH"
 if [ "${OS:-}" != "Windows_NT" ]; then
-  ./mongosh-dl --version 2.1.1 --out ${DOWNLOAD_DIR} --strip-path-components 1
+  ./mongosh-dl --version 2.1.1 --out ${DOWNLOAD_DIR} --strip-path-components 1 --retries 5
   chmod +x ./mongodl_test/bin/mongosh
   ./mongodl_test/bin/mongosh --version
 else
-  ./mongosh-dl --version 2.1.1 --out ${DOWNLOAD_DIR} --strip-path-components 1
+  ./mongosh-dl --version 2.1.1 --out ${DOWNLOAD_DIR} --strip-path-components 1 --retries 5
 fi
 
 # Ensure that we can use a downloaded mongodb directory.
 rm -rf ${DOWNLOAD_DIR}
 bash install-cli.sh "$(pwd)/orchestration"
-./mongodl --edition enterprise --version 7.0 --component archive --out ${DOWNLOAD_DIR} --strip-path-components 2
+./mongodl --edition enterprise --version 7.0 --component archive --out ${DOWNLOAD_DIR} --strip-path-components 2 --retries 5
 ./orchestration/drivers-orchestration run --existing-binaries-dir=${DOWNLOAD_DIR}
 ${DOWNLOAD_DIR}/mongod --version | grep v7.0
 ./orchestration/drivers-orchestration stop
@@ -56,19 +56,19 @@ fi
 export VALIDATE_DISTROS=1
 ./mongodl --list
 ./mongodl --edition enterprise --version 7.0.6 --component archive --no-download
-./mongodl --edition enterprise --version 3.6 --component archive --test
-./mongodl --edition enterprise --version 4.0 --component archive --test
-./mongodl --edition enterprise --version 4.2 --component archive --test
-./mongodl --edition enterprise --version 4.4 --component archive --test
-./mongodl --edition enterprise --version 5.0 --component archive --test
-./mongodl --edition enterprise --version 6.0 --component crypt_shared --test
-./mongodl --edition enterprise --version 8.0 --component archive --test
-./mongodl --edition enterprise --version rapid --component archive --test
-./mongodl --edition enterprise --version latest --component archive --out ${DOWNLOAD_DIR}
-./mongodl --edition enterprise --version latest-build --component archive --test
-./mongodl --edition enterprise --version latest-release --component archive --test
-./mongodl --edition enterprise --version v6.0-perf --component cryptd --test
-./mongodl --edition enterprise --version v8.0-perf --component cryptd --test
+./mongodl --edition enterprise --version 3.6 --component archive --test --retries 5
+./mongodl --edition enterprise --version 4.0 --component archive --test --retries 5
+./mongodl --edition enterprise --version 4.2 --component archive --test --retries 5
+./mongodl --edition enterprise --version 4.4 --component archive --test --retries 5
+./mongodl --edition enterprise --version 5.0 --component archive --test --retries 5
+./mongodl --edition enterprise --version 6.0 --component crypt_shared --test --retries 5
+./mongodl --edition enterprise --version 8.0 --component archive --test --retries 5
+./mongodl --edition enterprise --version rapid --component archive --test --retries 5
+./mongodl --edition enterprise --version latest --component archive --out ${DOWNLOAD_DIR} --retries 5
+./mongodl --edition enterprise --version latest-build --component archive --test --retries 5
+./mongodl --edition enterprise --version latest-release --component archive --test --retries 5
+./mongodl --edition enterprise --version v6.0-perf --component cryptd --test --retries 5
+./mongodl --edition enterprise --version v8.0-perf --component cryptd --test --retries 5
 
 popd
 make -C ${DRIVERS_TOOLS} test