Fix contextvar issues with async generators

Author: cbrewster

replit_river/client.py

@@ -60,7 +60,7 @@ async def send_rpc(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> ResponseType:
-        with _trace_procedure("rpc", service_name, procedure_name):
+        with _trace_procedure("rpc", service_name, procedure_name) as span:
             session = await self._transport.get_or_create_session()
             return await session.send_rpc(
                 service_name,
@@ -69,6 +69,7 @@ async def send_rpc(
                 request_serializer,
                 response_deserializer,
                 error_deserializer,
+                span,
             )
 
     async def send_upload(
@@ -82,7 +83,7 @@ async def send_upload(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> ResponseType:
-        with _trace_procedure("upload", service_name, procedure_name):
+        with _trace_procedure("upload", service_name, procedure_name) as span:
             session = await self._transport.get_or_create_session()
             return await session.send_upload(
                 service_name,
@@ -93,6 +94,7 @@ async def send_upload(
                 request_serializer,
                 response_deserializer,
                 error_deserializer,
+                span,
             )
 
     async def send_subscription(
@@ -104,7 +106,7 @@ async def send_subscription(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
-        with _trace_procedure("subscription", service_name, procedure_name):
+        with _trace_procedure("subscription", service_name, procedure_name) as span:
             session = await self._transport.get_or_create_session()
             async for msg in session.send_subscription(
                 service_name,
@@ -113,6 +115,7 @@ async def send_subscription(
                 request_serializer,
                 response_deserializer,
                 error_deserializer,
+                span,
             ):
                 if isinstance(msg, RiverError):
                     _record_river_error(msg)
@@ -129,7 +132,7 @@ async def send_stream(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
-        with _trace_procedure("stream", service_name, procedure_name):
+        with _trace_procedure("stream", service_name, procedure_name) as span:
             session = await self._transport.get_or_create_session()
             async for msg in session.send_stream(
                 service_name,
@@ -140,6 +143,7 @@ async def send_stream(
                 request_serializer,
                 response_deserializer,
                 error_deserializer,
+                span,
             ):
                 if isinstance(msg, RiverError):
                     _record_river_error(msg)
@@ -158,13 +162,13 @@ def _trace_procedure(
     procedure_type: Literal["rpc", "upload", "subscription", "stream"],
     service_name: str,
     procedure_name: str,
-) -> Generator[None, None, None]:
-    with tracer.start_as_current_span(
+) -> Generator[trace.Span, None, None]:
+    with tracer.start_span(
         f"river.client.{procedure_type}.{service_name}.{procedure_name}",
         kind=trace.SpanKind.CLIENT,
     ) as span:
         try:
-            yield
+            yield span
         except RiverException as e:
             span.set_attribute("river.error_code", e.code)
             span.set_attribute("river.error_message", e.message)

replit_river/client_session.py

@@ -5,6 +5,7 @@
 import nanoid  # type: ignore
 from aiochannel import Channel
 from aiochannel.errors import ChannelClosed
+from opentelemetry.trace import Span
 
 from replit_river.error_schema import (
     ERROR_CODE_STREAM_CLOSED,
@@ -37,6 +38,7 @@ async def send_rpc(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
+        span: Span,
     ) -> ResponseType:
         """Sends a single RPC request to the server.
 
@@ -51,6 +53,7 @@ async def send_rpc(
             payload=request_serializer(request),
             service_name=service_name,
             procedure_name=procedure_name,
+            span=span,
         )
         # Handle potential errors during communication
         try:
@@ -89,6 +92,7 @@ async def send_upload(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
+        span: Span,
     ) -> ResponseType:
         """Sends an upload request to the server.
 
@@ -107,6 +111,7 @@ async def send_upload(
                     service_name=service_name,
                     procedure_name=procedure_name,
                     payload=init_serializer(init),
+                    span=span,
                 )
                 first_message = False
             # If this request is not closed and the session is killed, we should
@@ -122,6 +127,7 @@ async def send_upload(
                     procedure_name=procedure_name,
                     control_flags=control_flags,
                     payload=request_serializer(item),
+                    span=span,
                 )
         except Exception as e:
             raise RiverServiceException(
@@ -171,6 +177,7 @@ async def send_subscription(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
+        span: Span,
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
         """Sends a subscription request to the server.
 
@@ -185,6 +192,7 @@ async def send_subscription(
             stream_id=stream_id,
             control_flags=STREAM_OPEN_BIT,
             payload=request_serializer(request),
+            span=span,
         )
 
         # Handle potential errors during communication
@@ -221,6 +229,7 @@ async def send_stream(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
+        span: Span,
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
         """Sends a subscription request to the server.
 
@@ -239,6 +248,7 @@ async def send_stream(
                     stream_id=stream_id,
                     control_flags=STREAM_OPEN_BIT,
                     payload=init_serializer(init),
+                    span=span,
                 )
             else:
                 # Get the very first message to open the stream
@@ -250,6 +260,7 @@ async def send_stream(
                     stream_id=stream_id,
                     control_flags=STREAM_OPEN_BIT,
                     payload=request_serializer(first),
+                    span=span,
                 )
 
         except StopAsyncIteration:

replit_river/session.py

@@ -6,6 +6,7 @@
 import nanoid  # type: ignore
 import websockets
 from aiochannel import Channel, ChannelClosed
+from opentelemetry.trace import Span, use_span
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from websockets.exceptions import ConnectionClosed
 
@@ -365,6 +366,7 @@ async def send_message(
         control_flags: int = 0,
         service_name: str | None = None,
         procedure_name: str | None = None,
+        span: Span | None = None,
     ) -> None:
         """Send serialized messages to the websockets."""
         # if the session is not active, we should not do anything
@@ -382,9 +384,11 @@ async def send_message(
             serviceName=service_name,
             procedureName=procedure_name,
         )
-        TraceContextTextMapPropagator().inject(
-            msg, None, TransportMessageTracingSetter()
-        )
+        if span:
+            with use_span(span):
+                TraceContextTextMapPropagator().inject(
+                    msg, None, TransportMessageTracingSetter()
+                )
         try:
             # We need this lock to ensure the buffer order and message sending order
             # are the same.