Don't record GeneratorExit errors in stream RPC spans

Author: cbrewster

replit_river/client.py

@@ -1,8 +1,9 @@
 import logging
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
+from asyncio import CancelledError
+from collections.abc import AsyncIterable, Awaitable, Callable
 from contextlib import contextmanager
 from datetime import timedelta
-from typing import Any, Generator, Generic, Literal, Optional, Union
+from typing import Any, AsyncGenerator, Generator, Generic, Literal, Optional, Union
 
 from opentelemetry import trace
 from opentelemetry.trace import Span, SpanKind, StatusCode
@@ -109,7 +110,7 @@ async def send_subscription(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
-    ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
+    ) -> AsyncGenerator[Union[ResponseType, ErrorType], None]:
         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(
@@ -135,7 +136,7 @@ async def send_stream(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
-    ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
+    ) -> AsyncGenerator[Union[ResponseType, ErrorType], None]:
         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(
@@ -160,15 +161,27 @@ def _trace_procedure(
     service_name: str,
     procedure_name: str,
 ) -> Generator[Span, None, None]:
-    with tracer.start_span(
+    span = tracer.start_span(
         f"river.client.{procedure_type}.{service_name}.{procedure_name}",
         kind=SpanKind.CLIENT,
-    ) as span:
-        try:
-            yield span
-        except RiverException as e:
-            _record_river_error(span, RiverError(code=e.code, message=e.message))
-            raise e
+    )
+    set_status = False
+    try:
+        yield span
+    except RiverException as e:
+        span.record_exception(e)
+        _record_river_error(span, RiverError(code=e.code, message=e.message))
+        set_status = True
+        raise e
+    except (Exception, CancelledError) as e:
+        span.record_exception(e)
+        span.set_status(StatusCode.ERROR, f"{type(e).__name__}: {e}")
+        set_status = True
+        raise e
+    finally:
+        if not set_status:
+            span.set_status(StatusCode.OK)
+        span.end()
 
 
 def _record_river_error(span: Span, error: RiverError) -> None:

replit_river/client_session.py

@@ -1,8 +1,8 @@
 import asyncio
 import logging
-from collections.abc import AsyncIterable, AsyncIterator
+from collections.abc import AsyncIterable
 from datetime import timedelta
-from typing import Any, Callable, Optional, Union
+from typing import Any, AsyncGenerator, Callable, Optional, Union
 
 import nanoid  # type: ignore
 from aiochannel import Channel
@@ -194,7 +194,7 @@ async def send_subscription(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
         span: Span,
-    ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
+    ) -> AsyncGenerator[Union[ResponseType, ErrorType], None]:
         """Sends a subscription request to the server.
 
         Expects the input and output be messages that will be msgpacked.
@@ -246,7 +246,7 @@ async def send_stream(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
         span: Span,
-    ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
+    ) -> AsyncGenerator[Union[ResponseType, ErrorType], None]:
         """Sends a subscription request to the server.
 
         Expects the input and output be messages that will be msgpacked.

tests/test_opentelemetry.py

@@ -1,3 +1,4 @@
+import contextlib
 from datetime import timedelta
 from typing import AsyncGenerator, AsyncIterator, Iterator
 
@@ -182,3 +183,39 @@ async def stream_data() -> AsyncGenerator[str, None]:
     assert len(spans) == 1
     assert spans[0].name == "river.client.stream.test_service.stream_method_error"
     assert spans[0].status.status_code == StatusCode.ERROR
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("handlers", [{**basic_stream}])
+async def test_stream_method_span_generator_exit_not_recorded(
+    client: Client, span_exporter: InMemorySpanExporter
+) -> None:
+    async def stream_data() -> AsyncGenerator[str, None]:
+        yield "Stream 1"
+        yield "Stream 2"
+        yield "Stream 3"
+
+    responses = []
+    stream = client.send_stream(
+        "test_service",
+        "stream_method",
+        "Initial Stream Data",
+        stream_data(),
+        serialize_request,
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+    )
+    async with contextlib.aclosing(stream) as generator:
+        async for response in generator:
+            responses.append(response)
+            break
+
+    assert responses == [
+        "Stream response for Initial Stream Data",
+    ]
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "river.client.stream.test_service.stream_method"
+    assert spans[0].status.status_code == StatusCode.OK