fix(asm): simplify asm request context span registration [backport #5822 to 1.13] (#5826)

Backport of #5822 to 1.13

Remove the use of span context._meta to store information about context
registration.

update tests for asm request context to unit test the new registration
process.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] OPTIONAL: PR description includes explicit acknowledgement of the
performance implications of the change as reported in the benchmarks PR
comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

---------

Co-authored-by: Christophe Papazian <[email protected]>
Co-authored-by: Munir Abdinur <[email protected]>

Author: 3 people

ddtrace/appsec/_asm_request_context.py

@@ -53,6 +53,7 @@ class ASM_Environment:
     def __init__(self, active=False):  # type: (bool) -> None
         self.active = active
         self.span = None
+        self.span_asm_context = None  # type: None | contextlib.AbstractContextManager
         self.waf_addresses = {}  # type: dict[str, Any]
         self.callbacks = {}  # type: dict[str, Any]
         self.telemetry = {}  # type: dict[str, Any]
@@ -82,12 +83,19 @@ def is_blocked():  # type: () -> bool
         return False
 
 
-def register(span):
+def register(span, span_asm_context=None):
     env = _ASM.get()
     if not env.active:
         log.debug("registering a span with no active asm context")
         return
     env.span = span
+    env.span_asm_context = span_asm_context
+
+
+def unregister(span):
+    env = _ASM.get()
+    if env.span_asm_context is not None and env.span is span:
+        env.span_asm_context.__exit__(None, None, None)
 
 
 class _DataHandler:

ddtrace/appsec/processor.py

@@ -220,8 +220,7 @@ def on_span_start(self, span):
         else:
             new_asm_context = _asm_request_context.asm_request_context_manager()
             new_asm_context.__enter__()
-            span.context._meta["ASM_CONTEXT_%d" % id(span)] = new_asm_context  # type: ignore
-            _asm_request_context.register(span)
+            _asm_request_context.register(span, new_asm_context)
 
         ctx = self._ddwaf._at_request_start()
 
@@ -376,7 +375,6 @@ def _is_needed(self, address):
 
     def on_span_finish(self, span):
         # type: (Span) -> None
-        asm_context = span.context._meta.get("ASM_CONTEXT_%d" % id(span), None)
         try:
             if span.span_type == SpanTypes.WEB:
                 # Force to set respond headers at the end
@@ -392,6 +390,4 @@ def on_span_finish(self, span):
                 self._ddwaf._at_request_end()
         finally:
             # release asm context if it was created by the span
-            if asm_context is not None:
-                asm_context.__exit__(None, None, None)  # type: ignore
-                del span.context._meta["ASM_CONTEXT_%d" % id(span)]
+            _asm_request_context.unregister(span)

releasenotes/notes/fix-asm-request-context-partial-flushing-7c45aaf55e836261.yaml

@@ -0,0 +1,3 @@
+fixes:
+  - |
+    ASM: Fixes encoding error when using AppSec and a trace is partial flushed.

tests/appsec/test_processor.py

@@ -654,16 +654,16 @@ def test_ddwaf_run_contained_oserror(tracer_appsec, caplog):
     assert "OSError: ddwaf run failed" in caplog.text
 
 
-def test_asm_context_meta(tracer_appsec):
+def test_asm_context_registration(tracer_appsec):
     tracer = tracer_appsec
 
     # For a web type span, a context manager is added, but then removed
     with tracer.trace("test", span_type=SpanTypes.WEB) as span:
-        assert span.context._meta["ASM_CONTEXT_%d" % (id(span),)]
-    assert span.context._meta.get("ASM_CONTEXT_%d" % (id(span),)) is None
+        assert _asm_request_context._ASM.get().span_asm_context
+    assert _asm_request_context._ASM.get().span_asm_context is None
 
     # Regression test, if the span type changes after being created, we always removed
     with tracer.trace("test", span_type=SpanTypes.WEB) as span:
         span.span_type = SpanTypes.HTTP
-        assert span.context._meta["ASM_CONTEXT_%d" % (id(span),)]
-    assert span.context._meta.get("ASM_CONTEXT_%d" % (id(span),)) is None
+        assert _asm_request_context._ASM.get().span_asm_context
+    assert _asm_request_context._ASM.get().span_asm_context is None