⚡️ Speed up method OtelInterceptor.intercept_unary_unary by 49%

[codeflash-ai]

📄 49% (0.49x) speedup for OtelInterceptor.intercept_unary_unary in chromadb/telemetry/opentelemetry/grpc.py

⏱️ Runtime : 664 microseconds → 445 microseconds (best of 91 runs)

📝 Explanation and details

The optimization achieves a 49% speedup by eliminating repeated import overhead, which was the primary performance bottleneck.

Key optimizations:

1. Module-level imports: Moved tracer, _encode_span_id, and _encode_trace_id imports to the top of the file. The line profiler shows the original code spent 36.4% of execution time (895,930ns out of 2,459,270ns) on the repeated from chromadb.telemetry.opentelemetry import tracer import in every call.

2. Reduced attribute lookups: Cached span.get_span_context() in a local variable ctx to avoid calling it twice for trace_id and span_id extraction. Also cached result.code() to avoid repeated method calls.

3. Optimized metadata handling:

   • Used conditional list creation instead of always slicing with [:]
   • Used append() calls instead of extend() for adding the two metadata items
   • Preserved the original tuple conversion for gRPC compatibility

The optimization is most effective for scenarios with frequent tracer-disabled calls (46-99% speedup in tests) and high-volume call patterns (52% speedup for 500 calls), making it ideal for production environments where telemetry may be conditionally enabled and gRPC calls are frequent.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 546 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import binascii
import types

import grpc
# imports
import pytest
from chromadb.telemetry.opentelemetry.grpc import OtelInterceptor
from opentelemetry.trace import SpanKind, StatusCode


def _encode_span_id(span_id: int) -> str:
    return binascii.hexlify(span_id.to_bytes(8, "big")).decode()

def _encode_trace_id(trace_id: int) -> str:
    return binascii.hexlify(trace_id.to_bytes(16, "big")).decode()

class _ClientCallDetails:
    # Mimic grpc.ClientCallDetails for test purposes
    def __init__(self, method, timeout, metadata, credentials):
        self.method = method
        self.timeout = timeout
        self.metadata = metadata
        self.credentials = credentials

class DummySpanContext:
    def __init__(self, trace_id, span_id):
        self.trace_id = trace_id
        self.span_id = span_id

class DummySpan:
    def __init__(self, trace_id, span_id):
        self._context = DummySpanContext(trace_id, span_id)
        self._attributes = {}
        self._status = None
        self._status_desc = None
        self._entered = False
        self._exited = False

    def __enter__(self):
        self._entered = True
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._exited = True

    def get_span_context(self):
        return self._context

    def set_attribute(self, key, value):
        self._attributes[key] = value

    def set_status(self, status, description=None):
        self._status = status
        self._status_desc = description

class DummyTracer:
    def __init__(self, trace_id=0x123456789abcdef0123456789abcdef0, span_id=0x123456789abcdef0):
        self.trace_id = trace_id
        self.span_id = span_id
        self.last_span = None

    def start_as_current_span(self, name, kind=None):
        self.last_span = DummySpan(self.trace_id, self.span_id)
        return self.last_span

class DummyResult:
    def __init__(self, details=None, code=grpc.StatusCode.OK):
        self._details = details
        self._code = code

    def details(self):
        return self._details

    def code(self):
        return self._code

class DummyErrorResult(DummyResult):
    def __init__(self, details=None, code=grpc.StatusCode.UNKNOWN):
        super().__init__(details, code)

class DummyException(Exception):
    pass

# --- Unit Tests ---

# Setup: Patch global tracer for tests
tracer = None  # Will be set in tests

# --- Basic Test Cases ---

def test_basic_successful_call_adds_trace_and_span_metadata():
    """Test basic call: metadata is added, span attributes are set, result is returned."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    # Prepare a dummy continuation that returns a successful DummyResult
    def continuation(client_call_details, request):
        # Check that metadata contains chroma-traceid and chroma-spanid
        meta_keys = [k for k, v in client_call_details.metadata]
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=10,
        metadata=[("foo", "bar")],
        credentials=None,
    )
    request = "test-request"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 3.97μs -> 2.71μs (46.4% faster)
    # Check span attributes set
    span = tracer.last_span

def test_basic_no_metadata():
    """Test call with no initial metadata: metadata is created and added."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    def continuation(client_call_details, request):
        # Should still have chroma-traceid and chroma-spanid
        meta_keys = [k for k, v in client_call_details.metadata]
        return DummyResult(details=None, code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=5,
        metadata=None,
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output
    span = tracer.last_span

def test_basic_tracer_none_passthrough():
    """If tracer is None, interceptor should pass through unchanged."""
    global tracer
    tracer = None
    interceptor = OtelInterceptor()
    called = []
    def continuation(client_call_details, request):
        called.append(True)
        # Should not have chroma-traceid or chroma-spanid
        if client_call_details.metadata:
            meta_keys = [k for k, v in client_call_details.metadata]
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=1,
        metadata=[("foo", "bar")],
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 4.77μs -> 3.26μs (46.2% faster)

# --- Edge Test Cases ---

def test_error_status_sets_span_status_and_attributes():
    """If result code is not OK, span status should be set to ERROR."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    def continuation(client_call_details, request):
        return DummyErrorResult(details="fail", code=grpc.StatusCode.UNKNOWN)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=2,
        metadata=[("foo", "bar")],
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 3.46μs -> 3.00μs (15.3% faster)
    span = tracer.last_span

def test_exception_in_continuation_sets_span_status_and_reraises():
    """If continuation raises, span status should be set to ERROR and exception re-raised."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    def continuation(client_call_details, request):
        raise DummyException("test error")
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=2,
        metadata=[("foo", "bar")],
        credentials=None,
    )
    request = "req"
    with pytest.raises(DummyException):
        interceptor.intercept_unary_unary(continuation, client_call_details, request) # 2.46μs -> 1.47μs (67.7% faster)
    span = tracer.last_span

def test_metadata_is_tuple_and_preserves_original():
    """Ensure metadata is always tuple and original metadata is preserved."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    def continuation(client_call_details, request):
        meta_keys = [k for k, v in client_call_details.metadata]
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=1,
        metadata=[("foo", "bar")],
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 3.72μs -> 2.67μs (39.0% faster)

def test_span_id_and_trace_id_encoding():
    """Test that span and trace IDs are encoded as expected."""
    global tracer
    # Use known values
    tracer = DummyTracer(trace_id=0x000102030405060708090a0b0c0d0e0f, span_id=0x0102030405060708)
    interceptor = OtelInterceptor()
    def continuation(client_call_details, request):
        # Check encoded values
        meta = dict(client_call_details.metadata)
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=1,
        metadata=None,
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output

def test_handles_empty_metadata_list():
    """Test that empty metadata list is handled gracefully."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    def continuation(client_call_details, request):
        meta_keys = [k for k, v in client_call_details.metadata]
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=1,
        metadata=[],
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 4.31μs -> 2.91μs (48.1% faster)

# --- Large Scale Test Cases ---

def test_large_metadata_list():
    """Test with a large metadata list (up to 1000 elements)."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    initial_metadata = [(f"key{i}", f"value{i}") for i in range(995)]
    def continuation(client_call_details, request):
        # Should preserve all original keys and add two new ones
        meta_keys = [k for k, v in client_call_details.metadata]
        for i in range(995):
            pass
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method="/test.Service/Method",
        timeout=20,
        metadata=initial_metadata,
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 29.3μs -> 27.9μs (5.08% faster)

def test_large_scale_multiple_calls():
    """Test many calls in succession to check for state leaks or performance issues."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    call_count = 500
    for i in range(call_count):
        def continuation(client_call_details, request):
            meta_keys = [k for k, v in client_call_details.metadata]
            return DummyResult(details=str(i), code=grpc.StatusCode.OK)
        client_call_details = _ClientCallDetails(
            method=f"/test.Service/Method{i}",
            timeout=1,
            metadata=[(f"foo{i}", f"bar{i}")],
            credentials=None,
        )
        request = f"req{i}"
        codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 596μs -> 392μs (51.9% faster)
        span = tracer.last_span

def test_large_scale_long_method_name():
    """Test with a very long gRPC method name."""
    global tracer
    tracer = DummyTracer()
    interceptor = OtelInterceptor()
    long_method = "/test.Service/" + "A" * 900
    def continuation(client_call_details, request):
        meta_keys = [k for k, v in client_call_details.metadata]
        return DummyResult(details="ok", code=grpc.StatusCode.OK)
    client_call_details = _ClientCallDetails(
        method=long_method,
        timeout=1,
        metadata=None,
        credentials=None,
    )
    request = "req"
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output
    span = tracer.last_span

def test_large_scale_different_span_ids():
    """Test that different tracer span/trace IDs are propagated correctly."""
    global tracer
    interceptor = OtelInterceptor()
    for i in range(10):
        tracer = DummyTracer(trace_id=i+100, span_id=i+200)
        def continuation(client_call_details, request):
            meta = dict(client_call_details.metadata)
            return DummyResult(details="ok", code=grpc.StatusCode.OK)
        client_call_details = _ClientCallDetails(
            method="/test.Service/Method",
            timeout=1,
            metadata=None,
            credentials=None,
        )
        request = "req"
        codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import binascii
# Patch point for the module import
import sys
import types
import types as py_types

import grpc
# imports
import pytest  # used for our unit tests
from chromadb.telemetry.opentelemetry.grpc import OtelInterceptor
from opentelemetry.trace import SpanKind, StatusCode


# Minimal ClientCallDetails implementation for testing
class _ClientCallDetails:
    def __init__(self, method, timeout, metadata, credentials):
        self.method = method
        self.timeout = timeout
        self.metadata = metadata
        self.credentials = credentials

# Dummy tracer and span for testing
class DummySpanContext:
    def __init__(self, trace_id=12345678901234567890123456789012, span_id=9876543210987654):
        self.trace_id = trace_id
        self.span_id = span_id

class DummySpan:
    def __init__(self, context=None):
        self.context = context or DummySpanContext()
        self.ended = False
        self.attributes = {}
        self.status = None
        self.status_desc = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.ended = True

    def get_span_context(self):
        return self.context

    def set_attribute(self, key, value):
        self.attributes[key] = value

    def set_status(self, status, description=None):
        self.status = status
        self.status_desc = description

class DummyTracer:
    def __init__(self, span_context=None):
        self.span_context = span_context or DummySpanContext()
        self.started_spans = []

    def start_as_current_span(self, name, kind=None):
        span = DummySpan(self.span_context)
        self.started_spans.append((name, kind, span))
        return span

chromadb_telemetry_opentelemetry = py_types.ModuleType("chromadb.telemetry.opentelemetry")
chromadb_telemetry_opentelemetry.tracer = DummyTracer()

# ---- UNIT TESTS ----

# Dummy gRPC response object for testing
class DummyGRPCResponse:
    def __init__(self, code=grpc.StatusCode.OK, details=""):
        self._code = code
        self._details = details

    def code(self):
        return self._code

    def details(self):
        return self._details

# 1. Basic Test Cases

def test_basic_successful_call_adds_trace_and_span_id():
    """Test a basic successful unary-unary call with no metadata."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    client_call_details = _ClientCallDetails(method, 5.0, None, None)
    called = {}

    def continuation(details, req):
        # Should receive new metadata with trace and span ids
        called['details'] = details
        called['req'] = req
        return DummyGRPCResponse()

    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output
    meta_keys = [k for k, v in called['details'].metadata]

def test_basic_metadata_preserved_and_extended():
    """Test that existing metadata is preserved and extended with trace/span ids."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    original_metadata = [("foo", "bar")]
    client_call_details = _ClientCallDetails(method, 5.0, original_metadata, None)
    called = {}

    def continuation(details, req):
        called['details'] = details
        return DummyGRPCResponse()

    interceptor.intercept_unary_unary(continuation, client_call_details, request) # 2.98μs -> 1.69μs (76.2% faster)
    # Metadata should contain both original and new keys
    meta = dict(called['details'].metadata)

def test_basic_status_and_attributes_set_on_span():
    """Test that span attributes are set according to gRPC response."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    client_call_details = _ClientCallDetails(method, 5.0, None, None)

    def continuation(details, req):
        return DummyGRPCResponse(code=grpc.StatusCode.NOT_FOUND, details="not found")

    tracer = chromadb_telemetry_opentelemetry.tracer
    tracer.started_spans.clear()
    interceptor.intercept_unary_unary(continuation, client_call_details, request)
    name, kind, span = tracer.started_spans[-1]

# 2. Edge Test Cases

def test_no_tracer_returns_continuation_directly(monkeypatch):
    """If tracer is None, should call and return continuation without modification."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    client_call_details = _ClientCallDetails(method, 5.0, None, None)
    called = {}

    # Patch tracer to None
    monkeypatch.setattr(chromadb_telemetry_opentelemetry, "tracer", None)
    def continuation(details, req):
        called['details'] = details
        called['req'] = req
        return "direct_result"

    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.34μs -> 1.17μs (99.7% faster)
    # Restore tracer
    monkeypatch.setattr(chromadb_telemetry_opentelemetry, "tracer", DummyTracer())


def test_exception_in_continuation_sets_error_on_span():
    """Test that exception in continuation sets error attributes and re-raises."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    client_call_details = _ClientCallDetails(method, 5.0, None, None)

    class MyException(Exception):
        pass

    def continuation(details, req):
        raise MyException("boom")

    tracer = chromadb_telemetry_opentelemetry.tracer
    tracer.started_spans.clear()
    with pytest.raises(MyException):
        interceptor.intercept_unary_unary(continuation, client_call_details, request)
    name, kind, span = tracer.started_spans[-1]

def test_empty_metadata_and_details():
    """Test when metadata is empty and details() returns empty string."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    client_call_details = _ClientCallDetails(method, 5.0, [], None)

    def continuation(details, req):
        # details() returns empty string
        class Resp(DummyGRPCResponse):
            def details(self): return ""
        return Resp()

    tracer = chromadb_telemetry_opentelemetry.tracer
    tracer.started_spans.clear()
    interceptor.intercept_unary_unary(continuation, client_call_details, request)
    name, kind, span = tracer.started_spans[-1]

def test_nonstandard_metadata_types():
    """Test that metadata with non-string values is handled (should be stringified by grpc)."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    # Metadata with int value
    original_metadata = [("num", 123)]
    client_call_details = _ClientCallDetails(method, 5.0, original_metadata, None)
    called = {}

    def continuation(details, req):
        called['details'] = details
        return DummyGRPCResponse()

    interceptor.intercept_unary_unary(continuation, client_call_details, request) # 3.05μs -> 1.66μs (83.6% faster)
    meta = dict(called['details'].metadata)

def test_multiple_calls_unique_span_and_trace_ids():
    """Each call should generate new span/trace ids in metadata."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    request = "request"
    client_call_details = _ClientCallDetails(method, 5.0, None, None)
    results = []

    class CustomTracer(DummyTracer):
        def __init__(self):
            super().__init__()
            self.count = 0
        def start_as_current_span(self, name, kind=None):
            # Each call gets a new context
            self.count += 1
            ctx = DummySpanContext(trace_id=1000+self.count, span_id=2000+self.count)
            span = DummySpan(ctx)
            self.started_spans.append((name, kind, span))
            return span

    chromadb_telemetry_opentelemetry.tracer = CustomTracer()

    def continuation(details, req):
        results.append(dict(details.metadata))
        return DummyGRPCResponse()

    interceptor.intercept_unary_unary(continuation, client_call_details, request)
    interceptor.intercept_unary_unary(continuation, client_call_details, request)
    # Check that traceid/spanid are different for each call
    traceids = [meta["chroma-traceid"] for meta in results]
    spanids = [meta["chroma-spanid"] for meta in results]

# 3. Large Scale Test Cases

def test_large_metadata_and_large_request():
    """Test with large metadata and large request payload."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    # Large metadata (999 elements)
    large_metadata = [(f"key{i}", f"value{i}") for i in range(999)]
    client_call_details = _ClientCallDetails(method, 5.0, large_metadata, None)
    # Large request (string of 10,000 chars)
    request = "x" * 10000
    called = {}

    def continuation(details, req):
        called['details'] = details
        called['req'] = req
        return DummyGRPCResponse()

    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 3.37μs -> 1.87μs (80.6% faster)
    # Metadata should be preserved and extended
    meta = dict(called['details'].metadata)
    for i in range(999):
        pass

def test_high_volume_of_calls():
    """Test interceptor under high volume of calls (1000 calls)."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    client_call_details = _ClientCallDetails(method, 5.0, None, None)
    call_count = 1000
    seen_traceids = set()
    seen_spanids = set()

    class CustomTracer(DummyTracer):
        def __init__(self):
            super().__init__()
            self.counter = 0
        def start_as_current_span(self, name, kind=None):
            self.counter += 1
            ctx = DummySpanContext(trace_id=1000000+self.counter, span_id=2000000+self.counter)
            span = DummySpan(ctx)
            self.started_spans.append((name, kind, span))
            return span

    chromadb_telemetry_opentelemetry.tracer = CustomTracer()

    def continuation(details, req):
        meta = dict(details.metadata)
        seen_traceids.add(meta["chroma-traceid"])
        seen_spanids.add(meta["chroma-spanid"])
        return DummyGRPCResponse()

    for _ in range(call_count):
        interceptor.intercept_unary_unary(continuation, client_call_details, "req")

def test_large_metadata_with_edge_values():
    """Test with metadata containing empty strings and very long strings."""
    interceptor = OtelInterceptor()
    method = "/service/method"
    # Metadata with empty string key/value and long string
    long_str = "a" * 1000
    metadata = [("", ""), ("long", long_str)]
    client_call_details = _ClientCallDetails(method, 5.0, metadata, None)
    called = {}

    def continuation(details, req):
        called['details'] = details
        return DummyGRPCResponse()

    interceptor.intercept_unary_unary(continuation, client_call_details, "req") # 2.96μs -> 1.75μs (69.0% faster)
    meta = dict(called['details'].metadata)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-OtelInterceptor.intercept_unary_unary-mh1tobsp and push.