React: add flow control to reactive readers

rbt/v1alpha1/react.proto

@@ -18,6 +18,12 @@ message QueryRequest {
 
   // Authorization bearer token.
   optional string bearer_token = 3;
+
+  // ID of the `QueryResponse` that this request acknowledges; see
+  // `AcknowledgeQueryResponse`. 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.
+  string acknowledge_query_response_id = 4;
 }
 
 message QueryResponse {
@@ -40,8 +46,20 @@ message QueryResponse {
   //
   // TODO(benh): send these as bytes?
   repeated string idempotency_keys = 2;
+
+  // Unique ID of this response, used to acknowledge it.
+  string query_response_id = 4;
+}
+
+////////////////////////////////////////////////////////////////////////
+
+message AcknowledgeQueryResponseRequest {
+  // ID of the `QueryResponse` being acknowledged.
+  string query_response_id = 1;
 }
 
+message AcknowledgeQueryResponseResponse {}
+
 ////////////////////////////////////////////////////////////////////////
 
 // TODO(benh): support batch, e.g., create a `MutateRequests` which
@@ -90,8 +108,15 @@ 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. Every response must be acknowledged by the
+  // client before a next one is produced, and that next response
+  // reflects the _latest_ state, skipping any states that came and
+  // went while the client was busy.
+  //
+  // A client acknowledges either by calling
+  // `AcknowledgeQueryResponse`, or, when its responses arrive over a
+  // websocket, by sending a `QueryRequest` with
+  // `acknowledge_query_response_id` set on that same websocket.
   //
   // NOTE: the service name and state ID is expected to be part of the
   // gRPC metadata, exactly like any other reboot requests.
@@ -101,6 +126,16 @@ service React {
     };
   }
 
+  // Acknowledges that a client has processed a response from `Query`
+  // and is ready for a next one. Must be called on the same server
+  // that produced the response.
+  rpc AcknowledgeQueryResponse(AcknowledgeQueryResponseRequest)
+      returns (AcknowledgeQueryResponseResponse) {
+    option (google.api.http) = {
+      post: "/rbt.v1alpha1.React/AcknowledgeQueryResponse"
+    };
+  }
+
   // 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

@@ -474,7 +474,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,
@@ -498,25 +500,39 @@ async def loop():
                         assert task is not None
 
                         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
-
-                            self._responses[task] = asyncio.Future()
-                            self._responses[task].set_result(response)
-
-                            if not have_first_response.is_set():
-                                have_first_response.set()
-                            else:
-                                self._event.set()
-
-                            await self._used_response[task].wait()
+                            if query_response.HasField('response'):
+                                response = self._response_type()
+                                response.ParseFromString(
+                                    query_response.response
+                                )
+
+                                self._used_response[task].clear()
+
+                                self._calls[task] = call
+
+                                self._responses[task] = asyncio.Future()
+                                self._responses[task].set_result(response)
+
+                                if not have_first_response.is_set():
+                                    have_first_response.set()
+                                else:
+                                    self._event.set()
+
+                                await self._used_response[task].wait()
+
+                            # Only now that the response has been used
+                            # do we ask for a next one, so that it
+                            # reflects the latest state rather than a
+                            # state that has already been superseded.
+                            await stub.AcknowledgeQueryResponse(
+                                react_pb2.AcknowledgeQueryResponseRequest(
+                                    query_response_id=query_response.
+                                    query_response_id,
+                                ),
+                                # The same metadata ensures we're
+                                # routed to the same server.
+                                metadata=metadata,
+                            )
 
                         raise RuntimeError('React.Query should be infinite')