Bug Fix: Defensively copy context entities

Before this change, concurrent async tasks would all share the same
instance of the entities list. This change makes it so they each get
their own copy of the list. This matters because the recorder modifies
the list in place, which makes it so concurrent subtasks have the
wrong parent subsegment.

Author: tyler-dodge

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)]
+        results = await asyncio.gather(*tasks)
+        for result in results:
+            assert result == segment.id
+
 
 async def test_async_context_managers(loop):
     xray_recorder.configure(service='test', sampling=False, context=AsyncContext(loop=loop))