diff --git a/sdks/python/taskito/mixins/lifecycle.py b/sdks/python/taskito/mixins/lifecycle.py index 24b890f3..86705ba9 100644 --- a/sdks/python/taskito/mixins/lifecycle.py +++ b/sdks/python/taskito/mixins/lifecycle.py @@ -299,12 +299,15 @@ def sighup_handler(signum: int, frame: Any) -> None: EventType.WORKER_STOPPED, {"worker_id": worker_id}, ) + # Stop the consumer poll loops before waiting on anything else: their + # per-interval reads otherwise keep contending for the storage lock + # while the heartbeat makes its final round trip. + stop_log_consumers.set() stop_heartbeat.set() heartbeat_thread.join(timeout=6) - # Drain managed log consumers before tearing down resources a handler + # Join managed log consumers before tearing down resources a handler # might still hold: give them a shared deadline (the worker drain # timeout) to finish their in-flight message, not a fixed slice each. - stop_log_consumers.set() consumer_deadline = time.monotonic() + self._drain_timeout for consumer in log_consumer_threads: consumer.join(timeout=max(0.0, consumer_deadline - time.monotonic())) diff --git a/sdks/python/tests/core/test_pubsub_log.py b/sdks/python/tests/core/test_pubsub_log.py index 8a95270c..c313ae1a 100644 --- a/sdks/python/tests/core/test_pubsub_log.py +++ b/sdks/python/tests/core/test_pubsub_log.py @@ -177,7 +177,10 @@ def handle(n: int) -> None: lambda: handled == [1], timeout=30, message="consumer should run before shutdown" ) queue.shutdown() - thread.join(timeout=5) + # SQLite is opened with ``busy_timeout = 5000``, so one contended query + # during the drain can block for 5s on its own — a 5s budget can't bound + # a shutdown. Wait past that, well under the 30s drain timeout. + thread.join(timeout=20) assert not thread.is_alive()