feat(tracing) Implement an interceptor for grpc server

Implement gRPC ServerInterceptor; starts(or continues) a transaction for each grpc request

Author: hossein-raeisi

sentry_sdk/consts.py

@@ -79,6 +79,7 @@ class OP:
     WEBSOCKET_SERVER = "websocket.server"
     SOCKET_CONNECTION = "socket.connection"
     SOCKET_DNS = "socket.dns"
+    GRPC_SERVER = "grpc.server"
 
 
 # This type exists to trick mypy and PyCharm into thinking `init` and `Client`

sentry_sdk/integrations/grpc/__init__.py

@@ -0,0 +1 @@
+from .server import ServerInterceptor

sentry_sdk/integrations/grpc/server.py

@@ -0,0 +1,59 @@
+from sentry_sdk import Hub
+from sentry_sdk._types import MYPY
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_CUSTOM
+
+if MYPY:
+    from typing import Callable
+
+try:
+    import grpc
+except ImportError:
+    raise DidNotEnable("grpcio is not installed")
+
+
+class ServerInterceptor(grpc.ServerInterceptor):
+    def __init__(self, find_name=None):
+        # type: (ServerInterceptor, Callable | None) -> ServerInterceptor
+        if find_name:
+            self._find_method_name = find_name
+        super(ServerInterceptor, self).__init__()
+
+    def intercept_service(self, continuation, handler_call_details):
+        handler = continuation(handler_call_details)
+        if not handler or not handler.unary_unary:
+            return handler
+
+        def behavior(request, context):
+            hub = Hub(Hub.current)
+
+            name = self._find_method_name(context)
+
+            if name:
+                meta_data = dict(context.invocation_metadata())
+
+                transaction = Transaction.continue_from_headers(
+                    meta_data,
+                    op=OP.GRPC_SERVER,
+                    name=name,
+                    source=TRANSACTION_SOURCE_CUSTOM,
+                )
+
+                with hub.start_transaction(transaction=transaction):
+                    try:
+                        return handler.unary_unary(request, context)
+                    except BaseException as e:
+                        raise e
+            else:
+                return handler.unary_unary(request, context)
+
+        return grpc.unary_unary_rpc_method_handler(
+            behavior,
+            request_deserializer=handler.request_deserializer,
+            response_serializer=handler.response_serializer,
+        )
+
+    @staticmethod
+    def _find_name(context):
+        return context._rpc_event.call_details.method.decode()

setup.py

@@ -66,6 +66,7 @@ def get_file_text(file_name):
         "fastapi": ["fastapi>=0.79.0"],
         "pymongo": ["pymongo>=3.1"],
         "opentelemetry": ["opentelemetry-distro>=0.35b0"],
+        "grpcio": ["grpcio>=1.21.1"]
     },
     classifiers=[
         "Development Status :: 5 - Production/Stable",

test-requirements.txt

@@ -13,3 +13,7 @@ executing
 asttokens
 responses
 ipdb
+grpcio-tools
+protobuf
+mypy-protobuf
+types-protobuf

tests/conftest.py

@@ -298,20 +298,21 @@ def flush(timeout=None, callback=None):
         monkeypatch.setattr(test_client.transport, "capture_event", append)
         monkeypatch.setattr(test_client, "flush", flush)
 
-        return EventStreamReader(events_r)
+        return EventStreamReader(events_r, events_w)
 
     return inner
 
 
 class EventStreamReader(object):
-    def __init__(self, file):
-        self.file = file
+    def __init__(self, read_file, write_file):
+        self.read_file = read_file
+        self.write_file = write_file
 
     def read_event(self):
-        return json.loads(self.file.readline().decode("utf-8"))
+        return json.loads(self.read_file.readline().decode("utf-8"))
 
     def read_flush(self):
-        assert self.file.readline() == b"flush\n"
+        assert self.read_file.readline() == b"flush\n"
 
 
 # scope=session ensures that fixture is run earlier

tests/integrations/grpc/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("grpc")

tests/integrations/grpc/test_server.py

@@ -0,0 +1,121 @@
+from __future__ import absolute_import
+
+from concurrent import futures
+
+import grpc
+import pytest
+
+from sentry_sdk import Hub, start_transaction
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations.grpc.server import ServerInterceptor
+from tests.integrations.grpc.test_service_pb2 import TestMessage
+from tests.integrations.grpc.test_service_pb2_grpc import (
+    TestServiceServicer,
+    add_TestServiceServicer_to_server,
+    TestServiceStub,
+)
+
+PORT = 50051
+
+
+@pytest.mark.forked
+def test_grpc_server_starts_transaction(sentry_init, capture_events_forksafe):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events_forksafe()
+
+    server = _set_up()
+
+    with grpc.insecure_channel(f"localhost:{PORT}") as channel:
+        stub = TestServiceStub(channel)
+        stub.TestServe(TestMessage(text="test"))
+
+    _tear_down(server=server)
+
+    events.write_file.close()
+    event = events.read_event()
+    span = event["spans"][0]
+
+    assert event["type"] == "transaction"
+    assert event["transaction_info"] == {
+        "source": "custom",
+    }
+    assert event["contexts"]["trace"]["op"] == OP.GRPC_SERVER
+    assert span["op"] == "test"
+
+
+@pytest.mark.forked
+def test_grpc_server_continues_transaction(sentry_init, capture_events_forksafe):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events_forksafe()
+
+    server = _set_up()
+
+    with grpc.insecure_channel(f"localhost:{PORT}") as channel:
+        stub = TestServiceStub(channel)
+
+        with start_transaction() as transaction:
+            metadata = (
+                (
+                    "baggage",
+                    "sentry-trace_id={trace_id},sentry-environment=test,"
+                    "sentry-transaction=test-transaction,sentry-sample_rate=1.0".format(
+                        trace_id=transaction.trace_id
+                    ),
+                ),
+                (
+                    "sentry-trace",
+                    "{trace_id}-{parent_span_id}-{sampled}".format(
+                        trace_id=transaction.trace_id,
+                        parent_span_id=transaction.span_id,
+                        sampled=1,
+                    )
+                )
+            )
+            stub.TestServe(TestMessage(text="test"), metadata=metadata)
+
+    _tear_down(server=server)
+
+    events.write_file.close()
+    event = events.read_event()
+    span = event["spans"][0]
+
+    assert event["type"] == "transaction"
+    assert event["transaction_info"] == {
+        "source": "custom",
+    }
+    assert event["contexts"]["trace"]["op"] == OP.GRPC_SERVER
+    assert event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+    assert span["op"] == "test"
+
+
+def _set_up():
+    server = grpc.server(
+        futures.ThreadPoolExecutor(max_workers=2),
+        interceptors=[ServerInterceptor(find_name=_find_name)],
+    )
+
+    add_TestServiceServicer_to_server(TestService, server)
+    server.add_insecure_port(f"[::]:{PORT}")
+    server.start()
+
+    return server
+
+
+def _tear_down(server: grpc.Server):
+    server.stop(None)
+
+
+def _find_name(request):
+    return request.__class__
+
+
+class TestService(TestServiceServicer):
+    events = []
+
+    @staticmethod
+    def TestServe(request, context):
+        hub = Hub.current
+        with hub.start_span(op="test", description="test"):
+            pass
+
+        return TestMessage(text=request.text)

tests/integrations/grpc/test_service.proto

@@ -0,0 +1,11 @@
+syntax = "proto3";
+
+package test_grpc_server;
+
+service TestService{
+  rpc TestServe(TestMessage) returns (TestMessage);
+}
+
+message TestMessage {
+  string text = 1;
+}