wip

Author: quinna-h

ddtrace/_trace/trace_handlers.py

@@ -140,7 +140,7 @@ def _start_span(ctx: core.ExecutionContext, call_trace: bool = True, **kwargs) -
 
     if config._inferred_proxy_services_enabled:
         # dispatch event for checking headers and possibly making an inferred proxy span
-        core.dispatch("inferred_proxy.start", (ctx, tracer, span_kwargs, call_trace, integration_config))
+        core.dispatch("inferred_proxy.start", (ctx, tracer, span_kwargs, call_trace))
         # re-get span_kwargs in case an inferred span was created and we have a new span_kwargs.child_of field
         span_kwargs = ctx.get_item("span_kwargs", span_kwargs)
 
@@ -246,14 +246,15 @@ def _set_inferred_proxy_tags(span, status_code):
                 inferred_span.set_tag(ERROR_STACK, span.get_tag(ERROR_STACK))
 
 
-def _on_inferred_proxy_start(ctx, tracer, span_kwargs, call_trace, integration_config):
+def _on_inferred_proxy_start(ctx, tracer, span_kwargs, call_trace):
     # Skip creating another inferred span if one has already been created for this request
     if ctx.get_item("inferred_proxy_span"):
         return
 
     # some integrations like Flask / WSGI store headers from environ in 'distributed_headers'
     # and normalized headers in 'headers'
     headers = ctx.get_item("headers", ctx.get_item("distributed_headers", None))
+    integration_config = ctx.get_item("integration_config")
 
     # Inferred Proxy Spans
     if integration_config and headers is not None:
@@ -995,14 +996,15 @@ 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 _on_asgi_websocket_receive_message(ctx, scope, message, integration_config):
+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")
 
     span.set_tag_str(COMPONENT, integration_config.integration_name)
     span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
@@ -1025,14 +1027,15 @@ def _on_asgi_websocket_receive_message(ctx, scope, message, integration_config):
         _copy_trace_level_tags(span, ctx.parent.span)
 
 
-def _on_asgi_websocket_send_message(ctx, scope, message, integration_config):
+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")
 
     span.set_tag_str(COMPONENT, integration_config.integration_name)
     span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
@@ -1049,14 +1052,15 @@ def _on_asgi_websocket_send_message(ctx, scope, message, integration_config):
         )
 
 
-def _on_asgi_websocket_close_message(ctx, scope, message, integration_config):
+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")
 
     span.set_tag_str(COMPONENT, integration_config.integration_name)
     span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
@@ -1077,14 +1081,15 @@ def _on_asgi_websocket_close_message(ctx, scope, message, integration_config):
         _copy_trace_level_tags(span, ctx.parent.span)
 
 
-def _on_asgi_websocket_disconnect_message(ctx, scope, message, integration_config):
+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")
 
     span.set_tag_str(COMPONENT, integration_config.integration_name)
     span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)

ddtrace/contrib/internal/asgi/middleware.py

@@ -494,7 +494,7 @@ def _handle_websocket_send_message(self, scope: Mapping[str, Any], message: Mapp
             activate=True,
             tags={COMPONENT: self.integration_config.integration_name, SPAN_KIND: SpanKind.PRODUCER},
         ) as ctx:  # noqa: F841
-            core.dispatch("asgi.websocket.send.message", (ctx, scope, message, self.integration_config))
+            core.dispatch("asgi.websocket.send.message", (ctx, scope, message))
 
     def _handle_websocket_close_message(self, scope: Mapping[str, Any], message: Mapping[str, Any], request_span: Span):
         current_receive_span = scope.get("datadog", {}).get("current_receive_span")
@@ -517,7 +517,7 @@ def _handle_websocket_close_message(self, scope: Mapping[str, Any], message: Map
             activate=True,
             tags={COMPONENT: self.integration_config.integration_name, SPAN_KIND: SpanKind.PRODUCER},
         ) as ctx:  # noqa: F841
-            core.dispatch("asgi.websocket.close.message", (ctx, scope, message, self.integration_config))
+            core.dispatch("asgi.websocket.close.message", (ctx, scope, message))
 
         _cleanup_previous_receive(scope)
 
@@ -553,7 +553,7 @@ def _handle_websocket_receive_message(
         request_span: Span,
     ):
         _cleanup_previous_receive(scope)
-        core.dispatch("asgi.websocket.receive.message", (ctx, scope, message, self.integration_config))
+        core.dispatch("asgi.websocket.receive.message", (ctx, scope, message))
 
         scope["datadog"]["current_receive_span"] = recv_span
 
@@ -576,7 +576,7 @@ def _handle_websocket_disconnect_message(
             activate=True,
             tags={COMPONENT: self.integration_config.integration_name, SPAN_KIND: SpanKind.CONSUMER},
         ) as ctx:
-            core.dispatch("asgi.websocket.disconnect.message", (ctx, scope, message, self.integration_config))
+            core.dispatch("asgi.websocket.disconnect.message", (ctx, scope, message))
 
         if request_span and request_span.error == 0:
             request_span.finish()