fix(profiling): clear finished thread ids from ThreadSpanLinks (#11235)

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Emmett Butler <[email protected]>

Author: taegyunkim

ddtrace/internal/datadog/profiling/stack_v2/include/constants.hpp

@@ -1,3 +1,5 @@
+#pragma once
+
 #include "dd_wrapper/include/constants.hpp"
 
 // Default sampling frequency in microseconds.  This will almost certainly be overridden by dynamic sampling.

ddtrace/internal/datadog/profiling/stack_v2/include/thread_span_links.hpp

@@ -21,6 +21,13 @@ struct Span
       , span_type(_span_type)
     {
     }
+
+    // for testing
+    bool operator==(const Span& other) const
+    {
+        return span_id == other.span_id && local_root_span_id == other.local_root_span_id &&
+               span_type == other.span_type;
+    }
 };
 
 class ThreadSpanLinks
@@ -38,6 +45,7 @@ class ThreadSpanLinks
 
     void link_span(uint64_t thread_id, uint64_t span_id, uint64_t local_root_span_id, std::string span_type);
     const std::optional<Span> get_active_span_from_thread_id(uint64_t thread_id);
+    void unlink_span(uint64_t thread_id);
     void reset();
 
     static void postfork_child();

ddtrace/internal/datadog/profiling/stack_v2/src/stack_v2.cpp

@@ -86,6 +86,7 @@ stack_v2_thread_unregister(PyObject* self, PyObject* args)
     }
 
     Sampler::get().unregister_thread(id);
+    ThreadSpanLinks::get_instance().unlink_span(id);
     Py_RETURN_NONE;
 }
 

ddtrace/internal/datadog/profiling/stack_v2/src/thread_span_links.cpp

@@ -12,12 +12,14 @@ ThreadSpanLinks::link_span(uint64_t thread_id, uint64_t span_id, uint64_t local_
 {
     std::lock_guard<std::mutex> lock(mtx);
 
-    if (thread_id_to_span.find(thread_id) == thread_id_to_span.end()) {
+    auto it = thread_id_to_span.find(thread_id);
+    if (it == thread_id_to_span.end()) {
         thread_id_to_span[thread_id] = std::make_unique<Span>(span_id, local_root_span_id, span_type);
+    } else {
+        it->second->span_id = span_id;
+        it->second->local_root_span_id = local_root_span_id;
+        it->second->span_type = span_type;
     }
-    thread_id_to_span[thread_id]->span_id = span_id;
-    thread_id_to_span[thread_id]->local_root_span_id = local_root_span_id;
-    thread_id_to_span[thread_id]->span_type = span_type;
 }
 
 const std::optional<Span>
@@ -33,6 +35,14 @@ ThreadSpanLinks::get_active_span_from_thread_id(uint64_t thread_id)
     return span;
 }
 
+void
+ThreadSpanLinks::unlink_span(uint64_t thread_id)
+{
+    std::lock_guard<std::mutex> lock(mtx);
+
+    thread_id_to_span.erase(thread_id); // This is a no-op if the key is not found
+}
+
 void
 ThreadSpanLinks::reset()
 {

ddtrace/internal/datadog/profiling/stack_v2/test/thread_span_links.cpp

@@ -1,9 +1,13 @@
+#include "thread_span_links.hpp"
+
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include <optional>
+#include <random>
 #include <string>
 #include <thread>
-
-#include "thread_span_links.hpp"
+#include <unordered_set>
 
 static void
 get()
@@ -39,6 +43,48 @@ TEST(ThreadSpanLinksConcurrency, GetSetRace)
     t2.join();
 }
 
+TEST(ThreadSpanLinks, ClearFinished)
+{
+    unsigned int num_thread_ids = 100;
+    std::unordered_set<uint64_t> thread_ids;
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<uint64_t> dis(0, UINT64_MAX);
+
+    // Generate random 100 native thread ids
+    for (unsigned int i = 0; i < num_thread_ids; i++) {
+        thread_ids.insert(dis(gen));
+    }
+
+    // Call link_span with the thread ids
+    for (auto thread_id : thread_ids) {
+        Datadog::ThreadSpanLinks::get_instance().link_span(thread_id, thread_id, thread_id, "test");
+    }
+
+    std::unordered_set<uint64_t> finished_threads;
+    std::uniform_real_distribution<double> real_dis(0, 1);
+
+    for (auto thread_id : thread_ids) {
+        if (real_dis(gen) < 0.5) {
+            finished_threads.insert(thread_id);
+            Datadog::ThreadSpanLinks::get_instance().unlink_span(thread_id);
+        }
+    }
+
+    // Check that the unseen ids are removed
+    for (auto thread_id : thread_ids) {
+        std::optional<Datadog::Span> span_opt =
+          Datadog::ThreadSpanLinks::get_instance().get_active_span_from_thread_id(thread_id);
+        if (finished_threads.find(thread_id) == finished_threads.end()) {
+            EXPECT_EQ(span_opt, Datadog::Span(thread_id, thread_id, "test"));
+
+        } else {
+            EXPECT_EQ(span_opt, std::nullopt);
+        }
+    }
+}
+
 int
 main(int argc, char** argv)
 {

releasenotes/notes/profiling-span-link-leak-e7eb1aca3362cd2e.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    profiling: when a Python thread finishes, this change frees memory used for mapping
+    its thread id to ``Span``. The mapping is populated and used when
+    `DD_PROFILING_ENDPOINT_COLLECTION_ENABLED`` and
+    ``DD_PROFILING_STACK_V2_ENABLED`` were set to enable grouping of profiles
+    for endpoints.

tests/profiling_v2/collector/test_stack.py

@@ -1,6 +1,8 @@
 import os
 import sys
+import threading
 import time
+from unittest.mock import patch
 import uuid
 
 import pytest
@@ -9,6 +11,7 @@
 from ddtrace import tracer
 from ddtrace.internal.datadog.profiling import ddup
 from ddtrace.profiling.collector import stack
+from ddtrace.settings.profiling import config
 from tests.profiling.collector import pprof_utils
 
 
@@ -111,6 +114,63 @@ def test_push_span(stack_v2_enabled, tmp_path):
         )
 
 
+def test_push_span_unregister_thread(tmp_path, monkeypatch):
+    if sys.version_info[:2] == (3, 7):
+        pytest.skip("stack_v2 is not supported on Python 3.7")
+
+    with patch("ddtrace.internal.datadog.profiling.stack_v2.unregister_thread") as unregister_thread:
+        monkeypatch.setattr(config.stack, "v2_enabled", True)
+        tracer._endpoint_call_counter_span_processor.enable()
+
+        test_name = "test_push_span_unregister_thread"
+        pprof_prefix = str(tmp_path / test_name)
+        output_filename = pprof_prefix + "." + str(os.getpid())
+
+        assert ddup.is_available
+        ddup.config(env="test", service=test_name, version="my_version", output_filename=pprof_prefix)
+        ddup.start()
+
+        resource = str(uuid.uuid4())
+        span_type = ext.SpanTypes.WEB
+
+        def target_fun():
+            for _ in range(5):
+                time.sleep(0.01)
+
+        with stack.StackCollector(
+            None,
+            tracer=tracer,
+            endpoint_collection_enabled=True,
+            ignore_profiler=True,  # this is not necessary, but it's here to trim samples
+            _stack_collector_v2_enabled=True,
+        ):
+            with tracer.trace("foobar", resource=resource, span_type=span_type) as span:
+                span_id = span.span_id
+                local_root_span_id = span._local_root.span_id
+                t = threading.Thread(target=target_fun)
+                t.start()
+                t.join()
+                thread_id = t.ident
+        ddup.upload()
+
+        profile = pprof_utils.parse_profile(output_filename)
+        samples = pprof_utils.get_samples_with_label_key(profile, "span id")
+        assert len(samples) > 0
+        for sample in samples:
+            pprof_utils.assert_stack_event(
+                profile,
+                sample,
+                expected_event=pprof_utils.StackEvent(
+                    span_id=span_id,
+                    local_root_span_id=local_root_span_id,
+                    trace_type=span_type,
+                    trace_endpoint=resource,
+                ),
+            )
+
+        unregister_thread.assert_called_once_with(thread_id)
+
+
 @pytest.mark.parametrize("stack_v2_enabled", [True, False])
 def test_push_non_web_span(stack_v2_enabled, tmp_path):
     if sys.version_info[:2] == (3, 7) and stack_v2_enabled: