From 182aa58f6c8a1c48682fda2a1c1ae70fd8ffb207 Mon Sep 17 00:00:00 2001
From: sriram veeraghanta <veeraghanta.sriram@gmail.com>
Date: Tue, 26 Nov 2024 15:11:54 +0530
Subject: [PATCH] fix: tracer init fixes

---
 apiserver/plane/license/bgtasks/tracer.py |  7 ++--
 apiserver/plane/settings/common.py        | 21 ------------
 apiserver/plane/utils/telemetry.py        | 40 +++++++++++++++++++++++
 3 files changed, 43 insertions(+), 25 deletions(-)
 create mode 100644 apiserver/plane/utils/telemetry.py

diff --git a/apiserver/plane/license/bgtasks/tracer.py b/apiserver/plane/license/bgtasks/tracer.py
index 60588f5b7b..d857aff7e8 100644
--- a/apiserver/plane/license/bgtasks/tracer.py
+++ b/apiserver/plane/license/bgtasks/tracer.py
@@ -1,6 +1,5 @@
 # Third party imports
 from celery import shared_task
-from opentelemetry import trace
 
 # Module imports
 from plane.license.models import Instance
@@ -16,13 +15,11 @@ from plane.db.models import (
     Page,
     WorkspaceMember,
 )
+from plane.utils.telemetry import init_telemetry
 
 
 @shared_task
 def instance_traces():
-    # Get the tracer
-    tracer = trace.get_tracer(__name__)
-
     # Check if the instance is registered
     instance = Instance.objects.first()
 
@@ -31,6 +28,8 @@ def instance_traces():
         return
 
     if instance.is_telemetry_enabled:
+        # Get the tracer
+        tracer = init_telemetry()
         # Instance details
         with tracer.start_as_current_span("instance_details") as span:
             # Count of all models
diff --git a/apiserver/plane/settings/common.py b/apiserver/plane/settings/common.py
index 864010d3c7..ed42dfe19c 100644
--- a/apiserver/plane/settings/common.py
+++ b/apiserver/plane/settings/common.py
@@ -16,14 +16,6 @@ from sentry_sdk.integrations.django import DjangoIntegration
 from sentry_sdk.integrations.redis import RedisIntegration
 from corsheaders.defaults import default_headers
 
-# OpenTelemetry
-from opentelemetry import trace
-from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from opentelemetry.sdk.resources import Resource
-from opentelemetry.instrumentation.django import DjangoInstrumentor
-
 
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
@@ -33,19 +25,6 @@ SECRET_KEY = os.environ.get("SECRET_KEY", get_random_secret_key())
 # SECURITY WARNING: don't run with debug turned on in production!
 DEBUG = int(os.environ.get("DEBUG", "0"))
 
-# Configure the tracer provider
-service_name = os.environ.get("SERVICE_NAME", "plane-ce-api")
-resource = Resource.create({"service.name": service_name})
-trace.set_tracer_provider(TracerProvider(resource=resource))
-# Configure the OTLP exporter
-otel_endpoint = os.environ.get("OTLP_ENDPOINT", "https://telemetry.plane.so")
-otlp_exporter = OTLPSpanExporter(endpoint=otel_endpoint)
-span_processor = BatchSpanProcessor(otlp_exporter)
-trace.get_tracer_provider().add_span_processor(span_processor)
-# Initialize Django instrumentation
-DjangoInstrumentor().instrument()
-
-
 # Allowed Hosts
 ALLOWED_HOSTS = ["*"]
 
diff --git a/apiserver/plane/utils/telemetry.py b/apiserver/plane/utils/telemetry.py
new file mode 100644
index 0000000000..527bbb1cbe
--- /dev/null
+++ b/apiserver/plane/utils/telemetry.py
@@ -0,0 +1,40 @@
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.instrumentation.django import DjangoInstrumentor
+import os
+
+
+def init_telemetry():
+    """Initialize OpenTelemetry with proper shutdown handling"""
+
+    # Check if already initialized to prevent double initialization
+    if trace.get_tracer_provider().__class__.__name__ == "TracerProvider":
+        return
+
+    # Configure the tracer provider
+    service_name = os.environ.get("SERVICE_NAME", "plane-ce-api")
+    resource = Resource.create({"service.name": service_name})
+    tracer_provider = TracerProvider(resource=resource)
+    trace.set_tracer_provider(tracer_provider)
+
+    # Configure the OTLP exporter
+    otel_endpoint = os.environ.get("OTLP_ENDPOINT", "https://telemetry.plane.so")
+    otlp_exporter = OTLPSpanExporter(endpoint=otel_endpoint)
+    span_processor = BatchSpanProcessor(otlp_exporter)
+    tracer_provider.add_span_processor(span_processor)
+
+    # Initialize Django instrumentation
+    DjangoInstrumentor().instrument()
+
+    return tracer_provider
+
+
+def shutdown_telemetry():
+    """Shutdown OpenTelemetry tracers and processors"""
+    provider = trace.get_tracer_provider()
+
+    if hasattr(provider, "shutdown"):
+        provider.shutdown()