DRIVERS-3032 Add retries to download attempts

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
@@ -832,6 +833,7 @@ def _dl_component(
     test: bool,
     no_download: bool,
     latest_build_branch: "str|None",
+    retry: int,
 ) -> ExpandResult:
     LOGGER.info(f"Download {component} {version}-{edition} for {target}-{arch}")
     if version in ("latest-build", "latest"):
@@ -863,7 +865,20 @@ def _dl_component(
         # This must go to stdout to be consumed by the calling program.
         print(dl_url)
         return None
-    cached = cache.download_file(dl_url).path
+    retries = retry
+    while True:
+        try:
+            cached = cache.download_file(dl_url).path
+            break
+        except ConnectionResetError:
+            retries -= 1
+            if retries == 0:
+                raise
+            attempt = retry - retries
+            LOGGER.warning(
+                f"Download attempt failed, retry attempt {attempt} of {retry}"
+            )
+            time.sleep(attempt**2)
     return _expand_archive(cached, out_dir, pattern, strip_components, test=test)
 
 
@@ -1145,6 +1160,7 @@ def main(argv=None):
         'download the with "--version=latest-build"',
         metavar="BRANCH_NAME",
     )
+    dl_grp.add_argument("--retry", 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 +1200,7 @@ def main(argv=None):
         test=args.test,
         no_download=args.no_download,
         latest_build_branch=args.latest_build_branch,
+        retry=int(args.retry),
     )
     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,
+    retry: int,
 ) -> int:
     LOGGER.info(f"Download {version} mongosh for {target}-{arch}")
     if version == "latest":
@@ -101,7 +103,20 @@ def _download(
         return ExpandResult.Okay
 
     req = urllib.request.Request(dl_url)
-    resp = urllib.request.urlopen(req)
+    retries = retry
+    while True:
+        try:
+            resp = urllib.request.urlopen(req)
+            break
+        except ConnectionResetError:
+            retries -= 1
+            if retries == 0:
+                raise
+            attempt = retry - retries
+            LOGGER.warning(
+                f"Download attempt failed, retry attempt {attempt} of {retry}"
+            )
+            time.sleep(attempt**2)
 
     with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as fp:
         buf = resp.read(1024 * 1024 * 4)
@@ -189,6 +204,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("--retry", help="The number of times to retry", default=0)
     args = parser.parse_args(argv)
 
     target = args.target
@@ -214,6 +230,7 @@ def main(argv=None):
         strip_components=args.strip_components,
         test=args.test,
         no_download=args.no_download,
+        retry=int(args.retry),
     )
     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: