move span aggregator implementation details to reset method

Author: mabdinur

ddtrace/_trace/processor/__init__.py

@@ -493,23 +493,48 @@ def _queue_span_count_metrics(self, metric_name: str, tag_name: str, min_count:
                 )
             self._span_metrics[metric_name] = defaultdict(int)
 
-    def _reset(self):
+    def _reset(
+        self,
+        user_processors: Optional[List[TraceProcessor]] = None,
+        compute_stats: Optional[bool] = None,
+        apm_opt_out: Optional[bool] = None,
+        appsec_enabled: Optional[bool] = None,
+        reset_buffer: bool = True,
+    ) -> None:
         try:
             self.writer.stop()
         except ServiceStatusError:
             # Some writers (ex: AgentWriter), start when the first trace chunk is encoded. Stopping
             # the writer before that point will raise a ServiceStatusError.
             pass
+
+        if isinstance(self.writer, AgentWriter) and appsec_enabled:
+            # If appsec is enabled, we need to reset the API version to v0.4
+            # to ensure that the appsec meta_struct data is encoded.
+            self.writer._api_version = "v0.4"
         # Re-create the background writer thread
         self.writer = self.writer.recreate()
+        # Re-create the sampling processor with the compute_stats and apm_opt_out values
+        # Avoids overriding the agent sampling rules
+        if compute_stats is None:
+            compute_stats = self.sampling_processor._compute_stats_enabled
+        if apm_opt_out is None:
+            apm_opt_out = self.sampling_processor.apm_opt_out
         self.sampling_processor = TraceSamplingProcessor(
-            config._trace_compute_stats,
+            compute_stats,
             get_span_sampling_rules(),
-            asm_config._apm_opt_out,
+            apm_opt_out,
             self.sampling_processor.sampler._agent_based_samplers,
         )
-        self._traces = defaultdict(lambda: _Trace())
-        self._span_metrics = {
-            "spans_created": defaultdict(int),
-            "spans_finished": defaultdict(int),
-        }
+        # Override previously set user processors
+        if user_processors is not None:
+            self.user_processors = user_processors
+        # Reset the trace buffer and span metrics. This is useful when a process forks and we would like to
+        # avoid sending traces from both the parent and child processes.
+        if reset_buffer:
+            self._traces = defaultdict(lambda: _Trace())
+            self._span_metrics = {
+                "spans_created": defaultdict(int),
+                "spans_finished": defaultdict(int),
+            }
+            self._lock = RLock()

ddtrace/_trace/tracer.py

@@ -368,13 +368,6 @@ def configure(
         if compute_stats_enabled is not None:
             config._trace_compute_stats = compute_stats_enabled
 
-        if isinstance(self._span_aggregator.writer, AgentWriter):
-            if appsec_enabled:
-                self._span_aggregator.writer._api_version = "v0.4"
-
-        if trace_processors:
-            self._span_aggregator.user_processors = trace_processors
-
         if any(
             x is not None
             for x in [
@@ -384,7 +377,7 @@ def configure(
                 iast_enabled,
             ]
         ):
-            self._recreate()
+            self._recreate(trace_processors, compute_stats_enabled, appsec_enabled, reset_buffer=False)
 
         if context_provider is not None:
             self.context_provider = context_provider
@@ -412,15 +405,26 @@ def _generate_diagnostic_logs(self):
 
     def _child_after_fork(self):
         self._pid = getpid()
-        self._recreate()
+        self._recreate(reset_buffer=True)
         self._new_process = True
 
-    def _recreate(self):
+    def _recreate(
+        self,
+        trace_processors: Optional[List[TraceProcessor]] = None,
+        compute_stats_enabled: Optional[bool] = None,
+        appsec_enabled: Optional[bool] = None,
+        reset_buffer: bool = True,
+    ) -> None:
         """Re-initialize the tracer's processors and trace writer"""
         # Stop the writer.
         # This will stop the periodic thread in HTTPWriters, preventing memory leaks and unnecessary I/O.
         self.enabled = config._tracing_enabled
-        self._span_aggregator._reset()
+        self._span_aggregator._reset(
+            user_processors=trace_processors,
+            compute_stats=compute_stats_enabled,
+            appsec_enabled=appsec_enabled,
+            reset_buffer=reset_buffer,
+        )
         self._span_processors, self._appsec_processor = _default_span_processors_factory(
             self._endpoint_call_counter_span_processor,
         )

tests/tracer/test_processors.py

@@ -24,6 +24,7 @@
 from ddtrace.internal.sampling import SamplingMechanism
 from ddtrace.internal.sampling import SpanSamplingRule
 from ddtrace.internal.telemetry.constants import TELEMETRY_NAMESPACE
+from ddtrace.internal.writer import AgentWriter
 from ddtrace.trace import Context
 from ddtrace.trace import Span
 from tests.utils import DummyTracer
@@ -116,7 +117,9 @@ def process_trace(self, trace):
 
 
 def test_aggregator_reset():
-    """Test that the aggregator can reset trace buffers, sampling processor and trace writer"""
+    """Test that on reset, the aggregator clears user processors,
+    recreates the sampling processor, trace writer, and trace buffer
+    """
 
     class DDProc(TraceProcessor):
         def process_trace(self, trace):
@@ -155,6 +158,38 @@ def process_trace(self, trace):
     assert len(aggr._span_metrics["spans_created"]) == 0
 
 
+def test_aggregator_reset_with_args():
+    """Test that the span aggregator can reset trace buffers, sampling processor and trace api version"""
+
+    class DDProc(TraceProcessor):
+        def process_trace(self, trace):
+            return trace
+
+    class UserProc(TraceProcessor):
+        def process_trace(self, trace):
+            return trace
+
+    dd_proc = DDProc()
+    user_proc = UserProc()
+    aggr = SpanAggregator(
+        partial_flush_enabled=False,
+        partial_flush_min_spans=1,
+        dd_processors=[dd_proc],
+        user_processors=[user_proc],
+    )
+
+    aggr.writer = AgentWriter("", api_version="v0.5")
+    span = Span("span", on_finish=[aggr.on_span_finish])
+    aggr.on_span_start(span)
+    aggr._reset(user_processors=[], apm_opt_out=False, compute_stats=False, appsec_enabled=True, reset_buffer=False)
+    assert aggr.user_processors == []
+    assert aggr.sampling_processor.apm_opt_out is False
+    assert aggr.sampling_processor._compute_stats_enabled is False
+    assert aggr.writer._api_version == "v0.4"
+    assert span.trace_id in aggr._traces
+    assert len(aggr._span_metrics["spans_created"]) == 1
+
+
 def test_aggregator_bad_processor():
     class Proc(TraceProcessor):
         def process_trace(self, trace):