Repository: DataDog/dd-trace-py

Issue:
Decreased success rate of integration_agent CircleCI job
The `integration_agent` CircleCI job has been failing more often this past week. This job runs the integration tests with the Datadog Agent. 

We see the same two tests failing for 3.11 in two failing pipelines on 1.x ([[1]](https://app.circleci.com/pipelines/github/DataDog/dd-trace-py/31440/workflows/392068c8-7b66-4ea3-b5ae-21879d9820b1/jobs/2129132/parallel-runs/0/steps/0-106), [[2]](https://app.circleci.com/pipelines/github/DataDog/dd-trace-py/31771/workflows/80a433f7-4d5f-42e2-92aa-4439df620fc6/jobs/2150184/parallel-runs/0/steps/0-106)):

```
FAILED tests/integration/test_integration.py::test_regression_logging_in_context[True-True-True]
FAILED tests/integration/test_integration.py::test_regression_logging_in_context[True-True-False]
```



Discussion:

Resolving patch:
diff --git a/ddtrace/appsec/_asm_request_context.py b/ddtrace/appsec/_asm_request_context.py
--- a/ddtrace/appsec/_asm_request_context.py
+++ b/ddtrace/appsec/_asm_request_context.py
@@ -3,6 +3,7 @@
 
 from ddtrace import config
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
+from ddtrace.internal import _context
 from ddtrace.internal.logger import get_logger
 
 
@@ -33,7 +34,22 @@
 # contexts without spans
 
 
+_WAF_ADDRESSES = "waf_addresses"
+_CALLBACKS = "callbacks"
+_TELEMETRY = "telemetry"
+_CONTEXT_CALL = "context"
+_WAF_CALL = "waf_run"
+_BLOCK_CALL = "block"
+_WAF_RESULTS = "waf_results"
+
+
 class ASM_Environment:
+    """
+    an object of this class contains all asm data (waf and telemetry)
+    for a single request. It is bound to a single asm request context.
+    It is contained into a ContextVar.
+    """
+
     def __init__(self, active=False):  # type: (bool) -> None
         self.active = active
         self.span = None
@@ -60,6 +76,12 @@ def register(span):
 
 
 class _DataHandler:
+    """
+    An object of this class is created by each asm request context.
+    It handles the creation and destruction of ASM_Environment object.
+    It allows the ASM context to be reentrant.
+    """
+
     main_id = 0
 
     def __init__(self):
@@ -70,17 +92,18 @@ def __init__(self):
         self.active = True
         self.token = _ASM.set(env)
 
-        env.telemetry["waf_results"] = [], [], []
-        env.callbacks["context"] = []
+        env.telemetry[_WAF_RESULTS] = [], [], []
+        env.callbacks[_CONTEXT_CALL] = []
 
     def finalise(self):
         if self.active:
             env = _ASM.get()
             # assert _CONTEXT_ID.get() == self._id
-            callbacks = env.callbacks.get("context", [])
-            for function in callbacks:
-                function()
-            _ASM.reset(self.token)
+            callbacks = env.callbacks.get(_CONTEXT_CALL)
+            if callbacks is not None:
+                for function in callbacks:
+                    function()
+                _ASM.reset(self.token)
             self.active = False
 
 
@@ -94,6 +117,12 @@ def set_value(category, address, value):  # type: (str, str, Any) -> None
         asm_context_attr[address] = value
 
 
+def set_waf_address(address, value, span):  # type: (str, Any, Any) -> None
+    set_value(_WAF_ADDRESSES, address, value)
+    if span:
+        _context.set_item(address, value, span=span)
+
+
 def get_value(category, address, default=None):  # type: (str, str, Any) -> Any
     env = _ASM.get()
     if not env.active:
@@ -105,21 +134,25 @@ def get_value(category, address, default=None):  # type: (str, str, Any) -> Any
     return default
 
 
+def get_waf_address(address, default=None):  # type: (str, Any) -> Any
+    return get_value(_WAF_ADDRESSES, address, default=default)
+
+
 def add_context_callback(function):  # type: (Any) -> None
-    callbacks = get_value("callbacks", "context")
+    callbacks = get_value(_CALLBACKS, _CONTEXT_CALL)
     if callbacks is not None:
         callbacks.append(function)
 
 
 def set_waf_callback(value):  # type: (Any) -> None
-    set_value("callbacks", "waf_run", value)
+    set_value(_CALLBACKS, _WAF_CALL, value)
 
 
 def call_waf_callback(custom_data=None):
     # type: (dict[str, Any] | None) -> None
     if not config._appsec_enabled:
         return
-    callback = get_value("callbacks", "waf_run")
+    callback = get_value(_CALLBACKS, _WAF_CALL)
     if callback:
         return callback(custom_data)
     else:
@@ -128,11 +161,11 @@ def call_waf_callback(custom_data=None):
 
 def set_ip(ip):  # type: (Optional[str]) -> None
     if ip is not None:
-        set_value("waf_addresses", SPAN_DATA_NAMES.REQUEST_HTTP_IP, ip)
+        set_waf_address(SPAN_DATA_NAMES.REQUEST_HTTP_IP, ip, _ASM.get().span)
 
 
 def get_ip():  # type: () -> Optional[str]
-    return get_value("waf_addresses", SPAN_DATA_NAMES.REQUEST_HTTP_IP)
+    return get_value(_WAF_ADDRESSES, SPAN_DATA_NAMES.REQUEST_HTTP_IP)
 
 
 # Note: get/set headers use Any since we just carry the headers here without changing or using them
@@ -142,19 +175,19 @@ def get_ip():  # type: () -> Optional[str]
 
 def set_headers(headers):  # type: (Any) -> None
     if headers is not None:
-        set_value("waf_addresses", SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, headers)
+        set_waf_address(SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, headers, _ASM.get().span)
 
 
 def get_headers():  # type: () -> Optional[Any]
-    return get_value("waf_addresses", SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, {})
+    return get_value(_WAF_ADDRESSES, SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, {})
 
 
 def set_headers_case_sensitive(case_sensitive):  # type: (bool) -> None
-    set_value("waf_addresses", SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES_CASE, case_sensitive)
+    set_waf_address(SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES_CASE, case_sensitive, _ASM.get().span)
 
 
 def get_headers_case_sensitive():  # type: () -> bool
-    return get_value("waf_addresses", SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES_CASE, False)  # type : ignore
+    return get_value(_WAF_ADDRESSES, SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES_CASE, False)  # type : ignore
 
 
 def set_block_request_callable(_callable):  # type: (Optional[Callable]) -> None
@@ -164,14 +197,14 @@ def set_block_request_callable(_callable):  # type: (Optional[Callable]) -> None
     functools.partial.
     """
     if _callable:
-        set_value("callbacks", "block", _callable)
+        set_value(_CALLBACKS, _BLOCK_CALL, _callable)
 
 
 def block_request():  # type: () -> None
     """
     Calls or returns the stored block request callable, if set.
     """
-    _callable = get_value("callbacks", "block")
+    _callable = get_value(_CALLBACKS, _BLOCK_CALL)
     if _callable:
         _callable()
     else:
@@ -204,11 +237,11 @@ def set_waf_results(result_data, result_info, is_blocked):  # type: (Any, Any, b
 
 
 def get_waf_results():  # type: () -> Tuple[List[Any], List[Any], List[bool]] | None
-    return get_value("telemetry", "waf_results")
+    return get_value(_TELEMETRY, _WAF_RESULTS)
 
 
 def reset_waf_results():  # type: () -> None
-    set_value("telemetry", "waf_results", ([], [], []))
+    set_value(_TELEMETRY, _WAF_RESULTS, ([], [], []))
 
 
 @contextlib.contextmanager
@@ -216,6 +249,9 @@ def asm_request_context_manager(
     remote_ip=None, headers=None, headers_case_sensitive=False, block_request_callable=None
 ):
     # type: (Optional[str], Any, bool, Optional[Callable]) -> Generator[_DataHandler|None, None, None]
+    """
+    The ASM context manager
+    """
     if True or config._appsec_enabled:
         resources = _DataHandler()
         asm_request_context_set(remote_ip, headers, headers_case_sensitive, block_request_callable)
diff --git a/ddtrace/appsec/processor.py b/ddtrace/appsec/processor.py
--- a/ddtrace/appsec/processor.py
+++ b/ddtrace/appsec/processor.py
@@ -238,19 +238,16 @@ def waf_callable(custom_data=None):
         _asm_request_context.set_waf_callback(waf_callable)
         _asm_request_context.add_context_callback(_set_waf_request_metrics)
         if headers is not None:
-            _context.set_items(
-                {
-                    SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES: headers,
-                    SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES_CASE: headers_case_sensitive,
-                },
-                span=span,
+            _asm_request_context.set_waf_address(SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, headers, span)
+            _asm_request_context.set_waf_address(
+                SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES_CASE, headers_case_sensitive, span
             )
             if not peer_ip:
                 return
 
             ip = trace_utils._get_request_header_client_ip(headers, peer_ip, headers_case_sensitive)
             # Save the IP and headers in the context so the retrieval can be skipped later
-            _context.set_item("http.request.remote_ip", ip, span=span)
+            _asm_request_context.set_waf_address(SPAN_DATA_NAMES.REQUEST_HTTP_IP, ip, span)
             if ip and self._is_needed(WAF_DATA_NAMES.REQUEST_HTTP_IP):
                 log.debug("[DDAS-001-00] Executing ASM WAF for checking IP block")
                 # _asm_request_context.call_callback()
@@ -289,7 +286,7 @@ def _waf_action(self, span, ctx, custom_data=None):
                 if custom_data is not None and custom_data.get(key) is not None:
                     value = custom_data.get(key)
                 else:
-                    value = _context.get_item(SPAN_DATA_NAMES[key], span=span)
+                    value = _asm_request_context.get_value("waf_addresses", SPAN_DATA_NAMES[key])
 
                 if value:
                     data[waf_name] = _transform_headers(value) if key.endswith("HEADERS_NO_COOKIES") else value
@@ -344,7 +341,7 @@ def update_metric(name, value):
                 (SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, "request"),
                 (SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES, "response"),
             ]:
-                headers_req = _context.get_item(id_tag, span=span)
+                headers_req = _asm_request_context.get_waf_address(id_tag)
                 if headers_req:
                     _set_headers(span, headers_req, kind=kind)
 
@@ -357,7 +354,7 @@ def update_metric(name, value):
             # Partial DDAS-011-00
             span.set_tag_str(APPSEC.EVENT, "true")
 
-            remote_ip = _context.get_item(SPAN_DATA_NAMES.REQUEST_HTTP_IP, span=span)
+            remote_ip = _asm_request_context.get_waf_address(SPAN_DATA_NAMES.REQUEST_HTTP_IP)
             if remote_ip:
                 # Note that if the ip collection is disabled by the env var
                 # DD_TRACE_CLIENT_IP_HEADER_DISABLED actor.ip won't be sent
diff --git a/ddtrace/contrib/trace_utils.py b/ddtrace/contrib/trace_utils.py
--- a/ddtrace/contrib/trace_utils.py
+++ b/ddtrace/contrib/trace_utils.py
@@ -19,6 +19,7 @@
 
 from ddtrace import Pin
 from ddtrace import config
+from ddtrace.ext import SpanTypes
 from ddtrace.ext import http
 from ddtrace.ext import user
 from ddtrace.internal import _context
@@ -507,28 +508,30 @@ def set_http_meta(
     if retries_remain is not None:
         span.set_tag_str(http.RETRIES_REMAIN, str(retries_remain))
 
-    if config._appsec_enabled:
+    if span.span_type == SpanTypes.WEB and config._appsec_enabled:
+        from ddtrace.appsec._asm_request_context import set_waf_address
+        from ddtrace.appsec._constants import SPAN_DATA_NAMES
+
         status_code = str(status_code) if status_code is not None else None
 
-        _context.set_items(
-            {
-                k: v
-                for k, v in [
-                    ("http.request.uri", raw_uri),
-                    ("http.request.method", method),
-                    ("http.request.cookies", request_cookies),
-                    ("http.request.query", parsed_query),
-                    ("http.request.headers", request_headers),
-                    ("http.response.headers", response_headers),
-                    ("http.response.status", status_code),
-                    ("http.request.path_params", request_path_params),
-                    ("http.request.body", request_body),
-                    ("http.request.remote_ip", request_ip),
-                ]
-                if v is not None
-            },
-            span=span,
-        )
+        addresses = {
+            k: v
+            for k, v in [
+                (SPAN_DATA_NAMES.REQUEST_URI_RAW, raw_uri),
+                (SPAN_DATA_NAMES.REQUEST_METHOD, method),
+                (SPAN_DATA_NAMES.REQUEST_COOKIES, request_cookies),
+                (SPAN_DATA_NAMES.REQUEST_QUERY, parsed_query),
+                (SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, request_headers),
+                (SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES, response_headers),
+                (SPAN_DATA_NAMES.RESPONSE_STATUS, status_code),
+                (SPAN_DATA_NAMES.REQUEST_PATH_PARAMS, request_path_params),
+                (SPAN_DATA_NAMES.REQUEST_BODY, request_body),
+                (SPAN_DATA_NAMES.REQUEST_HTTP_IP, request_ip),
+            ]
+            if v is not None
+        }
+        for k, v in addresses.items():
+            set_waf_address(k, v, span)
 
     if route is not None:
         span.set_tag_str(http.ROUTE, route)


Test patch:
diff --git a/tests/tracer/test_trace_utils.py b/tests/tracer/test_trace_utils.py
--- a/tests/tracer/test_trace_utils.py
+++ b/tests/tracer/test_trace_utils.py
@@ -21,6 +21,7 @@
 from ddtrace.context import Context
 from ddtrace.contrib import trace_utils
 from ddtrace.contrib.trace_utils import _get_request_header_client_ip
+from ddtrace.ext import SpanTypes
 from ddtrace.ext import http
 from ddtrace.internal import _context
 from ddtrace.internal.compat import six
@@ -306,6 +307,7 @@ def test_ext_service(int_config, pin, config_val, default, expected):
 
 
 @pytest.mark.parametrize("appsec_enabled", [False, True])
+@pytest.mark.parametrize("span_type", [SpanTypes.WEB, SpanTypes.HTTP, None])
 @pytest.mark.parametrize(
     "method,url,status_code,status_msg,query,request_headers,response_headers,uri,path_params,cookies",
     [
@@ -357,9 +359,11 @@ def test_set_http_meta(
     path_params,
     cookies,
     appsec_enabled,
+    span_type,
 ):
     int_config.http.trace_headers(["my-header"])
     int_config.trace_query_string = True
+    span.span_type = span_type
     with override_global_config({"_appsec_enabled": appsec_enabled}):
         trace_utils.set_http_meta(
             span,
@@ -414,7 +418,7 @@ def test_set_http_meta(
             tag = "http.request.headers." + header
             assert span.get_tag(tag) == value
 
-    if appsec_enabled:
+    if appsec_enabled and span.span_type == SpanTypes.WEB:
         if uri is not None:
             assert _context.get_item("http.request.uri", span=span) == uri
         if method is not None: