Merge branch '2.18' into backport-11759-to-2.18

Author: mabdinur

ddtrace/contrib/internal/asgi/middleware.py

@@ -91,6 +91,18 @@ async def _blocked_asgi_app(scope, receive, send):
     await send({"type": "http.response.body", "body": b""})
 
 
+def _parse_response_cookies(response_headers):
+    cookies = {}
+    try:
+        result = response_headers.get("set-cookie", "").split("=", maxsplit=1)
+        if len(result) == 2:
+            cookie_key, cookie_value = result
+            cookies[cookie_key] = cookie_value
+    except Exception:
+        log.debug("failed to extract response cookies", exc_info=True)
+    return cookies
+
+
 class TraceMiddleware:
     """
     ASGI application middleware that traces the requests.
@@ -211,7 +223,6 @@ async def __call__(self, scope, receive, send):
                 peer_ip = client[0]
             else:
                 peer_ip = None
-
             trace_utils.set_http_meta(
                 span,
                 self.integration_config,
@@ -234,15 +245,8 @@ async def wrapped_send(message):
                 except Exception:
                     log.warning("failed to extract response headers", exc_info=True)
                     response_headers = None
-
                 if span and message.get("type") == "http.response.start" and "status" in message:
-                    cookies = {}
-                    try:
-                        cookie_key, cookie_value = response_headers.get("set-cookie", "").split("=", maxsplit=1)
-                        cookies[cookie_key] = cookie_value
-                    except Exception:
-                        log.debug("failed to extract response cookies", exc_info=True)
-
+                    cookies = _parse_response_cookies(response_headers)
                     status_code = message["status"]
                     trace_utils.set_http_meta(
                         span,

ddtrace/contrib/internal/celery/app.py

@@ -133,10 +133,6 @@ def _traced_apply_async_inner(func, instance, args, kwargs):
                 if task_span:
                     task_span.set_exc_info(*sys.exc_info())
 
-                prerun_span = core.get_item("prerun_span")
-                if prerun_span:
-                    prerun_span.set_exc_info(*sys.exc_info())
-
                 raise
             finally:
                 task_span = core.get_item("task_span")
@@ -147,11 +143,4 @@ def _traced_apply_async_inner(func, instance, args, kwargs):
                     )
                     task_span.finish()
 
-                prerun_span = core.get_item("prerun_span")
-                if prerun_span:
-                    log.debug(
-                        "The task_postrun signal was not called, so manually closing span: %s", prerun_span._pprint()
-                    )
-                    prerun_span.finish()
-
     return _traced_apply_async_inner

ddtrace/contrib/internal/celery/signals.py

@@ -54,9 +54,6 @@ def trace_prerun(*args, **kwargs):
     service = config.celery["worker_service_name"]
     span = pin.tracer.trace(c.WORKER_ROOT_SPAN, service=service, resource=task.name, span_type=SpanTypes.WORKER)
 
-    # Store an item called "prerun span" in case task_postrun doesn't get called
-    core.set_item("prerun_span", span)
-
     # set span.kind to the type of request being performed
     span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
 

ddtrace/internal/datadog/profiling/dd_wrapper/include/uploader_builder.hpp

@@ -43,6 +43,8 @@ class UploaderBuilder
     static void set_output_filename(std::string_view _output_filename);
 
     static std::variant<Uploader, std::string> build();
+
+    static void postfork_child();
 };
 
 } // namespace Datadog

ddtrace/internal/datadog/profiling/dd_wrapper/src/code_provenance.cpp

@@ -14,9 +14,8 @@ namespace Datadog {
 void
 Datadog::CodeProvenance::postfork_child()
 {
-    get_instance().mtx.~mutex();            // Destroy the mutex
+    // NB placement-new to re-init and leak the mutex because doing anything else is UB
     new (&get_instance().mtx) std::mutex(); // Recreate the mutex
-    get_instance().reset();                 // Reset the state
 }
 
 void

ddtrace/internal/datadog/profiling/dd_wrapper/src/ddup_interface.cpp

@@ -24,6 +24,7 @@ ddup_postfork_child()
     Datadog::Uploader::postfork_child();
     Datadog::SampleManager::postfork_child();
     Datadog::CodeProvenance::postfork_child();
+    Datadog::UploaderBuilder::postfork_child();
 }
 
 void

ddtrace/internal/datadog/profiling/dd_wrapper/src/profile.cpp

@@ -186,6 +186,6 @@ Datadog::Profile::collect(const ddog_prof_Sample& sample, int64_t endtime_ns)
 void
 Datadog::Profile::postfork_child()
 {
-    profile_mtx.unlock();
+    new (&profile_mtx) std::mutex();
     cycle_buffers();
 }

ddtrace/internal/datadog/profiling/dd_wrapper/src/uploader.cpp

@@ -175,5 +175,6 @@ Datadog::Uploader::postfork_parent()
 void
 Datadog::Uploader::postfork_child()
 {
-    unlock();
+    // NB placement-new to re-init and leak the mutex because doing anything else is UB
+    new (&upload_lock) std::mutex();
 }

ddtrace/internal/datadog/profiling/dd_wrapper/src/uploader_builder.cpp

@@ -186,3 +186,10 @@ Datadog::UploaderBuilder::build()
 
     return Datadog::Uploader{ output_filename, ddog_exporter };
 }
+
+void
+Datadog::UploaderBuilder::postfork_child()
+{
+    // NB placement-new to re-init and leak the mutex because doing anything else is UB
+    new (&tag_mutex) std::mutex();
+}

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

@@ -67,6 +67,11 @@ _stack_v2_atfork_child()
     // so we don't even reveal this function to the user
     _set_pid(getpid());
     ThreadSpanLinks::postfork_child();
+
+    // `thread_info_map_lock` and `task_link_map_lock` are global locks held in echion
+    // NB placement-new to re-init and leak the mutex because doing anything else is UB
+    new (&thread_info_map_lock) std::mutex;
+    new (&task_link_map_lock) std::mutex;
 }
 
 __attribute__((constructor)) void