⚡️ Speed up method OtelInterceptor._intercept_call by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for OtelInterceptor._intercept_call in chromadb/telemetry/opentelemetry/grpc.py

⏱️ Runtime : 170 microseconds → 158 microseconds (best of 120 runs)

📝 Explanation and details

The optimization achieves a 7% speedup through three key improvements:

1. Replaced binascii.hexlify().decode() with .hex()

• Changed binascii.hexlify(span_id.to_bytes(8, "big")).decode() to span_id.to_bytes(8, "big").hex()
• The .hex() method is a direct, native operation that's faster than the two-step hexlify+decode process
• Line profiler shows encoding functions run ~30% faster (13.3μs → 9.3μs for span_id, 24.6μs → 16.3μs for trace_id)

2. Eliminated redundant get_span_context() calls

• Original code called span.get_span_context() twice; optimized version stores it in span_ctx variable
• Reduces method call overhead in the hot path where trace/span IDs are extracted

3. Minimized result.code() calls and improved conditional logic

• Cached result.code() in a code variable instead of calling it 3 times
• Restructured the details check to avoid the redundant and result.details() call when hasattr fails
• Optimized metadata handling to avoid unnecessary slicing with client_call_details.metadata[:]

These optimizations are particularly effective for high-frequency gRPC calls where the encoding functions and span context access happen repeatedly. The test results show consistent 6-13% improvements across various scenarios, with the largest gains (13.7%) occurring when handling large trace/span IDs where the encoding optimization has maximum impact.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import binascii
import types

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


# Minimal stub for grpc and opentelemetry objects
class DummyStatusCode:
    OK = "OK"
    ERROR = "ERROR"
    CANCELLED = "CANCELLED"
    UNKNOWN = "UNKNOWN"
    # Simulate grpc.StatusCode enums
    def __init__(self, name, value):
        self.name = name
        self.value = (value,)

class DummyResult:
    def __init__(self, code_name="OK", code_value=0, details_val=None):
        self._code = DummyStatusCode(code_name, code_value)
        self._details_val = details_val

    def code(self):
        return self._code

    def details(self):
        return self._details_val

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

    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, description)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return False

class DummyTracer:
    def __init__(self, trace_id=1, span_id=2):
        self.trace_id = trace_id
        self.span_id = span_id

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

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

# The function to test, adapted to be testable (tracer is injected)
def _intercept_call(continuation, client_call_details, request_or_iterator, tracer):
    # If tracer is None, just call through
    if tracer is None:
        return continuation(client_call_details, request_or_iterator)
    with tracer.start_as_current_span(
        f"RPC {client_call_details.method}", kind="CLIENT"
    ) as span:
        # Prepare metadata for propagation
        metadata = (
            client_call_details.metadata[:] if client_call_details.metadata else []
        )
        metadata.extend(
            [
                (
                    "chroma-traceid",
                    _encode_trace_id(span.get_span_context().trace_id),
                ),
                ("chroma-spanid", _encode_span_id(span.get_span_context().span_id)),
            ]
        )
        # Update client call details with new metadata
        new_client_details = _ClientCallDetails(
            client_call_details.method,
            client_call_details.timeout,
            tuple(metadata),  # Ensure metadata is a tuple
            client_call_details.credentials,
        )
        try:
            result = continuation(new_client_details, request_or_iterator)
            # Set attributes based on the result
            if hasattr(result, "details") and result.details():
                span.set_attribute("rpc.detail", result.details())
            span.set_attribute("rpc.status_code", result.code().name.lower())
            span.set_attribute("rpc.status_code_value", result.code().value[0])
            # Set span status based on gRPC call result
            if result.code().name != "OK":
                span.set_status("ERROR", description=str(result.code().name))
            return result
        except Exception as e:
            # Log exception details and re-raise
            span.set_attribute("rpc.error", str(e))
            span.set_status("ERROR", description=str(e))
            raise

# ------------------------
# Unit tests start here
# ------------------------

# 1. Basic Test Cases















#------------------------------------------------
import binascii
# Patch module import for chromadb.telemetry.opentelemetry.tracer
import sys
import types
import types as pytypes

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


class _ClientCallDetails:
    # Minimal implementation for testing
    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._ctx = DummySpanContext(trace_id, span_id)
        self._attrs = {}
        self._status = None
        self._desc = None
        self._ended = False

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        self._ended = True

    def get_span_context(self):
        return self._ctx

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

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

class DummyTracer:
    def __init__(self, trace_id=1, span_id=2):
        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

chromadb_mod = pytypes.ModuleType("chromadb")
telemetry_mod = pytypes.ModuleType("chromadb.telemetry")
opentelemetry_mod = pytypes.ModuleType("chromadb.telemetry.opentelemetry")
opentelemetry_mod.tracer = None  # default, can be set in tests
telemetry_mod.opentelemetry = opentelemetry_mod
chromadb_mod.telemetry = telemetry_mod
sys.modules["chromadb"] = chromadb_mod
sys.modules["chromadb.telemetry"] = telemetry_mod
sys.modules["chromadb.telemetry.opentelemetry"] = opentelemetry_mod

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

    def code(self):
        return self._code

    def details(self):
        return self._details

class DummyContinuation:
    def __init__(self, result=None, raise_exc=None):
        self.result = result
        self.raise_exc = raise_exc
        self.called_with = []

    def __call__(self, client_call_details, request_or_iterator):
        self.called_with.append((client_call_details, request_or_iterator))
        if self.raise_exc:
            raise self.raise_exc
        return self.result

class DummyClientCallDetails:
    def __init__(self, method, timeout=None, metadata=None, credentials=None):
        self.method = method
        self.timeout = timeout
        self.metadata = metadata
        self.credentials = credentials
from chromadb.telemetry.opentelemetry.grpc import OtelInterceptor

# --- Unit Tests ---

@pytest.fixture
def interceptor():
    return OtelInterceptor()

# --- Basic Test Cases ---

def test_no_tracer_returns_continuation_result(interceptor):
    """
    If tracer is None, _intercept_call should just call and return the result of continuation.
    """
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = None
    result = DummyResult(grpc.StatusCode.OK, details="ok")
    continuation = DummyContinuation(result=result)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 2.96μs -> 2.95μs (0.339% faster)

def test_tracer_sets_metadata_and_span_attributes_ok(interceptor):
    """
    When tracer is set, _intercept_call should inject trace/span IDs into metadata,
    and set span attributes for an OK result.
    """
    tracer = DummyTracer(trace_id=123, span_id=456)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK, details="All good")
    continuation = DummyContinuation(result=result)
    orig_metadata = [("foo", "bar")]
    call_details = DummyClientCallDetails(method="TestMethod", metadata=orig_metadata)
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 13.9μs -> 12.5μs (11.6% faster)
    # Check metadata is extended
    new_call_details = continuation.called_with[0][0]
    # Check span attributes
    span = tracer.last_span

def test_tracer_sets_error_status_on_non_ok(interceptor):
    """
    Should set error status on the span if result code is not OK.
    """
    tracer = DummyTracer(trace_id=1, span_id=2)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.INVALID_ARGUMENT, details="bad arg")
    continuation = DummyContinuation(result=result)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 12.9μs -> 12.3μs (4.26% faster)
    span = tracer.last_span

def test_tracer_handles_no_details(interceptor):
    """
    Should not fail if result.details() is None.
    """
    tracer = DummyTracer()
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK, details=None)
    continuation = DummyContinuation(result=result)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 10.2μs -> 9.48μs (7.81% faster)
    span = tracer.last_span

def test_tracer_handles_missing_details_method(interceptor):
    """
    Should not fail if result does not have a details method.
    """
    tracer = DummyTracer()
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer

    class ResultNoDetails:
        def code(self):
            return grpc.StatusCode.OK

    result = ResultNoDetails()
    continuation = DummyContinuation(result=result)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 10.9μs -> 10.1μs (8.37% faster)
    span = tracer.last_span

# --- Edge Test Cases ---

def test_metadata_none_and_empty(interceptor):
    """
    Should handle None and empty metadata gracefully.
    """
    tracer = DummyTracer(trace_id=999, span_id=888)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK)
    continuation = DummyContinuation(result=result)
    # None metadata
    call_details_none = DummyClientCallDetails(method="TestMethod", metadata=None)
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details_none, req); ret_none = codeflash_output # 10.2μs -> 9.43μs (8.15% faster)
    new_call_details_none = continuation.called_with[0][0]
    # Empty metadata
    call_details_empty = DummyClientCallDetails(method="TestMethod", metadata=[])
    continuation.called_with.clear()
    codeflash_output = interceptor._intercept_call(continuation, call_details_empty, req); ret_empty = codeflash_output # 5.88μs -> 5.51μs (6.75% faster)
    new_call_details_empty = continuation.called_with[0][0]

def test_handles_exception_in_continuation(interceptor):
    """
    Should set error attributes and status if continuation raises.
    """
    tracer = DummyTracer()
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    exc = ValueError("fail!")
    continuation = DummyContinuation(raise_exc=exc)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    with pytest.raises(ValueError):
        interceptor._intercept_call(continuation, call_details, req) # 9.68μs -> 8.84μs (9.58% faster)
    span = tracer.last_span

def test_handles_unusual_method_name_and_metadata(interceptor):
    """
    Should handle odd method names and metadata types.
    """
    tracer = DummyTracer(trace_id=0xDEADBEEF, span_id=0xFEEDFACE)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK)
    # Unusual method name and metadata
    call_details = DummyClientCallDetails(
        method="!@#$%^&*()_+",
        metadata=[("x", ""), ("y", None), ("z", "🚀")],
    )
    continuation = DummyContinuation(result=result)
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 10.5μs -> 9.80μs (6.92% faster)
    new_call_details = continuation.called_with[0][0]
    # Check injected trace/span IDs are hex strings
    chroma_traceid = [v for k, v in new_call_details.metadata if k == "chroma-traceid"][0]
    chroma_spanid = [v for k, v in new_call_details.metadata if k == "chroma-spanid"][0]

def test_handles_timeout_and_credentials(interceptor):
    """
    Should propagate timeout and credentials.
    """
    tracer = DummyTracer()
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK)
    credentials = object()
    call_details = DummyClientCallDetails(
        method="TestMethod", timeout=123.45, metadata=[], credentials=credentials
    )
    continuation = DummyContinuation(result=result)
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 9.70μs -> 9.10μs (6.66% faster)
    new_call_details = continuation.called_with[0][0]

# --- Large Scale Test Cases ---

def test_large_metadata_list(interceptor):
    """
    Should handle large metadata lists efficiently and correctly.
    """
    tracer = DummyTracer(trace_id=0x123456789ABCDEF0, span_id=0x0FEDCBA987654321)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    large_metadata = [(f"key{i}", f"value{i}") for i in range(1000)]
    result = DummyResult(grpc.StatusCode.OK)
    call_details = DummyClientCallDetails(method="TestMethod", metadata=large_metadata)
    continuation = DummyContinuation(result=result)
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 15.8μs -> 15.4μs (2.78% faster)
    new_call_details = continuation.called_with[0][0]

def test_many_calls_in_sequence(interceptor):
    """
    Should not leak state between calls; each call gets its own span.
    """
    tracer = DummyTracer(trace_id=42, span_id=99)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK)
    continuation = DummyContinuation(result=result)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    for i in range(10):
        codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 47.7μs -> 44.2μs (7.96% faster)
        span = tracer.last_span

def test_handles_large_trace_and_span_ids(interceptor):
    """
    Should encode large trace/span IDs correctly.
    """
    big_trace_id = 2**128 - 1
    big_span_id = 2**64 - 1
    tracer = DummyTracer(trace_id=big_trace_id, span_id=big_span_id)
    sys.modules["chromadb.telemetry.opentelemetry"].tracer = tracer
    result = DummyResult(grpc.StatusCode.OK)
    continuation = DummyContinuation(result=result)
    call_details = DummyClientCallDetails(method="TestMethod")
    req = object()
    codeflash_output = interceptor._intercept_call(continuation, call_details, req); ret = codeflash_output # 9.27μs -> 8.16μs (13.7% faster)
    new_call_details = continuation.called_with[0][0]
    chroma_traceid = [v for k, v in new_call_details.metadata if k == "chroma-traceid"][0]
    chroma_spanid = [v for k, v in new_call_details.metadata if k == "chroma-spanid"][0]
# 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_call-mh1tgx7e and push.