Tracing: nested subsegments using `trace_method` doesn't work across threads, even when setting the trace context

[tibbe]

Expected Behaviour

I tried to adapt the recipe for correct tracing in presence of a ThreadPoolExecutor as documented in https://github.com/aws/aws-xray-sdk-python#trace-threadpoolexecutor. I'd expect a function marked with trace_method running on the thread pool to be nested under the function that submited it to the pool, assuming the trace context was set as in the linked docs.

See the code snippet further below for a complete Lambda handler that demonstrates the problem.

Current Behaviour

I'd expect the trace for h() below to nest under f() but it instead shows up as a sibling of the handler itself (i.e. it's technically "above" f()):

Code snippet

from concurrent.futures import ThreadPoolExecutor

from aws_lambda_powertools import Tracer
from aws_xray_sdk.core import xray_recorder

tracer = Tracer()


@tracer.capture_lambda_handler
def handler(event, context):
    del event, context  # Unused.

    with ThreadPoolExecutor() as pool:
        @tracer.capture_method
        def f():
            @tracer.capture_method
            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)
        f()

Possible Solution

No response

Steps to Reproduce

Create a Lambda containing the above snippet and send it any event.

AWS Lambda Powertools for Python version

latest

AWS Lambda function runtime

3.9

Packaging format used

PyPi

Debugging logs

START RequestId: 210251f5-ee0d-432a-a085-3a8f7c9d529d Version: $LATEST
[WARNING]	2023-03-24T11:17:10.071Z	210251f5-ee0d-432a-a085-3a8f7c9d529d	Removing Segment/Subsugment Name invalid characters from ## benetics.repro.handler.<locals>.f.
[WARNING]	2023-03-24T11:17:10.072Z	210251f5-ee0d-432a-a085-3a8f7c9d529d	Removing Segment/Subsugment Name invalid characters from ## benetics.repro.handler.<locals>.f.<locals>.h.
END RequestId: 210251f5-ee0d-432a-a085-3a8f7c9d529d
REPORT RequestId: 210251f5-ee0d-432a-a085-3a8f7c9d529d	Duration: 6.67 ms	Billed Duration: 2049 ms	Memory Size: 1024 MB	Max Memory Used: 58 MB	Init Duration: 2041.79 ms	
XRAY TraceId: 1-641d86b3-323a174c2c681c082b002ce6	SegmentId: 7058f8f456770f5c	Sampled: true	

[tibbe]

Here's the raw trace:

{
    "Id": "1-641d86b3-323a174c2c681c082b002ce6",
    "Duration": 2.342,
    "LimitExceeded": false,
    "Segments": [
        {
            "Id": "7383857efa7c1d04",
            "Document": {
                "id": "7383857efa7c1d04",
                "name": "backend-tibbe-Repro",
                "start_time": 1679656627.739,
                "trace_id": "1-641d86b3-323a174c2c681c082b002ce6",
                "end_time": 1679656630.081,
                "http": {
                    "response": {
                        "status": 200
                    }
                },
                "aws": {
                    "request_id": "210251f5-ee0d-432a-a085-3a8f7c9d529d"
                },
                "origin": "AWS::Lambda",
                "resource_arn": "arn:aws:lambda:eu-central-1:xxxxxxxxxxxx:function:backend-tibbe-Repro"
            }
        },
        {
            "Id": "7058f8f456770f5c",
            "Document": {
                "id": "7058f8f456770f5c",
                "name": "backend-tibbe-Repro",
                "start_time": 1679656630.0707042,
                "trace_id": "1-641d86b3-323a174c2c681c082b002ce6",
                "end_time": 1679656630.0771244,
                "parent_id": "7383857efa7c1d04",
                "aws": {
                    "account_id": "xxxxxxxxxxxx",
                    "function_arn": "arn:aws:lambda:eu-central-1:xxxxxxxxxxxx:function:backend-tibbe-Repro",
                    "resource_names": [
                        "backend-tibbe-Repro"
                    ]
                },
                "origin": "AWS::Lambda::Function",
                "subsegments": [
                    {
                        "id": "2c19ff749967fb32",
                        "name": "Overhead",
                        "start_time": 1679656630.0767124,
                        "end_time": 1679656630.0769603,
                        "aws": {
                            "function_arn": "arn:aws:lambda:eu-central-1:xxxxxxxxxxxx:function:backend-tibbe-Repro"
                        }
                    },
                    {
                        "id": "f30867ce48ab7aaf",
                        "name": "Invocation",
                        "start_time": 1679656630.0709977,
                        "end_time": 1679656630.0766628,
                        "aws": {
                            "function_arn": "arn:aws:lambda:eu-central-1:xxxxxxxxxxxx:function:backend-tibbe-Repro"
                        },
                        "subsegments": [
                            {
                                "id": "8df51c43ef63b540",
                                "name": "## benetics.repro.handler.locals.f.locals.h",
                                "start_time": 1679656630.0723202,
                                "end_time": 1679656630.0727339,
                                "in_progress": false,
                                "namespace": "local"
                            },
                            {
                                "id": "1fef83373b9e62d5",
                                "name": "## handler",
                                "start_time": 1679656630.0714564,
                                "end_time": 1679656630.0730584,
                                "in_progress": false,
                                "annotations": {
                                    "ColdStart": true,
                                    "Service": "Repro"
                                },
                                "namespace": "local",
                                "subsegments": [
                                    {
                                        "id": "ce6b35994cefd0da",
                                        "name": "## benetics.repro.handler.locals.f",
                                        "start_time": 1679656630.0717008,
                                        "end_time": 1679656630.0724592,
                                        "in_progress": false,
                                        "namespace": "local"
                                    }
                                ]
                            }
                        ]
                    },
                    {
                        "id": "448b42c019df6b09",
                        "name": "Initialization",
                        "start_time": 1679656628.0294325,
                        "end_time": 1679656630.069455,
                        "aws": {
                            "function_arn": "arn:aws:lambda:eu-central-1:xxxxxxxxxxxx:function:backend-tibbe-Repro"
                        }
                    }
                ]
            }
        }
    ]
}

[heitorlessa]

thank you so much @tibbe! much appreciated. We'll investigate next week, it seems to be the same in X-Ray SDK so we need to spend time creating a repro on GitHub to do this properly. Leaving triage label so it shows up on our daily updates.

[rubenfonseca]

Looking at it now

[rubenfonseca]

I've looked at this extensively today, and I agree with Heitor that using the XRay SDK directly leads us to the same problem. Example code:

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)
        f()

This is producing the wrong traces:

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

So my next step is to prepare a minimal repository with the sample above and submit a bug report upstream. I'll update this issue with the link.

[rubenfonseca]

I've opened a new issue upstream here: aws/aws-xray-sdk-python#388

[heitorlessa]

Putting this on hold until Ruben is back from PTO mid-next week (we also need to hear from X-Ray).