feat: add span pointer attributes for websocket context propagation (#15153)

## Description

<!-- Provide an overview of the change and motivation for the change -->


From
[RFC](https://docs.google.com/document/d/1RhIoUJhY7nawH9pkdbtbNy_G-7P02IR76Ka-SyJQ4f8/edit?tab=t.0):

> There is no direct context extraction on the incoming messages traced
since they already hold a relationship (a link) to the trace where the
context propagation happens (the handshake). The same applies to the
send message spans (for context injection).

This PR implements span pointers to connect outgoing and incoming
messages over websocket.

Span pointer attributes:

- link.name: span-pointer-down (if outgoing) / span-pointer-up (if
incoming)
- dd.kind: span-pointer
- ptr.kind: websocket
- ptr.dir: d (if outgoing) / u (if incoming)
- ptr.hash: S<128 bit hex handshake trace id><64 bit hex handshake
parent id><32 bit hex counter> (if outgoing on server or incoming on
client) / C<128 bit hex handshake trace id><64 bit hex handshake parent
id><32 bit hex counter> (if outgoing on client or incoming on server)




## Testing

<!-- Describe your testing strategy or note what tests are included -->

See `test_websocket_context_propagation` generated snapshot file

`test_websocket_context_propagation` flame graph:

`websocket.receive` parent and `websocket.send` / `websocket.close`

<img width="929" height="204" alt="Screenshot 2025-11-10 at 3 28 26 PM"
src="https://github.com/user-attachments/assets/4c2b6656-c38c-4953-b569-ece83ed1db72"
/>

<img width="929" height="247" alt="Screenshot 2025-11-10 at 3 28 58 PM"
src="https://github.com/user-attachments/assets/f4cada27-e508-4444-9341-267a37ecac5a"
/>

## Risks

<!-- Note any risks associated with this change, or "None" if no risks
-->

## Additional Notes

<!-- Any other information that would be helpful for reviewers -->

---------

Co-authored-by: Brett Langdon <brett.langdon@datadoghq.com>

Author: quinna-h

ddtrace/_trace/_span_pointer.py

@@ -24,6 +24,11 @@ class _SpanPointerDirection(Enum):
     DOWNSTREAM = "d"
 
 
+class _SpanPointerDirectionName(Enum):
+    UPSTREAM = "span-pointer-up"
+    DOWNSTREAM = "span-pointer-down"
+
+
 class _SpanPointerDescription(NamedTuple):
     # Not to be confused with _SpanPointer. This class describes the parameters
     # required to attach a span pointer to a Span. It lets us decouple code

ddtrace/_trace/trace_handlers.py

@@ -15,7 +15,10 @@
 import ddtrace
 from ddtrace import config
 from ddtrace._trace._inferred_proxy import create_inferred_proxy_span_if_headers_exist
+from ddtrace._trace._span_link import SpanLinkKind as _SpanLinkKind
 from ddtrace._trace._span_pointer import _SpanPointerDescription
+from ddtrace._trace._span_pointer import _SpanPointerDirection
+from ddtrace._trace._span_pointer import _SpanPointerDirectionName
 from ddtrace._trace.span import Span
 from ddtrace._trace.utils import extract_DD_context_from_messages
 from ddtrace.constants import _SPAN_MEASURED_KEY
@@ -31,6 +34,7 @@
 from ddtrace.contrib.internal.trace_utils import _set_url_tag
 from ddtrace.ext import SpanKind
 from ddtrace.ext import SpanLinkKind
+from ddtrace.ext import SpanTypes
 from ddtrace.ext import db
 from ddtrace.ext import http
 from ddtrace.ext import net
@@ -43,6 +47,7 @@
 from ddtrace.internal.constants import FLASK_ENDPOINT
 from ddtrace.internal.constants import FLASK_URL_RULE
 from ddtrace.internal.constants import FLASK_VIEW_ARGS
+from ddtrace.internal.constants import HTTP_REQUEST_UPGRADED
 from ddtrace.internal.constants import MESSAGING_BATCH_COUNT
 from ddtrace.internal.constants import MESSAGING_DESTINATION_NAME
 from ddtrace.internal.constants import MESSAGING_MESSAGE_ID
@@ -57,6 +62,9 @@
 
 log = get_logger(__name__)
 
+_WEBSOCKET_LINK_ATTRS_EXECUTED = {SPAN_LINK_KIND: SpanLinkKind.EXECUTED}
+_WEBSOCKET_LINK_ATTRS_RESUMING = {SPAN_LINK_KIND: SpanLinkKind.RESUMING}
+
 
 class _TracedIterable(wrapt.ObjectProxy):
     def __init__(self, wrapped, span, parent_span, wrapped_is_iterator=False):
@@ -992,12 +1000,109 @@ def _set_client_ip_tags(scope: Mapping[str, Any], span: Span):
             log.debug("Could not validate client IP address for websocket send message: %s", str(e))
 
 
+def _init_websocket_message_counters(scope: Dict[str, Any]) -> None:
+    if "datadog" not in scope:
+        scope["datadog"] = {}
+    if "websocket_receive_counter" not in scope["datadog"]:
+        scope["datadog"]["websocket_receive_counter"] = 0
+    if "websocket_send_counter" not in scope["datadog"]:
+        scope["datadog"]["websocket_send_counter"] = 0
+
+
+def _increment_websocket_counter(scope: Dict[str, Any], counter_type: str) -> int:
+    """
+    Increment and return websocket message counter (either websocket_receive_counter or websocket_send_counter)
+    """
+    scope["datadog"][counter_type] += 1
+    return scope["datadog"][counter_type]
+
+
+def _build_websocket_span_pointer_hash(
+    handshake_trace_id: int,
+    handshake_span_id: int,
+    counter: int,
+    is_server: bool,
+    is_incoming: bool,
+) -> str:
+    """
+    Build websocket span pointer hash.
+
+    Format: <prefix><128 bit hex trace id><64 bit hex span id><32 bit hex counter>
+    Prefix: 'S' for server outgoing or client incoming, 'C' for server incoming or client outgoing
+    """
+    if (is_server and not is_incoming) or (not is_server and is_incoming):
+        prefix = "S"
+    else:
+        prefix = "C"
+
+    trace_id_hex = f"{handshake_trace_id:032x}"
+    span_id_hex = f"{handshake_span_id:016x}"
+    counter_hex = f"{counter:08x}"
+
+    return f"{prefix}{trace_id_hex}{span_id_hex}{counter_hex}"
+
+
+def _has_distributed_tracing_context(span: Span) -> bool:
+    """
+    Check if the handshake span has extracted distributed tracing context.
+
+    A websocket server must not set the span pointer if the handshake has not extracted a context
+
+    A span has distributed tracing context if it has a parent context that was
+    extracted from headers.
+    """
+    if not span or not span._parent_context:
+        return False
+    return span._parent_context._is_remote
+
+
+def _add_websocket_span_pointer_attributes(
+    link_attributes: Dict[str, Any],
+    integration_config: Any,
+    handshake_span: Span,
+    scope: Dict[str, Any],
+    is_incoming: bool,
+) -> None:
+    """
+    Add span pointer attributes to link_attributes for websocket message correlation.
+    """
+
+    if not integration_config.distributed_tracing or not _has_distributed_tracing_context(handshake_span):
+        return
+
+    # Increment counter based on message direction
+    counter_type = "websocket_receive_counter" if is_incoming else "websocket_send_counter"
+    counter = _increment_websocket_counter(scope, counter_type)
+
+    ptr_hash = _build_websocket_span_pointer_hash(
+        handshake_trace_id=handshake_span.trace_id,
+        handshake_span_id=handshake_span.span_id,
+        counter=counter,
+        is_server=True,
+        is_incoming=is_incoming,
+    )
+
+    if is_incoming:
+        link_name = _SpanPointerDirectionName.UPSTREAM
+        ptr_direction = _SpanPointerDirection.UPSTREAM
+    else:
+        link_name = _SpanPointerDirectionName.DOWNSTREAM
+        ptr_direction = _SpanPointerDirection.DOWNSTREAM
+
+    link_attributes.update(
+        {
+            "link.name": link_name,
+            "dd.kind": _SpanLinkKind.SPAN_POINTER.value,
+            "ptr.kind": SpanTypes.WEBSOCKET,
+            "ptr.dir": ptr_direction,
+            "ptr.hash": ptr_hash,
+        }
+    )
+
+
 def _on_asgi_websocket_receive_message(ctx, scope, message):
     """
     Handle websocket receive message events.
-
-    This handler is called when a websocket receive message event is dispatched.
-    It sets up the span with appropriate tags, metrics, and links.
     """
     span = ctx.span
     integration_config = ctx.get_item("integration_config")
@@ -1011,24 +1116,24 @@ def _on_asgi_websocket_receive_message(ctx, scope, message):
     span.set_metric(websocket.MESSAGE_FRAMES, 1)
 
     if hasattr(ctx, "parent") and ctx.parent.span:
-        span.set_link(
-            trace_id=ctx.parent.span.trace_id,
-            span_id=ctx.parent.span.span_id,
-            attributes={SPAN_LINK_KIND: SpanLinkKind.EXECUTED},
+        handshake_span = ctx.parent.span
+        link_attributes = _WEBSOCKET_LINK_ATTRS_EXECUTED.copy()
+
+        _add_websocket_span_pointer_attributes(
+            link_attributes, integration_config, handshake_span, scope, is_incoming=True
         )
 
+        span.link_span(handshake_span.context, link_attributes)
+
         if getattr(integration_config, "asgi_websocket_messages_inherit_sampling", True):
-            _inherit_sampling_tags(span, ctx.parent.span._local_root)
+            _inherit_sampling_tags(span, handshake_span._local_root)
 
-        _copy_trace_level_tags(span, ctx.parent.span)
+        _copy_trace_level_tags(span, handshake_span)
 
 
 def _on_asgi_websocket_send_message(ctx, scope, message):
     """
     Handle websocket send message events.
-
-    This handler is called when a websocket send message event is dispatched.
-    It sets up the span with appropriate tags, metrics, and links.
     """
     span = ctx.span
     integration_config = ctx.get_item("integration_config")
@@ -1041,19 +1146,19 @@ def _on_asgi_websocket_send_message(ctx, scope, message):
     span.set_metric(websocket.MESSAGE_FRAMES, 1)
 
     if hasattr(ctx, "parent") and ctx.parent.span:
-        span.set_link(
-            trace_id=ctx.parent.span.trace_id,
-            span_id=ctx.parent.span.span_id,
-            attributes={SPAN_LINK_KIND: SpanLinkKind.RESUMING},
+        handshake_span = ctx.parent.span
+        link_attributes = _WEBSOCKET_LINK_ATTRS_RESUMING.copy()
+
+        _add_websocket_span_pointer_attributes(
+            link_attributes, integration_config, handshake_span, scope, is_incoming=False
         )
 
+        span.link_span(handshake_span.context, link_attributes)
+
 
 def _on_asgi_websocket_close_message(ctx, scope, message):
     """
     Handle websocket close message events.
-
-    This handler is called when a websocket close message event is dispatched.
-    It sets up the span with appropriate tags, metrics, and links.
     """
     span = ctx.span
     integration_config = ctx.get_item("integration_config")
@@ -1068,21 +1173,21 @@ def _on_asgi_websocket_close_message(ctx, scope, message):
     _set_websocket_close_tags(span, message)
 
     if hasattr(ctx, "parent") and ctx.parent.span:
-        span.set_link(
-            trace_id=ctx.parent.span.trace_id,
-            span_id=ctx.parent.span.span_id,
-            attributes={SPAN_LINK_KIND: SpanLinkKind.RESUMING},
+        handshake_span = ctx.parent.span
+        link_attributes = _WEBSOCKET_LINK_ATTRS_RESUMING.copy()
+
+        _add_websocket_span_pointer_attributes(
+            link_attributes, integration_config, handshake_span, scope, is_incoming=False
         )
 
-        _copy_trace_level_tags(span, ctx.parent.span)
+        span.link_span(handshake_span.context, link_attributes)
+
+        _copy_trace_level_tags(span, handshake_span)
 
 
 def _on_asgi_websocket_disconnect_message(ctx, scope, message):
     """
     Handle websocket disconnect message events.
-
-    This handler is called when a websocket disconnect message event is dispatched.
-    It sets up the span with appropriate tags, metrics, and links.
     """
     span = ctx.span
     integration_config = ctx.get_item("integration_config")
@@ -1093,16 +1198,19 @@ def _on_asgi_websocket_disconnect_message(ctx, scope, message):
     _set_websocket_close_tags(span, message)
 
     if hasattr(ctx, "parent") and ctx.parent.span:
-        span.set_link(
-            trace_id=ctx.parent_span.trace_id,
-            span_id=ctx.parent_span.span_id,
-            attributes={SPAN_LINK_KIND: SpanLinkKind.EXECUTED},
+        handshake_span = ctx.parent.span
+        link_attributes = _WEBSOCKET_LINK_ATTRS_EXECUTED.copy()
+
+        _add_websocket_span_pointer_attributes(
+            link_attributes, integration_config, handshake_span, scope, is_incoming=True
         )
 
+        span.link_span(handshake_span.context, link_attributes)
+
         if getattr(integration_config, "asgi_websocket_messages_inherit_sampling", True):
-            _inherit_sampling_tags(span, ctx.parent.span._local_root)
+            _inherit_sampling_tags(span, handshake_span._local_root)
 
-        _copy_trace_level_tags(span, ctx.parent.span)
+        _copy_trace_level_tags(span, handshake_span)
 
 
 def _on_asgi_request(ctx: core.ExecutionContext) -> None:
@@ -1115,14 +1223,15 @@ def _on_asgi_request(ctx: core.ExecutionContext) -> None:
     span = _start_span(ctx)
     ctx.set_item("req_span", span)
 
-    if scope["type"] == "websocket":
-        span._set_tag_str("http.upgraded", "websocket")
-
     if "datadog" not in scope:
         scope["datadog"] = {"request_spans": [span]}
     else:
         scope["datadog"]["request_spans"].append(span)
 
+    if scope["type"] == "websocket":
+        span._set_tag_str(HTTP_REQUEST_UPGRADED, SpanTypes.WEBSOCKET)
+        _init_websocket_message_counters(scope)
+
 
 def listen():
     core.on("wsgi.request.prepare", _on_request_prepare)

ddtrace/internal/constants.py

@@ -67,6 +67,7 @@
 HTTP_REQUEST_HEADER = "http.request.header"
 HTTP_REQUEST_PARAMETER = "http.request.parameter"
 HTTP_REQUEST_BODY = "http.request.body"
+HTTP_REQUEST_UPGRADED = "http.upgraded"
 HTTP_REQUEST_PATH_PARAMETER = "http.request.path.parameter"
 REQUEST_PATH_PARAMS = "http.request.path_params"
 STATUS_403_TYPE_AUTO = {"status_code": 403, "type": "auto"}

releasenotes/notes/websocket-span-pointers-25e07939aa75527a.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+     feat(asgi): Enable context propagation between websocket message spans.