Adding gRPC instrumentor

ext/opentelemetry-ext-grpc/setup.cfg

@@ -50,3 +50,6 @@ test =
 
 [options.packages.find]
 where = src
+
+opentelemetry_instrumentor =
+grpc = opentelemetry.ext.grpc:GrpcInstrumentor

ext/opentelemetry-ext-grpc/src/opentelemetry/ext/grpc/__init__.py

@@ -16,9 +16,60 @@
 # pylint:disable=import-self
 # pylint:disable=no-name-in-module
 # pylint:disable=relative-beyond-top-level
+from contextlib import contextmanager
 
+import grpc
+from concurrent import futures
 from opentelemetry import trace
 from opentelemetry.ext.grpc.version import __version__
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.ext.grpc.grpcext import intercept_channel, intercept_server
+
+
+class GrpcInstrumentorServer(BaseInstrumentor):
+    trace.set_tracer_provider(TracerProvider())
+    tracer = trace.get_tracer_provider()
+    server = None
+
+    def _instrument(self, **kwargs):
+        self.server = grpc.server(futures.ThreadPoolExecutor())
+        self.server = intercept_server(self.server, server_interceptor())
+
+    def _uninstrument(self, **kwargs):
+        if hasattr(self.server, 'stop'):
+            return self.server.stop(kwargs.get('grace'))
+
+
+class GrpcInstrumentorClient(BaseInstrumentor):
+    trace.set_tracer_provider(TracerProvider())
+    tracer = trace.get_tracer_provider()
+    channel = None
+
+    def _instrument(self, **kwargs):
+        hostport = kwargs.get('hostport')
+
+        if kwargs.get('channel_type') == 'secure':
+            self.channel = secure_channel_wrapper(hostport, kwargs.get("credentials"))
+
+        else:
+            self.channel = insecure_channel_wrapper(hostport)
+
+    def _uninstrument(self, **kwargs):
+        if hasattr(self.channel, 'close'):
+            return self.channel.close()
+
+
+@contextmanager
+def insecure_channel_wrapper(hostport):
+    with grpc.insecure_channel(hostport) as channel:
+        yield intercept_channel(channel, client_interceptor())
+
+
+@contextmanager
+def secure_channel_wrapper(hostport, credentials):
+    with grpc.secure_channel(hostport, credentials) as channel:
+        yield intercept_channel(channel, client_interceptor())
 
 
 def client_interceptor(tracer_provider=None):

ext/opentelemetry-ext-grpc/tests/test_server_interceptor.py

@@ -26,6 +26,8 @@
 from opentelemetry.ext.grpc.grpcext import intercept_server
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.test.test_base import TestBase
+from opentelemetry.ext.grpc import GrpcInstrumentorServer, GrpcInstrumentorClient
+
 
 
 class UnaryUnaryMethodHandler(grpc.RpcMethodHandler):
@@ -49,6 +51,39 @@ def service(self, handler_call_details):
 
 
 class TestOpenTelemetryServerInterceptor(TestBase):
+
+    def test_instrumentor(self):
+
+        def handler(request, context):
+            return b""
+
+        grpc_server_instrumentor = GrpcInstrumentorServer()
+        grpc_server_instrumentor.instrument()
+
+        server = grpc_server_instrumentor.server
+
+        server.add_generic_rpc_handlers((UnaryUnaryRpcHandler(handler),))
+
+        port = server.add_insecure_port("[::]:0")
+        grpc_client_instrumentor = GrpcInstrumentorClient()
+        grpc_client_instrumentor.instrument(hostport=port)
+        channel = grpc.insecure_channel("localhost:{:d}".format(port))
+
+        try:
+            server.start()
+            channel.unary_unary("")(b"")
+        finally:
+            server.stop(None)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        self.assertEqual(span.name, "")
+        self.assertIs(span.kind, trace.SpanKind.SERVER)
+        self.check_span_instrumentation_info(span, opentelemetry.ext.grpc)
+
+
     def test_create_span(self):
         """Check that the interceptor wraps calls with spans server-side."""