⚡️ Speed up function _temporal_context by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for _temporal_context in temporalio/nexus/_operation_context.py

⏱️ Runtime : 27.8 microseconds → 24.5 microseconds (best of 66 runs)

📝 Explanation and details

The optimization inlines the context retrieval logic directly into _temporal_context, eliminating a function call overhead.

Key changes:

• Moved the start_ctx and cancel_ctx retrieval from _try_temporal_context() directly into _temporal_context()
• Removed the intermediate function call that was consuming 87.5% of the execution time in the original version

Why this is faster:

• Eliminates function call overhead: The original version had _temporal_context() calling _try_temporal_context(), which involved stack frame creation, parameter passing, and return value handling
• Reduces instruction count: Direct execution of context retrieval logic removes the call/return instructions
• Better CPU cache locality: Fewer jumps between code locations improves instruction cache efficiency

Performance characteristics:
The optimization is particularly effective for error cases (when no context is set) as shown in the test results - achieving 15.3-18.7% speedup when _temporal_context() needs to raise the "Not in Nexus operation context" error. This makes sense because these are likely the most common code paths where the function is called to validate context existence.

The _try_temporal_context() function remains unchanged to preserve the existing API for other callers who need the optional return behavior.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

from contextvars import ContextVar
from typing import Optional, Union

# imports
import pytest  # used for our unit tests
from temporalio.nexus._operation_context import _temporal_context


# Dummy context classes for testing
class _TemporalStartOperationContext:
    def __init__(self, value):
        self.value = value

class _TemporalCancelOperationContext:
    def __init__(self, value):
        self.value = value

_temporal_start_operation_context: ContextVar[_TemporalStartOperationContext] = (
    ContextVar("temporal-start-operation-context")
)

_temporal_cancel_operation_context: ContextVar[_TemporalCancelOperationContext] = (
    ContextVar("temporal-cancel-operation-context")
)
from temporalio.nexus._operation_context import _temporal_context

# unit tests

# ----- BASIC TEST CASES -----





def test_both_contexts_set_raises():
    """Test that setting both contexts raises a RuntimeError."""
    start_ctx = _TemporalStartOperationContext("start")
    cancel_ctx = _TemporalCancelOperationContext("cancel")
    token1 = _temporal_start_operation_context.set(start_ctx)
    token2 = _temporal_cancel_operation_context.set(cancel_ctx)
    try:
        with pytest.raises(RuntimeError) as excinfo:
            _temporal_context()
    finally:
        _temporal_start_operation_context.reset(token1)
        _temporal_cancel_operation_context.reset(token2)

def test_neither_context_set_raises():
    """Test that not setting any context raises a RuntimeError."""
    with pytest.raises(RuntimeError) as excinfo:
        _temporal_context() # 625ns -> 542ns (15.3% faster)

def test_try_temporal_context_none():
    """Test that _try_temporal_context returns None if neither context is set."""

def test_context_reset_behavior():
    """Test that resetting contextvars properly removes context."""
    ctx = _TemporalStartOperationContext("start")
    token = _temporal_start_operation_context.set(ctx)
    _temporal_start_operation_context.reset(token)





def test_contextvars_isolation():
    """Test that contextvars are isolated between test runs (no leakage)."""
    # Set context, reset, and ensure next test sees nothing
    ctx = _TemporalStartOperationContext("isolate")
    token = _temporal_start_operation_context.set(ctx)
    _temporal_start_operation_context.reset(token)


#------------------------------------------------
from __future__ import annotations

from contextvars import ContextVar, Token
from typing import Optional, Union

# imports
import pytest  # used for our unit tests
from temporalio.nexus._operation_context import _temporal_context


# Dummy context classes for testing
class _TemporalStartOperationContext:
    def __init__(self, val: str = "start"):
        self.val = val

class _TemporalCancelOperationContext:
    def __init__(self, val: str = "cancel"):
        self.val = val

# ContextVars as used in the function
_temporal_start_operation_context: ContextVar[_TemporalStartOperationContext] = (
    ContextVar("temporal-start-operation-context")
)
_temporal_cancel_operation_context: ContextVar[_TemporalCancelOperationContext] = (
    ContextVar("temporal-cancel-operation-context")
)
from temporalio.nexus._operation_context import _temporal_context

# unit tests

# ----------- BASIC TEST CASES -----------




def test_temporal_context_none_set():
    """Test when neither context is set, should raise error."""
    with pytest.raises(RuntimeError, match="Not in Nexus operation context."):
        _temporal_context() # 792ns -> 667ns (18.7% faster)

# ----------- EDGE TEST CASES -----------


def test_temporal_context_contextvar_isolation():
    """Test that contextvars are isolated between tests (no residual state)."""



def test_temporal_context_contextvar_token_stack():
    """Test nested contextvar set/reset using tokens (contextvar stack)."""
    token1 = _temporal_start_operation_context.set(_TemporalStartOperationContext("first"))
    token2 = _temporal_start_operation_context.set(_TemporalStartOperationContext("second"))
    try:
        codeflash_output = _temporal_context(); ctx = codeflash_output
    finally:
        _temporal_start_operation_context.reset(token2)
        codeflash_output = _temporal_context(); ctx = codeflash_output
        _temporal_start_operation_context.reset(token1)

def test_temporal_context_contextvar_none_explicitly():
    """Test explicitly setting contextvar to None."""
    token = _temporal_start_operation_context.set(None)
    try:
        with pytest.raises(RuntimeError, match="Not in Nexus operation context."):
            _temporal_context()
    finally:
        _temporal_start_operation_context.reset(token)

# ----------- LARGE SCALE TEST CASES -----------




def test_temporal_context_large_scale_none_contexts():
    """Test that with no context set, always raises error, repeated."""
    for _ in range(100):
        with pytest.raises(RuntimeError, match="Not in Nexus operation context."):
            _temporal_context()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from temporalio.nexus._operation_context import _temporal_context
import pytest

def test__temporal_context():
    with pytest.raises(RuntimeError, match='Not\\ in\\ Nexus\\ operation\\ context\\.'):
        _temporal_context()

To edit these changes git checkout codeflash/optimize-_temporal_context-mgib5dhm and push.