ThreadPoolExecutor Functions Do Not Appear as Child Spans in New Relic Traces

[akash-kola-test]

Is your feature request related to a problem? Please describe.

When using ThreadPoolExecutor to execute functions in a multi-threaded environment, the traces for these executor functions do not appear as child spans of the parent span. Instead, they are either missing or appear disconnected from the parent span, making it difficult to track the execution flow properly.

Feature Description

The New Relic Python Agent should correctly propagate trace context when using ThreadPoolExecutor, ensuring that functions executed in the thread pool are captured as child spans of the initiating parent span. This would allow better visibility into multi-threaded applications and maintain trace hierarchy.

Describe Alternatives

Explicitly wrapping executor functions with the New Relic tracing context before submitting them to ThreadPoolExecutor.

Additional context

Consider the following example Flask application using ThreadPoolExecutor:

from concurrent.futures import ThreadPoolExecutor
import time

import newrelic.agent
from flask import Flask

app = Flask(__name__)

# Initialize New Relic agent
newrelic.agent.initialize("newrelic.ini")

# Global ThreadPoolExecutor
executor = ThreadPoolExecutor(max_workers=3)

@newrelic.agent.function_trace()
def function_a():
    time.sleep(3)

@newrelic.agent.function_trace()
def demo():
    futures = [executor.submit(function_a) for _ in range(3)]
    for future in futures:
        try:
            future.result()
        except Exception as e:
            print(f"Thread execution error: {e}")
            newrelic.agent.notice_error(e)

@app.get("/test")
def hello_world():
    start_time = time.time()
    demo()
    end_time = time.time()
    print("Elapsed time:", end_time - start_time)
    return "Threads Executed!"

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)

Expected Behavior:

The spans for function_a should appear as child spans of the demo function in the New Relic UI.

Current Behavior:

The spans for function_a do not appear as child spans under demo, leading to broken or missing trace relationships.

Priority

Please help us better understand this feature request by choosing a priority from the following options:
Really Want

I've searched, cloned this repo, but I haven't found this feature. Very happy to be corrected if it does exist.

[workato-integration]

https://new-relic.atlassian.net/browse/NR-378461

[akash-kola-test]

Hi Team,

I spent some time understanding the New Relic Python Agent’s code, particularly how it handles context propagation with trace wrappers. Based on that, I implemented the following approach, which successfully produces the expected traces for functions executed within ThreadPoolExecutor.

from concurrent.futures import Future, ThreadPoolExecutor
from contextvars import copy_context
from functools import partial
from typing import Any, Callable, Iterable, Iterator, ParamSpec, TypeVar

from newrelic.api.time_trace import current_trace
from newrelic.common.object_wrapper import wrap_function_wrapper
from newrelic.core.context import context_wrapper

def bind_submit(func, *args, **kwargs):
    return {"func": func, "args": args, "kwargs": kwargs}

def wrap_ContextThreadPoolExecutor_submit(wrapped, instance, args, kwargs):
    trace = current_trace()
    if not trace:
        return wrapped(*args, **kwargs)

    # Use hardened function signature bind so we have safety net catchall of args and kwargs.
    bound_args = bind_submit(*args, **kwargs)
    bound_args["func"] = context_wrapper(bound_args["func"], trace=trace, strict=True)
    return wrapped(bound_args["func"], *bound_args["args"], **bound_args["kwargs"])

P = ParamSpec("P")
T = TypeVar("T")

class ContextThreadPoolExecutor(ThreadPoolExecutor):
    """ThreadPoolExecutor that copies the current execution context to worker threads."""

    def submit(
        self, func: Callable[P, T], *args: P.args, **kwargs: P.kwargs
    ) -> Future[T]:
        """Submit a function to the executor, ensuring context is copied."""
        return super().submit(partial(copy_context().run, func, *args, **kwargs))

    def map(
        self,
        fn: Callable[..., T],
        *iterables: Iterable[Any],
        timeout: float | None = None,
        chunksize: int = 1,
    ) -> Iterator[T]:
        """Map a function to multiple iterables while maintaining execution context."""
        contexts = [copy_context() for _ in range(len(iterables[0]))]  # type: ignore[arg-type]

        def _wrapped_fn(*args: Any) -> T:
            return contexts.pop().run(fn, *args)

        return super().map(
            _wrapped_fn, *iterables, timeout=timeout, chunksize=chunksize
        )

wrap_function_wrapper(
    __name__, "ContextThreadPoolExecutor.submit", wrap_ContextThreadPoolExecutor_submit
)

Since this approach successfully produces the traces, I wanted to check if there’s anything we can take from this and implement within the New Relic Python Agent itself.

Would this approach align with the agent's internals, or is there another preferred way to handle trace propagation in ThreadPoolExecutor?

If there’s a possible direction to explore, I’d be happy to try implementing it from my side with some guidance. Looking forward to any thoughts or suggestions!

Thanks.

EDIT:
I’ve also attached an image showing the traces generated using this approach. You can see how the spans are structured within the trace.