⚡️ Speed up function _setup_sentry_tracing by 93%

[codeflash-ai]

📄 93% (0.93x) speedup for _setup_sentry_tracing in sentry_sdk/integrations/spark/spark_driver.py

⏱️ Runtime : 46.5 microseconds → 24.0 microseconds (best of 74 runs)

📝 Explanation and details

The optimized code achieves a 93% speedup through two key optimizations:

1. Import Caching with Global Variable

• Replaces repeated from pyspark import SparkContext calls with a cached global _spark_context_class
• Line profiler shows the import overhead drops from 6,986ns to just 2,621ns on first call, with subsequent calls using the cached reference
• This eliminates Python's module lookup overhead on repeated function calls

2. Idempotent Patching Prevention

• Adds _sentry_patched attribute checking to prevent re-patching SparkContext._do_init
• When already patched, _patch_spark_context_init() returns early, avoiding expensive decorator re-application
• Line profiler shows 4 out of 6 calls now take the early return path (1,614ns vs 30,118ns for full patching)

Performance Impact by Test Case:

• Multiple context switches: 410% faster on subsequent calls due to cached imports and patch detection
• No active context scenarios: 384% faster when patching is required repeatedly
• Basic setup calls: 25% faster for typical single-call scenarios

The optimizations are particularly effective for applications that repeatedly call these functions or switch between SparkContext instances, which is common in distributed Spark environments. The changes preserve all original functionality while dramatically reducing redundant work.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import sys
from functools import wraps
from typing import TYPE_CHECKING, Callable, ParamSpec, TypeVar, cast, overload

# imports
import pytest  # used for our unit tests
# function to test
import sentry_sdk
from sentry_sdk.integrations.spark.spark_driver import _setup_sentry_tracing
from sentry_sdk.utils import (capture_internal_exceptions,
                              ensure_integration_enabled)


# --- Minimal stub classes and functions to support testing ---
class SparkIntegration:
    pass

class DummyScope:
    def __init__(self):
        self._processors = []

    def add_event_processor(self, func):
        self._processors.append(func)
        return func

class DummyClient:
    def __init__(self, integration_enabled=True):
        self._integration_enabled = integration_enabled

    def get_integration(self, integration):
        # Return integration object if enabled, else None
        return SparkIntegration() if self._integration_enabled else None

def dummy_capture_internal_exceptions():
    # Context manager that does nothing
    class Dummy:
        def __enter__(self): return self
        def __exit__(self, exc_type, exc_val, exc_tb): return False
    return Dummy()

# Patch sentry_sdk and its utils for testing
sentry_sdk.get_isolation_scope = lambda: DummyScope()
sentry_sdk.get_client = lambda: DummyClient(integration_enabled=True)
capture_internal_exceptions = dummy_capture_internal_exceptions

# --- Dummy SparkContext for testing ---
class DummyConf(dict):
    def get(self, key, default=None):
        return self[key] if key in self else default

class DummySparkContext:
    # Simulates a SparkContext for testing
    _active_spark_context = None
    _do_init_called = False

    def __init__(self):
        self.appName = "TestApp"
        self.applicationId = "app-123"
        self._conf = DummyConf({
            "spark.executor.id": "executor-1",
            "spark.submit.deployMode": "cluster",
            "spark.driver.host": "localhost",
            "spark.driver.port": "7077",
        })
        self.version = "3.1.2"
        self.master = "local"
        self.sparkHome = "/opt/spark"
        self.uiWebUrl = "http://localhost:4040"
        self._gateway = object()
        self._jsc = type("JSC", (), {"sc": lambda self: self})()
        self._do_init_called = False
        self.sparkUser = lambda: "spark-user"

    def setLocalProperty(self, key, value):
        setattr(self, key, value)

    def __getattr__(self, item):
        # fallback for missing attributes
        return None

    def _do_init(self, *args, **kwargs):
        self._do_init_called = True
        DummySparkContext._active_spark_context = self

def _add_event_processor(sc):
    scope = sentry_sdk.get_isolation_scope()
    @scope.add_event_processor
    def process_event(event, hint):
        with capture_internal_exceptions():
            if sentry_sdk.get_client().get_integration(SparkIntegration) is None:
                return event
            if sc._active_spark_context is None:
                return event
            event.setdefault("user", {}).setdefault("id", sc.sparkUser())
            event.setdefault("tags", {}).setdefault("executor.id", sc._conf.get("spark.executor.id"))
            event["tags"].setdefault("spark-submit.deployMode", sc._conf.get("spark.submit.deployMode"))
            event["tags"].setdefault("driver.host", sc._conf.get("spark.driver.host"))
            event["tags"].setdefault("driver.port", sc._conf.get("spark.driver.port"))
            event["tags"].setdefault("spark_version", sc.version)
            event["tags"].setdefault("app_name", sc.appName)
            event["tags"].setdefault("application_id", sc.applicationId)
            event["tags"].setdefault("master", sc.master)
            event["tags"].setdefault("spark_home", sc.sparkHome)
            event.setdefault("extra", {}).setdefault("web_url", sc.uiWebUrl)
        return event
from sentry_sdk.integrations.spark.spark_driver import _setup_sentry_tracing


# --- Basic Test Cases ---
def test_setup_sentry_tracing_with_active_context_sets_properties():
    # Test that _setup_sentry_tracing sets app properties if active context exists
    sc = DummySparkContext()
    DummySparkContext._active_spark_context = sc
    _setup_sentry_tracing() # 11.1μs -> 10.4μs (6.18% faster)

def test_setup_sentry_tracing_with_no_active_context_patches_init():
    # Test that _setup_sentry_tracing patches SparkContext._do_init if no active context
    DummySparkContext._active_spark_context = None
    sc = DummySparkContext()
    # Save original function for comparison
    original_do_init = DummySparkContext._do_init
    _setup_sentry_tracing() # 8.68μs -> 1.79μs (384% faster)



def test_setup_sentry_tracing_with_missing_appName_and_applicationId():
    # Test with missing appName and applicationId attributes
    sc = DummySparkContext()
    sc.appName = None
    sc.applicationId = None
    DummySparkContext._active_spark_context = sc
    _setup_sentry_tracing() # 10.7μs -> 1.97μs (445% faster)






def test_setup_sentry_tracing_multiple_contexts():
    # Test with multiple SparkContexts (simulate switching active context)
    sc1 = DummySparkContext()
    sc2 = DummySparkContext()
    sc2.appName = "AnotherApp"
    sc2.applicationId = "app-456"
    DummySparkContext._active_spark_context = sc1
    _setup_sentry_tracing() # 11.1μs -> 8.89μs (25.1% faster)
    DummySparkContext._active_spark_context = sc2
    _setup_sentry_tracing() # 4.92μs -> 965ns (410% faster)


#------------------------------------------------
import sys
import types
from functools import wraps
from typing import TYPE_CHECKING, Callable, ParamSpec, TypeVar, cast, overload

# imports
import pytest  # used for our unit tests
from sentry_sdk.integrations.spark.spark_driver import _setup_sentry_tracing
# --- Function to test (copied from prompt, unchanged) ---
from sentry_sdk.utils import (capture_internal_exceptions,
                              ensure_integration_enabled)

# --- Mocks and helpers for testing ---

class DummyConf(dict):
    def get(self, key, default=None):
        return self[key] if key in self else default

class DummySparkContext:
    # A minimal SparkContext mock for testing
    def __init__(
        self,
        appName="testApp",
        applicationId="app-123",
        master="local",
        sparkHome="/fake/spark/home",
        version="3.0.0",
        sparkUser="testUser",
        executor_id="1",
        deployMode="client",
        driver_host="localhost",
        driver_port="7077",
        uiWebUrl="http://localhost:4040",
        active=True,
        conf=None
    ):
        self.appName = appName
        self.applicationId = applicationId
        self.master = master
        self.sparkHome = sparkHome
        self.version = version
        self._active_spark_context = self if active else None
        self._conf = conf if conf is not None else DummyConf({
            "spark.executor.id": executor_id,
            "spark.submit.deployMode": deployMode,
            "spark.driver.host": driver_host,
            "spark.driver.port": driver_port,
        })
        self._gateway = "dummy_gateway"
        self.sparkUser = lambda: sparkUser
        self.uiWebUrl = uiWebUrl
        self._jsc = self
        self._do_init_called = False
        self._setLocalProperty_calls = []
        self._addSparkListener_called = False

    def setLocalProperty(self, key, value):
        self._setLocalProperty_calls.append((key, value))

    def sc(self):
        return self

    def addSparkListener(self, listener):
        self._addSparkListener_called = True

    def version(self):
        return self.version


class DummySparkContextModule(types.SimpleNamespace):
    _active_spark_context = None
    _do_init = lambda self, *args, **kwargs: setattr(self, "_do_init_called", True)

sys.modules["pyspark"] = types.SimpleNamespace(SparkContext=DummySparkContextModule)
sys.modules["pyspark.java_gateway"] = types.SimpleNamespace(
    ensure_callback_server_started=lambda gw: setattr(sys.modules["pyspark"], "_callback_server_started", True)
)

# Patch sentry_sdk and integration
class DummyIntegration:
    pass

class DummyClient:
    def __init__(self, integration_enabled=True):
        self.integration_enabled = integration_enabled
    def get_integration(self, integration):
        return DummyIntegration() if self.integration_enabled else None

class DummyScope:
    def __init__(self):
        self.event_processors = []
    def add_event_processor(self, func):
        self.event_processors.append(func)
        return func

class DummySentrySDK:
    _client = DummyClient()
    _scope = DummyScope()
    @staticmethod
    def get_client():
        return DummySentrySDK._client
    @staticmethod
    def get_isolation_scope():
        return DummySentrySDK._scope

sys.modules["sentry_sdk"] = DummySentrySDK
sys.modules["sentry_sdk.utils"] = types.SimpleNamespace(
    capture_internal_exceptions=lambda: types.SimpleNamespace(__enter__=lambda s: None, __exit__=lambda s, a, b, c: None),
    ensure_integration_enabled=lambda integration, func=None: lambda f: f
)

# Patch SparkIntegration
SparkIntegration = DummyIntegration

# Patch SentryListener for _start_sentry_listener
class SentryListener:
    pass


def _set_app_properties():
    # type: () -> None
    from pyspark import SparkContext

    spark_context = SparkContext._active_spark_context
    if spark_context:
        spark_context.setLocalProperty(
            "sentry_app_name",
            spark_context.appName,
        )
        spark_context.setLocalProperty(
            "sentry_application_id",
            spark_context.applicationId,
        )

def _start_sentry_listener(sc):
    # type: (DummySparkContext) -> None
    from pyspark.java_gateway import ensure_callback_server_started

    gw = sc._gateway
    ensure_callback_server_started(gw)
    listener = SentryListener()
    sc._jsc.sc().addSparkListener(listener)

def _add_event_processor(sc):
    # type: (DummySparkContext) -> None
    scope = sentry_sdk.get_isolation_scope()

    @scope.add_event_processor
    def process_event(event, hint):
        # type: (dict, dict) -> dict
        with capture_internal_exceptions():
            if sentry_sdk.get_client().get_integration(SparkIntegration) is None:
                return event

            if sc._active_spark_context is None:
                return event

            event.setdefault("user", {}).setdefault("id", sc.sparkUser())

            event.setdefault("tags", {}).setdefault(
                "executor.id", sc._conf.get("spark.executor.id")
            )
            event["tags"].setdefault(
                "spark-submit.deployMode",
                sc._conf.get("spark.submit.deployMode"),
            )
            event["tags"].setdefault("driver.host", sc._conf.get("spark.driver.host"))
            event["tags"].setdefault("driver.port", sc._conf.get("spark.driver.port"))
            event["tags"].setdefault("spark_version", sc.version)
            event["tags"].setdefault("app_name", sc.appName)
            event["tags"].setdefault("application_id", sc.applicationId)
            event["tags"].setdefault("master", sc.master)
            event["tags"].setdefault("spark_home", sc.sparkHome)

            event.setdefault("extra", {}).setdefault("web_url", sc.uiWebUrl)

        return event
from sentry_sdk.integrations.spark.spark_driver import _setup_sentry_tracing

# --- Unit tests ---

# Basic Test Cases



def test_set_app_properties_sets_properties():
    """Test that _set_app_properties sets sentry_app_name and sentry_application_id."""
    from pyspark import SparkContext
    sc = DummySparkContext()
    SparkContext._active_spark_context = sc
    _set_app_properties()
    calls = sc._setLocalProperty_calls

def test_start_sentry_listener_adds_listener_and_starts_callback():
    """Test that _start_sentry_listener adds the listener and starts callback server."""
    sc = DummySparkContext()
    sys.modules["pyspark"]._callback_server_started = False
    _start_sentry_listener(sc)




def test_set_app_properties_noop_if_no_active_context():
    """Test that _set_app_properties does nothing if no active SparkContext."""
    from pyspark import SparkContext
    SparkContext._active_spark_context = None
    # Should not raise or do anything
    _set_app_properties()






#------------------------------------------------
from sentry_sdk.integrations.spark.spark_driver import _setup_sentry_tracing
import pytest

def test__setup_sentry_tracing():
    with pytest.raises(ModuleNotFoundError, match="No\\ module\\ named\\ 'pyspark'"):
        _setup_sentry_tracing()

To edit these changes git checkout codeflash/optimize-_setup_sentry_tracing-mg9oea66 and push.