fix(asgi): fix websocket spans duration and close span [backport 3.12] (#14426)

Backport 358bfd5 from #14388 to 3.12.

- Fix the duration of websocket.receive spans to close when the next
websocket.receive span is started; previously the
`cleanup_previous_receive` was in a handler method, which led to some
delay
- Fix the duration of the `fastapi.request` handshake span to close when
the connection is upgraded, not after the connection is closed
- Minor fix where websocket.close parent should be the request span if
`DD_TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES` is disabled

Trace (when `DD_TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES` is disabled):

<img width="1062" height="487" alt="Screenshot 2025-08-21 at 1 00 39 PM"
src="https://github.com/user-attachments/assets/42199884-a712-45fd-a888-d564b21cadf2"
/>


## 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: Quinna Halim <[email protected]>

Author: dd-octo-sts[bot]

ddtrace/contrib/internal/asgi/middleware.py

@@ -359,12 +359,14 @@ async def wrapped_receive():
 
                 Note:
                     - Spans are linked to the handshake span
-                    - Receive spans remain open until the next receive operation or connection close
+                    - Receive spans are finished exactly when the next receive operation starts
                 """
 
                 if not self.integration_config.trace_asgi_websocket_messages:
                     return await receive()
 
+                current_receive_span = scope.get("datadog", {}).get("current_receive_span")
+
                 with core.context_with_data(
                     "asgi.websocket.receive_message",
                     tracer=self.tracer,
@@ -377,6 +379,10 @@ async def wrapped_receive():
                     call_trace=False,
                     activate=True,
                 ) as ctx:
+                    if current_receive_span:
+                        current_receive_span.finish()
+                        scope["datadog"].pop("current_receive_span", None)
+
                     recv_span = ctx.span
                     try:
                         message = await receive()
@@ -421,6 +427,16 @@ async def wrapped_send(message: Mapping[str, Any]):
                 """
                 try:
                     if (
+                        scope["type"] == "websocket"
+                        and self.integration_config.trace_asgi_websocket_messages
+                        and message.get("type") == "websocket.accept"
+                    ):
+                        # Close handshake span once connection is upgraded
+                        if span and span.error == 0:
+                            span.finish()
+                        return await send(message)
+
+                    elif (
                         scope["type"] == "websocket"
                         and self.integration_config.trace_asgi_websocket_messages
                         and message.get("type") == "websocket.send"
@@ -549,7 +565,7 @@ def _handle_websocket_close_message(self, scope: Mapping[str, Any], message: Map
         current_receive_span = scope.get("datadog", {}).get("current_receive_span")
 
         # Use receive span as parent if it exists, otherwise use main span
-        if current_receive_span:
+        if self.integration_config.websocket_messages_separate_traces and current_receive_span:
             parent_span = current_receive_span
         else:
             parent_span = request_span
@@ -607,8 +623,6 @@ def _handle_http_response(
     def _handle_websocket_receive_message(
         self, scope: Mapping[str, Any], message: Mapping[str, Any], recv_span: Span, request_span: Span
     ):
-        _cleanup_previous_receive(scope)
-
         recv_span.set_tag_str(COMPONENT, self.integration_config.integration_name)
         recv_span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
         recv_span.set_tag_str(websocket.RECEIVE_DURATION_TYPE, "blocking")
@@ -660,8 +674,3 @@ def _handle_websocket_disconnect_message(
 
             _copy_trace_level_tags(disconnect_span, request_span)
             _set_websocket_close_tags(disconnect_span, message)
-
-        # The context system automatically finishes the disconnect_span
-        # Only finish the main span if needed
-        if request_span and request_span.error == 0:
-            request_span.finish()

releasenotes/notes/fix-websocket-durations-0e0b07916912294d.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    tracing: Fixes issue with ``websocket.receive`` span not closing exactly when another ``websocket.receive`` span was opened. 
+  - |
+    tracing: Fixes duration of websocket handshake span such that the handshake span closes when the connection is upgraded. 
+  - |
+    tracing: Fix where ``websocket.close`` parent should be the handshake span when configuration is disabled.