[cross-repo from server#265] Conformance: signals/queries still time out on server 0.2.125 (#79)

Author: durable-workflow-ops

src/durable_workflow/client.py

@@ -1171,6 +1171,7 @@ def __init__(
         self.control_token = control_token
         self.worker_token = worker_token
         self.namespace = namespace
+        self.timeout = timeout
         self.retry_policy = retry_policy or TransportRetryPolicy()
         self.metrics = metrics or NOOP_METRICS
         self.payload_size_warning_config = (

src/durable_workflow/worker.py

@@ -21,6 +21,7 @@
 import inspect
 import logging
 import sys
+import threading
 import time
 import traceback
 import uuid
@@ -347,6 +348,8 @@ def __init__(
         self._shutdown_timeout = shutdown_timeout
         self._in_flight: set[asyncio.Task[Any]] = set()
         self._query_tasks_supported = False
+        self._query_thread_stop = threading.Event()
+        self._query_thread: threading.Thread | None = None
         # In-flight slot accounting feeds the periodic heartbeat so operators
         # see free-slot counts without the worker having to re-derive them at
         # shutdown. Counters are bumped/decremented around dispatch.
@@ -897,15 +900,15 @@ async def call_interceptor(
 
         return await handler(context)
 
-    async def _run_query_task(self, task: dict[str, Any]) -> str:
+    async def _run_query_task(self, task: dict[str, Any], *, client: Client | None = None) -> str:
         context = QueryTaskInterceptorContext(
             worker_id=self.worker_id,
             task_queue=self.task_queue,
             task=task,
         )
 
         async def call_core(ctx: QueryTaskInterceptorContext) -> str:
-            return await self._run_query_task_core(ctx.task)
+            return await self._run_query_task_core(ctx.task, client=client)
 
         handler = call_core
         for interceptor in reversed(self.interceptors):
@@ -923,7 +926,8 @@ async def call_interceptor(
 
         return await handler(context)
 
-    async def _run_query_task_core(self, task: dict[str, Any]) -> str:
+    async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | None = None) -> str:
+        client = client or self.client
         query_task_id: str = task["query_task_id"]
         attempt: int = task.get("query_task_attempt", 1)
         wf_type: str = task.get("workflow_type", "")
@@ -939,6 +943,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason="query_payload_decode_failed",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
 
@@ -953,6 +958,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 f"no workflow registered for type {wf_type!r}",
                 reason="query_workflow_type_not_registered",
                 failure_type="WorkflowTypeNotRegistered",
+                client=client,
             )
             return "failed"
 
@@ -992,6 +998,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason="query_payload_decode_failed",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
         except QueryFailed as e:
@@ -1003,6 +1010,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason=reason,
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
         except Exception as e:
@@ -1013,11 +1021,12 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason="query_rejected",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
 
         try:
-            await self.client.complete_query_task(
+            await client.complete_query_task(
                 query_task_id=query_task_id,
                 lease_owner=self.worker_id,
                 query_task_attempt=attempt,
@@ -1044,6 +1053,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason=server_reason if server_reason else "query_result_completion_failed",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
         except AvroNotInstalledError as e:
@@ -1057,6 +1067,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason="query_result_encode_failed",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
         except (TypeError, ValueError) as e:
@@ -1067,6 +1078,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason="query_result_encode_failed",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
         except Exception as e:
@@ -1077,6 +1089,7 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
                 reason="query_result_completion_failed",
                 failure_type=type(e).__name__,
                 stack_trace=traceback.format_exc(),
+                client=client,
             )
             return "failed"
 
@@ -1107,9 +1120,11 @@ async def _fail_query_task(
         reason: str,
         failure_type: str | None = None,
         stack_trace: str | None = None,
+        client: Client | None = None,
     ) -> None:
+        client = client or self.client
         try:
-            await self.client.fail_query_task(
+            await client.fail_query_task(
                 query_task_id=query_task_id,
                 lease_owner=self.worker_id,
                 query_task_attempt=attempt,
@@ -1217,11 +1232,12 @@ async def _dispatch_activity_task(self, task: dict[str, Any]) -> None:
             self._record_task_metrics("activity", outcome, time.perf_counter() - task_start)
             self._act_semaphore.release()
 
-    async def _poll_query_tasks(self) -> None:
-        while not self._stop.is_set():
+    async def _poll_query_tasks(self, *, client: Client | None = None, track_tasks: bool = True) -> None:
+        client = client or self.client
+        while not self._stop.is_set() and not self._query_thread_stop.is_set():
             try:
                 poll_start = time.perf_counter()
-                task = await self.client.poll_query_task(
+                task = await client.poll_query_task(
                     worker_id=self.worker_id,
                     task_queue=self.task_queue,
                     timeout=self._poll_timeout,
@@ -1236,18 +1252,86 @@ async def _poll_query_tasks(self) -> None:
                 await asyncio.sleep(0)
                 continue
             self._record_poll_metrics("query", "task", time.perf_counter() - poll_start)
-            self._track(self._dispatch_query_task(task))
+            if track_tasks:
+                self._track(self._dispatch_query_task(task, client=client))
+            else:
+                await self._dispatch_query_task(task, client=client)
 
-    async def _dispatch_query_task(self, task: dict[str, Any]) -> None:
+    async def _dispatch_query_task(self, task: dict[str, Any], *, client: Client | None = None) -> None:
         task_start = time.perf_counter()
         outcome = "error"
         try:
-            outcome = await self._run_query_task(task)
+            outcome = await self._run_query_task(task, client=client)
         except Exception:
             log.exception("unhandled error in query task execution")
         finally:
             self._record_task_metrics("query", outcome, time.perf_counter() - task_start)
 
+    def _clone_client_for_query_tasks(self) -> Client:
+        warning_config = self._payload_size_warning_config()
+
+        return Client(
+            self.client.base_url,
+            token=self.client.token,
+            control_token=self.client.control_token,
+            worker_token=self.client.worker_token,
+            namespace=self.client.namespace,
+            timeout=getattr(self.client, "timeout", 60.0),
+            retry_policy=self.client.retry_policy,
+            metrics=self.metrics,
+            payload_size_limit_bytes=(
+                warning_config.limit_bytes
+                if warning_config is not None
+                else serializer.DEFAULT_PAYLOAD_SIZE_BYTES
+            ),
+            payload_size_warning_threshold_percent=(
+                warning_config.threshold_percent
+                if warning_config is not None
+                else serializer.DEFAULT_WARNING_THRESHOLD_PERCENT
+            ),
+            payload_size_warnings=warning_config is not None,
+            external_storage=self.external_storage,
+            external_storage_threshold_bytes=self.external_storage_threshold_bytes,
+            external_storage_cache=self.external_storage_cache,
+        )
+
+    def _can_clone_client_for_query_tasks(self) -> bool:
+        return type(self.client) is Client
+
+    def _start_query_task_thread(self) -> None:
+        if self._query_thread is not None and self._query_thread.is_alive():
+            return
+
+        self._query_thread_stop.clear()
+        self._query_thread = threading.Thread(
+            target=self._run_query_task_thread,
+            name=f"durable-workflow-query-poller-{self.worker_id}",
+            daemon=True,
+        )
+        self._query_thread.start()
+
+    def _run_query_task_thread(self) -> None:
+        try:
+            asyncio.run(self._query_task_thread_main())
+        except Exception:
+            log.exception("query task poller thread stopped unexpectedly")
+
+    async def _query_task_thread_main(self) -> None:
+        async with self._clone_client_for_query_tasks() as client:
+            await self._poll_query_tasks(client=client, track_tasks=False)
+
+    async def _stop_query_task_thread(self) -> None:
+        self._query_thread_stop.set()
+        thread = self._query_thread
+
+        if thread is None or not thread.is_alive():
+            return
+
+        await asyncio.to_thread(
+            thread.join,
+            min(self._shutdown_timeout, self._poll_timeout + 1),
+        )
+
     async def run(self) -> None:
         """Register the worker and poll until `stop()` is called or the task is cancelled."""
         await self._register()
@@ -1256,9 +1340,15 @@ async def run(self) -> None:
         hb_loop = asyncio.create_task(self._heartbeat_loop())
         loops = [wf_loop, act_loop, hb_loop]
         if self._query_tasks_supported:
-            loops.append(asyncio.create_task(self._poll_query_tasks()))
-        with contextlib.suppress(asyncio.CancelledError):
-            await asyncio.gather(*loops)
+            if self._can_clone_client_for_query_tasks():
+                self._start_query_task_thread()
+            else:
+                loops.append(asyncio.create_task(self._poll_query_tasks()))
+        try:
+            with contextlib.suppress(asyncio.CancelledError):
+                await asyncio.gather(*loops)
+        finally:
+            await self._stop_query_task_thread()
 
     async def _heartbeat_loop(self) -> None:
         """Periodically refresh the server-side worker registration.
@@ -1386,7 +1476,10 @@ async def run_until(
             await self._register()
             background_tasks.append(asyncio.create_task(self._heartbeat_loop()))
             if self._query_tasks_supported:
-                background_tasks.append(asyncio.create_task(self._poll_query_tasks()))
+                if self._can_clone_client_for_query_tasks():
+                    self._start_query_task_thread()
+                else:
+                    background_tasks.append(asyncio.create_task(self._poll_query_tasks()))
 
             while True:
                 desc = await self.client.describe_workflow(workflow_id)
@@ -1465,6 +1558,7 @@ async def run_until(
     async def stop(self) -> None:
         """Stop polling and drain in-flight tasks up to the configured shutdown timeout."""
         self._stop.set()
+        await self._stop_query_task_thread()
         if self._in_flight:
             log.info("draining %d in-flight task(s)…", len(self._in_flight))
             done, pending = await asyncio.wait(

tests/test_worker.py

@@ -4,6 +4,7 @@
 import contextlib
 import logging
 import sys
+import threading
 from unittest.mock import AsyncMock
 
 import pytest
@@ -1476,6 +1477,70 @@ async def test_run_skips_query_loop_without_query_task_capability(self, mock_cli
         assert mock_client.poll_activity_task.call_count >= 1
         mock_client.poll_query_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_query_thread_processes_tasks_while_event_loop_is_blocked(
+        self, mock_client: AsyncMock
+    ) -> None:
+        completed = threading.Event()
+        query_task = {
+            "query_task_id": "qt-thread",
+            "query_task_attempt": 1,
+            "workflow_type": "query-wf",
+            "workflow_id": "wf-1",
+            "run_id": "run-1",
+            "query_name": "status",
+            "payload_codec": "json",
+            "workflow_arguments": serializer.envelope([], codec="json"),
+            "query_arguments": serializer.envelope([], codec="json"),
+            "history_events": [],
+        }
+
+        class QueryThreadClient:
+            def __init__(self) -> None:
+                self.polled = False
+                self.completed_kwargs: dict[str, object] | None = None
+
+            async def __aenter__(self) -> "QueryThreadClient":
+                return self
+
+            async def __aexit__(self, *_: object) -> None:
+                return None
+
+            async def poll_query_task(self, **_: object) -> dict[str, object] | None:
+                if not self.polled:
+                    self.polled = True
+                    return query_task
+                await asyncio.sleep(0.01)
+                return None
+
+            async def complete_query_task(self, **kwargs: object) -> dict[str, str]:
+                self.completed_kwargs = kwargs
+                completed.set()
+                return {"outcome": "completed"}
+
+            async def fail_query_task(self, **_: object) -> dict[str, str]:
+                raise AssertionError("query task should complete")
+
+        query_client = QueryThreadClient()
+        worker = Worker(
+            mock_client,
+            task_queue="q1",
+            workflows=[QueryWorkflow],
+            activities=[],
+            poll_timeout=0.01,
+            shutdown_timeout=0.2,
+        )
+        worker._clone_client_for_query_tasks = lambda: query_client  # type: ignore[method-assign]
+
+        worker._start_query_task_thread()
+
+        assert completed.wait(timeout=1.0)
+        await worker.stop()
+
+        assert query_client.completed_kwargs is not None
+        assert query_client.completed_kwargs["query_task_id"] == "qt-thread"
+        assert query_client.completed_kwargs["result"] == {"status": "ready"}
+
 
 class TestWorkerHeartbeats:
     @pytest.mark.asyncio