From ac8cce112341bb31a575cb1e61a55acb196fc600 Mon Sep 17 00:00:00 2001
From: Roger Yang <80478925+RogerHYang@users.noreply.github.com>
Date: Tue, 10 Sep 2024 09:23:17 -0700
Subject: [PATCH] feat: id generator with separate source of randomness (#1010)

---
 .../openinference/instrumentation/config.py   | 37 ++++++++++++++++++-
 .../tests/test_id_generator.py                | 25 +++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 python/openinference-instrumentation/tests/test_id_generator.py

diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py
index dd20abefc..1d441f6f8 100644
--- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py
+++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py
@@ -1,6 +1,7 @@
 import os
 from contextlib import contextmanager
 from dataclasses import dataclass, field, fields
+from secrets import randbits
 from typing import (
     Any,
     Callable,
@@ -21,7 +22,16 @@
     detach,
     set_value,
 )
-from opentelemetry.trace import Link, Span, SpanKind, Tracer, use_span
+from opentelemetry.sdk.trace import IdGenerator
+from opentelemetry.trace import (
+    INVALID_SPAN_ID,
+    INVALID_TRACE_ID,
+    Link,
+    Span,
+    SpanKind,
+    Tracer,
+    use_span,
+)
 from opentelemetry.util.types import Attributes, AttributeValue
 
 from openinference.semconv.trace import (
@@ -343,10 +353,32 @@ def end(self, end_time: Optional[int] = None) -> None:
         span.end(end_time)
 
 
+class _IdGenerator(IdGenerator):
+    """
+    An IdGenerator that uses a different source of randomness to
+    avoid being affected by seeds set by user application.
+    """
+
+    def generate_span_id(self) -> int:
+        while (span_id := randbits(64)) == INVALID_SPAN_ID:
+            continue
+        return span_id
+
+    def generate_trace_id(self) -> int:
+        while (trace_id := randbits(128)) == INVALID_TRACE_ID:
+            continue
+        return trace_id
+
+
 class OITracer(wrapt.ObjectProxy):  # type: ignore[misc]
     def __init__(self, wrapped: Tracer, config: TraceConfig) -> None:
         super().__init__(wrapped)
         self._self_config = config
+        self._self_id_generator = _IdGenerator()
+
+    @property
+    def id_generator(self) -> IdGenerator:
+        return self._self_id_generator
 
     @contextmanager
     def start_as_current_span(
@@ -391,7 +423,8 @@ def start_span(
         set_status_on_exception: bool = True,
     ) -> Span:
         tracer = cast(Tracer, self.__wrapped__)
-        span = tracer.start_span(
+        span = tracer.__class__.start_span(
+            self,
             name=name,
             context=context,
             kind=kind,
diff --git a/python/openinference-instrumentation/tests/test_id_generator.py b/python/openinference-instrumentation/tests/test_id_generator.py
new file mode 100644
index 000000000..7d81574e9
--- /dev/null
+++ b/python/openinference-instrumentation/tests/test_id_generator.py
@@ -0,0 +1,25 @@
+from random import seed
+
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from openinference.instrumentation import OITracer, TraceConfig
+
+
+def test_id_generator_is_unaffected_by_seed() -> None:
+    in_memory_span_exporter = InMemorySpanExporter()
+    tracer_provider = TracerProvider()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(in_memory_span_exporter))
+    tracer = tracer_provider.get_tracer(__name__)
+    oi_tracer = OITracer(tracer, TraceConfig())
+    n = 10
+    for tr in (tracer, oi_tracer):
+        for _ in range(n):
+            seed(42)
+            with tr.start_as_current_span("parent"):
+                tr.start_span("child").end()
+    spans = in_memory_span_exporter.get_finished_spans()
+    assert len(spans) == n * 2 * 2
+    assert len(set(span.context.trace_id for span in spans)) == (n + 1)
+    assert len(set(span.context.span_id for span in spans)) == (n + 1) * 2