flo-python 0.1.0.dev4__py3-none-any.whl → 0.1.2__py3-none-any.whl

flo/worker.py

@@ -22,6 +22,7 @@ Example::
 
 import asyncio
 import contextlib
+import dataclasses
 import json
 import logging
 import secrets
@@ -31,9 +32,14 @@ from dataclasses import dataclass, field
 from typing import Any
 
 from .client import FloClient
-from .exceptions import NonRetryableError, is_connection_error
+from .exceptions import NonRetryableError, NotFoundError, ServerError, is_connection_error
+from .retry import _sleep_unless_stopped, _worker_backoff, is_refused_unrun
 from .types import (
+    _DEFAULT_WORKER_BLOCK_MS,
+    ActionRegisterOptions,
     ActionType,
+    Reason,
+    StatusCode,
     StreamGroupAckOptions,
     StreamGroupNackOptions,
     StreamGroupReadOptions,
@@ -41,8 +47,10 @@ from .types import (
     StreamRecord,
     TaskAssignment,
     WorkerAwaitOptions,
+    WorkerRegisterOptions,
     WorkerTouchOptions,
 )
+from .wire import _validate_block_ms
 
 logger = logging.getLogger("flo.worker")
 
@@ -68,6 +76,164 @@ class ActionResult:
 ActionHandler = Callable[["ActionContext"], Awaitable[bytes | dict[str, Any] | ActionResult]]
 
 
+class _EmptyPollBackoff:
+    """Paces a poll loop whose blocking reads come back empty early.
+
+    A blocking read can come back empty early for two reasons that look the
+    same on the wire. The server may have no room to park it (re-polling at
+    once would spin), or an append woke it so the consumer re-reads. A full
+    server answers within a round trip, so only answers under
+    min(250 ms, block_ms/2) count as early. Even then, one may be an append
+    wake, so the first in a row is re-polled at once.
+    """
+
+    _EARLY_S = 0.25
+    _FIRST_S = 0.05
+    _MAX_S = 1.0
+
+    def __init__(self, block_ms: int, stop: asyncio.Event) -> None:
+        self._early_s = min(self._EARLY_S, block_ms / 2000.0)
+        self._stop = stop
+        self._streak = 0
+
+    def reset(self) -> None:
+        self._streak = 0
+
+    def _next_pause(self, elapsed_s: float) -> float:
+        if elapsed_s >= self._early_s:
+            self._streak = 0
+            return 0.0
+        self._streak += 1
+        if self._streak == 1:
+            return 0.0
+        # The exponent is capped: the float power overflows on a long streak,
+        # and the pause reaches _MAX_S long before the cap.
+        return min(self._FIRST_S * 2.0 ** min(self._streak - 2, 16), self._MAX_S)
+
+    async def empty(self, elapsed_s: float) -> None:
+        pause = self._next_pause(elapsed_s)
+        if pause:
+            with contextlib.suppress(asyncio.TimeoutError):
+                await asyncio.wait_for(self._stop.wait(), timeout=pause)
+
+
+# A result (complete, fail, ack, nack) is sent at most this many times. Past
+# that it is dropped: the lease or ack timeout runs out on the server and the
+# task or record is delivered again, so delivery stays at-least-once.
+_RESULT_SENDS = 3
+
+# A result to send: the op name, the send on a given client, and whether the
+# handler was cancelled.
+_Report = tuple[str, Callable[[FloClient], Awaitable[None]], bool]
+
+# How long a stopping worker keeps reconnecting to resend a lost result, so
+# stop() can't hang on a server that is gone.
+_STOP_RECONNECT_S = 2.0
+
+# auto_touch never touches more often than this, whatever the timeout.
+_MIN_TOUCH_INTERVAL_MS = 100
+
+# The longest lease a touch may ask for; the server refuses more.
+_MAX_TOUCH_EXTEND_MS = 3_600_000
+
+# How often an action worker heartbeats. The server refuses work to a worker
+# silent past its heartbeat timeout (90 s by default) and later removes it.
+_HEARTBEAT_INTERVAL_S = 30.0
+
+
+def _touch_pace(timeout_ms: int) -> tuple[float, int]:
+    """How often (in seconds) to touch a run of an action with timeout_ms,
+    and the lease each touch asks for: the timeout, at most an hour, touched
+    every third of that (at least _MIN_TOUCH_INTERVAL_MS apart) so two can
+    fail in a row."""
+    extend_ms = min(timeout_ms, _MAX_TOUCH_EXTEND_MS)
+    return max(extend_ms / 3, _MIN_TOUCH_INTERVAL_MS) / 1000, extend_ms
+
+
+# Touch refusals that mean the run moved on: touching again can't help.
+_TOUCH_FINAL = (StatusCode.CONFLICT, StatusCode.NOT_FOUND, StatusCode.BAD_REQUEST)
+
+
+async def _send_result(
+    what: str,
+    send: Callable[[], Awaitable[object]],
+    generation: Callable[[], int],
+    reconnect: Callable[[int], Awaitable[None]],
+    stop: asyncio.Event,
+) -> bool:
+    """Send a result, resending the identical request after a reconnect when
+    the connection drops or the request times out.
+
+    A refusal is not resent: the server read the request and said no. Once
+    the worker is stopping, reconnects get _STOP_RECONNECT_S in all. Nothing
+    here raises; a result that can't be sent is logged and dropped.
+    """
+    loop = asyncio.get_running_loop()
+    stop_deadline: float | None = None
+
+    async def reconnect_bounded(seen: int) -> bool:
+        """Reconnect; False if the stop allowance ran out first."""
+        nonlocal stop_deadline
+        task = asyncio.ensure_future(reconnect(seen))
+        try:
+            if not stop.is_set():
+                waiter = asyncio.ensure_future(stop.wait())
+                try:
+                    await asyncio.wait({task, waiter}, return_when=asyncio.FIRST_COMPLETED)
+                finally:
+                    waiter.cancel()
+            if not task.done():
+                if stop_deadline is None:
+                    stop_deadline = loop.time() + _STOP_RECONNECT_S
+                await asyncio.wait({task}, timeout=max(stop_deadline - loop.time(), 0.0))
+        finally:
+            if not task.done():
+                task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+        if task.cancelled():
+            return False
+        task.result()
+        return True
+
+    for attempt in range(1, _RESULT_SENDS + 1):
+        seen = generation()
+        try:
+            await send()
+            return True
+        except Exception as e:
+            if not (is_connection_error(e) or isinstance(e, asyncio.TimeoutError)):
+                logger.error(f"{what} was refused, not resending: {e}")
+                return False
+            if attempt == _RESULT_SENDS:
+                logger.error(f"Dropped {what} after {_RESULT_SENDS} sends: {e}")
+                return False
+            if stop.is_set() and stop_deadline is None:
+                stop_deadline = loop.time() + _STOP_RECONNECT_S
+            if stop_deadline is not None and loop.time() >= stop_deadline:
+                logger.error(f"Dropped {what}: connection lost while the worker stops: {e}")
+                return False
+            logger.warning(
+                f"Lost {what} (send {attempt}/{_RESULT_SENDS}): {e}; reconnecting to resend"
+            )
+            try:
+                if not await reconnect_bounded(seen):
+                    logger.error(f"Dropped {what}: no connection before the worker stopped")
+                    return False
+            except Exception as rc_err:
+                # The next send fails fast on the closed connection and
+                # counts as one of the attempts.
+                logger.warning(f"Reconnect to resend {what} failed: {rc_err}")
+    return False
+
+
+def _worker_block_ms(block_ms: int | None) -> int:
+    if not block_ms:
+        return _DEFAULT_WORKER_BLOCK_MS
+    _validate_block_ms(block_ms)
+    return block_ms
+
+
 @dataclass
 class ActionWorkerOptions:
     """Configuration for a Flo action worker.
@@ -78,7 +244,13 @@ class ActionWorkerOptions:
     worker_id: str = ""
     concurrency: int = 10
     action_timeout: float = 300.0  # 5 minutes
-    block_ms: int = 30000
+    block_ms: int = _DEFAULT_WORKER_BLOCK_MS
+    # Touch each run every third of its action's timeout while the handler
+    # runs, so a long handler isn't taken back and run again elsewhere.
+    auto_touch: bool = False
+
+    def __post_init__(self) -> None:
+        self.block_ms = _worker_block_ms(self.block_ms)
 
 
 @dataclass
@@ -190,7 +362,8 @@ class ActionWorker:
         worker_id: str | None = None,
         concurrency: int = 10,
         action_timeout: float = 300.0,
-        block_ms: int = 30000,
+        block_ms: int = _DEFAULT_WORKER_BLOCK_MS,
+        auto_touch: bool = False,
     ):
         """Initialize a Flo worker from a connected client.
 
@@ -199,7 +372,11 @@ class ActionWorker:
             worker_id: Unique worker identifier (auto-generated if not provided).
             concurrency: Maximum number of concurrent actions.
             action_timeout: Timeout for action handlers in seconds.
-            block_ms: Timeout for blocking dequeue in milliseconds.
+            block_ms: Long-poll wait per await in milliseconds, at most 300000.
+                0 means the 30000 default.
+            auto_touch: Touch each run every third of its action's timeout
+                while its handler runs. Off by default: a handler that runs
+                past the timeout without touching is taken back and run again.
         """
         self._parent_client = parent_client
         self.config = ActionWorkerOptions(
@@ -207,16 +384,27 @@ class ActionWorker:
             concurrency=concurrency,
             action_timeout=action_timeout,
             block_ms=block_ms,
+            auto_touch=auto_touch,
         )
 
         self._client: FloClient | None = None
         self._result_client: FloClient | None = None
+        # Concurrent results that fail on the same dead connection reconnect
+        # it once: the generation tells a waiter someone already did.
+        self._result_reconnect_lock = asyncio.Lock()
+        self._result_generation = 0
         self._handlers: dict[str, ActionHandler] = {}
+        self._action_options: dict[str, ActionRegisterOptions] = {}
         self._running = False
         self._stop_event = asyncio.Event()
         self._tasks: set[asyncio.Task[None]] = set()
         self._semaphore: asyncio.Semaphore | None = None
         self._heartbeat_task: asyncio.Task[None] | None = None
+        self._action_names: list[str] = []
+        # The load heartbeats report. Not the semaphore, which also counts the
+        # slot an outstanding await holds, nor _tasks, which also counts
+        # results still being sent.
+        self._running_handlers = 0
 
     @staticmethod
     def _generate_worker_id() -> str:
@@ -227,11 +415,14 @@ class ActionWorker:
             hostname = "unknown"
         return f"{hostname}-{secrets.token_hex(4)}"
 
-    def action(self, name: str) -> Callable[[ActionHandler], ActionHandler]:
-        """Decorator to register an action handler.
+    def action(
+        self, name: str, options: ActionRegisterOptions | None = None
+    ) -> Callable[[ActionHandler], ActionHandler]:
+        """Decorator to register an action handler (see register_action).
 
         Args:
             name: The action name to register.
+            options: The action's timeout, attempts and owner.
 
         Returns:
             Decorator function.
@@ -243,17 +434,30 @@ class ActionWorker:
         """
 
         def decorator(handler: ActionHandler) -> ActionHandler:
-            self.register_action(name, handler)
+            self.register_action(name, handler, options)
             return handler
 
         return decorator
 
-    def register_action(self, name: str, handler: ActionHandler) -> None:
+    def register_action(
+        self,
+        name: str,
+        handler: ActionHandler,
+        options: ActionRegisterOptions | None = None,
+    ) -> None:
         """Register an action handler.
 
+        The worker registers the action with the server when it starts, with
+        exactly these options (the server's defaults when none: a 30 s
+        timeout, 3 attempts). Registering replaces the action's settings, so
+        workers registering one action with different options leave it with
+        whichever registered last. The options' namespace is ignored: the
+        action goes in the worker's.
+
         Args:
             name: The action name.
             handler: Async function that handles the action.
+            options: The action's timeout, attempts and owner.
 
         Raises:
             ValueError: If action is already registered.
@@ -261,6 +465,9 @@ class ActionWorker:
         if name in self._handlers:
             raise ValueError(f"Action '{name}' is already registered")
         self._handlers[name] = handler
+        self._action_options[name] = dataclasses.replace(
+            options or ActionRegisterOptions(), namespace=None
+        )
         logger.info(f"Registered action: {name}")
 
     async def start(self) -> None:
@@ -281,17 +488,13 @@ class ActionWorker:
         )
 
         # Create a dedicated connection using the parent client's endpoint and namespace.
-        # Timeout must accommodate block_ms + action_timeout so blocking reads
-        # (ACTION_AWAIT with block_ms) don't get killed by socket-level timeout.
-        worker_timeout_ms = max(
-            self.config.block_ms + 5000,
-            int(self.config.action_timeout * 1000),
-        )
+        worker_timeout_ms = int(self._parent_client._timeout * 1000)
         self._client = FloClient(
             self._parent_client._endpoint,
             namespace=self._parent_client.namespace,
             debug=self._parent_client._debug,
             timeout_ms=worker_timeout_ms,
+            retry=self._parent_client._retry,
         )
         await self._client.connect()
 
@@ -303,27 +506,21 @@ class ActionWorker:
             namespace=self._parent_client.namespace,
             debug=self._parent_client._debug,
             timeout_ms=worker_timeout_ms,
+            retry=self._parent_client._retry,
         )
         await self._result_client.connect()
 
         try:
             # Register actions with the server
             action_names = list(self._handlers.keys())
+            self._action_names = action_names
             for action_name in action_names:
-                await self._client.action.register(action_name, ActionType.USER)
+                await self._client.action.register(
+                    action_name, ActionType.USER, self._action_options[action_name]
+                )
                 logger.debug(f"Registered action with server: {action_name}")
 
-            # Register worker
-            from .types import WorkerRegisterOptions
-
-            await self._client.worker.register(
-                self.config.worker_id,
-                action_names,
-                WorkerRegisterOptions(
-                    concurrency=self.config.concurrency,
-                    machine_id=socket.gethostname(),
-                ),
-            )
+            await self._register(self._client)
             logger.info(f"Worker registered with {len(action_names)} actions")
 
             # Initialize concurrency control
@@ -358,20 +555,51 @@ class ActionWorker:
             self._running = False
             logger.info("Worker stopped")
 
+    async def _register(self, client: FloClient) -> None:
+        """Register this worker, with the same options every time."""
+        await client.worker.register(
+            self.config.worker_id,
+            self._action_names,
+            WorkerRegisterOptions(
+                concurrency=self.config.concurrency,
+                machine_id=socket.gethostname(),
+            ),
+        )
+
+    async def _heartbeat(self) -> None:
+        """Send one heartbeat, registering again if the server removed this
+        worker: unregistered, it would go on taking work the server can't
+        see it holding."""
+        # On the result connection: on the poll connection it would wait
+        # behind a blocking await, up to block_ms, past the server's timeout.
+        client = self._result_client or self._client
+        assert client is not None
+        load = self._running_handlers
+        seen = self._result_generation
+        try:
+            await client.worker.heartbeat(self.config.worker_id, current_load=load)
+        except NotFoundError:
+            logger.warning("Worker was removed by the server; registering again")
+            await self._register(client)
+            return
+        except Exception as e:
+            # Nothing else reconnects this connection for an idle worker, which
+            # would then miss every heartbeat and be refused work for good.
+            lost = is_connection_error(e) or isinstance(e, asyncio.TimeoutError)
+            if self._result_client is not None and lost:
+                with contextlib.suppress(Exception):
+                    await self._reconnect_result_client(seen)
+            raise
+        logger.debug(f"Heartbeat sent (load={load})")
+
     async def _heartbeat_loop(self) -> None:
         """Send periodic heartbeats to keep the worker registration alive."""
-        assert self._client is not None
         while self._running and not self._stop_event.is_set():
             try:
-                await asyncio.sleep(30)
+                await asyncio.sleep(_HEARTBEAT_INTERVAL_S)
                 if not self._running:
                     break
-                current_load = len(self._tasks)
-                await self._client.worker.heartbeat(
-                    self.config.worker_id,
-                    current_load=current_load,
-                )
-                logger.debug(f"Heartbeat sent (load={current_load})")
+                await self._heartbeat()
             except asyncio.CancelledError:
                 break
             except Exception as e:
@@ -381,6 +609,13 @@ class ActionWorker:
         """Main polling loop for tasks."""
         assert self._client is not None
         assert self._semaphore is not None
+        # Pause before each reconnect after the first in a row: a server that
+        # accepts and then drops the connection would otherwise be hit in a
+        # tight loop. A poll that gets an answer resets it.
+        reconnect_pause = 0.0
+        loop = asyncio.get_running_loop()
+        backoff = _EmptyPollBackoff(self.config.block_ms, self._stop_event)
+        failed = _worker_backoff()
         while self._running and not self._stop_event.is_set():
             try:
                 # Wait for semaphore slot
@@ -392,16 +627,21 @@ class ActionWorker:
                     break
 
                 # Await task from server
+                started = loop.time()
                 result = await self._client.worker.await_task(
                     self.config.worker_id,
                     action_names,
                     WorkerAwaitOptions(block_ms=self.config.block_ms),
                 )
+                reconnect_pause = 0.0
+                failed.reset()
 
                 if result.task is None:
                     # No task available, release semaphore and continue
                     self._semaphore.release()
+                    await backoff.empty(loop.time() - started)
                     continue
+                backoff.reset()
 
                 # Execute task in background
                 task = asyncio.create_task(self._execute_task(result.task))
@@ -412,213 +652,252 @@ class ActionWorker:
                 break
             except Exception as e:
                 self._semaphore.release()
-                if is_connection_error(e):
+                reason = e.reason if isinstance(e, ServerError) else None
+                if reason == Reason.WORKER_DRAINING:
+                    # The server drained this worker: take no more work, and
+                    # let the running handlers finish.
+                    logger.info("Worker is draining; no longer taking work")
+                    break
+                if reason == Reason.WORKER_UNHEALTHY:
+                    # It missed its heartbeats: one now lets it take work
+                    # again; the backoff keeps a heartbeat that fails from
+                    # turning this into a hot loop.
+                    try:
+                        await self._heartbeat()
+                    except Exception as hb_err:
+                        logger.warning(f"Heartbeat failed: {hb_err}")
+                    d = failed.next()
+                    logger.warning(f"Await refused, worker unhealthy; retrying in {d:.3f}s")
+                    await _sleep_unless_stopped(self._stop_event, d)
+                elif is_connection_error(e):
+                    # stop() interrupts the poll connection: that is not a
+                    # drop to recover from.
+                    if self._stop_event.is_set():
+                        break
+                    if reconnect_pause:
+                        with contextlib.suppress(asyncio.TimeoutError):
+                            await asyncio.wait_for(self._stop_event.wait(), reconnect_pause)
+                        if self._stop_event.is_set():
+                            break
+                    reconnect_pause = min(max(reconnect_pause * 2, 1.0), 30.0)
                     logger.warning("Connection lost, reconnecting...")
                     try:
                         await self._client.reconnect()
                         # Also reconnect result client
                         if self._result_client is not None:
                             try:
-                                await self._result_client.reconnect()
+                                await self._reconnect_result_client(only_if_down=True)
                             except Exception as rc_err:
                                 logger.warning(f"Failed to reconnect result client: {rc_err}")
                         # Re-register worker after reconnect
                         try:
-                            from .types import WorkerRegisterOptions
-
-                            await self._client.worker.register(
-                                self.config.worker_id,
-                                action_names,
-                                WorkerRegisterOptions(
-                                    concurrency=self.config.concurrency,
-                                    machine_id=socket.gethostname(),
-                                ),
-                            )
+                            await self._register(self._client)
                         except Exception as reg_err:
                             logger.warning(f"Failed to re-register worker: {reg_err}")
                         logger.info("Reconnected, resuming work")
                     except Exception as recon_err:
                         logger.error(f"Reconnect failed: {recon_err}, retrying...")
-                        await asyncio.sleep(1)
                 else:
-                    logger.error(f"Await error: {e}, retrying...")
-                    await asyncio.sleep(1)
+                    d = failed.next()
+                    logger.error(f"Await error: {e}, retrying in {d:.3f}s...")
+                    await _sleep_unless_stopped(self._stop_event, d)
 
-    async def _send_with_retry(self, op: str, fn: Callable[[], Awaitable[None]]) -> None:
-        """Attempt to send a result (Complete/Fail), reconnecting on connection error."""
-        max_attempts = 3
-        for attempt in range(1, max_attempts + 1):
-            try:
-                await fn()
+    async def _reconnect_result_client(
+        self, failed_generation: int | None = None, *, only_if_down: bool = False
+    ) -> None:
+        """Reconnect the result connection, unless it was reconnected since
+        `failed_generation` was observed, or, with `only_if_down`, it is up.
+
+        The poll connection dropping says nothing certain about this one, and
+        replacing a healthy connection would cut off results in flight on it.
+        """
+        assert self._result_client is not None
+        async with self._result_reconnect_lock:
+            if failed_generation is not None and failed_generation != self._result_generation:
                 return
-            except Exception as e:
-                if not is_connection_error(e) or not self._running:
-                    raise
-                logger.warning(
-                    f"Connection lost while sending {op} result "
-                    f"(attempt {attempt}/{max_attempts}), reconnecting..."
-                )
-                if self._result_client is not None:
-                    await self._result_client.reconnect()
-        raise RuntimeError(f"Failed to send {op} result after {max_attempts} attempts")
+            if only_if_down and self._result_client.is_connected:
+                return
+            await self._result_client.reconnect()
+            self._result_generation += 1
 
     async def _execute_task(self, task: TaskAssignment) -> None:
-        """Execute a task with error handling."""
-        assert self._client is not None
-        assert self._semaphore is not None
-        rc = self._result_client or self._client
+        """Run a task's handler, then report its outcome."""
+        rc = self._result_client
+        assert rc is not None
         try:
-            logger.info(
-                f"Executing action: {task.task_type} (task={task.task_id}, attempt={task.attempt})"
-            )
-
-            # Get handler
-            handler = self._handlers.get(task.task_type)
-            if handler is None:
-                logger.error(f"No handler registered for action: {task.task_type}")
-                await self._send_with_retry(
-                    "fail",
-                    lambda: rc.worker.fail(
-                        self.config.worker_id,
-                        task.task_type,
-                        task.task_id,
-                        f"No handler for: {task.task_type}",
-                    ),
-                )
-                return
-
-            # Create action context
-            ctx = ActionContext(
-                task_id=task.task_id,
-                action_name=task.task_type,
-                payload=task.payload,
-                attempt=task.attempt,
-                created_at=task.created_at,
-                namespace=self._parent_client.namespace,
-                _worker=self,
-            )
-
-            # Execute with timeout
+            self._running_handlers += 1
             try:
-                result = await asyncio.wait_for(
-                    handler(ctx),
-                    timeout=self.config.action_timeout,
-                )
+                op, send, cancelled = await self._run_handler(task)
+            finally:
+                self._running_handlers -= 1
+            # Reporting stays outside the handler's error handling: a complete
+            # that can't be delivered must not turn into a fail.
+            sent = await _send_result(
+                f"{op} for task {task.task_id}",
+                lambda: send(rc),
+                lambda: self._result_generation,
+                self._reconnect_result_client,
+                self._stop_event,
+            )
+            if sent and op == "complete":
+                logger.info(f"Action completed: {task.task_type} (task={task.task_id})")
+            if cancelled:
+                raise asyncio.CancelledError()
+        except Exception as e:
+            logger.error(f"Failed to report result for task {task.task_id}: {e}")
+        finally:
+            if self._semaphore is not None:
+                self._semaphore.release()
 
-                # 3-way dispatch based on result type
-                from .types import WorkerCompleteOptions
-
-                if isinstance(result, ActionResult):
-                    # Named outcome
-                    await self._send_with_retry(
-                        "complete",
-                        lambda: rc.worker.complete(
-                            self.config.worker_id,
-                            task.task_type,
-                            task.task_id,
-                            result.data,
-                            WorkerCompleteOptions(outcome=result.outcome),
-                        ),
-                    )
-                elif isinstance(result, dict):
-                    # Plain dict → JSON serialize
-                    result_bytes = json.dumps(result).encode("utf-8")
-                    await self._send_with_retry(
-                        "complete",
-                        lambda: rc.worker.complete(
-                            self.config.worker_id,
-                            task.task_type,
-                            task.task_id,
-                            result_bytes,
-                        ),
-                    )
-                else:
-                    # bytes or other → pass through
-                    await self._send_with_retry(
-                        "complete",
-                        lambda: rc.worker.complete(
-                            self.config.worker_id,
-                            task.task_type,
-                            task.task_id,
-                            result if isinstance(result, bytes) else b"",
-                        ),
-                    )
-                logger.info(f"Action completed: {task.task_type}")
+    async def _run_handler(self, task: TaskAssignment) -> "_Report":
+        """Run the task's handler. Its outcome alone decides complete or fail.
 
-            except asyncio.TimeoutError:
-                logger.error(f"Action timed out: {task.task_type}")
-                await self._send_with_retry(
-                    "fail",
-                    lambda: rc.worker.fail(
-                        self.config.worker_id,
-                        task.task_type,
-                        task.task_id,
-                        "Action timed out",
-                    ),
-                )
+        Returns the result to send, and whether the handler was cancelled, in
+        which case the caller re-raises the cancellation after reporting.
+        """
+        from .types import WorkerCompleteOptions, WorkerFailOptions
 
-            except asyncio.CancelledError:
-                logger.warning(f"Action cancelled: {task.task_type}")
-                await self._send_with_retry(
-                    "fail",
-                    lambda: rc.worker.fail(
-                        self.config.worker_id,
-                        task.task_type,
-                        task.task_id,
-                        "Action cancelled",
-                    ),
-                )
+        wid = self.config.worker_id
 
-            except Exception as exc:
-                logger.error(f"Action failed: {task.task_type} - {exc}")
-                from .types import WorkerFailOptions
-
-                retry = not isinstance(exc, NonRetryableError)
-                err_msg = str(exc)
-                await self._send_with_retry(
-                    "fail",
-                    lambda: rc.worker.fail(
-                        self.config.worker_id,
-                        task.task_type,
-                        task.task_id,
-                        err_msg,
-                        WorkerFailOptions(retry=retry),
-                    ),
-                )
+        def fail(message: str, retry: bool = True, cancelled: bool = False) -> _Report:
+            options = WorkerFailOptions(retry=retry)
+            return (
+                "fail",
+                lambda c: c.worker.fail(wid, task.task_type, task.task_id, message, options),
+                cancelled,
+            )
 
-        except Exception as e:
-            logger.error(f"Failed to report task result: {e}")
+        def complete(data: bytes, options: WorkerCompleteOptions | None = None) -> _Report:
+            return (
+                "complete",
+                lambda c: c.worker.complete(wid, task.task_type, task.task_id, data, options),
+                False,
+            )
 
+        logger.info(
+            f"Executing action: {task.task_type} (task={task.task_id}, attempt={task.attempt})"
+        )
+        handler = self._handlers.get(task.task_type)
+        if handler is None:
+            logger.error(f"No handler registered for action: {task.task_type}")
+            return fail(f"No handler for: {task.task_type}")
+
+        ctx = ActionContext(
+            task_id=task.task_id,
+            action_name=task.task_type,
+            payload=task.payload,
+            attempt=task.attempt,
+            created_at=task.created_at,
+            namespace=self._parent_client.namespace,
+            _worker=self,
+        )
+        # Not wait_for: a TimeoutError the handler raises itself would look
+        # the same as the deadline passing. Only the deadline is final.
+        run = asyncio.ensure_future(handler(ctx))
+        timeout_ms = self._action_options[task.task_type].timeout_ms
+        stop_touching = asyncio.Event()
+        touching = (
+            asyncio.ensure_future(self._keep_touching(task, timeout_ms, stop_touching))
+            if self.config.auto_touch and timeout_ms > 0
+            else None
+        )
+        try:
+            await asyncio.wait({run}, timeout=self.config.action_timeout)
+        except asyncio.CancelledError:
+            run.cancel()
+            logger.warning(f"Action cancelled: {task.task_type}")
+            return fail("Action cancelled", cancelled=True)
         finally:
-            if self._semaphore is not None:
-                self._semaphore.release()
+            if touching is not None:
+                # Let a touch in flight finish: cancelling it mid-request
+                # would close the result connection the outcome goes on.
+                stop_touching.set()
+                with contextlib.suppress(Exception):
+                    await asyncio.shield(touching)
+        if not run.done():
+            run.cancel()
+            with contextlib.suppress(BaseException):
+                await run
+            logger.error(f"Action timed out: {task.task_type}")
+            return fail("Action timed out", retry=False)
+        try:
+            result = run.result()
+        except asyncio.CancelledError:
+            logger.warning(f"Action cancelled: {task.task_type}")
+            return fail("Action cancelled", cancelled=True)
+        except Exception as exc:
+            logger.error(f"Action failed: {task.task_type} - {exc!r}")
+            message = str(exc) or type(exc).__name__
+            return fail(message, retry=not isinstance(exc, NonRetryableError))
+
+        if isinstance(result, ActionResult):
+            return complete(result.data, WorkerCompleteOptions(outcome=result.outcome))
+        if isinstance(result, dict):
+            return complete(json.dumps(result).encode("utf-8"))
+        return complete(result if isinstance(result, bytes) else b"")
 
     async def _touch_task(self, action_name: str, task_id: str, extend_ms: int) -> None:
-        """Extend lease on a task (internal method)."""
-        if self._client is None:
+        """Extend lease on a task (internal method).
+
+        On the result connection: the poll connection may be held by a
+        blocking await for up to block_ms, longer than a lease can wait.
+        """
+        client = self._result_client or self._client
+        if client is None:
             raise RuntimeError("Worker not connected")
-        await self._client.worker.touch(
+        await client.worker.touch(
             self.config.worker_id,
             action_name,
             task_id,
             WorkerTouchOptions(extend_ms=extend_ms),
         )
 
+    async def _keep_touching(
+        self, task: TaskAssignment, timeout_ms: int, stop: asyncio.Event
+    ) -> None:
+        """Touch a run at _touch_pace(timeout_ms) until stop is set.
+
+        Each touch asks for this worker's timeout (at most an hour), so the
+        lease covers it even if another worker's registration shortened the
+        action's. Conflict, not_found or bad_request means the run moved on
+        (finished, cancelled, taken back) and touching stops; any other
+        refusal, an overloaded leader say, is tried again next interval, and
+        a lost connection is reconnected for it.
+        """
+        interval_s, extend_ms = _touch_pace(timeout_ms)
+        while True:
+            try:
+                await asyncio.wait_for(stop.wait(), interval_s)
+                return
+            except asyncio.TimeoutError:
+                pass
+            seen = self._result_generation
+            try:
+                await self._touch_task(task.task_type, task.task_id, extend_ms)
+            except ServerError as e:
+                if e.status_code in _TOUCH_FINAL:
+                    logger.warning(f"Touch for task {task.task_id} was refused, stopping: {e}")
+                    return
+                logger.warning(f"Touch for task {task.task_id} was refused: {e}; trying again")
+            except Exception as e:
+                logger.warning(f"Touch for task {task.task_id} lost: {e!r}")
+                if is_connection_error(e) or isinstance(e, asyncio.TimeoutError):
+                    with contextlib.suppress(Exception):
+                        await self._reconnect_result_client(seen)
+
     def stop(self) -> None:
         """Signal the worker to stop.
 
-        This sets a flag that will cause the polling loop to exit
-        after the current iteration completes. Also interrupts in-flight
-        connections to unblock any blocking Await call immediately.
+        Ends the polling loop, interrupting a blocking Await. Running
+        handlers finish and their results are sent before start() returns.
         """
         logger.info("Stopping worker...")
         self._running = False
         self._stop_event.set()
-        # Interrupt connections to unblock any blocking Await
+        # Only the poll connection is interrupted: the result connection
+        # still has to carry the results of the handlers that are finishing.
         if self._client:
             self._client.interrupt()
-        if self._result_client:
-            self._result_client.interrupt()
 
     async def close(self) -> None:
         """Stop and close the worker."""
@@ -664,8 +943,20 @@ class StreamWorkerOptions:
     worker_id: str = ""
     concurrency: int = 10
     batch_size: int = 10
-    block_ms: int = 30000
+    block_ms: int = _DEFAULT_WORKER_BLOCK_MS
     message_timeout: float = 300.0  # 5 minutes
+    # Drain this consumer's pending (delivered-but-unacked) entries via
+    # group_claim after a reconnect, before resuming normal group_read.
+    # Required for at-least-once across reconnects: group_read alone only
+    # returns records past last_delivered_id. Set False for at-most-once.
+    redeliver_pending_on_reconnect: bool = True
+    # Minimum idle time (ms) an entry must have before the reconnect drain
+    # claims it. 0 = drain everything this consumer already owns; > 0 also
+    # steals entries abandoned by other (dead) consumers idle at least this long.
+    claim_min_idle_ms: int = 0
+
+    def __post_init__(self) -> None:
+        self.block_ms = _worker_block_ms(self.block_ms)
 
 
 @dataclass
@@ -747,8 +1038,10 @@ class StreamWorker:
         worker_id: str | None = None,
         concurrency: int = 10,
         batch_size: int = 10,
-        block_ms: int = 30000,
+        block_ms: int = _DEFAULT_WORKER_BLOCK_MS,
         message_timeout: float = 300.0,
+        redeliver_pending_on_reconnect: bool = True,
+        claim_min_idle_ms: int = 0,
     ):
         self._parent_client = parent_client
         self._handler = handler
@@ -761,9 +1054,16 @@ class StreamWorker:
             batch_size=batch_size,
             block_ms=block_ms,
             message_timeout=message_timeout,
+            redeliver_pending_on_reconnect=redeliver_pending_on_reconnect,
+            claim_min_idle_ms=claim_min_idle_ms,
         )
 
         self._client: FloClient | None = None
+        self._ack_client: FloClient | None = None
+        # Concurrent acks that fail on the same dead connection reconnect it
+        # once: the generation tells a waiter someone already did.
+        self._ack_reconnect_lock = asyncio.Lock()
+        self._ack_generation = 0
         self._running = False
         self._stop_event = asyncio.Event()
         self._tasks: set[asyncio.Task[None]] = set()
@@ -790,33 +1090,46 @@ class StreamWorker:
             f"concurrency={self.config.concurrency})"
         )
 
-        # Timeout must accommodate block_ms + message_timeout so blocking reads
-        # (group_read with block_ms) don't get killed by socket-level timeout.
-        worker_timeout_ms = max(
-            self.config.block_ms + 5000,
-            int(self.config.message_timeout * 1000),
-        )
+        worker_timeout_ms = int(self._parent_client._timeout * 1000)
         self._client = FloClient(
             self._parent_client._endpoint,
             namespace=self._parent_client.namespace,
             debug=self._parent_client._debug,
             timeout_ms=worker_timeout_ms,
+            retry=self._parent_client._retry,
         )
         await self._client.connect()
 
+        # Acks and nacks get their own connection: the polling one is held by
+        # a group_read for up to block_ms, and an ack queued behind it can
+        # arrive after the group's ack timeout and see the record redelivered.
+        self._ack_client = FloClient(
+            self._parent_client._endpoint,
+            namespace=self._parent_client.namespace,
+            debug=self._parent_client._debug,
+            timeout_ms=worker_timeout_ms,
+            retry=self._parent_client._retry,
+        )
         try:
-            # Join consumer group
-            await self._client.stream.group_join(
-                self.config.stream,
-                self.config.group,
-                self.config.consumer,
-            )
-            logger.info(f"Joined consumer group {self.config.group} on stream {self.config.stream}")
+            await self._ack_client.connect()
+        except BaseException:
+            await self._client.close()
+            raise
 
+        try:
+            # Semaphore and running state must be set up before the join: a
+            # connection error during the join triggers _handle_reconnect ->
+            # _drain_pending, which acquires the semaphore.
             self._semaphore = asyncio.Semaphore(self.config.concurrency)
             self._running = True
             self._stop_event.clear()
 
+            # Join consumer group. A connection error here (a stalled or dropped
+            # connection at startup) is retried via reconnect rather than
+            # killing the worker — mirroring the poll loop and ack retry.
+            await self._join_group()
+            logger.info(f"Joined consumer group {self.config.group} on stream {self.config.stream}")
+
             await self._poll_loop()
 
         finally:
@@ -839,6 +1152,9 @@ class StreamWorker:
 
                 await self._client.close()
                 self._client = None
+            if self._ack_client:
+                await self._ack_client.close()
+                self._ack_client = None
 
             self._running = False
             logger.info("Stream worker stopped")
@@ -847,6 +1163,9 @@ class StreamWorker:
         """Main polling loop for stream records."""
         assert self._client is not None
         assert self._semaphore is not None
+        loop = asyncio.get_running_loop()
+        backoff = _EmptyPollBackoff(self.config.block_ms, self._stop_event)
+        failed = _worker_backoff()
 
         while self._running and not self._stop_event.is_set():
             try:
@@ -857,6 +1176,7 @@ class StreamWorker:
                     self._semaphore.release()
                     break
 
+                started = loop.time()
                 result = await self._client.stream.group_read(
                     self.config.stream,
                     self.config.group,
@@ -866,10 +1186,13 @@ class StreamWorker:
                         block_ms=self.config.block_ms,
                     ),
                 )
+                failed.reset()
 
                 if not result.records:
                     self._semaphore.release()
+                    await backoff.empty(loop.time() - started)
                     continue
+                backoff.reset()
 
                 # Release semaphore before dispatching — each task will
                 # acquire its own slot via _process_record.
@@ -888,6 +1211,10 @@ class StreamWorker:
                 break
             except Exception as e:
                 self._semaphore.release()
+                if self._stop_event.is_set():
+                    # stop() interrupted the read. Reconnecting would also
+                    # reconnect the ack connection under the draining handlers.
+                    break
                 if is_connection_error(e):
                     logger.warning("Stream worker lost connection, reconnecting...")
                     try:
@@ -897,46 +1224,146 @@ class StreamWorker:
                         logger.error(f"Stream worker reconnect failed: {recon_err}, retrying...")
                         await asyncio.sleep(1)
                 else:
-                    logger.error(f"Stream read error: {e}, retrying...")
-                    await asyncio.sleep(1)
+                    d = failed.next()
+                    logger.error(f"Stream read error: {e}, retrying in {d:.3f}s...")
+                    await _sleep_unless_stopped(self._stop_event, d)
+
+    async def _join_group(self) -> None:
+        """Join the consumer group, retrying on connection errors.
+
+        A connection error (a stalled or dropped connection at startup) triggers
+        a reconnect-and-retry instead of failing the worker. A join the server
+        refused without running it (overloaded or unavailable, past the
+        client's own retries) is tried again with the worker backoff. Other
+        errors (e.g. an invalid request) remain fatal and are raised.
+        """
+        assert self._client is not None
+        failed = _worker_backoff()
+        while True:
+            if self._stop_event.is_set():
+                return
+            try:
+                await self._client.stream.group_join(
+                    self.config.stream,
+                    self.config.group,
+                    self.config.consumer,
+                )
+                return
+            except Exception as e:
+                if is_refused_unrun(e):
+                    d = failed.next()
+                    logger.warning(f"Group join refused: {e}, retrying in {d:.3f}s...")
+                    await _sleep_unless_stopped(self._stop_event, d)
+                    continue
+                if not is_connection_error(e):
+                    raise
+                logger.warning(f"Connection lost during group join: {e}, reconnecting...")
+                with contextlib.suppress(Exception):
+                    await self._handle_reconnect()
+                await asyncio.sleep(1)
 
     async def _handle_reconnect(self) -> None:
-        """Reconnect and re-join the consumer group."""
+        """Reconnect, re-join the consumer group, and drain pending entries.
+
+        After re-joining, this consumer's pending (delivered-but-unacked)
+        entries are re-processed via a group_claim cursor loop before normal
+        reads resume — group_read only returns records past last_delivered_id,
+        so without this an in-flight record at crash time is never re-surfaced.
+        See ``redeliver_pending_on_reconnect``.
+        """
         assert self._client is not None
         await self._client.reconnect()
+        # The ack connection went down with the poll connection (e.g. a server
+        # restart) but may not notice until an ack fails on it.
+        try:
+            await self._reconnect_ack_client()
+        except Exception as e:
+            logger.warning(f"Failed to reconnect ack client: {e}")
         await self._client.stream.group_join(
             self.config.stream,
             self.config.group,
             self.config.consumer,
         )
+        if self.config.redeliver_pending_on_reconnect:
+            await self._drain_pending()
 
-    async def _ack_with_retry(
-        self, record_ids: list[StreamID], options: StreamGroupAckOptions
-    ) -> None:
-        """Ack with retry on connection error."""
-        max_attempts = 3
-        for attempt in range(1, max_attempts + 1):
+    async def _drain_pending(self) -> None:
+        """Re-process this consumer's pending entries via a group_claim loop.
+
+        Records flow through the same _process_record path (and concurrency
+        gate) as live reads, so a handler that succeeds acks and clears the PEL
+        entry. Bounded per stream by a max page count so a pathological backlog
+        can't wedge the reconnect indefinitely; leftovers are reclaimed on the
+        next reconnect or by group_read once acked.
+        """
+        if self._semaphore is None or self._client is None:
+            # start() has not initialized the concurrency gate yet.
+            return
+
+        max_pages = 1000  # safety cap: max_pages * batch_size entries
+        cursor = StreamID()  # MIN — scan from the start of the PEL
+        pages = 0
+        while True:
+            if self._stop_event.is_set():
+                return
             try:
-                assert self._client is not None
-                await self._client.stream.group_ack(
+                result = await self._client.stream.group_claim(
                     self.config.stream,
                     self.config.group,
-                    record_ids,
-                    options,
+                    self.config.consumer,
+                    self.config.claim_min_idle_ms,
+                    cursor,
+                    self.config.batch_size,
                 )
-                return
             except Exception as e:
-                if not is_connection_error(e) or not self._running:
-                    raise
-                logger.warning(
-                    "Connection lost while acking "
-                    f"(attempt {attempt}/{max_attempts}), reconnecting..."
-                )
-                await self._handle_reconnect()
+                # Don't fail the reconnect on a drain error; normal reads resume.
+                logger.warning(f"Pending drain claim failed: {e}")
+                return
+
+            for record in result.records:
+                await self._semaphore.acquire()  # same gate as live reads
+                if self._stop_event.is_set():
+                    self._semaphore.release()
+                    return
+                task = asyncio.create_task(self._process_record(record))
+                self._tasks.add(task)
+                task.add_done_callback(self._tasks.discard)
+
+            pages += 1
+            if result.done or not result.records or pages >= max_pages:
+                return
+            cursor = result.next_cursor
+
+    async def _reconnect_ack_client(self, failed_generation: int | None = None) -> None:
+        """Reconnect the ack connection, unless it was reconnected since
+        `failed_generation` was observed."""
+        assert self._ack_client is not None
+        async with self._ack_reconnect_lock:
+            if failed_generation is not None and failed_generation != self._ack_generation:
+                return
+            await self._ack_client.reconnect()
+            self._ack_generation += 1
+
+    async def _on_ack_connection(
+        self,
+        op: str,
+        record_id: StreamID,
+        send: Callable[[FloClient], Awaitable[object]],
+    ) -> bool:
+        """Send an ack or nack on the ack connection, resending it after a
+        reconnect if the connection drops. Never raises; False if dropped."""
+        ack_client = self._ack_client
+        assert ack_client is not None
+        return await _send_result(
+            f"{op} for record {record_id}",
+            lambda: send(ack_client),
+            lambda: self._ack_generation,
+            self._reconnect_ack_client,
+            self._stop_event,
+        )
 
     async def _process_record(self, record: StreamRecord) -> None:
         """Process a single record: call handler, then ack or nack."""
-        assert self._client is not None
         assert self._semaphore is not None
         try:
             ctx = StreamContext(
@@ -946,36 +1373,41 @@ class StreamWorker:
                 group=self.config.group,
                 consumer=self.config.consumer,
             )
-
             try:
-                await asyncio.wait_for(
-                    self._handler(ctx),
-                    timeout=self.config.message_timeout,
-                )
-
-                # Success — ack with retry
-                await self._ack_with_retry(
-                    [record.id],
-                    StreamGroupAckOptions(consumer=self.config.consumer),
-                )
-
+                await asyncio.wait_for(self._handler(ctx), timeout=self.config.message_timeout)
+                handled = True
             except Exception as e:
                 logger.error(
                     f"Record processing failed (stream={self.config.stream}, id={record.id}): {e}"
                 )
-                try:
-                    await self._client.stream.group_nack(
+                handled = False
+
+            # Outside the handler's error handling: an ack that can't be
+            # delivered must not turn into a nack.
+            if handled:
+                await self._on_ack_connection(
+                    "ack",
+                    record.id,
+                    lambda c: c.stream.group_ack(
+                        self.config.stream,
+                        self.config.group,
+                        [record.id],
+                        StreamGroupAckOptions(consumer=self.config.consumer),
+                    ),
+                )
+            else:
+                await self._on_ack_connection(
+                    "nack",
+                    record.id,
+                    lambda c: c.stream.group_nack(
                         self.config.stream,
                         self.config.group,
                         [record.id],
                         StreamGroupNackOptions(consumer=self.config.consumer),
-                    )
-                except Exception as nack_err:
-                    logger.error(f"Failed to nack record: {nack_err}")
-
+                    ),
+                )
         except Exception as e:
             logger.error(f"Failed to process record: {e}")
-
         finally:
             self._semaphore.release()
 
@@ -992,6 +1424,8 @@ class StreamWorker:
         self.stop()
         if self._client:
             await self._client.close()
+        if self._ack_client:
+            await self._ack_client.close()
 
     async def __aenter__(self) -> "StreamWorker":
         """Async context manager entry."""