fix: cancellation large sweep

CHANGELOG.md

@@ -4,6 +4,8 @@ This contains the changes between releases.
 
 # Unreleased
 
+* Indexed and narrowed claim-time cancellation maintenance so worker polling avoids scanning large historical task tables for tasks without cancellation policies.
+
 # 0.5.0
 
 * Added a Habitat tasks filter for terminal failures and defaulted the failed status view to hide failed attempts that were retried later.

sdks/python/src/absurd_sdk/__init__.py

@@ -21,7 +21,7 @@
     Awaitable,
     Callable,
     Dict,
-    Iterator,
+    Generator,
     List,
     Literal,
     Mapping,
@@ -279,7 +279,7 @@ class TimeoutError(Exception):
 
 
 @contextmanager
-def _task_state_exception_handling() -> Iterator[None]:
+def _task_state_exception_handling() -> Generator[None, None, None]:
     try:
         yield
     except Exception as err:
@@ -1212,7 +1212,7 @@ def __init__(
     ) -> None:
         self._queue_name = _validate_queue_name(queue_name)
         self._default_max_attempts = default_max_attempts
-        self._hooks: AbsurdHooks = hooks if hooks is not None else cast(AbsurdHooks, {})
+        self._hooks: AbsurdHooks = hooks if hooks is not None else {}
         self._registry: Dict[str, Dict[str, Any]] = {}
         self._worker_running = False
 
@@ -2291,7 +2291,7 @@ async def execute() -> JsonValue:
 
                 hook_result = wrap_hook(ctx, execute)
                 if inspect.isawaitable(hook_result):
-                    result = cast(JsonValue, await cast(Awaitable[Any], hook_result))
+                    result = cast(JsonValue, await hook_result)
                 else:
                     result = cast(JsonValue, hook_result)
             else:

sdks/python/tests/test_events.py

@@ -3,7 +3,7 @@
 from datetime import datetime, timezone, timedelta
 from psycopg import sql
 
-from absurd_sdk import Absurd, TimeoutError
+from absurd_sdk import Absurd, JsonObject, TimeoutError
 
 
 def _set_fake_now(conn, fake_time):
@@ -64,7 +64,7 @@ def waiter_task(params, ctx):
     assert sleeping_run["wake_event"] == event_name
 
     # Emit event
-    payload = {"value": 42}
+    payload: JsonObject = {"value": 42}
     client.emit_event(event_name, payload)
 
     # Task should now be pending
@@ -86,7 +86,7 @@ def test_event_emitted_before_await_is_cached(conn, queue_name):
     client.create_queue()
 
     event_name = f"cached_event_{queue}"
-    payload = {"data": "pre-emitted"}
+    payload: JsonObject = {"data": "pre-emitted"}
 
     # Emit event before task even exists
     client.emit_event(event_name, payload)
@@ -113,8 +113,8 @@ def test_emit_event_is_first_write_wins(conn, queue_name):
     client.create_queue()
 
     event_name = f"stable_event_{queue}"
-    first_payload = {"value": 1}
-    second_payload = {"value": 2}
+    first_payload: JsonObject = {"value": 1}
+    second_payload: JsonObject = {"value": 2}
 
     first_emit_at = datetime(2024, 5, 1, 9, 0, 0, tzinfo=timezone.utc)
     _set_fake_now(conn, first_emit_at)
@@ -221,7 +221,7 @@ def multi_waiter_task(params, ctx):
         assert task["state"] == "sleeping"
 
     # Emit event once
-    payload = {"data": "broadcast"}
+    payload: JsonObject = {"data": "broadcast"}
     client.emit_event(event_name, payload)
 
     # All tasks should resume and complete

sdks/python/tests/test_task_context.py

@@ -5,7 +5,7 @@
 import pytest
 from psycopg import sql
 
-from absurd_sdk import Absurd, SuspendTask, _create_task_context
+from absurd_sdk import Absurd, JsonObject, SuspendTask, _create_task_context
 
 
 def _fetch_checkpoint(conn, queue, task_id, checkpoint_name):
@@ -106,7 +106,7 @@ def test_await_event_flow(conn, queue_name):
     assert wake_event == event_name
     assert available_at is None
 
-    payload = {"status": "ok"}
+    payload: JsonObject = {"status": "ok"}
     client.emit_event(event_name, payload, queue_name=queue)
 
     replay = client.claim_tasks(worker_id="worker")[0]

sql/absurd.sql

@@ -331,6 +331,13 @@ begin
     'r_' || p_queue_name
   );
 
+  execute format(
+    'create index if not exists %I on absurd.%I (task_id)
+      where cancellation is not null',
+    ('t_' || p_queue_name) || '_ci',
+    't_' || p_queue_name
+  );
+
   execute format(
     'create index if not exists %I on absurd.%I (event_name)',
     ('w_' || p_queue_name) || '_eni',
@@ -953,6 +960,7 @@ begin
       'select task_id
          from absurd.%I
         where state in (''pending'', ''sleeping'', ''running'')
+          and cancellation is not null
           and (
             (
               (cancellation->>''max_delay'')::bigint is not null