[Corpus] Use threads for datastore operations under targets feature flag

src/clusterfuzz/_internal/base/concurrency.py

@@ -22,13 +22,13 @@
 
 
 @contextlib.contextmanager
-def make_pool(pool_size=POOL_SIZE, max_pool_size=None):
+def make_pool(pool_size=POOL_SIZE, max_pool_size=None, use_threads=False):
   """Returns a pool that can (usually) execute tasks concurrently."""
   if max_pool_size is not None:
     pool_size = min(pool_size, max_pool_size)
 
   # Don't use processes on Windows and unittests to avoid hangs.
-  if (environment.get_value('PY_UNITTESTS') or
+  if (use_threads or environment.get_value('PY_UNITTESTS') or
       environment.platform() == 'WINDOWS'):
     yield futures.ThreadPoolExecutor(pool_size)
   else:

src/clusterfuzz/_internal/base/feature_flags.py

@@ -43,6 +43,7 @@ class FeatureFlags(Enum):
   SWARMING_MAX_PENDING_TASKS = 'swarming_max_pending_tasks'
 
   ENABLE_FUZZ_FOR_BOTS = 'enable_fuzz_for_bots'
+  STORAGE_THREADED_OPS_FUZZ_TARGETS = 'storage_threaded_ops_fuzz_targets'
 
   @property
   def flag(self):

src/clusterfuzz/_internal/bot/tasks/utasks/corpus_pruning_task.py

@@ -28,6 +28,7 @@
 from google.cloud import ndb
 from google.protobuf import timestamp_pb2
 
+from clusterfuzz._internal.base import feature_flags
 from clusterfuzz._internal.base import utils
 from clusterfuzz._internal.bot.fuzzers import engine_common
 from clusterfuzz._internal.bot.fuzzers import options
@@ -1213,6 +1214,23 @@ def _utask_preprocess(fuzzer_name, job_type, uworker_env):
     setup_input.global_blacklisted_functions.extend(
         leak_blacklist.get_global_blacklisted_functions())
 
+  if uworker_env is None:
+    uworker_env = {}
+
+  threaded_ops_flag = (
+      feature_flags.FeatureFlags.STORAGE_THREADED_OPS_FUZZ_TARGETS)
+  if threaded_ops_flag.enabled:
+    threaded_ops_targets = threaded_ops_flag.string_value
+    if threaded_ops_targets:
+      allowed_targets = [
+          t.strip() for t in threaded_ops_targets.split(',') if t.strip()
+      ]
+      if fuzzer_name in allowed_targets:
+        uworker_env['USE_THREADED_STORAGE_OPS'] = 'True'
+        # TODO(paulovlb): Remove this once fixed.
+        logs.info(f'[Corpus Fix] Enabled threaded storage ops for target: '
+                  f'{fuzzer_name}')
+
   logs.info('done preprocess')
   return uworker_msg_pb2.Input(  # pylint: disable=no-member
       job_type=job_type,

src/clusterfuzz/_internal/google_cloud_utils/storage.py

@@ -1361,11 +1361,27 @@ def _error_tolerant_delete_signed_url(url: str):
     logs.warning(f'Failed to delete: {url}')
 
 
+def _should_use_threads(urls):
+  """Returns True if threads should be used for operations on |urls|."""
+  if environment.get_value('USE_THREADED_STORAGE_OPS'):
+    if urls:
+      # TODO(paulovlb): Remove this once fixed.
+      logs.info(
+          f'[Corpus Fix] Using threads for GCS operations on URL: {urls[0]}')
+    return True
+  return False
+
+
 def upload_signed_urls(signed_urls: List[str], files: List[str]) -> List[bool]:
+  """Uploads files to signed URLs."""
   if not signed_urls:
     return []
   logs.info('Uploading URLs.')
-  with concurrency.make_pool(_POOL_SIZE) as pool:
+  use_threads = _should_use_threads(signed_urls)
+  if use_threads:
+    # TODO(paulovlb): Remove this once fixed.
+    logs.info('[Corpus Fix] Using thread pool for uploading URLs.')
+  with concurrency.make_pool(_POOL_SIZE, use_threads=use_threads) as pool:
     result = list(
         pool.map(_error_tolerant_upload_signed_url, zip(signed_urls, files)))
   logs.info('Done uploading URLs.')
@@ -1394,7 +1410,11 @@ def download_signed_urls(signed_urls: List[str],
   urls_and_filepaths = list(zip(signed_urls, filepaths))
 
   def synchronous_download_urls(urls_and_filepaths):
-    with concurrency.make_pool(_POOL_SIZE) as pool:
+    use_threads = _should_use_threads(signed_urls)
+    if use_threads:
+      # TODO(paulovlb): Remove this once fixed.
+      logs.info('[Corpus Fix] Using thread pool for downloading URLs.')
+    with concurrency.make_pool(_POOL_SIZE, use_threads=use_threads) as pool:
       return list(
           pool.map(_error_tolerant_download_signed_url_to_file,
                    urls_and_filepaths))
@@ -1413,10 +1433,15 @@ def synchronous_download_urls(urls_and_filepaths):
 
 
 def delete_signed_urls(urls):
+  """Deletes signed URLs."""
   if not urls:
     return
   logs.info('Deleting URLs.')
-  with concurrency.make_pool(_POOL_SIZE) as pool:
+  use_threads = _should_use_threads(urls)
+  if use_threads:
+    # TODO(paulovlb): Remove this once fixed.
+    logs.info('[Corpus Fix] Using thread pool for deleting URLs.')
+  with concurrency.make_pool(_POOL_SIZE, use_threads=use_threads) as pool:
     pool.map(_error_tolerant_delete_signed_url, urls)
   logs.info('Done deleting URLs.')