Add a field for TTL in WindowTaskRateLimit. (#4385)

Author: jonathanmetzman

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,

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."""

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))