Implement RPC timeouts

Author: blast-hardcheese

replit_river/client.py

@@ -2,6 +2,7 @@
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
 from contextlib import contextmanager
 from typing import Any, Generator, Generic, Literal, Optional, Union
+from datetime import timedelta
 
 from opentelemetry import trace
 from opentelemetry.trace import Span, SpanKind, StatusCode
@@ -60,6 +61,7 @@ async def send_rpc(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
+        timeout: timedelta,
     ) -> ResponseType:
         with _trace_procedure("rpc", service_name, procedure_name) as span:
             session = await self._transport.get_or_create_session()
@@ -71,6 +73,7 @@ async def send_rpc(
                 response_deserializer,
                 error_deserializer,
                 span,
+                timeout,
             )
 
     async def send_upload(

replit_river/client_session.py

@@ -1,3 +1,5 @@
+import asyncio
+from datetime import timedelta
 import logging
 from collections.abc import AsyncIterable, AsyncIterator
 from typing import Any, Callable, Optional, Union
@@ -8,6 +10,7 @@
 from opentelemetry.trace import Span
 
 from replit_river.error_schema import (
+    ERROR_CODE_CANCEL,
     ERROR_CODE_STREAM_CLOSED,
     RiverException,
     RiverServiceException,
@@ -39,6 +42,7 @@ async def send_rpc(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
         span: Span,
+        timeout: timedelta,
     ) -> ResponseType:
         """Sends a single RPC request to the server.
 
@@ -58,7 +62,19 @@ async def send_rpc(
         # Handle potential errors during communication
         try:
             try:
-                response = await output.get()
+                async with asyncio.timeout(int(timeout.total_seconds())):
+                    response = await output.get()
+            except asyncio.CancelledError as e:
+                # TODO(dstewart) After protocol v2, change this to STREAM_CANCEL_BIT
+                await self.send_message(
+                    stream_id=stream_id,
+                    control_flags=STREAM_CLOSED_BIT,
+                    payload={"type": "CLOSE"},
+                    service_name=service_name,
+                    procedure_name=procedure_name,
+                    span=span,
+                )
+                raise RiverException(ERROR_CODE_CANCEL, str(e)) from e
             except ChannelClosed as e:
                 raise RiverServiceException(
                     ERROR_CODE_STREAM_CLOSED,

replit_river/codegen/client.py

@@ -852,6 +852,7 @@ def __init__(self, client: river.Client[Any]):
             async def {name}(
               self,
               input: {render_type_expr(input_type)},
+              timeout: datetime.timedelta,
             ) -> {render_type_expr(output_type)}:
               return await self.client.send_rpc(
                 {repr(schema_name)},
@@ -860,6 +861,7 @@ async def {name}(
                 {reindent("                    ", render_input_method)},
                 {reindent("                    ", parse_output_method)},
                 {reindent("                    ", parse_error_method)},
+                timeout,
               )
                     """,
                     )

replit_river/rpc.py

@@ -51,7 +51,7 @@
 ]
 ACK_BIT = 0x0001
 STREAM_OPEN_BIT = 0x0002
-STREAM_CLOSED_BIT = 0x0004
+STREAM_CLOSED_BIT = 0x0004  # Synonymous with the cancel bit in v2
 
 # these codes are retriable
 # if the server sends a response with one of these codes,

tests/test_communication.py

@@ -1,4 +1,5 @@
 import asyncio
+from datetime import timedelta
 from typing import AsyncGenerator
 
 import pytest
@@ -18,6 +19,7 @@ async def test_rpc_method(client: Client) -> None:
         serialize_request,
         deserialize_response,
         deserialize_error,
+        timedelta(seconds=20),
     )
     assert response == "Hello, Alice!"
 

tests/test_opentelemetry.py

@@ -1,3 +1,4 @@
+from datetime import timedelta
 from typing import AsyncGenerator
 
 import pytest
@@ -21,6 +22,7 @@ async def test_rpc_method_span(
         serialize_request,
         deserialize_response,
         deserialize_error,
+        timedelta(seconds=20),
     )
     assert response == "Hello, Alice!"
     spans = span_exporter.get_finished_spans()