[CIFuzz] Retry HTTP requests on certain errors (#4516)

Skia has seen some unhandled connection reset errors.
When we see these errors when downloading old OSS-Fuzz build, retry downloading it, since the error is likely transient.

2 files changed, 77 insertions, 7 deletions

diff --git a/infra/cifuzz/fuzz_target.py b/infra/cifuzz/fuzz_target.py
index 7050ae2d..d469d855 100644
--- a/infra/cifuzz/fuzz_target.py
+++ b/infra/cifuzz/fuzz_target.py
@@ -20,6 +20,7 @@ import stat
 import subprocess
 import sys
 import tempfile
+import time
 import urllib.error
 import urllib.request
 import zipfile
@@ -367,11 +368,44 @@ class FuzzTarget:
     return download_and_unpack_zip(corpus_url, corpus_dir)
 
 
-def download_and_unpack_zip(http_url, out_dir):
+def download_url(url, filename, num_retries=3):
+  """Downloads the file located at |url|, using HTTP to |filename|.
+
+  Args:
+    url: A url to a file to download.
+    filename: The path the file should be downloaded to.
+    num_retries: The number of times to retry the download on
+       ConnectionResetError.
+
+  Returns:
+    True on success.
+  """
+  sleep_time = 1
+
+  for _ in range(num_retries):
+    try:
+      urllib.request.urlretrieve(url, filename)
+      return True
+    except urllib.error.HTTPError:
+      # In these cases, retrying probably wont work since the error probably
+      # means there is nothing at the URL to download.
+      logging.error('Unable to download from: %s.', url)
+      return False
+    except ConnectionResetError:
+      # These errors are more likely to be transient. Retry.
+      pass
+    time.sleep(sleep_time)
+
+  logging.error('Failed to download %s, %d times.', url, num_retries)
+
+  return False
+
+
+def download_and_unpack_zip(url, out_dir):
   """Downloads and unpacks a zip file from an http url.
 
   Args:
-    http_url: A url to the zip file to be downloaded and unpacked.
+    url: A url to the zip file to be downloaded and unpacked.
     out_dir: The path where the zip file should be extracted to.
 
   Returns:
@@ -384,17 +418,15 @@ def download_and_unpack_zip(http_url, out_dir):
   # Gives the temporary zip file a unique identifier in the case that
   # that download_and_unpack_zip is done in parallel.
   with tempfile.NamedTemporaryFile(suffix='.zip') as tmp_file:
-    try:
-      urllib.request.urlretrieve(http_url, tmp_file.name)
-    except urllib.error.HTTPError:
-      logging.error('Unable to download build from: %s.', http_url)
+    result = download_url(url, tmp_file.name)
+    if not result:
       return None
 
     try:
       with zipfile.ZipFile(tmp_file.name, 'r') as zip_file:
         zip_file.extractall(out_dir)
     except zipfile.BadZipFile:
-      logging.error('Error unpacking zip from %s. Bad Zipfile.', http_url)
+      logging.error('Error unpacking zip from %s. Bad Zipfile.', url)
       return None
   return out_dir
 
diff --git a/infra/cifuzz/fuzz_target_test.py b/infra/cifuzz/fuzz_target_test.py
index 04aac4c9..00422659 100644
--- a/infra/cifuzz/fuzz_target_test.py
+++ b/infra/cifuzz/fuzz_target_test.py
@@ -18,6 +18,7 @@ import sys
 import tempfile
 import unittest
 import unittest.mock
+import urllib.error
 
 import parameterized
 from pyfakefs import fake_filesystem_unittest
@@ -319,6 +320,43 @@ class DownloadOSSFuzzBuildDirIntegrationTests(unittest.TestCase):
     self.assertIsNone(test_target.download_oss_fuzz_build())
 
 
+class DownloadUrlTest(unittest.TestCase):
+  """Test that download_url works."""
+  URL = 'example.com/file'
+  FILE_PATH = '/tmp/file'
+
+  @unittest.mock.patch('time.sleep')
+  @unittest.mock.patch('urllib.request.urlretrieve', return_value=True)
+  def test_download_url_no_error(self, mocked_urlretrieve, _):
+    """Tests that download_url works when there is no error."""
+    self.assertTrue(fuzz_target.download_url(self.URL, self.FILE_PATH))
+    self.assertEqual(1, mocked_urlretrieve.call_count)
+
+  @unittest.mock.patch('time.sleep')
+  @unittest.mock.patch('logging.error')
+  @unittest.mock.patch('urllib.request.urlretrieve',
+                       side_effect=urllib.error.HTTPError(
+                           None, None, None, None, None))
+  def test_download_url_http_error(self, mocked_urlretrieve, mocked_error, _):
+    """Tests that download_url doesn't retry when there is an HTTP error."""
+    self.assertFalse(fuzz_target.download_url(self.URL, self.FILE_PATH))
+    mocked_error.assert_called_with('Unable to download from: %s.', self.URL)
+    self.assertEqual(1, mocked_urlretrieve.call_count)
+
+  @unittest.mock.patch('time.sleep')
+  @unittest.mock.patch('logging.error')
+  @unittest.mock.patch('urllib.request.urlretrieve',
+                       side_effect=ConnectionResetError)
+  def test_download_url_connection_error(self, mocked_urlretrieve, mocked_error,
+                                         mocked_sleep):
+    """Tests that download_url doesn't retry when there is an HTTP error."""
+    self.assertFalse(fuzz_target.download_url(self.URL, self.FILE_PATH))
+    self.assertEqual(3, mocked_urlretrieve.call_count)
+    self.assertEqual(3, mocked_sleep.call_count)
+    mocked_error.assert_called_with('Failed to download %s, %d times.',
+                                    self.URL, 3)
+
+
 class DownloadAndUnpackZipUnitTests(unittest.TestCase):
   """Test the download and unpack functionality in the fuzz_target module."""