fix(sampling): revert "refactor(writer): handle agent responses in the tracer (#6178)" [backport 2.0] (#7337)

Backport 31ef044 from #7333 to 2.0.

This reverts commit ec3e14d. Which
caused the tracer to ignore trace-agent sampling rates, often times
leading to much higher ingestion rates.

I tested this manually by looking at
`datadog.trace_agent.receiver.traces_priority` between <1.19 and >=1.19.

<1.19:  some P1 and a lot of P0 

>=1.19: a lot more P1 (and  a little more  P-1) and no P0

Basically before 1.19 there were a ton of traces getting dropped with P0
priority due to the trace-agent recommended rates. After 1.19 we ignored
the trace-agent recommended rates, leading to no P0 traces, and far more
being sampled overall.

## 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/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## 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.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

Co-authored-by: Zachary Groves <[email protected]>

Author: github-actions[bot]

ddtrace/internal/ci_visibility/writer.py

@@ -29,6 +29,8 @@
 
     from ddtrace.vendor.dogstatsd import DogStatsd
 
+    from ...sampler import BaseSampler
+
 
 class CIVisibilityEventClient(WriterClientBase):
     def __init__(self):
@@ -79,6 +81,7 @@ class CIVisibilityWriter(HTTPWriter):
     def __init__(
         self,
         intake_url="",  # type: str
+        sampler=None,  # type: Optional[BaseSampler]
         processing_interval=None,  # type: Optional[float]
         timeout=None,  # type: Optional[float]
         dogstatsd=None,  # type: Optional[DogStatsd]
@@ -126,6 +129,7 @@ def __init__(
         super(CIVisibilityWriter, self).__init__(
             intake_url=intake_url,
             clients=clients,
+            sampler=sampler,
             processing_interval=processing_interval,
             timeout=timeout,
             dogstatsd=dogstatsd,
@@ -142,6 +146,7 @@ def recreate(self):
         # type: () -> HTTPWriter
         return self.__class__(
             intake_url=self.intake_url,
+            sampler=self._sampler,
             processing_interval=self._interval,
             timeout=self._timeout,
             dogstatsd=self.dogstatsd,

ddtrace/internal/writer/__init__.py

@@ -1,4 +1,3 @@
-from .writer import AgentResponse  # noqa
 from .writer import AgentWriter  # noqa
 from .writer import DEFAULT_SMA_WINDOW  # noqa
 from .writer import HTTPWriter  # noqa

ddtrace/internal/writer/writer.py

@@ -5,8 +5,6 @@
 import os
 import sys
 import threading
-from typing import Any
-from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -28,6 +26,8 @@
 from ...internal.utils.formats import parse_tags_str
 from ...internal.utils.http import Response
 from ...internal.utils.time import StopWatch
+from ...sampler import BasePrioritySampler
+from ...sampler import BaseSampler
 from .._encoding import BufferFull
 from .._encoding import BufferItemTooLarge
 from ..agent import get_connection
@@ -103,8 +103,10 @@ class LogWriter(TraceWriter):
     def __init__(
         self,
         out=sys.stdout,  # type: TextIO
+        sampler=None,  # type: Optional[BaseSampler]
     ):
         # type: (...) -> None
+        self._sampler = sampler
         self.encoder = JSONEncoderV2()
         self.out = out
 
@@ -115,7 +117,7 @@ def recreate(self):
         :rtype: :class:`LogWriter`
         :returns: A new :class:`LogWriter` instance
         """
-        writer = self.__class__(out=self.out)
+        writer = self.__class__(out=self.out, sampler=self._sampler)
         return writer
 
     def stop(self, timeout=None):
@@ -147,6 +149,7 @@ def __init__(
         self,
         intake_url,  # type: str
         clients,  # type: List[WriterClientBase]
+        sampler=None,  # type: Optional[BaseSampler]
         processing_interval=None,  # type: Optional[float]
         # Match the payload size since there is no functionality
         # to flush dynamically.
@@ -168,6 +171,7 @@ def __init__(
         self.intake_url = intake_url
         self._buffer_size = buffer_size
         self._max_payload_size = max_payload_size
+        self._sampler = sampler
         self._headers = headers or {}
         self._timeout = timeout
 
@@ -295,7 +299,6 @@ def _get_finalized_headers(self, count, client):
         return headers
 
     def _send_payload(self, payload, count, client):
-        # type: (...) -> Response
         headers = self._get_finalized_headers(count, client)
 
         self._metrics_dist("http.requests")
@@ -313,15 +316,15 @@ def _send_payload(self, payload, count, client):
                 self._intake_endpoint(client),
                 response.status,
                 response.reason,
-            )  # type: Tuple[Any, Any, Any]
+            )
             # Append the payload if requested
             if config._trace_writer_log_err_payload:
                 msg += ", payload %s"
                 # If the payload is bytes then hex encode the value before logging
                 if isinstance(payload, six.binary_type):
-                    log_args += (binascii.hexlify(payload).decode(),)  # type: ignore
+                    log_args += (binascii.hexlify(payload).decode(),)
                 else:
-                    log_args += (payload,)  # type: ignore
+                    log_args += (payload,)
 
             log.error(msg, *log_args)
             self._metrics_dist("http.dropped.bytes", len(payload))
@@ -447,12 +450,6 @@ def on_shutdown(self):
             self._reset_connection()
 
 
-class AgentResponse(object):
-    def __init__(self, rate_by_service):
-        # type: (Dict[str, float]) -> None
-        self.rate_by_service = rate_by_service
-
-
 class AgentWriter(HTTPWriter):
     """
     The Datadog Agent supports (at the time of writing this) receiving trace
@@ -468,6 +465,7 @@ class AgentWriter(HTTPWriter):
     def __init__(
         self,
         agent_url,  # type: str
+        sampler=None,  # type: Optional[BaseSampler]
         priority_sampling=False,  # type: bool
         processing_interval=None,  # type: Optional[float]
         # Match the payload size since there is no functionality
@@ -481,7 +479,6 @@ def __init__(
         api_version=None,  # type: Optional[str]
         reuse_connections=None,  # type: Optional[bool]
         headers=None,  # type: Optional[Dict[str, str]]
-        response_callback=None,  # type: Optional[Callable[[AgentResponse], None]]
     ):
         # type: (...) -> None
         if processing_interval is None:
@@ -537,10 +534,10 @@ def __init__(
         additional_header_str = os.environ.get("_DD_TRACE_WRITER_ADDITIONAL_HEADERS")
         if additional_header_str is not None:
             _headers.update(parse_tags_str(additional_header_str))
-        self._response_cb = response_callback
         super(AgentWriter, self).__init__(
             intake_url=agent_url,
             clients=[client],
+            sampler=sampler,
             processing_interval=processing_interval,
             buffer_size=buffer_size,
             max_payload_size=max_payload_size,
@@ -555,6 +552,7 @@ def recreate(self):
         # type: () -> HTTPWriter
         return self.__class__(
             agent_url=self.agent_url,
+            sampler=self._sampler,
             processing_interval=self._interval,
             buffer_size=self._buffer_size,
             max_payload_size=self._max_payload_size,
@@ -592,7 +590,6 @@ def _downgrade(self, payload, response, client):
         raise ValueError()
 
     def _send_payload(self, payload, count, client):
-        # type: (...) -> Response
         response = super(AgentWriter, self)._send_payload(payload, count, client)
         if response.status in [404, 415]:
             log.debug("calling endpoint '%s' but received %s; downgrading API", client.ENDPOINT, response.status)
@@ -608,15 +605,16 @@ def _send_payload(self, payload, count, client):
             else:
                 if payload is not None:
                     self._send_payload(payload, count, client)
-        elif response.status < 400:
-            if self._response_cb:
-                raw_resp = response.get_json()
-                if raw_resp and "rate_by_service" in raw_resp:
-                    self._response_cb(
-                        AgentResponse(
-                            rate_by_service=raw_resp["rate_by_service"],
+        elif response.status < 400 and isinstance(self._sampler, BasePrioritySampler):
+            result_traces_json = response.get_json()
+            if result_traces_json and "rate_by_service" in result_traces_json:
+                try:
+                    if isinstance(self._sampler, BasePrioritySampler):
+                        self._sampler.update_rate_by_service_sample_rates(
+                            result_traces_json["rate_by_service"],
                         )
-                    )
+                except ValueError:
+                    log.error("sample_rate is negative, cannot update the rate samplers")
         return response
 
     def start(self):

ddtrace/tracer.py

@@ -52,12 +52,10 @@
 from .internal.serverless.mini_agent import maybe_start_serverless_mini_agent
 from .internal.service import ServiceStatusError
 from .internal.utils.http import verify_url
-from .internal.writer import AgentResponse
 from .internal.writer import AgentWriter
 from .internal.writer import LogWriter
 from .internal.writer import TraceWriter
 from .provider import DefaultContextProvider
-from .sampler import BasePrioritySampler
 from .sampler import BaseSampler
 from .sampler import DatadogSampler
 from .sampler import RateSampler
@@ -240,6 +238,7 @@ def __init__(
         else:
             writer = AgentWriter(
                 agent_url=self._agent_url,
+                sampler=self._sampler,
                 priority_sampling=config._priority_sampling,
                 dogstatsd=get_dogstatsd_client(self._dogstatsd_url),
                 sync_mode=self._use_sync_mode(),
@@ -458,12 +457,12 @@ def configure(
         elif any(x is not None for x in [new_url, api_version, sampler, dogstatsd_url]):
             self._writer = AgentWriter(
                 self._agent_url,
+                sampler=self._sampler,
                 priority_sampling=priority_sampling in (None, True) or config._priority_sampling,
                 dogstatsd=get_dogstatsd_client(self._dogstatsd_url),
                 sync_mode=self._use_sync_mode(),
                 api_version=api_version,
                 headers={"Datadog-Client-Computed-Stats": "yes"} if compute_stats_enabled else {},
-                response_callback=self._agent_response_callback,
             )
         elif writer is None and isinstance(self._writer, LogWriter):
             # No need to do anything for the LogWriter.
@@ -511,20 +510,6 @@ def configure(
 
         self._generate_diagnostic_logs()
 
-    def _agent_response_callback(self, resp):
-        # type: (AgentResponse) -> None
-        """Handle the response from the agent.
-
-        The agent can return updated sample rates for the priority sampler.
-        """
-        try:
-            if isinstance(self._sampler, BasePrioritySampler):
-                self._sampler.update_rate_by_service_sample_rates(
-                    resp.rate_by_service,
-                )
-        except ValueError:
-            log.error("sample_rate is negative, cannot update the rate samplers")
-
     def _generate_diagnostic_logs(self):
         if config._debug_mode or config._startup_logs_enabled:
             try:

releasenotes/notes/fix-make-tracer-acknowledge-agent-sampling-decisions-2919150951ed1b81.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    sampling: This fix reverts a refactor which affected how the tracer handled the trace-agent's
+    recommended trace sampling rates, leading to an unintended increase in traces sampled.