feat(instrumentation): add OpenTelemetry tracing and metrics with basic configurations

jina/serve/executors/__init__.py

@@ -18,6 +18,7 @@
 from jina.serve.helper import store_init_kwargs, wrap_func
 
 if TYPE_CHECKING:
+    from opentelemetry.context.context import Context
     from prometheus_client import Summary
 
 __dry_run_endpoint__ = '_jina_dry_run_'
@@ -315,7 +316,16 @@ async def __acall__(self, req_endpoint: str, **kwargs):
         elif __default_endpoint__ in self.requests:
             return await self.__acall_endpoint__(__default_endpoint__, **kwargs)
 
-    async def __acall_endpoint__(self, req_endpoint, **kwargs):
+    async def __acall_endpoint__(
+        self, req_endpoint, tracing_context: Optional['Context'], **kwargs
+    ):
+        async def exec_func(summary, tracing_context):
+            with summary:
+                if iscoroutinefunction(func):
+                    return await func(self, tracing_context=tracing_context, **kwargs)
+                else:
+                    return func(self, tracing_context=tracing_context, **kwargs)
+
         func = self.requests[req_endpoint]
 
         runtime_name = (
@@ -330,11 +340,18 @@ async def __acall_endpoint__(self, req_endpoint, **kwargs):
             else contextlib.nullcontext()
         )
 
-        with _summary:
-            if iscoroutinefunction(func):
-                return await func(self, **kwargs)
-            else:
-                return func(self, **kwargs)
+        if self.tracer:
+            with self.tracer.start_span(req_endpoint, context=tracing_context) as _:
+                from opentelemetry.propagate import extract
+                from opentelemetry.trace.propagation.tracecontext import (
+                    TraceContextTextMapPropagator,
+                )
+
+                tracing_carrier_context = {}
+                TraceContextTextMapPropagator().inject(tracing_carrier_context)
+                return await exec_func(_summary, extract(tracing_carrier_context))
+        else:
+            return await exec_func(_summary, None)
 
     @property
     def workspace(self) -> Optional[str]:

jina/serve/runtimes/request_handlers/data_request_handler.py

@@ -1,7 +1,6 @@
 from typing import TYPE_CHECKING, Dict, List, Optional
 
 from docarray import DocumentArray
-from opentelemetry.context.context import Context
 
 from jina import __default_endpoint__
 from jina.excepts import BadConfigSource
@@ -13,6 +12,7 @@
     import argparse
 
     from opentelemetry import metrics, trace
+    from opentelemetry.context.context import Context
     from prometheus_client import CollectorRegistry
 
     from jina.logging.logger import JinaLogger
@@ -147,7 +147,7 @@ def _parse_params(parameters: Dict, executor_name: str):
         return parsed_params
 
     async def handle(
-        self, requests: List['DataRequest'], tracing_context: Context = None
+        self, requests: List['DataRequest'], tracing_context: 'Context' = None
     ) -> DataRequest:
         """Initialize private parameters and execute private loading functions.
 

tests/integration/instrumentation/__init__.py

@@ -54,11 +54,14 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @requests(on='/index')
-    def empty(self, docs: 'DocumentArray', tracing_context: Context, **kwargs):
-        with self.tracer.start_as_current_span(
-            'dummy', context=tracing_context
-        ) as span:
-            span.set_attribute('len_docs', len(docs))
+    def empty(
+        self, docs: 'DocumentArray', tracing_context: Optional[Context], **kwargs
+    ):
+        if self.tracer:
+            with self.tracer.start_span('dummy', context=tracing_context) as span:
+                span.set_attribute('len_docs', len(docs))
+                return docs
+        else:
             return docs
 
 

tests/integration/instrumentation/test_flow_instrumentation.py

@@ -15,9 +15,9 @@
 @pytest.mark.parametrize(
     'protocol, client_type, num_internal_spans',
     [
-        ('grpc', 'GRPCClient', 1),
-        ('http', 'HTTPClient', 4),
-        ('websocket', 'WebSocketClient', 6),
+        ('grpc', 'GRPCClient', 2),
+        ('http', 'HTTPClient', 5),
+        ('websocket', 'WebSocketClient', 7),
     ],
 )
 def test_gateway_instrumentation(