feat: add DD_TRACE_API_VERSION environment variable (#2918)

* feat: add ``DD_TRACE_ENCODING`` environment variable

The new variable can be used to activate the v0.5 trace encoding on
demand.

* rename env var

* better docs and Tracer.configure API extension

Author: P403n1x87

benchmarks/encoder/config.yaml

@@ -5,6 +5,7 @@ one-trace: &base_variant
   ltags: 0
   nmetrics: 0
   dd_origin: false
+  encoding: "v0.4"
 many-traces:
   <<: *base_variant
   ntraces: 100

benchmarks/encoder/scenario.py

@@ -9,9 +9,10 @@ class Encoder(bm.Scenario):
     ltags = bm.var(type=int)
     nmetrics = bm.var(type=int)
     dd_origin = bm.var(type=bool)
+    encoding = bm.var(type=str)
 
     def run(self):
-        encoder = utils.init_encoder()
+        encoder = utils.init_encoder(self.encoding)
         traces = utils.gen_traces(self)
 
         def _(loops):

benchmarks/encoder/utils.py

@@ -1,7 +1,7 @@
 import random
 import string
 
-from ddtrace.internal.encoding import Encoder
+from ddtrace.internal.encoding import MSGPACK_ENCODERS
 from ddtrace.span import Span
 
 
@@ -10,14 +10,14 @@
     # see https://github.com/DataDog/dd-trace-py/pull/2422
     from ddtrace.internal._encoding import BufferedEncoder  # noqa: F401
 
-    def init_encoder(max_size=8 << 20, max_item_size=8 << 20):
-        return Encoder(max_size, max_item_size)
+    def init_encoder(encoding, max_size=8 << 20, max_item_size=8 << 20):
+        return MSGPACK_ENCODERS[encoding](max_size, max_item_size)
 
 
 except ImportError:
 
-    def init_encoder():
-        return Encoder()
+    def init_encoder(encoding):
+        return MSGPACK_ENCODERS[encoding]()
 
 
 def _rands(size=6, chars=string.ascii_uppercase + string.digits):

ddtrace/encoding.py

@@ -1,10 +1,12 @@
-from .internal.encoding import Encoder
 from .internal.encoding import JSONEncoder
 from .internal.encoding import JSONEncoderV2
 from .internal.encoding import MsgpackEncoderV03 as MsgpackEncoder
 from .utils.deprecation import deprecation
 
 
+Encoder = MsgpackEncoder
+
+
 __all__ = (
     "Encoder",
     "JSONEncoder",

ddtrace/internal/encoding.py

@@ -11,7 +11,7 @@
 from .logger import get_logger
 
 
-__all__ = ["MsgpackEncoderV03", "MsgpackEncoderV05", "ListStringTable", "Encoder"]
+__all__ = ["MsgpackEncoderV03", "MsgpackEncoderV05", "ListStringTable", "MSGPACK_ENCODERS"]
 
 
 if TYPE_CHECKING:
@@ -30,27 +30,31 @@ def encode_traces(self, traces):
         # type: (List[List[Span]]) -> str
         """
         Encodes a list of traces, expecting a list of items where each items
-        is a list of spans. Before dump the string in a serialized format all
-        traces are normalized, calling the ``to_dict()`` method. The traces
+        is a list of spans. Before dumping the string in a serialized format all
+        traces are normalized according to the encoding format. The trace
         nesting is not changed.
 
         :param traces: A list of traces that should be serialized
         """
-        normalized_traces = [[span.to_dict() for span in trace] for trace in traces]
-        return self.encode(normalized_traces)
+        raise NotImplementedError()
 
-    @staticmethod
-    def encode(obj):
+    def encode(self, obj):
+        # type: (List[List[Any]]) -> str
         """
         Defines the underlying format used during traces or services encoding.
-        This method must be implemented and should only be used by the internal functions.
+        This method must be implemented and should only be used by the internal
+        functions.
         """
-        raise NotImplementedError
+        raise NotImplementedError()
 
 
 class JSONEncoder(_EncoderBase):
     content_type = "application/json"
 
+    def encode_traces(self, traces):
+        normalized_traces = [[span.to_dict() for span in trace] for trace in traces]
+        return self.encode(normalized_traces)
+
     @staticmethod
     def encode(obj):
         # type: (Any) -> str
@@ -93,4 +97,8 @@ def _decode_id_to_hex(hex_id):
         return int(hex_id, 16)
 
 
-Encoder = MsgpackEncoderV03
+MSGPACK_ENCODERS = {
+    "v0.3": MsgpackEncoderV03,
+    "v0.4": MsgpackEncoderV03,
+    "v0.5": MsgpackEncoderV05,
+}

ddtrace/internal/writer.py

@@ -28,8 +28,8 @@
 from ._encoding import BufferFull
 from ._encoding import BufferItemTooLarge
 from .agent import get_connection
-from .encoding import Encoder
 from .encoding import JSONEncoderV2
+from .encoding import MSGPACK_ENCODERS
 from .logger import get_logger
 from .runtime import container
 from .sma import SimpleMovingAverage
@@ -234,6 +234,7 @@ def __init__(
         dogstatsd=None,  # type: Optional[DogStatsd]
         report_metrics=False,  # type: bool
         sync_mode=False,  # type: bool
+        api_version=None,  # type: Optional[str]
     ):
         # type: (...) -> None
         # Pre-conditions:
@@ -256,10 +257,18 @@ def __init__(
         }
         self._timeout = timeout
 
-        if priority_sampler is not None:
-            self._endpoint = "v0.4/traces"
-        else:
-            self._endpoint = "v0.3/traces"
+        encoding_version = (
+            api_version or os.getenv("DD_TRACE_API_VERSION") or ("v0.4" if priority_sampler is not None else "v0.3")
+        )
+        try:
+            Encoder = MSGPACK_ENCODERS[encoding_version]
+        except KeyError:
+            raise ValueError(
+                "Unsupported encoding version: '%s'. The supported versions are: %r"
+                % (encoding_version, ", ".join(sorted(MSGPACK_ENCODERS.keys())))
+            )
+
+        self._endpoint = "%s/traces" % encoding_version
 
         self._container_info = container.get_container_info()
         if self._container_info and self._container_info.container_id:
@@ -348,10 +357,27 @@ def _put(self, data, headers):
                 conn.close()
 
     def _downgrade(self, payload, response):
+        if self._endpoint == "v0.5/traces":
+            self._endpoint = "v0.4/traces"
+            self._encoder = MSGPACK_ENCODERS["v0.4"](
+                max_size=self._buffer_size,
+                max_item_size=self._max_payload_size,
+            )
+            # Since we have to change the encoding in this case, the payload
+            # would need to be converted to the downgraded encoding before
+            # sending it, but we chuck it away instead.
+            log.warning(
+                "Dropping trace payload due to the downgrade to an incompatible API version (from v0.5 to v0.4). To "
+                "avoid this from happening in the future, either ensure that the Datadog agent has a v0.5/traces "
+                "endpoint available, or explicitly set the trace API version to, e.g., v0.4."
+            )
+            return None
         if self._endpoint == "v0.4/traces":
             self._endpoint = "v0.3/traces"
+            # These endpoints share the same encoding, so we can try sending the
+            # same payload over the downgraded endpoint.
             return payload
-        raise ValueError
+        raise ValueError()
 
     def _send_payload(self, payload, count):
         headers = self._headers.copy()
@@ -378,7 +404,8 @@ def _send_payload(self, payload, count):
                     self.agent_url,
                 )
             else:
-                return self._send_payload(payload, count)
+                if payload is not None:
+                    self._send_payload(payload, count)
         elif response.status >= 400:
             log.error(
                 "failed to send traces to Datadog Agent at %s: HTTP error status %s, reason %s",

ddtrace/tracer.py

@@ -297,6 +297,7 @@ def configure(
         writer=None,  # type: Optional[TraceWriter]
         partial_flush_enabled=None,  # type: Optional[bool]
         partial_flush_min_spans=None,  # type: Optional[int]
+        api_version=None,  # type: Optional[str]
     ):
         # type: (...) -> None
         """
@@ -393,6 +394,7 @@ def configure(
                 dogstatsd=get_dogstatsd_client(self._dogstatsd_url),
                 report_metrics=config.health_metrics_enabled,
                 sync_mode=self._use_sync_mode(),
+                api_version=api_version,
             )
         elif writer is None and isinstance(self.writer, LogWriter):
             # No need to do anything for the LogWriter.

docs/configuration.rst

@@ -170,6 +170,13 @@ below:
      - A map of case-insensitive header keys to tag names. Automatically applies matching header values as tags on root spans.
        For example, ``User-Agent:http.useragent,content-type:http.content_type``.
 
+       .. _dd-trace-api-version:
+   * - ``DD_TRACE_API_VERSION``
+     - String
+     - ``v0.4`` if priority sampling is enabled, else ``v0.3``
+     - The trace API version to use when sending traces to the Datadog agent.
+       Currently, the supported versions are: ``v0.3``, ``v0.4`` and ``v0.5``.
+
        .. _dd-profiling-enabled:
    * - ``DD_PROFILING_ENABLED``
      - Boolean

releasenotes/notes/feat-encoding-env-var-59eedcbbcf065204.yaml

@@ -0,0 +1,11 @@
+---
+prelude: >
+    The ``v0.5`` trace API version has been made available and can be selected
+    via the newly introduced ``DD_TRACE_API_VERSION`` environment variable by
+    setting ``DD_TRACE_API_VERSION="v0.5"``. This newer API version generates
+    smaller payloads, thus increasing the trace throughput to the Datadog agent,
+    for better performance.
+features:
+  - |
+    Added the ``DD_TRACE_API_VERSION`` environment variable for requesting a
+    specific trace API version.

scripts/profiles/encoders/run.py

@@ -1,20 +1,10 @@
 import sys
 
-from ddtrace.internal.encoding import Encoder
-from ddtrace.internal.encoding import MsgpackEncoderV03
-from ddtrace.internal.encoding import MsgpackEncoderV05
+from ddtrace.internal.encoding import MSGPACK_ENCODERS as ENCODERS
 from tests.tracer.test_encoders import gen_trace
 
 
-ENCODERS = {
-    "v0.3": MsgpackEncoderV03,
-    "v0.5": MsgpackEncoderV05,
-}
-
-try:
-    encoder = ENCODERS.get(sys.argv[1], Encoder)(8 << 20, 8 << 20)
-except IndexError:
-    encoder = Encoder(8 << 20, 8 << 20)
+encoder = ENCODERS[sys.argv[1]](8 << 20, 8 << 20)
 
 
 trace = gen_trace(nspans=1000)