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

flo/client.py

@@ -10,19 +10,31 @@ from typing import TYPE_CHECKING
 
 from .exceptions import (
     ConnectionFailedError,
+    InvalidChecksumError,
     InvalidEndpointError,
     NotConnectedError,
+    ProtocolError,
+    RequestTimeoutError,
     UnexpectedEofError,
     raise_for_status,
 )
-from .types import HEADER_SIZE, OpCode, StatusCode
-from .wire import RawResponse, parse_response_header, serialize_request
+from .retry import DEFAULT_RETRY_POLICY, RetryPolicy, _policy_backoff, _refused_unrun
+from .types import _DEFAULT_WORKER_BLOCK_MS, HEADER_SIZE, OpCode, OptionTag, StatusCode
+from .wire import (
+    OptionsIterator,
+    RawResponse,
+    compute_crc32,
+    parse_response_header,
+    serialize_request,
+)
 
 if TYPE_CHECKING:
     from .worker import ActionWorker, StreamRecordHandler, StreamWorker
 
 logger = logging.getLogger("flo")
 
+_SERVER_AWAIT_DEFAULT_MS = 30_000
+
 
 class FloClient:
     """Async client for Flo server.
@@ -40,19 +52,28 @@ class FloClient:
         namespace: str = "default",
         timeout_ms: int = 5000,
         debug: bool = False,
+        retry: RetryPolicy = DEFAULT_RETRY_POLICY,
     ):
         """Initialize Flo client.
 
         Args:
             endpoint: Server endpoint in "host:port" format.
             namespace: Default namespace for operations.
-            timeout_ms: Connection and operation timeout in milliseconds.
+            timeout_ms: Connection and operation timeout in milliseconds. A
+                blocking call (block_ms) is allowed this long on
+                top of its wait. A call that times out or is cancelled drops
+                the connection; call reconnect() before reusing the client.
             debug: Enable debug logging.
+            retry: How a request the server refused without running it
+                (overloaded or unavailable, ran = no) is sent again; backoffs
+                are in milliseconds. RetryPolicy(max_attempts=1) turns
+                retrying off. Workers made from this client use it too.
         """
         self._endpoint = endpoint
         self._namespace = namespace
         self._timeout = timeout_ms / 1000.0  # Convert to seconds
         self._debug = debug
+        self._retry = retry
 
         self._reader: asyncio.StreamReader | None = None
         self._writer: asyncio.StreamWriter | None = None
@@ -83,7 +104,27 @@ class FloClient:
 
     @staticmethod
     def _parse_endpoint(endpoint: str) -> tuple[str, int]:
-        """Parse endpoint string into host and port."""
+        """Parse endpoint string into host and port.
+
+        Accepts an optional URI scheme: ``flo://`` (plaintext) is stripped and
+        bare ``host:port`` remains valid. ``flos://`` (TLS) is reserved but not
+        yet implemented, and any other scheme is rejected so it fails here
+        rather than as a confusing DNS error at connect time.
+        """
+        if "://" in endpoint:
+            scheme, _, rest = endpoint.partition("://")
+            scheme = scheme.lower()
+            if scheme == "flo":
+                # Drop scheme and any trailing path/query (reserved for future use).
+                endpoint = rest.split("/", 1)[0]
+            elif scheme == "flos":
+                raise InvalidEndpointError(f"flos:// (TLS) is not yet supported: {endpoint}")
+            else:
+                raise InvalidEndpointError(
+                    f"Unsupported scheme {scheme}:// (expected flo:// or bare "
+                    f"host:port): {endpoint}"
+                )
+
         parts = endpoint.rsplit(":", 1)
         if len(parts) != 2:
             raise InvalidEndpointError(f"Invalid endpoint format: {endpoint} (expected host:port)")
@@ -167,13 +208,16 @@ class FloClient:
             timeout: Maximum total time to keep retrying (default 5 min).
         """
         async with self._lock:
-            # Close existing connection
-            if self._writer:
-                self._writer.close()
-                with contextlib.suppress(Exception):
-                    await self._writer.wait_closed()
+            # Close existing connection. It is detached before the await so a
+            # concurrent reconnect that installs its connection meanwhile
+            # keeps it.
+            old = self._writer
+            if old:
                 self._writer = None
                 self._reader = None
+                old.close()
+                with contextlib.suppress(Exception):
+                    await old.wait_closed()
 
         backoff = 1.0
         max_backoff = 30.0
@@ -183,10 +227,16 @@ class FloClient:
         while True:
             attempt += 1
             try:
-                self._reader, self._writer = await asyncio.wait_for(
+                reader, writer = await asyncio.wait_for(
                     asyncio.open_connection(self._host, self._port),
                     timeout=self._timeout,
                 )
+                # A concurrent reconnect may have installed its connection
+                # first; close it rather than leak it.
+                replaced = self._writer
+                self._reader, self._writer = reader, writer
+                if replaced is not None:
+                    replaced.close()
                 if self._debug:
                     logger.debug(f"[flo] Reconnected to {self._endpoint} (attempt {attempt})")
                 return
@@ -224,6 +274,24 @@ class FloClient:
         self._request_id += 1
         return self._request_id
 
+    @staticmethod
+    def _server_wait_ms(op_code: OpCode, value: bytes, options: bytes) -> int:
+        """How long the server may hold this request before answering."""
+        if op_code == OpCode.KV_GET:
+            # kv.get's body opens with wait_ms: [present:u8][u64].
+            if len(value) >= 9 and value[0] == 1:
+                return int.from_bytes(value[1:9], "little")
+            return 0
+        waits = {
+            opt.tag: opt.as_u32() or 0
+            for opt in OptionsIterator(options)
+            if opt.tag == OptionTag.BLOCK_MS
+        }
+        if op_code == OpCode.ACTION_AWAIT:
+            # An await that names no block_ms waits the server's default.
+            waits.setdefault(OptionTag.BLOCK_MS, _SERVER_AWAIT_DEFAULT_MS)
+        return max(waits.values(), default=0)
+
     async def _send_request(
         self,
         op_code: OpCode,
@@ -232,7 +300,35 @@ class FloClient:
         value: bytes,
         options: bytes = b"",
     ) -> RawResponse:
-        """Send a request and wait for response.
+        """Send a request and wait for its response, sending it again as the
+        retry policy allows while the server refuses it without running it.
+
+        Each try is a new request with a new id. The wait between tries is
+        taken outside the connection lock, so other calls proceed meanwhile.
+        """
+        wait = _policy_backoff(self._retry)
+        attempt = 1
+        while True:
+            response = await self._send_once(op_code, namespace, key, value, options)
+            if attempt >= self._retry.max_attempts or not _refused_unrun(
+                response.status, response.data
+            ):
+                return response
+            d = wait.next()
+            if self._debug:
+                logger.debug(f"[flo] {response.status.name}, not run: retrying in {d:.3f}s")
+            await asyncio.sleep(d)
+            attempt += 1
+
+    async def _send_once(
+        self,
+        op_code: OpCode,
+        namespace: str,
+        key: bytes,
+        value: bytes,
+        options: bytes = b"",
+    ) -> RawResponse:
+        """Send a request, under a new request id, and wait for its response.
 
         Args:
             op_code: Operation code.
@@ -252,6 +348,10 @@ class FloClient:
             raise NotConnectedError("Not connected to server")
 
         async with self._lock:
+            # A call that held the lock before this one may have dropped the
+            # connection.
+            if self._writer is None or self._reader is None:
+                raise NotConnectedError("Not connected to server")
             request_id = self._next_request_id()
 
             # Serialize request
@@ -267,46 +367,74 @@ class FloClient:
             if self._debug:
                 logger.debug(f"[flo] -> {op_code.name} ns={namespace} key={key!r}")
 
-            # Send request
-            self._writer.write(request)
-            await self._writer.drain()
-
-            # Read response header
+            # The server may hold a blocking request for its whole wait, so it
+            # gets the operation timeout on top of that wait.
+            read_timeout = self._timeout + self._server_wait_ms(op_code, value, options) / 1000.0
+            writer, reader = self._writer, self._reader
             try:
-                header_data = await asyncio.wait_for(
-                    self._reader.readexactly(HEADER_SIZE),
-                    timeout=self._timeout,
-                )
-            except asyncio.IncompleteReadError as e:
-                raise UnexpectedEofError("Connection closed while reading response header") from e
-
-            # Parse header to get data length
-            status, data_len, resp_request_id, crc = parse_response_header(header_data)
+                writer.write(request)
+                await writer.drain()
 
-            # Read response data
-            if data_len > 0:
                 try:
-                    response_data = await asyncio.wait_for(
-                        self._reader.readexactly(data_len),
-                        timeout=self._timeout,
+                    header_data = await asyncio.wait_for(
+                        reader.readexactly(HEADER_SIZE),
+                        timeout=read_timeout,
                     )
                 except asyncio.IncompleteReadError as e:
-                    raise UnexpectedEofError("Connection closed while reading response data") from e
-            else:
-                response_data = b""
+                    raise UnexpectedEofError(
+                        "Connection closed while reading response header"
+                    ) from e
+                except asyncio.TimeoutError as e:
+                    raise RequestTimeoutError(f"No response to {op_code.name}") from e
+
+                status, data_len, resp_request_id, crc = parse_response_header(header_data)
+
+                if data_len > 0:
+                    try:
+                        response_data = await asyncio.wait_for(
+                            reader.readexactly(data_len),
+                            timeout=self._timeout,
+                        )
+                    except asyncio.IncompleteReadError as e:
+                        raise UnexpectedEofError(
+                            "Connection closed while reading response data"
+                        ) from e
+                    except asyncio.TimeoutError as e:
+                        raise RequestTimeoutError(
+                            f"Response to {op_code.name} stalled mid-frame"
+                        ) from e
+                else:
+                    response_data = b""
+
+                computed_crc = compute_crc32(header_data, response_data)
+                if computed_crc != crc:
+                    raise InvalidChecksumError(
+                        f"CRC32 mismatch: 0x{computed_crc:08X} != 0x{crc:08X}"
+                    )
+                if resp_request_id == 0:
+                    # The server answers a request it could not parse with
+                    # id 0 and then closes; surface its error.
+                    raise_for_status(status, response_data)
+                if resp_request_id != request_id:
+                    raise ProtocolError(
+                        f"Response is for request {resp_request_id}, expected {request_id}"
+                    )
+            except BaseException:
+                # The reply to an abandoned request (timed out, cancelled, cut
+                # off mid-frame) may still arrive and would be read as the
+                # answer to the next one, and after a corrupt or mismatched
+                # frame the stream position is unknown, so the connection is
+                # dropped; the caller reconnects.
+                writer.close()
+                # reconnect() swaps the connection outside the lock.
+                if self._writer is writer:
+                    self._writer = None
+                    self._reader = None
+                raise
 
             if self._debug:
                 logger.debug(f"[flo] <- {status.name} {len(response_data)} bytes")
 
-            # Verify CRC32
-            from .wire import compute_crc32
-
-            computed_crc = compute_crc32(header_data, response_data)
-            if computed_crc != crc:
-                from .exceptions import InvalidChecksumError
-
-                raise InvalidChecksumError(f"CRC32 mismatch: 0x{computed_crc:08X} != 0x{crc:08X}")
-
             return RawResponse(status=status, data=response_data, request_id=resp_request_id)
 
     async def _send_and_check(
@@ -356,7 +484,8 @@ class FloClient:
         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,
     ) -> "ActionWorker":
         """Create a new ActionWorker from this client.
 
@@ -367,7 +496,11 @@ class FloClient:
             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.
 
         Returns:
             A new ActionWorker instance ready to register actions and start.
@@ -388,6 +521,7 @@ class FloClient:
             concurrency=concurrency,
             action_timeout=action_timeout,
             block_ms=block_ms,
+            auto_touch=auto_touch,
         )
 
     def new_stream_worker(
@@ -400,8 +534,10 @@ class FloClient:
         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,
     ) -> "StreamWorker":
         """Create a new StreamWorker from this client.
 
@@ -417,8 +553,17 @@ class FloClient:
             worker_id: Unique worker identifier (auto-generated if not provided).
             concurrency: Maximum number of concurrent record handlers.
             batch_size: Number of records to fetch per poll.
-            block_ms: Timeout for blocking read in milliseconds.
+            block_ms: Long-poll wait per read in milliseconds, at most 300000.
+                0 means the 30000 default.
             message_timeout: Timeout for record handlers in seconds.
+            redeliver_pending_on_reconnect: Drain this consumer's pending
+                (delivered-but-unacked) entries via group_claim after a
+                reconnect, before resuming normal reads. Defaults to True
+                (required for at-least-once across reconnects). Set False for
+                at-most-once.
+            claim_min_idle_ms: Minimum idle time (ms) before the reconnect drain
+                claims an entry. 0 drains everything this consumer owns; > 0
+                also steals entries abandoned by dead consumers.
 
         Returns:
             A new StreamWorker instance ready to start.
@@ -446,4 +591,6 @@ class FloClient:
             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,
         )

flo/exceptions.py

@@ -3,9 +3,11 @@
 Exception classes for Flo client errors.
 """
 
+import asyncio
 import contextlib
+from collections.abc import Callable
 
-from .types import StatusCode
+from .types import TABLE_HASH, VERSION, Ran, Reason, StatusCode
 
 
 class FloError(Exception):
@@ -56,6 +58,12 @@ class UnexpectedEofError(FloError):
     pass
 
 
+class RequestTimeoutError(FloError, asyncio.TimeoutError):
+    """The server did not answer a request in time; the connection was dropped."""
+
+    pass
+
+
 # =============================================================================
 # Protocol Errors
 # =============================================================================
@@ -73,10 +81,23 @@ class InvalidMagicError(ProtocolError):
     pass
 
 
-class UnsupportedVersionError(ProtocolError):
-    """Unsupported protocol version."""
-
-    pass
+class TableMismatchError(ProtocolError):
+    """The server was built from another protocol version or op table, so
+    neither can read the other. Upgrade the SDK and the server together."""
+
+    def __init__(self, server_version: int, server_table: int) -> None:
+        self.server_version = server_version
+        self.server_table = server_table
+        if server_version != VERSION:
+            message = (
+                f"server protocol {server_version}, client protocol {VERSION}: upgrade the client"
+            )
+        else:
+            message = (
+                f"server table 0x{server_table:016x}, client table 0x{TABLE_HASH:016x}: "
+                "upgrade the client"
+            )
+        super().__init__(message)
 
 
 class InvalidChecksumError(ProtocolError):
@@ -85,12 +106,6 @@ class InvalidChecksumError(ProtocolError):
     pass
 
 
-class InvalidReservedFieldError(ProtocolError):
-    """Reserved field contains non-zero value."""
-
-    pass
-
-
 class PayloadTooLargeError(ProtocolError):
     """Payload exceeds maximum allowed size."""
 
@@ -132,6 +147,12 @@ class ValueTooLargeError(ValidationError):
     pass
 
 
+class BlockTooLongError(ValidationError):
+    """A blocking wait (block_ms) exceeds 300000 ms (5 minutes)."""
+
+    pass
+
+
 # =============================================================================
 # Server Response Errors
 # =============================================================================
@@ -141,6 +162,10 @@ class ServerError(FloError):
     """Base class for server-returned errors."""
 
     status_code: StatusCode
+    # Why; Reason.UNCLASSIFIED when the server didn't say.
+    reason: Reason = Reason.UNCLASSIFIED
+    # Whether the request took effect; unknown unless the server said.
+    ran: Ran = Ran.UNKNOWN
 
     def __init__(self, message: str, status_code: StatusCode):
         super().__init__(message)
@@ -189,26 +214,42 @@ class RateLimitedError(ServerError):
         super().__init__(message, StatusCode.RATE_LIMITED)
 
 
+class UnavailableError(ServerError):
+    """No leader, or the shard isn't taking writes or is offline. Retryable.
+
+    The message says why; an offline shard stays unavailable until an
+    operator acts, so whether and when to retry is the caller's call.
+    """
+
+    def __init__(self, message: str = StatusCode.UNAVAILABLE.message()):
+        super().__init__(message, StatusCode.UNAVAILABLE)
+
+
 class InternalServerError(ServerError):
-    """Internal server error."""
+    """Internal server error. Not retryable: the server uses it for a write
+    that committed but wasn't applied, which must not be resent."""
 
     def __init__(self, message: str = "Internal server error"):
         super().__init__(message, StatusCode.INTERNAL_ERROR)
 
 
 class GenericServerError(ServerError):
-    """Generic server error."""
+    """Generic server error, or a status with no dedicated error type."""
 
-    def __init__(self, message: str = "Generic error"):
-        super().__init__(message, StatusCode.ERROR_GENERIC)
+    def __init__(
+        self, message: str = "Generic error", status_code: StatusCode = StatusCode.ERROR_GENERIC
+    ):
+        super().__init__(message, status_code)
 
 
 def is_connection_error(exc: BaseException) -> bool:
     """Return True if the exception indicates a broken connection.
 
-    Connection errors may be resolved by reconnecting.
+    Connection errors may be resolved by reconnecting. A timed-out request
+    counts: the client drops its connection rather than risk reading the late
+    reply as the next answer.
     """
-    return isinstance(exc, (UnexpectedEofError, NotConnectedError))
+    return isinstance(exc, (UnexpectedEofError, NotConnectedError, RequestTimeoutError, OSError))
 
 
 def raise_for_status(status: StatusCode, data: bytes = b"") -> None:
@@ -224,22 +265,41 @@ def raise_for_status(status: StatusCode, data: bytes = b"") -> None:
     if status == StatusCode.OK:
         return
 
-    # Try to decode error message from data
+    # A refusal's body: [reason:u16][ran:u8][message].
+    if len(data) < 3:
+        raise IncompleteResponseError(f"a refusal body of {len(data)} bytes")
+    try:
+        reason = Reason(int.from_bytes(data[0:2], "little"))
+        ran = Ran(data[2])
+    except ValueError as e:
+        raise IncompleteResponseError(f"a refusal with reason and ran {data[0:3].hex()}") from e
     message = status.message()
-    if data:
+    if len(data) > 3:
         with contextlib.suppress(UnicodeDecodeError, ValueError):
-            message = data.decode("utf-8")
+            message = data[3:].decode("utf-8")
+    error = _error_for(status, message)
+    error.reason = reason
+    error.ran = ran
+    raise error
+
+
+def _error_for(status: StatusCode, message: str) -> ServerError:
 
-    error_map = {
+    error_map: dict[StatusCode, Callable[[str], ServerError]] = {
         StatusCode.NOT_FOUND: NotFoundError,
         StatusCode.BAD_REQUEST: BadRequestError,
         StatusCode.CONFLICT: ConflictError,
         StatusCode.UNAUTHORIZED: UnauthorizedError,
         StatusCode.OVERLOADED: OverloadedError,
+        StatusCode.UNAVAILABLE: UnavailableError,
         StatusCode.RATE_LIMITED: RateLimitedError,
         StatusCode.INTERNAL_ERROR: InternalServerError,
         StatusCode.ERROR_GENERIC: GenericServerError,
     }
 
-    error_class = error_map.get(status, GenericServerError)
-    raise error_class(message)
+    error_class = error_map.get(status)
+    if error_class is not None:
+        return error_class(message)
+    if status.name not in StatusCode.__members__:
+        message = f"Unknown status {int(status)}: {message}"
+    return GenericServerError(message, status)