Bug Fix: Defensively copy context entities

aws_xray_sdk/core/async_context.py

@@ -1,5 +1,6 @@
 import asyncio
 import sys
+import copy
 
 from .context import Context as _Context
 
@@ -108,6 +109,13 @@ def task_factory(loop, coro):
     else:
         current_task = asyncio.Task.current_task(loop=loop)
     if current_task is not None and hasattr(current_task, 'context'):
-        setattr(task, 'context', current_task.context)
+        if current_task.context.get('entities'):
+            # Defensively copying because recorder modifies the list in place.
+            new_context = copy.copy(current_task.context)
+            new_context['entities'] = [item for item in current_task.context['entities']]
+        else:
+            # no reason to copy if there's no entities list.
+            new_context = current_task.context
+        setattr(task, 'context', new_context)
 
     return task

tests/test_async_recorder.py

@@ -3,6 +3,7 @@
 from .util import get_new_stubbed_recorder
 from aws_xray_sdk.version import VERSION
 from aws_xray_sdk.core.async_context import AsyncContext
+import asyncio
 
 
 xray_recorder = get_new_stubbed_recorder()
@@ -43,6 +44,28 @@ async def test_capture(loop):
     assert platform.python_implementation() == service.get('runtime')
     assert platform.python_version() == service.get('runtime_version')
 
+async def test_concurrent_calls(loop):
+    xray_recorder.configure(service='test', sampling=False, context=AsyncContext(loop=loop))
+    async with xray_recorder.in_segment_async('segment') as segment:
+        global counter
+        counter = 0
+        total_tasks = 10
+        event = asyncio.Event()
+        async def assert_task():
+            async with xray_recorder.in_subsegment_async('segment') as subsegment:
+                global counter
+                counter += 1
+                # Ensure that the task subsegments overlap
+                if counter < total_tasks:
+                    await event.wait()
+                else:
+                    event.set()
+                return subsegment.parent_id
+        tasks = [assert_task() for task in range(total_tasks)]
+        subsegs_parent_ids = await asyncio.gather(*tasks)
+        for subseg_parent_id in subsegs_parent_ids:
+            assert subseg_parent_id == segment.id
+
 
 async def test_async_context_managers(loop):
     xray_recorder.configure(service='test', sampling=False, context=AsyncContext(loop=loop))