⚡️ Speed up function in_operation by 42%

[codeflash-ai]

📄 42% (0.42x) speedup for in_operation in temporalio/nexus/_operation_context.py

⏱️ Runtime : 334 microseconds → 236 microseconds (best of 280 runs)

📝 Explanation and details

The optimization inlines the _try_temporal_context() function call directly into in_operation() to eliminate function call overhead.

Key changes:

• The in_operation() function now directly executes the context variable lookups and validation logic instead of calling _try_temporal_context()
• Added is not None check to the return statement to maintain the boolean return type
• The _try_temporal_context() function is preserved unchanged for other potential callers

Why this speeds up execution:

• Eliminates function call overhead: Python function calls involve stack frame creation, parameter passing, and return value handling - all eliminated by inlining
• Reduces total execution path: The profiler shows in_operation() spent 100% of its time calling _try_temporal_context(), so removing this indirection directly improves performance
• Better CPU cache locality: All logic executes in a single function scope without jumping between stack frames

Performance characteristics based on tests:

• Best gains on simple cases: 33-50% speedup when contexts are unset or have falsy values, where the overhead of function calls dominates the simple logic
• Consistent improvements across scale: 35-45% speedup even with 100-1000 repeated calls, showing the optimization scales well
• Particularly effective for high-frequency usage: Since in_operation() appears to be called frequently in performance-critical paths, the cumulative benefit of removing function call overhead is substantial

✅ Correctness verification report:

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

🌀 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 in_operation


# Dummy context classes for testing
class _TemporalStartOperationContext:
    pass

class _TemporalCancelOperationContext:
    pass

_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 in_operation

# --- Basic Test Cases ---

def test_in_operation_none_set():
    """Test when no context is set: should return False."""
    codeflash_output = in_operation() # 333ns -> 250ns (33.2% faster)

def test_in_operation_start_set():
    """Test when only start context is set: should return True."""
    _temporal_start_operation_context.set(_TemporalStartOperationContext())
    codeflash_output = in_operation() # 292ns -> 250ns (16.8% faster)

def test_in_operation_cancel_set():
    """Test when only cancel context is set: should return True."""
    _temporal_cancel_operation_context.set(_TemporalCancelOperationContext())
    codeflash_output = in_operation() # 291ns -> 250ns (16.4% faster)



def test_in_operation_context_set_to_none_explicitly():
    """Test when context vars are explicitly set to None."""
    _temporal_start_operation_context.set(None)
    _temporal_cancel_operation_context.set(None)
    codeflash_output = in_operation() # 333ns -> 291ns (14.4% faster)

def test_in_operation_context_set_to_falsey_object():
    """Test when context vars are set to a falsey value (should not happen, but test anyway)."""
    # Simulate a falsey value (not an actual context object)
    _temporal_start_operation_context.set(0)
    _temporal_cancel_operation_context.set(None)
    codeflash_output = in_operation() # 292ns -> 250ns (16.8% faster)
    _temporal_start_operation_context.set(None)
    _temporal_cancel_operation_context.set("")
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_context_set_to_non_context_object():
    """Test when context vars are set to objects of wrong type."""
    _temporal_start_operation_context.set("not a context")
    codeflash_output = in_operation() # 291ns -> 250ns (16.4% faster)
    _temporal_start_operation_context.set(None)
    _temporal_cancel_operation_context.set(12345)
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_context_var_reset():
    """Test that resetting context var disables operation."""
    token = _temporal_start_operation_context.set(_TemporalStartOperationContext())
    codeflash_output = in_operation() # 292ns -> 250ns (16.8% faster)
    _temporal_start_operation_context.reset(token)
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_nested_contexts():
    """Test nested context var usage with contextvars."""
    token = _temporal_start_operation_context.set(_TemporalStartOperationContext())
    codeflash_output = in_operation()
    # Nested context: set cancel context (should raise)
    token2 = _temporal_cancel_operation_context.set(_TemporalCancelOperationContext())
    with pytest.raises(RuntimeError):
        in_operation()
    # Reset nested cancel context
    _temporal_cancel_operation_context.reset(token2)
    codeflash_output = in_operation()
    # Reset start context
    _temporal_start_operation_context.reset(token)
    codeflash_output = in_operation()

def test_in_operation_multiple_sets_same_type():
    """Test multiple sets on the same context var."""
    token1 = _temporal_start_operation_context.set(_TemporalStartOperationContext())
    token2 = _temporal_start_operation_context.set(_TemporalStartOperationContext())
    codeflash_output = in_operation() # 333ns -> 250ns (33.2% faster)
    _temporal_start_operation_context.reset(token2)
    codeflash_output = in_operation() # 125ns -> 125ns (0.000% faster)
    _temporal_start_operation_context.reset(token1)
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)

# --- Large Scale Test Cases ---

def test_in_operation_many_start_contexts():
    """Test with many sequential start contexts (should always be True until reset)."""
    tokens = []
    for _ in range(100):  # Large but <1000 for performance
        tokens.append(_temporal_start_operation_context.set(_TemporalStartOperationContext()))
        codeflash_output = in_operation() # 11.8μs -> 8.34μs (41.3% faster)
    # Reset all tokens
    for token in reversed(tokens):
        _temporal_start_operation_context.reset(token)
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_many_cancel_contexts():
    """Test with many sequential cancel contexts (should always be True until reset)."""
    tokens = []
    for _ in range(100):  # Large but <1000 for performance
        tokens.append(_temporal_cancel_operation_context.set(_TemporalCancelOperationContext()))
        codeflash_output = in_operation() # 11.7μs -> 8.39μs (39.6% faster)
    # Reset all tokens
    for token in reversed(tokens):
        _temporal_cancel_operation_context.reset(token)
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_interleaved_contexts():
    """Test interleaving start/cancel contexts, ensuring error is raised if both are set."""
    start_tokens = []
    cancel_tokens = []
    # Set 50 start contexts
    for _ in range(50):
        start_tokens.append(_temporal_start_operation_context.set(_TemporalStartOperationContext()))
        codeflash_output = in_operation()
    # Set 50 cancel contexts, should raise error on first
    for i in range(50):
        with pytest.raises(RuntimeError):
            cancel_tokens.append(_temporal_cancel_operation_context.set(_TemporalCancelOperationContext()))
            in_operation()
    # Reset all cancel contexts
    for token in reversed(cancel_tokens):
        _temporal_cancel_operation_context.reset(token)
    # Should still be in start context
    codeflash_output = in_operation()
    # Reset all start contexts
    for token in reversed(start_tokens):
        _temporal_start_operation_context.reset(token)
    codeflash_output = in_operation()

def test_in_operation_stress_context_switching():
    """Stress test rapid switching between start and cancel contexts."""
    for i in range(100):  # Large but <1000 for performance
        token_start = _temporal_start_operation_context.set(_TemporalStartOperationContext())
        codeflash_output = in_operation() # 12.0μs -> 8.30μs (45.0% faster)
        _temporal_start_operation_context.reset(token_start)
        codeflash_output = in_operation()
        token_cancel = _temporal_cancel_operation_context.set(_TemporalCancelOperationContext()) # 11.7μs -> 8.30μs (41.5% faster)
        codeflash_output = in_operation()
        _temporal_cancel_operation_context.reset(token_cancel)
        codeflash_output = in_operation() # 12.0μs -> 8.14μs (47.0% faster)

def test_in_operation_context_var_isolation():
    """Test that contextvar values are isolated between tests (pytest fixture)."""
    codeflash_output = in_operation() # 292ns -> 250ns (16.8% faster)
    _temporal_start_operation_context.set(_TemporalStartOperationContext())
    codeflash_output = in_operation() # 125ns -> 83ns (50.6% faster)
    # The fixture should reset this before the next test
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
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 in_operation


# Dummy context classes for testing
class _TemporalStartOperationContext:
    pass

class _TemporalCancelOperationContext:
    pass

_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 in_operation

# unit tests

# Helper functions for test setup/teardown
def set_start_context(ctx):
    return _temporal_start_operation_context.set(ctx)

def set_cancel_context(ctx):
    return _temporal_cancel_operation_context.set(ctx)

def reset_context(var, token):
    var.reset(token)

# ----------- BASIC TEST CASES -----------
def test_in_operation_default_false():
    """Test that in_operation returns False when no context is set."""
    codeflash_output = not in_operation() # 292ns -> 250ns (16.8% faster)

def test_in_operation_start_context_true():
    """Test that in_operation returns True when start context is set."""
    ctx = _TemporalStartOperationContext()
    token = set_start_context(ctx)
    try:
        codeflash_output = in_operation()
    finally:
        reset_context(_temporal_start_operation_context, token)

def test_in_operation_cancel_context_true():
    """Test that in_operation returns True when cancel context is set."""
    ctx = _TemporalCancelOperationContext()
    token = set_cancel_context(ctx)
    try:
        codeflash_output = in_operation()
    finally:
        reset_context(_temporal_cancel_operation_context, token)

def test_in_operation_start_then_cancel_context():
    """Test that in_operation returns True when only cancel context is set after resetting start."""
    start_ctx = _TemporalStartOperationContext()
    cancel_ctx = _TemporalCancelOperationContext()
    start_token = set_start_context(start_ctx)
    cancel_token = set_cancel_context(cancel_ctx)
    try:
        # Reset start context so only cancel context remains
        reset_context(_temporal_start_operation_context, start_token)
        codeflash_output = in_operation()
    finally:
        reset_context(_temporal_cancel_operation_context, cancel_token)

def test_in_operation_cancel_then_start_context():
    """Test that in_operation returns True when only start context is set after resetting cancel."""
    start_ctx = _TemporalStartOperationContext()
    cancel_ctx = _TemporalCancelOperationContext()
    cancel_token = set_cancel_context(cancel_ctx)
    start_token = set_start_context(start_ctx)
    try:
        # Reset cancel context so only start context remains
        reset_context(_temporal_cancel_operation_context, cancel_token)
        codeflash_output = in_operation()
    finally:
        reset_context(_temporal_start_operation_context, start_token)

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

def test_in_operation_context_reset_to_none():
    """Test that in_operation returns False after contexts are reset to None."""
    start_ctx = _TemporalStartOperationContext()
    cancel_ctx = _TemporalCancelOperationContext()
    start_token = set_start_context(start_ctx)
    cancel_token = set_cancel_context(cancel_ctx)
    # Reset both contexts to None
    reset_context(_temporal_start_operation_context, start_token)
    reset_context(_temporal_cancel_operation_context, cancel_token)
    codeflash_output = not in_operation() # 334ns -> 291ns (14.8% faster)

def test_in_operation_context_set_to_falsey():
    """Test that in_operation returns False if context is set to a falsey value."""
    # ContextVars can be set to None, which is falsey
    start_token = set_start_context(None)
    cancel_token = set_cancel_context(None)
    try:
        codeflash_output = not in_operation()
    finally:
        reset_context(_temporal_start_operation_context, start_token)
        reset_context(_temporal_cancel_operation_context, cancel_token)

def test_in_operation_context_set_to_unexpected_type():
    """Test that in_operation returns True if context is set to a non-None, non-expected type."""
    # If context is set to a non-None value, even if not the expected type, it should return True
    start_token = set_start_context("unexpected")
    try:
        codeflash_output = in_operation()
    finally:
        reset_context(_temporal_start_operation_context, start_token)

def test_in_operation_context_set_to_zero():
    """Test that in_operation returns False if context is set to zero (falsey value)."""
    start_token = set_start_context(0)
    try:
        codeflash_output = not in_operation()
    finally:
        reset_context(_temporal_start_operation_context, start_token)

# ----------- LARGE SCALE TEST CASES -----------
def test_in_operation_many_context_changes():
    """Test that in_operation behaves correctly under many rapid context changes."""
    tokens = []
    # Rapidly set and reset start context
    for i in range(100):
        ctx = _TemporalStartOperationContext()
        token = set_start_context(ctx)
        tokens.append(token)
        codeflash_output = in_operation() # 11.5μs -> 8.35μs (38.3% faster)
    # Reset all contexts
    for token in reversed(tokens):
        reset_context(_temporal_start_operation_context, token)
    codeflash_output = not in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_many_cancel_context_changes():
    """Test that in_operation behaves correctly under many rapid cancel context changes."""
    tokens = []
    for i in range(100):
        ctx = _TemporalCancelOperationContext()
        token = set_cancel_context(ctx)
        tokens.append(token)
        codeflash_output = in_operation() # 11.3μs -> 8.35μs (35.8% faster)
    for token in reversed(tokens):
        reset_context(_temporal_cancel_operation_context, token)
    codeflash_output = not in_operation() # 125ns -> 83ns (50.6% faster)

def test_in_operation_context_stack_behavior():
    """Test that context stack behaves correctly (nested context changes)."""
    start_ctx1 = _TemporalStartOperationContext()
    start_ctx2 = _TemporalStartOperationContext()
    token1 = set_start_context(start_ctx1)
    token2 = set_start_context(start_ctx2)
    try:
        codeflash_output = in_operation()
        # Reset the top context, should still be in operation due to the lower one
        reset_context(_temporal_start_operation_context, token2)
        codeflash_output = in_operation()
        # Reset the last context, should not be in operation anymore
        reset_context(_temporal_start_operation_context, token1)
        codeflash_output = not in_operation()
    finally:
        # Ensure no context remains
        try:
            reset_context(_temporal_start_operation_context, token2)
        except ValueError:
            pass
        try:
            reset_context(_temporal_start_operation_context, token1)
        except ValueError:
            pass

def test_in_operation_large_scale_contexts():
    """Test scalability: set and reset up to 1000 contexts."""
    tokens = []
    for i in range(1000):
        ctx = _TemporalStartOperationContext()
        token = set_start_context(ctx)
        tokens.append(token)
        codeflash_output = in_operation() # 116μs -> 81.7μs (42.1% faster)
    for token in reversed(tokens):
        reset_context(_temporal_start_operation_context, token)
    codeflash_output = not in_operation() # 125ns -> 84ns (48.8% faster)

def test_in_operation_alternating_contexts():
    """Test alternating between start and cancel contexts many times."""
    start_tokens = []
    cancel_tokens = []
    for i in range(500):
        start_token = set_start_context(_TemporalStartOperationContext())
        start_tokens.append(start_token)
        codeflash_output = in_operation() # 58.9μs -> 41.2μs (42.9% faster)
        reset_context(_temporal_start_operation_context, start_token)
        cancel_token = set_cancel_context(_TemporalCancelOperationContext())
        cancel_tokens.append(cancel_token)
        codeflash_output = in_operation()
        reset_context(_temporal_cancel_operation_context, cancel_token) # 58.2μs -> 41.0μs (42.1% faster)
    codeflash_output = not in_operation() # 125ns -> 83ns (50.6% faster)
# 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 in_operation

def test_in_operation():
    in_operation()

To edit these changes git checkout codeflash/optimize-in_operation-mgib0ny1 and push.