DRIVERS-3032 Add retries to download attempts

.evergreen/mongodl.py

@@ -28,6 +28,7 @@
 import sys
 import tarfile
 import textwrap
+import time
 import urllib.error
 import urllib.request
 import warnings
@@ -832,6 +833,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 +861,29 @@ 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)
+
+    remaining = retries
+    while True:
+        try:
+            cached = cache.download_file(dl_url).path
+            return _expand_archive(
+                cached, out_dir, pattern, strip_components, test=test
+            )
+        except Exception:
+            remaining -= 1
+            if remaining < 1:
+                raise
+            attempt = retries - remaining
+            LOGGER.warning(
+                f"Download attempt failed, retry attempt {attempt} of {retries}"
+            )
+            time.sleep(attempt**2)
 
 
 def _pathjoin(items: "Iterable[str]") -> PurePath:
@@ -1145,6 +1164,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 +1204,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

@@ -14,6 +14,7 @@
 import subprocess
 import sys
 import tempfile
+import time
 import urllib.request
 from pathlib import Path
 
@@ -75,6 +76,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 +98,36 @@ 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
+    remaining = retries
+    while True:
+        try:
+            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
+        except Exception:
+            remaining -= 1
+            if remaining < 1:
+                raise
+            attempt = retries - remaining
+            LOGGER.warning(
+                f"Download attempt failed, retry attempt {attempt} of {retries}"
+            )
+            time.sleep(attempt**2)
 
 
 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 3"
     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 3"
     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 3
+./mongodl --edition enterprise --version 7.0 --component cryptd --out ${DOWNLOAD_DIR} --strip-path-components 1 --retries 3
 ./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 3
   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 3
 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 3
 ./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 3
+./mongodl --edition enterprise --version 4.0 --component archive --test --retries 3
+./mongodl --edition enterprise --version 4.2 --component archive --test --retries 3
+./mongodl --edition enterprise --version 4.4 --component archive --test --retries 3
+./mongodl --edition enterprise --version 5.0 --component archive --test --retries 3
+./mongodl --edition enterprise --version 6.0 --component crypt_shared --test --retries 3
+./mongodl --edition enterprise --version 8.0 --component archive --test --retries 3
+./mongodl --edition enterprise --version rapid --component archive --test --retries 3
+./mongodl --edition enterprise --version latest --component archive --out ${DOWNLOAD_DIR} --retries 3
+./mongodl --edition enterprise --version latest-build --component archive --test --retries 3
+./mongodl --edition enterprise --version latest-release --component archive --test --retries 3
+./mongodl --edition enterprise --version v6.0-perf --component cryptd --test --retries 3
+./mongodl --edition enterprise --version v8.0-perf --component cryptd --test --retries 3
 
 popd
 make -C ${DRIVERS_TOOLS} test