diff --git a/modelq/app/base.py b/modelq/app/base.py index b786361..cbf1648 100644 --- a/modelq/app/base.py +++ b/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): diff --git a/tests/test_base.py b/tests/test_base.py index fa07f0c..b4cf75b 100644 --- a/tests/test_base.py +++ b/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