From 3f7ae5381c8d50cdab560235fd3250aa467390d6 Mon Sep 17 00:00:00 2001 From: "T. Kowalski" Date: Tue, 9 Dec 2025 15:42:41 +0100 Subject: [PATCH] fix(profiling): do not remove links for just-created Tasks (#15552) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Echion PR: https://github.com/P403n1x87/echion/pull/210 This PR fixes a race condition happening around linking Parent and Children Tasks when using utilities like `asyncio.gather` or `asyncio.as_completed`. The race condition would manifest in the following way 1. User calls e.g. `asyncio.as_completed` from `parent_task` 2. Our patch creates Tasks for each awaitable passed to the function and inserts resulting Task objects into the Task Links Map (one of each is `new_task`) 3. A Sample is taken exactly at that moment, Task Link Map cleanup logic kicks in. a. `new_task` was created only just now, so was not part of the `all_tasks` `MirrorSet` snapshotted by the Sampler b. The cleanup logic sees that there's a link from `new_task` to `parent_task` but `new_task` isn't in `all_tasks` c. The cleanup logic assumes that means `new_task` previously existed but has completed, deduces the link must be removed d. Link is removed 4. `new_task` starts executing, but isn't linked to `parent_task` because we've cleaned it up before it got a chance to start The fix I propose is to introduce a new `previous_task_objects` set which contains the addresses of all existing Task objects at the time we previously sampled. If we see a Task object in the Task Links Map that doesn't exist in `all_tasks`, then two things are possible: either the Task used to exist and the link should be removed, or the Task was just created and the link should stay. To determine which case we're in, we look at `previous_task_objects` – if the Task is there then it previously existed and just completed (⇒ need to remove); if not, it was just created (⇒ shouldn't remove). I think this comes with one caveat: if we sample very rarely (which may happen in certain circumstances with adaptive sampling), then we could miss certain Tasks and never get rid of their links: ```py async def super_quick(): return None async def parent(): await asyncio.sleep(1.0) asyncio.gather(asyncio.create_task(super_quick())) ``` Given that super_quick will be… super quick, the following could happen: * We sample; we see parent ; `all_tasks = { parent }; task_links = {} ; previous_all_tasks = { parent }` * Parent continues, calls `asyncio.gather(...)` creates another Task ; `all_tasks = { parent, super_quick }; task_links = { super_quick -> parent }` * (We don’t sample yet because we rarely sample…) * Task execution continues, `super_quick` runs super quickly and finishes ; `all_tasks = { parent }; task_links = { super_quick -> parent }` * We sample; we see only parent (only Task in existence at that point); `current all_tasks = { parent }`; * `previous_all_tasks = { parent }` (still from the first sample we took) and `task_links = { super_quick -> parent }` * Our “debouncing” logic considers that `super_quick` is a new Task [because it’s not in `previous_all_tasks`] and doesn’t remove the link ⇒ The link will never go away because we’ll keep repeating this (potentially adding/removing other Tasks and Links, but this specific one will stay forever) The PR also fixes an issue with our patch for `as_completed` that would crash if `parent` was `None`. I don't think this can ever happen (since `as_completed` is necessarily called from a Task), but better safe than sorry. As this change fixes a race condition, it is pretty hard to strictly speaking test it. `dd-trace-py` already has a test that flaked in the past due to it (and that flaked in Echion as well): [`test_asyncio_as_completed`](https://github.com/datadog/dd-trace-py/blob/5b8add2c81604dafaa2b6c772c18efed71a6ae19/tests/profiling/collector/test_asyncio_as_completed.py#L62) --- .../profiling/stack_v2/echion/echion/threads.h | 16 ++++++++++++++-- ...ask-link-race-condition-891ac5c86af8a7c1.yaml | 4 ++++ 2 files changed, 18 insertions(+), 2 deletions(-) create mode 100644 releasenotes/notes/profiling-fix-task-link-race-condition-891ac5c86af8a7c1.yaml diff --git a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h index 0a993a8803d..e9541ecd5c8 100644 --- a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h +++ b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h @@ -209,6 +209,7 @@ ThreadInfo::unwind_tasks() std::unordered_set parent_tasks; std::unordered_map waitee_map; // Indexed by task origin std::unordered_map origin_map; // Indexed by task origin + static std::unordered_set previous_task_objects; auto maybe_all_tasks = get_all_tasks(reinterpret_cast(asyncio_loop)); if (!maybe_all_tasks) { @@ -232,14 +233,25 @@ ThreadInfo::unwind_tasks() if (all_task_origins.find(kv.first) == all_task_origins.end()) to_remove.push_back(kv.first); } - for (auto key : to_remove) - task_link_map.erase(key); + for (auto key : to_remove) { + // Only remove the link if the Child Task previously existed; otherwise it's a Task that + // has just been created and that wasn't in all_tasks when we took the snapshot. + if (previous_task_objects.find(key) != previous_task_objects.end()) { + task_link_map.erase(key); + } + } // Determine the parent tasks from the gather links. std::transform(task_link_map.cbegin(), task_link_map.cend(), std::inserter(parent_tasks, parent_tasks.begin()), [](const std::pair& kv) { return kv.second; }); + + // Copy all Task object pointers into previous_task_objects + previous_task_objects.clear(); + for (const auto& task : all_tasks) { + previous_task_objects.insert(task->origin); + } } for (auto& task : all_tasks) { diff --git a/releasenotes/notes/profiling-fix-task-link-race-condition-891ac5c86af8a7c1.yaml b/releasenotes/notes/profiling-fix-task-link-race-condition-891ac5c86af8a7c1.yaml new file mode 100644 index 00000000000..0caab9b4044 --- /dev/null +++ b/releasenotes/notes/profiling-fix-task-link-race-condition-891ac5c86af8a7c1.yaml @@ -0,0 +1,4 @@ +fixes: + - | + profiling: This fix resolves a race condition leading to incorrect stacks being reported + for asyncio parent/child Tasks (e.g. when using ``asyncio.gather``).