perf: stop reading every task_result payload to prune keys Redis already expires

modelq/app/base.py

@@ -386,44 +386,104 @@ def requeue_stuck_processing_tasks(self, threshold: float = 180.0):
                     # Remove from processing set
                     self.redis_client.srem("processing_tasks", task_id)
 
-    def prune_old_task_results(self, older_than_seconds: int = None):
+    SCAN_BATCH = 500
+
+    def prune_old_task_results(self, older_than_seconds: int = None) -> int:
         """
-        Deletes task result keys (stored with the prefix 'task_result:') whose
-        finished_at (or started_at if finished_at is not available) timestamp is older
-        than `older_than_seconds`. In addition, it also removes the corresponding
-        task key (stored with the prefix 'task:').
+        Deletes `task_result:*` keys (and their `task:*` twin) that have lost
+        their TTL and are older than `older_than_seconds`. Returns the count.
+
+        Expiry is Redis's job. Every task_result key is written with a TTL, so a
+        key that still has one needs nothing from us -- this only has to catch
+        keys whose TTL went missing (a write path that forgot `ex=`, a
+        RENAME/RESTORE that dropped it), which would otherwise live forever.
+
+        The cost is in deciding *which* keys to read, not in the delete. TTL is
+        an 8-byte reply and is pipelined, so a healthy keyspace is walked without
+        transferring a single payload; only keys already known to be broken get
+        read. The previous version GET the JSON of every key to compare one
+        timestamp: on a production shard that was 196.7M SCAN + 228.8M GET over
+        26 days -- ~2.2 hours of blocked event loop and 13.5TB of network output
+        -- to issue 2 deletes.
+
+        Bulk-reading with MGET would make this worse, not better. task_result
+        payloads reach 7MB, so a batched read builds one huge client output
+        buffer, and that is what pushes RSS past the container limit and gets
+        redis-server OOM-killed.
         """
         if older_than_seconds is None:
             older_than_seconds = self.TASK_RESULT_RETENTION
 
         now = time.time()
-        keys_deleted = 0
-
-        # Use scan_iter to avoid blocking Redis
-        for key in self.redis_client.scan_iter("task_result:*"):
-            try:
-                task_json = self.redis_client.get(key)
-                if not task_json:
+        pruned = 0
+        batch = []
+
+        def flush(batch):
+            """TTL the batch; only keys missing one are worth reading."""
+            if not batch:
+                return 0
+            pipe = self.redis_client.pipeline(transaction=False)
+            for key in batch:
+                pipe.ttl(key)
+            ttls = pipe.execute()
+
+            # -1 == exists with no expiry (leaked). -2 == already gone.
+            # Anything with a TTL is Redis's problem, not ours.
+            orphans = [key for key, ttl in zip(batch, ttls) if ttl == -1]
+            if not orphans:
+                return 0
+
+            # Only now do we read values, and only for the broken keys.
+            pipe = self.redis_client.pipeline(transaction=False)
+            for key in orphans:
+                pipe.get(key)
+            blobs = pipe.execute()
+
+            expired, keep = [], []
+            for key, blob in zip(orphans, blobs):
+                key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+                try:
+                    task_data = json.loads(blob) if blob else {}
+                except Exception as e:
+                    logger.error(f"Error parsing {key_str}: {e}")
+                    keep.append(key)
                     continue
-                task_data = json.loads(task_json)
-                # Use finished_at if available; otherwise fallback to started_at
-                timestamp = task_data.get("finished_at") or task_data.get("started_at")
-                if timestamp and (now - timestamp > older_than_seconds):
-                    # Delete the task_result key
-                    self.redis_client.delete(key)
-                    # Extract the task id from the key and delete the corresponding task key.
-                    key_str = key.decode("utf-8") if isinstance(key, bytes) else key
-                    task_id = key_str.split("task_result:")[-1]
-                    task_key = f"task:{task_id}"
-                    self.redis_client.delete(task_key)
-                    keys_deleted += 1
-                    logger.info(f"Deleted old keys: {key_str} and {task_key}")
-            except Exception as e:
+                stamp = task_data.get("finished_at") or task_data.get("started_at")
+                if stamp and (now - stamp > older_than_seconds):
+                    expired.append((key, key_str.split("task_result:")[-1]))
+                else:
+                    keep.append(key)
+
+            pipe = self.redis_client.pipeline(transaction=False)
+            for key, task_id in expired:
+                # UNLINK, not DELETE: frees multi-MB payloads off the main thread.
+                pipe.unlink(key)
+                pipe.unlink(f"task:{task_id}")
+            for key in keep:
+                # Not old enough to drop, but it must not live forever.
+                pipe.expire(key, older_than_seconds)
+            pipe.execute()
+
+            for _, task_id in expired:
+                logger.info(f"Pruned untracked task_result:{task_id} and task:{task_id}")
+            for key in keep:
                 key_str = key.decode("utf-8") if isinstance(key, bytes) else key
-                logger.error(f"Error processing key {key_str}: {e}")
+                logger.warning(f"task_result key had no TTL, set to {older_than_seconds}s: {key_str}")
+            return len(expired)
+
+        try:
+            for key in self.redis_client.scan_iter("task_result:*", count=self.SCAN_BATCH):
+                batch.append(key)
+                if len(batch) >= self.SCAN_BATCH:
+                    pruned += flush(batch)
+                    batch = []
+            pruned += flush(batch)
+        except Exception as e:
+            logger.error(f"Error scanning task_result keys: {e}")
 
-        if keys_deleted:
-            logger.info(f"Pruned {keys_deleted} task(s) older than {older_than_seconds} seconds.")
+        if pruned:
+            logger.info(f"Pruned {pruned} task(s) older than {older_than_seconds} seconds.")
+        return pruned
 
     def update_server_status(self, status: str):
         """
@@ -915,7 +975,12 @@ def _pruning_loop(self):
             with self._guarded_iteration("pruning"):
                 self.prune_inactive_servers(timeout_seconds=self.PRUNE_TIMEOUT)
                 self.requeue_stuck_processing_tasks(threshold=180)
-                self.prune_old_task_results(older_than_seconds=self.TASK_RESULT_RETENTION)
+                # prune_old_task_results() is deliberately NOT called here. Every
+                # task_result key is written with a TTL, so Redis expires it on
+                # its own; running the scan every PRUNE_CHECK_INTERVAL only
+                # re-read the whole keyspace. Measured on a production shard over
+                # 26 days: 196.7M SCAN + 228.8M GET to issue 2 DELs. Call it
+                # manually if you ever need to repair keys that lost their TTL.
             time.sleep(self.PRUNE_CHECK_INTERVAL)
 
     def check_middleware(self, middleware_event: str,task: Optional[Task] = None, error: Optional[Exception] = None):

tests/test_base.py

@@ -1047,3 +1047,97 @@ def test_task():
     queued_tasks = mq.get_all_queued_tasks()
     assert len(queued_tasks) == 1
     assert queued_tasks[0]["task_id"] == task.task_id
+
+
+# ---------------------------------------------------------------------------
+# prune_old_task_results: healthy keys must never be read
+# ---------------------------------------------------------------------------
+
+class _CountingRedis:
+    """Wraps a redis client and counts value reads, pipelined ones included."""
+
+    def __init__(self, inner):
+        self._inner = inner
+        self.get_calls = 0
+
+    def __getattr__(self, name):
+        return getattr(self._inner, name)
+
+    def get(self, *a, **kw):
+        self.get_calls += 1
+        return self._inner.get(*a, **kw)
+
+    def pipeline(self, *a, **kw):
+        return _CountingPipeline(self._inner.pipeline(*a, **kw), self)
+
+
+class _CountingPipeline:
+    def __init__(self, inner, parent):
+        self._inner = inner
+        self._parent = parent
+
+    def __getattr__(self, name):
+        return getattr(self._inner, name)
+
+    def get(self, *a, **kw):
+        self._parent.get_calls += 1
+        return self._inner.get(*a, **kw)
+
+
+def test_prune_does_not_read_keys_that_have_a_ttl(modelq_instance):
+    """The control case: a keyspace of healthy keys costs zero value reads.
+
+    This is the whole point of the rewrite. The old implementation GET every
+    key on every pass; production showed 228.8M GETs to issue 2 deletes.
+    """
+    raw = modelq_instance.redis_client
+    for i in range(50):
+        raw.set(f"task_result:h{i}", json.dumps({"finished_at": time.time()}), ex=3600)
+
+    counting = _CountingRedis(raw)
+    modelq_instance.redis_client = counting
+    pruned = modelq_instance.prune_old_task_results(older_than_seconds=86_400)
+
+    assert pruned == 0
+    assert counting.get_calls == 0, "healthy keys must never have their value read"
+    for i in range(50):
+        assert raw.get(f"task_result:h{i}") is not None
+        assert raw.ttl(f"task_result:h{i}") > 0
+
+
+def test_prune_reads_only_the_key_that_lost_its_ttl(modelq_instance):
+    """One broken key among many healthy ones costs exactly one read."""
+    raw = modelq_instance.redis_client
+    for i in range(50):
+        raw.set(f"task_result:h{i}", json.dumps({"finished_at": time.time()}), ex=3600)
+    raw.set("task_result:leaked", json.dumps({"finished_at": time.time() - 90_000}))
+
+    counting = _CountingRedis(raw)
+    modelq_instance.redis_client = counting
+    pruned = modelq_instance.prune_old_task_results(older_than_seconds=86_400)
+
+    assert pruned == 1
+    assert counting.get_calls == 1, "only the TTL-less key should be read"
+    assert raw.get("task_result:leaked") is None
+
+
+def test_prune_bounds_a_recent_orphan_instead_of_deleting_it(modelq_instance):
+    """A key with no TTL but not yet old is kept, and stops living forever."""
+    raw = modelq_instance.redis_client
+    raw.set("task_result:fresh", json.dumps({"finished_at": time.time()}))
+    assert raw.ttl("task_result:fresh") == -1
+
+    pruned = modelq_instance.prune_old_task_results(older_than_seconds=86_400)
+
+    assert pruned == 0
+    assert raw.get("task_result:fresh") is not None
+    assert raw.ttl("task_result:fresh") > 0
+
+
+def test_pruning_loop_does_not_scan_task_results(modelq_instance):
+    """The 60s loop must not call the scan; Redis TTLs handle expiry."""
+    import inspect
+
+    src = inspect.getsource(type(modelq_instance)._pruning_loop)
+    body = "\n".join(l for l in src.splitlines() if not l.strip().startswith("#"))
+    assert "prune_old_task_results" not in body