#402: add opentelemetry tracing and collector (#403)

* #402: add opentelemetry tracing and collector

* fix style

Author: Kraysent

app/commands/adminapi/command.py

@@ -1,14 +1,16 @@
 from pathlib import Path
 from typing import final
 
+import pydantic
 import pydantic_settings as settings
 import structlog
 import yaml
 
 from app.data import repositories
 from app.domain import adminapi as domain
-from app.lib import auth, clients, commands, config
+from app.lib import auth, clients, commands, config, tracing
 from app.lib.storage import postgres
+from app.lib.tracing import TracingConfig
 from app.lib.web import middlewares, server
 from app.presentation import adminapi as presentation
 
@@ -27,6 +29,8 @@ def __init__(self, config_path: str):
     def prepare(self):
         cfg = parse_config(self.config_path)
 
+        tracing.setup_tracing("adminapi", cfg.tracing)
+
         self.pg_storage = postgres.PgStorage(cfg.storage, log)
         self.pg_storage.connect()
 
@@ -65,6 +69,9 @@ class Config(config.ConfigSettings):
     storage: postgres.PgStorageConfig
     clients: ClientsConfig
     auth_enabled: bool
+    tracing: TracingConfig = pydantic.Field(
+        default_factory=lambda: TracingConfig(endpoint="localhost:4317", enabled=False)
+    )
 
 
 def parse_config(path: str) -> Config:

app/commands/dataapi/command.py

@@ -1,14 +1,16 @@
 from pathlib import Path
 from typing import final
 
+import pydantic
 import structlog
 import yaml
 
 from app.data import repositories
 from app.domain import dataapi as domain
 from app.domain import responders
-from app.lib import commands, config
+from app.lib import commands, config, tracing
 from app.lib.storage import postgres
+from app.lib.tracing import TracingConfig
 from app.lib.web import server
 from app.presentation import dataapi as presentation
 
@@ -28,6 +30,8 @@ def __init__(self, config_path: str) -> None:
     def prepare(self):
         self.config = parse_config(self.config_path)
 
+        tracing.setup_tracing("dataapi", self.config.tracing)
+
         self.pg_storage = postgres.PgStorage(self.config.storage, log)
         self.pg_storage.connect()
 
@@ -50,6 +54,9 @@ class Config(config.ConfigSettings):
     server: server.ServerConfig
     storage: postgres.PgStorageConfig
     catalogs: responders.CatalogConfig
+    tracing: TracingConfig = pydantic.Field(
+        default_factory=lambda: TracingConfig(endpoint="localhost:4317", enabled=False)
+    )
 
 
 def parse_config(path: str) -> Config:

app/lib/concurrency/errorgroup.py

@@ -1,3 +1,4 @@
+import contextvars
 import threading
 from collections.abc import Callable
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -23,14 +24,21 @@ def __init__(self) -> None:
         self._error_lock = threading.Lock()
 
     def run[T, **P](self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> TaskResult[T]:
-        future = self._executor.submit(self._run_with_error_handling, fn, *args, **kwargs)
+        ctx = contextvars.copy_context()
+        future = self._executor.submit(self._run_with_error_handling, ctx, fn, *args, **kwargs)
         self._futures.append(future)
 
         return TaskResult[T](future)
 
-    def _run_with_error_handling[T, **P](self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T | None:
+    def _run_with_error_handling[T, **P](
+        self,
+        ctx: contextvars.Context,
+        fn: Callable[P, T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T | None:
         try:
-            return fn(*args, **kwargs)
+            return ctx.run(fn, *args, **kwargs)
         except Exception as e:
             with self._error_lock:
                 if self._error is None:

app/lib/tracing.py

@@ -0,0 +1,27 @@
+import pydantic
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.grpc import trace_exporter
+from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
+from opentelemetry.sdk import trace as sdk_trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+
+class TracingConfig(pydantic.BaseModel):
+    endpoint: str
+    enabled: bool = False
+
+
+def setup_tracing(service_name: str, cfg: TracingConfig) -> None:
+    if not cfg.enabled:
+        return
+
+    provider = sdk_trace.TracerProvider(
+        resource=Resource.create({"service.name": service_name}),
+    )
+    provider.add_span_processor(
+        BatchSpanProcessor(trace_exporter.OTLPSpanExporter(endpoint=cfg.endpoint, insecure=True))
+    )
+    trace.set_tracer_provider(provider)
+
+    PsycopgInstrumentor().instrument()

app/lib/web/server/server.py

@@ -9,6 +9,7 @@
 import uvicorn
 from fastapi import exceptions, responses
 from fastapi.middleware import cors
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from starlette.middleware import base as smiddlewares
 
 from app.lib.web import errors, middlewares
@@ -88,6 +89,8 @@ def __init__(
 
         app.add_exception_handler(exceptions.RequestValidationError, validation_exception_handler)
 
+        FastAPIInstrumentor.instrument_app(app)
+
         self.app = app
         self.config = cfg
         self.logger = logger

configs/dev/adminapi.yaml

@@ -15,3 +15,7 @@ storage:
 
 clients:
   ads_token: token
+
+tracing:
+  endpoint: localhost:4317
+  enabled: true

configs/dev/dataapi.yaml

@@ -9,6 +9,10 @@ storage:
   user: hyperleda
   password: password
 
+tracing:
+  endpoint: localhost:4317
+  enabled: true
+
 catalogs:
   velocity:
     apexes:

configs/prod/adminapi.yaml

@@ -14,3 +14,7 @@ storage:
 
 clients:
   ads_token: fake
+
+tracing:
+  endpoint: tracing-collector:4317
+  enabled: true

configs/prod/dataapi.yaml

@@ -9,6 +9,10 @@ storage:
   user: hyperleda
   password: fake
 
+tracing:
+  endpoint: tracing-collector:4317
+  enabled: true
+
 catalogs:
   velocity:
     apexes:

configs/test/adminapi.yaml

@@ -14,3 +14,7 @@ storage:
 
 clients:
   ads_token: fake
+
+tracing:
+  endpoint: tracing-collector:4317
+  enabled: true