Skip to content

Commit

Permalink
feat: id generator with separate source of randomness (#1010)
Browse files Browse the repository at this point in the history
  • Loading branch information
RogerHYang authored Sep 10, 2024
1 parent 5bc4100 commit ac8cce1
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 2 deletions.
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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 (
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
25 changes: 25 additions & 0 deletions python/openinference-instrumentation/tests/test_id_generator.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit ac8cce1

Please sign in to comment.