diff --git a/src/clusterfuzz/_internal/base/concurrency.py b/src/clusterfuzz/_internal/base/concurrency.py index 52ce2d4b788..3edb6463bc8 100644 --- a/src/clusterfuzz/_internal/base/concurrency.py +++ b/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: diff --git a/src/clusterfuzz/_internal/base/feature_flags.py b/src/clusterfuzz/_internal/base/feature_flags.py index 83323f5ad64..5dfeccae1c1 100644 --- a/src/clusterfuzz/_internal/base/feature_flags.py +++ b/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): diff --git a/src/clusterfuzz/_internal/bot/tasks/utasks/corpus_pruning_task.py b/src/clusterfuzz/_internal/bot/tasks/utasks/corpus_pruning_task.py index f76cc3f5fc9..bd803079963 100644 --- a/src/clusterfuzz/_internal/bot/tasks/utasks/corpus_pruning_task.py +++ b/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, diff --git a/src/clusterfuzz/_internal/google_cloud_utils/storage.py b/src/clusterfuzz/_internal/google_cloud_utils/storage.py index 6ab7abb6118..aad16899c4b 100644 --- a/src/clusterfuzz/_internal/google_cloud_utils/storage.py +++ b/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.')