diff --git a/src/clusterfuzz/_internal/base/tasks/task_rate_limiting.py b/src/clusterfuzz/_internal/base/tasks/task_rate_limiting.py index f12e775f74..4d68a43436 100644 --- a/src/clusterfuzz/_internal/base/tasks/task_rate_limiting.py +++ b/src/clusterfuzz/_internal/base/tasks/task_rate_limiting.py @@ -37,7 +37,6 @@ class TaskRateLimiter: """Rate limiter for tasks. This limits tasks to 100 erroneous runs or 2000 succesful runs in 6 hours. It keeps track of task completion when record_task is called at the end of every task.""" - TASK_RATE_LIMIT_WINDOW = datetime.timedelta(hours=6) TASK_RATE_LIMIT_MAX_ERRORS = 100 # TODO(metzman): Reevaluate this number, it's probably too high. TASK_RATE_LIMIT_MAX_COMPLETIONS = 2000 @@ -74,7 +73,9 @@ def is_rate_limited(self) -> bool: if environment.get_value('COMMAND_OVERRIDE'): # A user wants to run this task. return False - window_start = _get_datetime_now() - self.TASK_RATE_LIMIT_WINDOW + window_start = ( + _get_datetime_now() - + data_types.WindowRateLimitTask.TASK_RATE_LIMIT_WINDOW) query = data_types.WindowRateLimitTask.query( data_types.WindowRateLimitTask.task_name == self.task_name, data_types.WindowRateLimitTask.task_argument == self.task_argument, diff --git a/src/clusterfuzz/_internal/datastore/data_types.py b/src/clusterfuzz/_internal/datastore/data_types.py index 28f765b606..fdb16d16ba 100644 --- a/src/clusterfuzz/_internal/datastore/data_types.py +++ b/src/clusterfuzz/_internal/datastore/data_types.py @@ -13,6 +13,7 @@ # limitations under the License. """Classes for objects stored in the datastore.""" +import datetime import re from google.cloud import ndb @@ -1119,13 +1120,21 @@ class WindowRateLimitTask(Model): it will have a different lifecycle (it's not needed after the window completes). This should have a TTL as TASK_RATE_LIMIT_WINDOW in task_rate_limiting.py (6 hours).""" - # TODO(metzman): Consider using task_id. + TASK_RATE_LIMIT_WINDOW = datetime.timedelta(hours=6) + timestamp = ndb.DateTimeProperty(auto_now_add=True, indexed=True) + # Only use this for TTL. It should only be saved to by ClusterFuzz, not read. + ttl_expiry_timestamp = ndb.DateTimeProperty() + # TODO(metzman): Consider using task_id. task_name = ndb.StringProperty(indexed=True) task_argument = ndb.StringProperty(indexed=True) job_name = ndb.StringProperty(indexed=True) status = ndb.StringProperty(choices=[TaskState.ERROR, TaskState.FINISHED]) + def _pre_put_hook(self): + self.ttl_expiry_timestamp = ( + datetime.datetime.now() + self.TASK_RATE_LIMIT_WINDOW) + class BuildMetadata(Model): """Metadata associated with a particular archived build.""" diff --git a/src/clusterfuzz/_internal/tests/core/base/tasks/task_rate_limiting_test.py b/src/clusterfuzz/_internal/tests/core/base/tasks/task_rate_limiting_test.py index 5b82c3f091..9f2ecde40f 100644 --- a/src/clusterfuzz/_internal/tests/core/base/tasks/task_rate_limiting_test.py +++ b/src/clusterfuzz/_internal/tests/core/base/tasks/task_rate_limiting_test.py @@ -91,7 +91,7 @@ def test_is_rate_limited_old_tasks(self): """Test is_rate_limited() with old tasks outside the window.""" # Add tasks outside the time window. window_start = ( - self.now - task_rate_limiting.TaskRateLimiter.TASK_RATE_LIMIT_WINDOW) + self.now - data_types.WindowRateLimitTask.TASK_RATE_LIMIT_WINDOW) self._create_n_tasks( task_rate_limiting.TaskRateLimiter.TASK_RATE_LIMIT_MAX_COMPLETIONS + 1, timestamp=window_start - datetime.timedelta(minutes=10))