1515
1616import datetime
1717
18+ from clusterfuzz ._internal .base import memoize
19+ from clusterfuzz ._internal .config import local_config
1820from clusterfuzz ._internal .datastore import data_types
1921from clusterfuzz ._internal .datastore import ndb_utils
2022from clusterfuzz ._internal .metrics import logs
@@ -33,6 +35,13 @@ def _get_datetime_now():
3335_UTASK_PSEUDO_TASKS = {'uworker_main' , 'postprocess' , 'preprocess' }
3436
3537
38+ @memoize .wrap (memoize .FifoInMemory (1 ))
39+ def optin_to_task_rate_limiting ():
40+ enabled = local_config .ProjectConfig ().get ('rate_limit_tasks.enabled' , False )
41+ logs .info (f'Using async HTTP: { enabled } .' )
42+ return enabled
43+
44+
3645class TaskRateLimiter :
3746 """Rate limiter for tasks. This limits tasks to 100 erroneous runs or 2000
3847 succesful runs in 6 hours. It keeps track of task completion when record_task
@@ -51,6 +60,8 @@ def __str__(self):
5160
5261 def record_task (self , success : bool ) -> None :
5362 """Records a task and whether it completed succesfully."""
63+ if not optin_to_task_rate_limiting ():
64+ return
5465 if self .task_name in _UTASK_PSEUDO_TASKS :
5566 # Don't rate limit these fake uworker tasks.
5667 return
@@ -67,6 +78,8 @@ def record_task(self, success: bool) -> None:
6778
6879 def is_rate_limited (self ) -> bool :
6980 """Checks if the given task is rate limited."""
81+ if not optin_to_task_rate_limiting ():
82+ return False
7083 if self .task_name in _UTASK_PSEUDO_TASKS :
7184 # Don't rate limit these fake tasks.
7285 return False
0 commit comments