Flow control for reactive readers

.claude/rules/python-annotate-every-parameter.md

@@ -0,0 +1,21 @@
+# Annotate every Python parameter and return type
+
+Give every parameter of every Python function and method a type
+annotation, and every function a return type — including `self`-less
+helpers, nested functions, and parameters whose type feels obvious
+from the name. Leaving one off because `mypy` doesn't demand it is
+not a reason to leave it off.
+
+**Why:** `mypy` only checks a function's body once its signature is
+annotated, so a single un-annotated parameter silently switches off
+checking for everything that flows through it. An annotation is also
+the cheapest documentation a reader gets: `stub` says nothing,
+`stub: react_pb2_grpc.ReactStub` says where to look next.
+
+**How to apply:** When writing or editing a `def`, annotate all of
+its parameters and its return type in the same edit. When a
+parameter's type is a generated protobuf message or gRPC stub, name
+that type rather than falling back to `Any` — import it if the module
+doesn't already. `self` and `cls` need no annotation, and neither do
+`*args`/`**kwargs` when they are forwarded verbatim to an already
+typed callee.

rbt/v1alpha1/options.proto

@@ -17,6 +17,13 @@ message ReaderMethodOptions {
     STREAMING = 2;
   }
   State state = 3;
+
+  // Whether the server warns, in its log, every time a client of a
+  // reactive read of this method falls behind and has updates skipped
+  // for it. Unset means yes. A method whose readers only ever want
+  // its latest state, so that skipping the states in between is the
+  // point, sets this to `false`.
+  optional bool warn_on_flow_control = 4;
 }
 
 message WriterMethodOptions {

rbt/v1alpha1/react.proto

@@ -18,6 +18,25 @@ message QueryRequest {
 
   // Authorization bearer token.
   optional string bearer_token = 3;
+
+  // Whether this client tells us which responses it has processed,
+  // so that we can stay at most a bounded number of responses ahead
+  // of it. A client that does not set this gets responses at whatever
+  // rate we produce them, whether or not it can consume them that
+  // fast.
+  bool client_continues_query = 4;
+
+  // Sequence number of the last `QueryResponse` this client has fully
+  // processed, which returns to the server the right to send that
+  // many more; see `Query`. This is how a client whose responses
+  // arrive over a websocket continues its query, on that same
+  // websocket, and it is the only way that works for such a client.
+  // Only meaningful on transports where a client can send more than
+  // one message per query, i.e. websockets. When this is set the
+  // other fields are ignored. A heartbeat is an otherwise empty
+  // `QueryRequest`, which is why presence matters here: sequence
+  // number zero is a real response.
+  optional uint64 continue_query_sequence_number = 5;
 }
 
 message QueryResponse {
@@ -40,8 +59,46 @@ message QueryResponse {
   //
   // TODO(benh): send these as bytes?
   repeated string idempotency_keys = 2;
+
+  // Unique ID of the query this response belongs to, which a client
+  // names when it continues that query. The same for every response
+  // of one `Query`. Only set when the client set
+  // `client_continues_query` on its `QueryRequest`; when unset, the
+  // client is not expected to continue anything.
+  optional string query_id = 4;
+
+  // Position of this response within its query, counting from zero.
+  // A client names the last one it has fully processed to continue
+  // the query; see `Query`.
+  uint64 sequence_number = 5;
+
+  // How long this response waited for the client to make room for it,
+  // and how many updates were merged into it while it waited. Both
+  // are zero for a client that kept up. A client that finds them set
+  // is one whose user saw a stale screen for that long, and which
+  // missed that many updates on the way to the state this response
+  // carries; see `Query`.
+  uint32 stall_milliseconds = 6;
+  uint32 skipped_updates = 7;
+}
+
+////////////////////////////////////////////////////////////////////////
+
+message ContinueQueryRequest {
+  // ID of the query being continued, from `QueryResponse.query_id`.
+  string query_id = 1;
+
+  // Sequence number of the last `QueryResponse` the client has fully
+  // processed. The server returns to itself the right to send one
+  // more response for each response between the last number this
+  // client named and this one, so a client that processed several
+  // while a previous `ContinueQuery` was in flight names only the
+  // newest and is credited for all of them.
+  uint64 sequence_number = 2;
 }
 
+message ContinueQueryResponse {}
+
 ////////////////////////////////////////////////////////////////////////
 
 // TODO(benh): support batch, e.g., create a `MutateRequests` which
@@ -90,8 +147,42 @@ message WebSocketsConnectionResponse {}
 service React {
   // Allows users to "watch" the response of a unary reader method for
   // changes. The current response of the method is sent when the
-  // stream is opened, then a new response is sent whenever the
-  // state changes in a way that causes the method's response to change.
+  // stream is opened.
+  //
+  // A client that sets `client_continues_query` is sent responses
+  // under a window: the server may run up to a fixed number of
+  // responses ahead of what that client has continued past, and each
+  // response it sends spends one of them. While the window has
+  // room the server sends every state change as it happens, so a
+  // client that keeps up sees every update at the rate they occur.
+  // Once the window is empty the server holds the response it has
+  // ready and merges every state that follows into it, so when the
+  // window reopens the client is sent a single response carrying the
+  // _latest_ state, with `stall_milliseconds` and `skipped_updates`
+  // saying how long it waited and how many updates that one response
+  // stood in for. So a client that keeps up loses nothing, and a
+  // client that cannot is at most a window behind rather than
+  // arbitrarily behind. A client that does not set
+  // `client_continues_query` gets responses at whatever rate the
+  // server produces them, whether or not it can consume them
+  // that fast.
+  //
+  // To continue a query a client names the sequence number of the
+  // last response it has fully processed, and the server credits it
+  // for every response up to that one. A client may therefore process
+  // several responses while a single `ContinueQuery` is in flight and
+  // name only the newest.
+  //
+  // How a client continues its query is determined by the transport
+  // its responses arrive on, and only that way works: a client
+  // reading over a websocket sends a `QueryRequest` with
+  // `continue_query_sequence_number` set on that same websocket, and
+  // a client reading over any other transport calls `ContinueQuery`.
+  // The server tracks each query's window per transport, so one that
+  // arrives the other way names a query the server has no record of,
+  // which it ignores: the request is accepted, no room is returned,
+  // and the query stops producing responses once its window
+  // empties.
   //
   // NOTE: the service name and state ID is expected to be part of the
   // gRPC metadata, exactly like any other reboot requests.
@@ -101,6 +192,20 @@ service React {
     };
   }
 
+  // Tells the server which responses from `Query` a client has fully
+  // processed, returning that much room to the query's window. Must
+  // be called on the same server that produced those responses.
+  //
+  // For a client whose responses arrive over a websocket this is the
+  // wrong mechanism; such a client continues its query on the
+  // websocket that carries its responses, by sending a `QueryRequest`
+  // with `continue_query_sequence_number` set.
+  rpc ContinueQuery(ContinueQueryRequest) returns (ContinueQueryResponse) {
+    option (google.api.http) = {
+      post: "/rbt.v1alpha1.React/ContinueQuery"
+    };
+  }
+
   // All connections must be HTTP/2 (h2), but the WebSocket protocol
   // is inherently HTTP/1.1. RFC 8441
   // (https://datatracker.ietf.org/doc/html/rfc8441) introduced

reboot/aio/contexts.py

@@ -342,6 +342,60 @@ def from_grpc_metadata(cls, metadata: GrpcMetadata) -> 'Participants':
         return participants
 
 
+class QueryContinuations:
+    """Calls `React.ContinueQuery` for the responses of a `React.Query`
+    that a client has fully processed, which is what returns room to
+    that query's window.
+
+    Keeps at most one call in flight. While one is on its way the
+    client keeps processing, so the next call names only the newest
+    response it has handled; the server credits every response up to
+    that one, so nothing is lost by not naming each individually.
+    """
+
+    def __init__(
+        self,
+        stub: react_pb2_grpc.ReactStub,
+        metadata: GrpcMetadata,
+    ) -> None:
+        self._stub = stub
+        self._metadata = metadata
+        self._continuation: Optional[asyncio.Task] = None
+
+    async def continue_past(
+        self,
+        query_response: react_pb2.QueryResponse,
+    ) -> None:
+        """Continues the query past `query_response`, if there is no
+        call already in flight. Raises whatever a previous call raised,
+        so that a failure to continue surfaces on the read."""
+        if not query_response.HasField('query_id'):
+            # An older backend doesn't send a query ID and doesn't
+            # expect to be told.
+            return
+
+        if self._continuation is not None and self._continuation.done():
+            await self._continuation
+            self._continuation = None
+
+        if self._continuation is None:
+            self._continuation = asyncio.ensure_future(
+                self._stub.ContinueQuery(
+                    react_pb2.ContinueQueryRequest(
+                        query_id=query_response.query_id,
+                        sequence_number=query_response.sequence_number,
+                    ),
+                    # The same metadata ensures we're routed to the
+                    # same server.
+                    metadata=self._metadata,
+                )
+            )
+
+    async def stop(self) -> None:
+        await wait_for_tasks([self._continuation], cancel=True)
+        self._continuation = None
+
+
 class React:
     """Encapsulates machinery necessary for contexts that are "reactive",
     aka, those that are initiated from calls to `React.Query` and who
@@ -474,7 +528,9 @@ async def query():
                         self._state_type_name, self._state_ref
                     )
 
-                    call = react_pb2_grpc.ReactStub(channel).Query(
+                    stub = react_pb2_grpc.ReactStub(channel)
+
+                    call = stub.Query(
                         react_pb2.QueryRequest(
                             method=self._method,
                             request=serialized_request,
@@ -497,26 +553,39 @@ async def query():
                     async def loop():
                         assert task is not None
 
+                        # Keep consuming however fast they arrive,
+                        # letting each response replace the one before
+                        # it. A reader that reads faster than this one
+                        # can re-run therefore has its updates
+                        # accumulated into the latest state, which is
+                        # the one waiting here when this reader next
+                        # asks for it. See
+                        # https://github.com/reboot-dev/mono/issues/4754.
+                        #
+                        # Replacing drops the idempotency keys the
+                        # response carried, which is deliberate: the
+                        # keys a reactive read reports upward are the
+                        # ones for its own state's mutations, which
+                        # `reactively()` aggregates for it, and a
+                        # browser observes a mutation on the query of
+                        # the state it mutated. Keys kept here would
+                        # reach no one.
                         async for query_response in call:
-                            if not query_response.HasField('response'):
-                                continue
-
-                            response = self._response_type()
-                            response.ParseFromString(query_response.response)
-
-                            self._used_response[task].clear()
-
-                            self._calls[task] = call
+                            if query_response.HasField('response'):
+                                response = self._response_type()
+                                response.ParseFromString(
+                                    query_response.response
+                                )
 
-                            self._responses[task] = asyncio.Future()
-                            self._responses[task].set_result(response)
+                                self._calls[task] = call
 
-                            if not have_first_response.is_set():
-                                have_first_response.set()
-                            else:
-                                self._event.set()
+                                self._responses[task] = asyncio.Future()
+                                self._responses[task].set_result(response)
 
-                            await self._used_response[task].wait()
+                                if not have_first_response.is_set():
+                                    have_first_response.set()
+                                else:
+                                    self._event.set()
 
                         raise RuntimeError('React.Query should be infinite')
 

reboot/aio/internals/middleware.py

@@ -6,6 +6,7 @@
 import uuid
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from google.protobuf import descriptor_pool
 from google.protobuf.message import Message
 from logging import Logger
 from reboot.aio.auth.authorizers import Authorizer
@@ -37,6 +38,7 @@
     StateRef,
     StateTypeName,
 )
+from reboot.options import get_method_options
 from reboot.settings import DOCS_BASE_URL
 from reboot.time import DateTimeWithTimeZone
 from typing import (
@@ -135,6 +137,33 @@ def state_type_name(self) -> StateTypeName:
     def service_names(self) -> list[ServiceName]:
         return self._service_names
 
+    def warns_on_flow_control(self, method: str) -> bool:
+        """Whether the server warns every time a client of a reactive
+        read of `method` falls behind and has updates skipped for it,
+        which is what the method's `warn_on_flow_control` reader
+        option decides; a method that has not set it warns."""
+        pool = descriptor_pool.Default()
+
+        for service_name in self._service_names:
+            try:
+                service = pool.FindServiceByName(service_name)
+            except KeyError:
+                continue
+
+            descriptor = service.methods_by_name.get(method)
+
+            if descriptor is None:
+                continue
+
+            reader = get_method_options(descriptor).reader
+
+            if reader.HasField('warn_on_flow_control'):
+                return reader.warn_on_flow_control
+
+            return True
+
+        return True
+
     def create_context(
         self,
         *,