fix(sdk): sampling support (#3027)

Author: galkleinman

packages/traceloop-sdk/tests/test_sampler_initialization.py

@@ -0,0 +1,35 @@
+"""Tests for OpenTelemetry sampler initialization functionality."""
+
+from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from traceloop.sdk import Traceloop
+from traceloop.sdk.tracing.tracing import TracerWrapper
+
+
+class TestSamplerInitialization:
+    """Test class for sampler initialization functionality."""
+
+    def test_init_with_rate_based_sampler(self):
+        """Test initialization with TraceIdRatioBased sampler."""
+        if hasattr(TracerWrapper, "instance"):
+            _trace_wrapper_instance = TracerWrapper.instance
+            del TracerWrapper.instance
+
+        try:
+            exporter = InMemorySpanExporter()
+            sampler = TraceIdRatioBased(0.5)  # 50% sampling rate
+
+            client = Traceloop.init(
+                app_name="test-rate-based-sampler",
+                sampler=sampler,
+                exporter=exporter,
+                disable_batch=True
+            )
+
+            assert client is None
+            assert hasattr(TracerWrapper, "instance")
+            assert TracerWrapper.instance is not None
+
+        finally:
+            if '_trace_wrapper_instance' in locals():
+                TracerWrapper.instance = _trace_wrapper_instance

packages/traceloop-sdk/traceloop/sdk/__init__.py

@@ -5,6 +5,7 @@
 from typing import Callable, Optional, Set
 from colorama import Fore
 from opentelemetry.sdk.trace import SpanProcessor, ReadableSpan
+from opentelemetry.sdk.trace.sampling import Sampler
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.sdk.metrics.export import MetricExporter
 from opentelemetry.sdk._logs.export import LogExporter
@@ -60,6 +61,7 @@ def init(
         logging_headers: Dict[str, str] = None,
         processor: Optional[SpanProcessor] = None,
         propagator: TextMapPropagator = None,
+        sampler: Optional[Sampler] = None,
         traceloop_sync_enabled: bool = False,
         should_enrich_metrics: bool = True,
         resource_attributes: dict = {},
@@ -144,6 +146,7 @@ def init(
             processor=processor,
             propagator=propagator,
             exporter=exporter,
+            sampler=sampler,
             should_enrich_metrics=should_enrich_metrics,
             image_uploader=image_uploader or ImageUploader(api_endpoint, api_key),
             instruments=instruments,

packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py

@@ -13,6 +13,7 @@
 )
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, SpanProcessor, ReadableSpan
+from opentelemetry.sdk.trace.sampling import Sampler
 from opentelemetry.propagators.textmap import TextMapPropagator
 from opentelemetry.propagate import set_global_textmap
 from opentelemetry.sdk.trace.export import (
@@ -69,6 +70,7 @@ def __new__(
         processor: SpanProcessor = None,
         propagator: TextMapPropagator = None,
         exporter: SpanExporter = None,
+        sampler: Optional[Sampler] = None,
         should_enrich_metrics: bool = True,
         instruments: Optional[Set[Instruments]] = None,
         block_instruments: Optional[Set[Instruments]] = None,
@@ -82,7 +84,7 @@ def __new__(
 
             obj.__image_uploader = image_uploader
             obj.__resource = Resource(attributes=TracerWrapper.resource_attributes)
-            obj.__tracer_provider = init_tracer_provider(resource=obj.__resource)
+            obj.__tracer_provider = init_tracer_provider(resource=obj.__resource, sampler=sampler)
             if processor:
                 Telemetry().capture("tracer:init", {"processor": "custom"})
                 obj.__spans_processor: SpanProcessor = processor
@@ -339,12 +341,15 @@ def init_spans_exporter(api_endpoint: str, headers: Dict[str, str]) -> SpanExpor
         return GRPCExporter(endpoint=f"{api_endpoint}", headers=headers)
 
 
-def init_tracer_provider(resource: Resource) -> TracerProvider:
+def init_tracer_provider(resource: Resource, sampler: Optional[Sampler] = None) -> TracerProvider:
     provider: TracerProvider = None
     default_provider: TracerProvider = get_tracer_provider()
 
     if isinstance(default_provider, ProxyTracerProvider):
-        provider = TracerProvider(resource=resource)
+        if sampler is not None:
+            provider = TracerProvider(resource=resource, sampler=sampler)
+        else:
+            provider = TracerProvider(resource=resource)
         trace.set_tracer_provider(provider)
     elif not hasattr(default_provider, "add_span_processor"):
         logging.error(