Add opentelemetry instrumentation

cbrewster · cbrewster · commit 177916d18ac8 · 2024-11-21T15:36:42.000-06:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -21,16 +21,18 @@ classifiers = [
 ]
 dependencies = [
 	"pydantic==2.9.2",
-    "aiochannel>=1.2.1",
-    "black>=23.11,<25.0",
-    "grpcio-tools>=1.59.3",
-    "grpcio>=1.59.3",
-    "msgpack-types>=0.3.0",
-    "msgpack>=1.0.7",
-    "nanoid>=2.0.0",
-    "protobuf>=4.24.4",
-    "pydantic-core>=2.20.1",
-    "websockets>=12.0",
+	"aiochannel>=1.2.1",
+	"black>=23.11,<25.0",
+	"grpcio-tools>=1.59.3",
+	"grpcio>=1.59.3",
+	"msgpack-types>=0.3.0",
+	"msgpack>=1.0.7",
+	"nanoid>=2.0.0",
+	"protobuf>=4.24.4",
+	"pydantic-core>=2.20.1",
+	"websockets>=12.0",
+	"opentelemetry-sdk>=1.28.2",
+	"opentelemetry-api>=1.28.2",
 ]
 
 [tool.uv]
diff --git a/replit_river/client.py b/replit_river/client.py
@@ -1,8 +1,12 @@
 import logging
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
-from typing import Any, Generic, Optional, Union
+from contextlib import contextmanager
+from typing import Any, Generator, Generic, Literal, Optional, Union
+
+from opentelemetry import trace
 
 from replit_river.client_transport import ClientTransport
+from replit_river.error_schema import RiverException
 from replit_river.transport_options import (
     HandshakeMetadataType,
     TransportOptions,
@@ -17,6 +21,7 @@
 )
 
 logger = logging.getLogger(__name__)
+tracer = trace.get_tracer(__name__)
 
 
 class Client(Generic[HandshakeMetadataType]):
@@ -55,15 +60,16 @@ async def send_rpc(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> ResponseType:
-        session = await self._transport.get_or_create_session()
-        return await session.send_rpc(
-            service_name,
-            procedure_name,
-            request,
-            request_serializer,
-            response_deserializer,
-            error_deserializer,
-        )
+        with _trace_procedure("rpc", service_name, procedure_name):
+            session = await self._transport.get_or_create_session()
+            return await session.send_rpc(
+                service_name,
+                procedure_name,
+                request,
+                request_serializer,
+                response_deserializer,
+                error_deserializer,
+            )
 
     async def send_upload(
         self,
@@ -76,17 +82,18 @@ async def send_upload(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> ResponseType:
-        session = await self._transport.get_or_create_session()
-        return await session.send_upload(
-            service_name,
-            procedure_name,
-            init,
-            request,
-            init_serializer,
-            request_serializer,
-            response_deserializer,
-            error_deserializer,
-        )
+        with _trace_procedure("upload", service_name, procedure_name):
+            session = await self._transport.get_or_create_session()
+            return await session.send_upload(
+                service_name,
+                procedure_name,
+                init,
+                request,
+                init_serializer,
+                request_serializer,
+                response_deserializer,
+                error_deserializer,
+            )
 
     async def send_subscription(
         self,
@@ -97,15 +104,16 @@ async def send_subscription(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
-        session = await self._transport.get_or_create_session()
-        return session.send_subscription(
-            service_name,
-            procedure_name,
-            request,
-            request_serializer,
-            response_deserializer,
-            error_deserializer,
-        )
+        with _trace_procedure("subscription", service_name, procedure_name):
+            session = await self._transport.get_or_create_session()
+            return session.send_subscription(
+                service_name,
+                procedure_name,
+                request,
+                request_serializer,
+                response_deserializer,
+                error_deserializer,
+            )
 
     async def send_stream(
         self,
@@ -118,14 +126,33 @@ async def send_stream(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
-        session = await self._transport.get_or_create_session()
-        return session.send_stream(
-            service_name,
-            procedure_name,
-            init,
-            request,
-            init_serializer,
-            request_serializer,
-            response_deserializer,
-            error_deserializer,
-        )
+        with _trace_procedure("stream", service_name, procedure_name):
+            session = await self._transport.get_or_create_session()
+            return session.send_stream(
+                service_name,
+                procedure_name,
+                init,
+                request,
+                init_serializer,
+                request_serializer,
+                response_deserializer,
+                error_deserializer,
+            )
+
+
+@contextmanager
+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(
+        f"river.client.{procedure_type}.{service_name}.{procedure_name}",
+        kind=trace.SpanKind.CLIENT,
+    ) as span:
+        try:
+            yield
+        except RiverException as e:
+            span.set_attribute("river.error_code", e.code)
+            span.set_attribute("river.error_message", e.message)
+            raise e
diff --git a/replit_river/rpc.py b/replit_river/rpc.py
@@ -21,6 +21,7 @@
 
 import grpc
 from aiochannel import Channel, ChannelClosed
+from opentelemetry.propagators.textmap import Setter
 from pydantic import BaseModel, ConfigDict, Field
 
 from replit_river.error_schema import (
@@ -85,6 +86,11 @@ class ControlMessageHandshakeResponse(BaseModel):
     status: HandShakeStatus
 
 
+class PropagationContext(BaseModel):
+    traceparent: str
+    tracestate: str
+
+
 class TransportMessage(BaseModel):
     id: str
     # from_ is used instead of from because from is a reserved keyword in Python
@@ -96,12 +102,30 @@ class TransportMessage(BaseModel):
     procedureName: Optional[str] = None
     streamId: str
     controlFlags: int
+    tracing: Optional[PropagationContext] = None
     payload: Any
     model_config = ConfigDict(populate_by_name=True)
     # need this because we create TransportMessage objects with destructuring
     # where the key is "from"
 
 
+class TransportMessageTracingSetter(Setter[TransportMessage]):
+    """
+    Handles propagating tracing context to the recipient of the message.
+    """
+
+    def set(self, carrier: TransportMessage, key: str, value: str) -> None:
+        if not carrier.tracing:
+            carrier.tracing = PropagationContext(traceparent="", tracestate="")
+        match key:
+            case "traceparent":
+                carrier.tracing.traceparent = value
+            case "tracestate":
+                carrier.tracing.tracestate = value
+            case _:
+                logger.warning("unknown trace propagation key", extra={"key": key})
+
+
 class GrpcContext(grpc.aio.ServicerContext):
     """Represents a gRPC-compatible ServicerContext for River interop."""
 
diff --git a/replit_river/session.py b/replit_river/session.py
@@ -6,6 +6,7 @@
 import nanoid  # type: ignore
 import websockets
 from aiochannel import Channel, ChannelClosed
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from websockets.exceptions import ConnectionClosed
 
 from replit_river.message_buffer import MessageBuffer, MessageBufferClosedError
@@ -31,6 +32,7 @@
     STREAM_OPEN_BIT,
     GenericRpcHandler,
     TransportMessage,
+    TransportMessageTracingSetter,
 )
 
 logger = logging.getLogger(__name__)
@@ -380,6 +382,9 @@ async def send_message(
             serviceName=service_name,
             procedureName=procedure_name,
         )
+        TraceContextTextMapPropagator().inject(
+            msg, None, TransportMessageTracingSetter()
+        )
         try:
             # We need this lock to ensure the buffer order and message sending order
             # are the same.
diff --git a/scripts/lint.sh b/scripts/lint.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
 set -ex
 
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-#!/bin/bash`
	`1`	`+#!/usr/bin/env bash`
`2`	`2`
`3`	`3`	`set -ex`
`4`	`4`