add config variables and sampling rules

Author: quinna-h

ddtrace/contrib/internal/asgi/middleware.py

@@ -1,3 +1,4 @@
+import ipaddress
 import os
 import sys
 from typing import Any
@@ -23,6 +24,7 @@
 from ddtrace.internal.schema.span_attribute_schema import SpanDirection
 from ddtrace.internal.utils import get_blocked
 from ddtrace.internal.utils import set_blocked
+from ddtrace.internal.utils.formats import asbool
 from ddtrace.trace import Span
 
 
@@ -34,7 +36,12 @@
         service_name=config._get_service(default="asgi"),
         request_span_name="asgi.request",
         distributed_tracing=True,
-        _trace_asgi_websocket=os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False),
+        _trace_asgi_websocket=asbool(os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False)),
+        # TODO: set as initially false until we gradually release feature
+        _trace_asgi_websocket_messages=asbool(os.getenv("DD_TRACE_WEBSOCKET_MESSAGES_ENABLED", default=True)),
+        _asgi_websockets_inherit_sampling=asbool(
+            os.getenv("DD_TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING", default=True)
+        ),
     ),
 )
 
@@ -135,7 +142,6 @@ async def __call__(self, scope, receive, send):
                 return await self.app(scope, receive, send)
 
             method = "WEBSOCKET"
-
         else:
             return await self.app(scope, receive, send)
         try:
@@ -177,11 +183,12 @@ async def __call__(self, scope, receive, send):
             ctx.set_item("req_span", span)
 
             # set span.kind to the type of request being performed
+            # question: should this be network.client.ip or
+
             span.set_tag_str(SPAN_KIND, SpanKind.SERVER)
 
             if scope["type"] == "websocket":
                 span.set_tag_str("http.upgraded", "websocket")
-                # TODO: add websocket session id
 
             if "datadog" not in scope:
                 scope["datadog"] = {"request_spans": [span]}
@@ -255,13 +262,14 @@ async def __call__(self, scope, receive, send):
             async def wrapped_send(message):
                 """
                 websocket.message.type
-                websocket.session.id*
                 websocket.message.length
-                websocket.message.frames*
-                websocket.message.send_time*
+                websocket.message.frames
                 """
                 try:
-                    if message.get("type") == "websocket.send":
+                    if (
+                        self.integration_config._trace_asgi_websocket_messages
+                        and message.get("type") == "websocket.send"
+                    ):
                         with self.tracer.trace(
                             "websocket.send",
                             service=span.service,
@@ -270,15 +278,38 @@ async def wrapped_send(message):
                         ) as send_span:
                             send_span.set_tag_str(COMPONENT, self.integration_config.integration_name)
                             send_span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
+
+                            # set tags related to peer.hostname
+                            client = scope.get("client")
+                            if len(client) >= 1:
+                                client_ip = client[0]
+                                span.set_tag_str("out.host", client_ip)
+                                try:
+                                    ipaddress.ip_address(client_ip)  # validate ip address
+                                    span.set_tag_str("network.client.ip", client_ip)
+                                except ValueError:
+                                    pass
                             # set link to http handshake span
                             # The link should have the attribute dd.kind set to resuming.
                             # If the instrumented library supports multicast or broadcast sending,
                             # there must be a link to the handshake span of every affected connection;
-                            # any websocket.session.id tag then should be added to the span link instead.
-                    elif message.get("type") == "websocket.close":
+                            send_span.set_link(
+                                trace_id=span.trace_id, span_id=span.span_id, attributes={"dd.kind": "resuming"}
+                            )
+                            send_span.set_metric("websocket.message.frames", 1)
+                            if "text" in message:
+                                send_span.set_tag_str("websocket.message.type", "text")
+                                send_span.set_metric("websocket.message.length", len(message["text"].encode("utf-8")))
+                            elif "binary" in message:
+                                send_span.set_tag_str("websocket.message.type", "binary")
+                                send_span.set_metric("websocket.message.length", len(message["bytes"]))
+
+                    elif (
+                        self.integration_config._trace_asgi_websocket_messages
+                        and message.get("type") == "websocket.close"
+                    ):
                         """
                         tags:
-                        websocket.session.id*
                         websocket.close.code
                         websocket.close.reason
 
@@ -291,7 +322,15 @@ async def wrapped_send(message):
                         ) as close_span:
                             close_span.set_tag_str(COMPONENT, self.integration_config.integration_name)
                             close_span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
-                            # TODO: add span links?
+                            client = scope.get("client")
+                            if len(client) >= 1:
+                                client_ip = client[0]
+                                span.set_tag_str("out.host", client_ip)
+                                try:
+                                    ipaddress.ip_address(client_ip)  # validate ip address
+                                    span.set_tag_str("network.client.ip", client_ip)
+                                except ValueError:
+                                    pass
                             if "text" in message:
                                 close_span.set_tag_str("websocket.message.type", "text")
                                 close_span.set_metric("websocket.message.length", len(message["text"].encode("utf-8")))
@@ -343,27 +382,35 @@ async def wrapped_receive():
                 """
                 tags:
                 websocket.message.type
-                websocket.session.id*
                 websocket.message.length
-                websocket.message.frames*
-                websocket.message.receive_time*
+                websocket.message.frames
+
                 _dd.dm.inherited
                 _dd.dm.service
                 _dd.dm.resource
                 """
 
                 try:
                     message = await receive()
-                    if message["type"] == "websocket.receive":
-                        with self.tracer.trace(
-                            "websocket.receive",
+                    if (
+                        self.integration_config._trace_asgi_websocket_messages
+                        and message["type"] == "websocket.receive"
+                    ):
+                        with self.tracer.start_span(
+                            name="websocket.receive",
                             service=span.service,
                             resource=f"websocket {scope.get('path', '')}",
                             span_type="websocket",
+                            child_of=None,
                         ) as ws_span:
                             ws_span.set_tag_str(COMPONENT, self.integration_config.integration_name)
                             ws_span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
                             core.dispatch("asgi.websocket.receive", (message,))
+                            # the span should have link to http handshake span
+                            # the link should have attribute dd.kind set to executed_by
+                            ws_span.set_link(
+                                trace_id=span.trace_id, span_id=span.span_id, attributes={"dd.kind": "executed_by"}
+                            )
 
                             if "text" in message:
                                 ws_span.set_tag_str("websocket.message.type", "text")
@@ -372,27 +419,34 @@ async def wrapped_receive():
                                 ws_span.set_tag_str("websocket.message.type", "binary")
                                 ws_span.set_metric("websocket.message.length", len(message["bytes"]))
 
-                            ws_span.set_metric("_dd.dm.inherited", 1)
-                            parent_span = self.tracer.current_root_span()
-                            if parent_span is not None:
-                                ws_span.set_tag_str("_dd.dm.service", parent_span.service)
-                                ws_span.set_tag_str("_dd.dm.resource", parent_span.resource)
+                            # since asgi is a high level framework, frames is always 1
+                            ws_span.set_metric("websocket.message.frames", 1)
 
-                    elif message["type"] == "websocket.disconnect":
+                            if self.integration_config._asgi_websockets_inherit_sampling:
+                                ws_span.set_metric("_dd.dm.inherited", 1)
+                                ws_span.set_tag_str("_dd.dm.service", span.service)
+                                ws_span.set_tag_str("_dd.dm.resource", span.resource)
+
+                    elif (
+                        self.integration_config._trace_asgi_websocket_messages
+                        and message["type"] == "websocket.disconnect"
+                    ):
                         # peer closes the connection
-                        with self.tracer.trace(
-                            "websocket.close",
+                        # in this case the span will be trace root (will behave like the websocket.receive use case)
+                        with self.tracer.start_span(
+                            name="websocket.close",
                             service=span.service,
                             resource=f"websocket {scope.get('path', '')}",
                             span_type="websocket",
+                            child_of=None,
                         ) as close_span:
                             close_span.set_tag_str(COMPONENT, self.integration_config.integration_name)
                             close_span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
-                            close_span.set_metric("_dd.dm.inherited", 1)
-                            parent_span = self.tracer.current_root_span()
-                            if parent_span is not None:
-                                close_span.set_tag_str("_dd.dm.service", parent_span.service)
-                                close_span.set_tag_str("_dd.dm.resource", parent_span.resource)
+
+                            if self.integration_config._asgi_websockets_inherit_sampling:
+                                close_span.set_metric("_dd.dm.inherited", 1)
+                                close_span.set_tag_str("_dd.dm.service", span.service)
+                                close_span.set_tag_str("_dd.dm.resource", span.resource)
 
                             code = message.get("code")
                             reason = message.get("reason")
@@ -427,7 +481,7 @@ async def wrapped_blocked_send(message):
                     message["more_body"] = False
                     core.dispatch("asgi.finalize_response", (content, None))
                 try:
-                    return await send(message)
+                    result = await send(message)
                 finally:
                     trace_utils.set_http_meta(
                         span, self.integration_config, status_code=status, response_headers=headers
@@ -440,7 +494,6 @@ async def wrapped_blocked_send(message):
             try:
                 core.dispatch("asgi.start_request", ("asgi",))
                 # Do not block right here. Wait for route to be resolved in starlette/patch.py
-                # TODO: do I need to change to wrapped_receive?
                 return await self.app(scope, wrapped_recv, wrapped_send)
             except BlockingException as e:
                 set_blocked(e.args[0])

ddtrace/contrib/internal/django/patch.py

@@ -80,7 +80,7 @@
         ),
         use_handler_resource_format=asbool(os.getenv("DD_DJANGO_USE_HANDLER_RESOURCE_FORMAT", default=False)),
         use_legacy_resource_format=asbool(os.getenv("DD_DJANGO_USE_LEGACY_RESOURCE_FORMAT", default=False)),
-        _trace_asgi_websocket=os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False),
+        _trace_asgi_websocket=asbool(os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False)),
     ),
 )
 

ddtrace/contrib/internal/fastapi/patch.py

@@ -12,6 +12,7 @@
 from ddtrace.contrib.internal.starlette.patch import traced_route_init
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.schema import schematize_service_name
+from ddtrace.internal.utils.formats import asbool
 from ddtrace.internal.utils.wrappers import unwrap as _u
 from ddtrace.settings.asm import config as asm_config
 from ddtrace.trace import Pin
@@ -26,7 +27,11 @@
         request_span_name="fastapi.request",
         distributed_tracing=True,
         trace_query_string=None,  # Default to global config
-        _trace_asgi_websocket=os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False),
+        _trace_asgi_websocket=asbool(os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False)),
+        _trace_asgi_websocket_messages=asbool(os.getenv("DD_TRACE_WEBSOCKET_MESSAGES_ENABLED", default=True)),
+        _asgi_websockets_inherit_sampling=asbool(
+            os.getenv("DD_TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING", default=True)
+        ),
     ),
 )
 

ddtrace/contrib/internal/starlette/patch.py

@@ -24,6 +24,7 @@
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils import get_blocked
 from ddtrace.internal.utils import set_argument_value
+from ddtrace.internal.utils.formats import asbool
 from ddtrace.internal.utils.wrappers import unwrap as _u
 from ddtrace.settings.asm import config as asm_config
 from ddtrace.trace import Pin
@@ -39,7 +40,8 @@
         _default_service=schematize_service_name("starlette"),
         request_span_name="starlette.request",
         distributed_tracing=True,
-        _trace_asgi_websocket=os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False),
+        _trace_asgi_websocket=asbool(os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False)),
+        _trace_asgi_websocket_messages=asbool(os.getenv("DD_TRACE_WEBSOCKET_MESSAGES_ENABLED", default=True)),
     ),
 )
 

tests/contrib/fastapi/test_fastapi.py

@@ -540,7 +540,8 @@ def test_table_query_snapshot(snapshot_client):
     }
 
 
-@snapshot()
+# @pytest.mark.subprocess(env=dict(DD_TRACE_WEBSOCKET_MESSAGES_ENABLED="true"))
+@snapshot(ignores=["meta._dd.span_links", "metrics.websocket.message.length"])
 def test_traced_websocket(test_spans, snapshot_app):
     client = TestClient(snapshot_app)
     with override_config("fastapi", dict(_trace_asgi_websocket=True)):
@@ -550,6 +551,16 @@ def test_traced_websocket(test_spans, snapshot_app):
             websocket.send_text("ping")
 
 
+# @snapshot()
+# def test_websocket_trace_sampling_priority(test_spans, snapshot_app):
+#     from ddtrace import tracer
+#     with tracer.trace("test-root", service="mytest", resource="/ws", _sampling_priority=0):
+#         client
+#         with TestClient(app).websocket_connect("/ws") as ws:
+#             ws.send_text("ping")
+#             ws.receive_text()
+
+
 def test_dont_trace_websocket_by_default(client, test_spans):
     initial_event_count = len(test_spans.pop_traces())
     with client.websocket_connect("/ws") as websocket: