Bug: nested subsegments don't work across threads

[rubenfonseca]

We're tracking a bug where nested subsegments don't work across threads, even when using the official suggestion in the documentation.

We've managed to reproduce the issue with some minimal code:

Code snippet

from concurrent.futures import ThreadPoolExecutor

from aws_xray_sdk.core import xray_recorder

@xray_recorder.capture("## lambda handler")
def lambda_handler(event, context):
    with ThreadPoolExecutor() as pool:
        @xray_recorder.capture("## f")
        def f():
            @xray_recorder.capture("## h")
            def h():
                pass

            def g(trace_entity):
                xray_recorder.set_trace_entity(trace_entity)
                h()
                xray_recorder.clear_trace_entities()

            curr = xray_recorder.get_trace_entity()
            pool.submit(g, curr)

This produces the following XRay traces:

XRay trace

XRay Event at (2023-04-05T16:32:52.012000) with id (1-642d8694-731f0a1a6d204264677086db) and duration (0.388s)
 - 0.388s - lambda-function-url-HelloWorldFunction-sRQB00cHvXOP [HTTP: 200]
 - 0.021s - lambda-function-url-HelloWorldFunction-sRQB00cHvXOP
   - 0.223s - Initialization
   - 0.020s - Invocation
     - 0.019s - ## lambda handler
       - 0.018s - ## f
     - 0.000s - ## h
   - 0.000s - Overhead

The expected output would be for the "h" trace to be nested under "f":

Expected output

XRay Event at (2023-04-05T16:32:52.012000) with id (1-642d8694-731f0a1a6d204264677086db) and duration (0.388s)
 - 0.388s - lambda-function-url-HelloWorldFunction-sRQB00cHvXOP [HTTP: 200]
 - 0.021s - lambda-function-url-HelloWorldFunction-sRQB00cHvXOP
   - 0.223s - Initialization
   - 0.020s - Invocation
     - 0.019s - ## lambda handler
       - 0.018s - ## f
         - 0.000s - ## h
   - 0.000s - Overhead

[srprash]

Hi @rubenfonseca
Thanks for providing the code snippet. I was able to run it and confirm that the ## h subsegment indeed becomes the child of the Invocation subsegment. I debugged and found that the current entity when the h() method is being invoked is actually the ## f subsegment, however it doesn't become the parent when the method executes on the thread pool.

I noticed similar behavior in the code for official suggestion in the documentation.

I am still investigating the root cause but thought I would provide an update of my findings in the meantime. Apologies for the delay.

[srprash]

I suspect this _refresh_context method is the source of problem.
It tries to get a segment object from the current thread local(code). If this segment is not present, then the context is reinitialized (code) making the current entity to be the facade segment, which is essentially the Invocation subsegment in the final trace.

The question is why isn't the segment present on thread local attribute when the ## h subsegment is created in a thread pool execution?

[rubenfonseca]

Thank you for the update @srprash. When I dug into the problem, I come to basically the same conclusion. However, my python skills are not strong enough to understand why the thread local attribute is not present.

[srprash]

I did some more investigation and found that the thread local state is not being propagated between the main thread and the thread on the thread pool even though they share the same thread local instance. That explains why the segment stored in main thread is not present when ##h subsegment is created.
In a nutshell, multi-threading is not working as I expected. Found this SO post that is related to the behavior: https://stackoverflow.com/questions/72374802/shared-memory-in-threads-confusion