Repository: celery/celery
File: t/integration/test_mem_leak_in_exception_handling.py (Python)

Code before:
                raise ValueError("Deep nested exception at level 5")

            def deep_level_4():
                local_data = {"level": 4, "nested": {"data": list(range(50))}}  # noqa: F841
                deep_level_5()

            def deep_level_3():
                local_data = [1, 2, 3, {"nested": True}]  # noqa: F841
                deep_level_4()

            def deep_level_2():
                deep_level_3()

            def deep_level_1():
                deep_level_2()

            deep_level_1()

        self.task_success = task_success
        self.task_unhandled_exception = task_unhandled_exception
        self.task_retry_then_fail = task_retry_then_fail
        self.task_nested_exception_stack = task_nested_exception_stack


def get_memory_usage():
    """Get current memory usage in bytes.

    Returns RSS (total process memory) if psutil is available,
    otherwise returns Python heap allocations via tracemalloc.
    Note: These measurements are not directly comparable.
    """
    try:
        import psutil
        process = psutil.Process(os.getpid())
        return process.memory_info().rss
    except ImportError:
        # Fallback to tracemalloc if psutil not available
        current, peak = tracemalloc.get_traced_memory()
        return current


def test_mem_leak_unhandled_exceptions():
    """Test that reproduces the memory leak when tasks raise unhandled exceptions."""

    # Setup
    test_instance = MemoryLeakUnhandledExceptionsTest()

    # Enable memory tracing
    tracemalloc.start()

    # Warm up - run some successful tasks first

Review comment (suggestion):
```suggestion
    """
    Get current memory usage in bytes.
```

Code after:
            """Task that raises exception through deeply nested function calls."""
            def deep_level_5():
                local_data = {"level": 5, "data": list(range(100))}  # noqa: F841
                raise ValueError("Deep nested exception at level 5")

            def deep_level_4():
                local_data = {"level": 4, "nested": {"data": list(range(50))}}  # noqa: F841
                deep_level_5()

            def deep_level_3():
                local_data = [1, 2, 3, {"nested": True}]  # noqa: F841
                deep_level_4()

            def deep_level_2():
                deep_level_3()

            def deep_level_1():
                deep_level_2()

            deep_level_1()

        self.task_success = task_success
        self.task_unhandled_exception = task_unhandled_exception
        self.task_retry_then_fail = task_retry_then_fail
        self.task_nested_exception_stack = task_nested_exception_stack


def get_memory_usage():
    """
    Get current memory usage in bytes.

    Returns RSS (total process memory) if psutil is available,
    otherwise returns Python heap allocations via tracemalloc.
    Note: These measurements are not directly comparable.
    """
    try:
        import psutil
        process = psutil.Process(os.getpid())
        return process.memory_info().rss
    except ImportError:
        # Fallback to tracemalloc if psutil not available
        current, peak = tracemalloc.get_traced_memory()
        return current


def test_mem_leak_unhandled_exceptions():
    """Test that reproduces the memory leak when tasks raise unhandled exceptions."""

    # Setup
    test_instance = MemoryLeakUnhandledExceptionsTest()