concurrent asyncio tasks and xray are not working with Lambda

Hello, I am opening this issue on request of @NathanielRN after a discussion with him in https://github.com/aws/aws-xray-sdk-python/issues/203 .

The problem is that when using aioboto3, asyncio (with gather or wait) and xray in Lambda we either get an exception or no subsegments showing calls to S3 or SQS. Here is the last code snippet from that issue:

import asyncio
from io import BytesIO

import aioboto3
from aws_xray_sdk.core import xray_recorder, patch_all
from aws_xray_sdk.core.async_context import AsyncContext

# here I tried adding your solution (not added when solution to use AsyncContext() is used below)
# xray_recorder.configure(service='repro_xray_issue')  # uncomment this for second solution trial

patch_all()


def lambda_handler(a,b):
    asyncio.run(main())


async def main():
    # here I tried adding your solution (not added when solution to remove AsyncContext() is used above)
    xray_recorder.configure(service='repro_xray_issue', context=AsyncContext())    # comment this for second solution trial
    async with xray_recorder.in_segment_async('my_segment_name') as segment:  # comment this for second solution trial
        filelike1 = BytesIO()
        filelike2 = BytesIO()
        res1, res2 = await asyncio.gather(s3_get(filelike1), s3_get(filelike2))  # .wait() doesn't work either


@xray_recorder.capture_async('s3_get')  # tried with this and without also
async def s3_get(filelike):
    async with aioboto3.Session().client('s3') as s3:
        return await s3.download_fileobj('s3-validation-files-003', 'test.txt', filelike)

[abivolmv]

@NathanielRN this, propagation, happens automatically when I don’t use

xray_recorder.configure(context=AsyncContext())
    async with xray_recorder.in_segment_async('my_segment_name') as segment:

Something happens here that detaches the API incoming trace from Lambda custom trace. We can see that the trace is reaching the Lambda handler in the other parallel trace :

I ended up doing this, though it feels counter-intuitive, it shows one full trace:

    trace_id = xray_recorder.get_trace_entity().trace_id
    parent_id = xray_recorder.current_segment().id
    # enable xray to work with async
    # even though Lambda has its own segment we need to create one async segment
    xray_recorder.configure(context=AsyncContext())
    async with xray_recorder.in_segment_async('add_brochure', traceid=trace_id, parent_id=parent_id) as segment:
       ...

[NathanielRN]

Hey @abivolmv, thanks for your patience! Can you see release 2.10.0 now? https://github.com/aws/aws-xray-sdk-python/releases/tag/2.10.0

Let us know if it works for you 🙂