Add internal setting for additional writer headers (#2689)

For testing it's desirable to be able to pass additional headers with
the trace requests made to the agent. With this ability we can propagate
tokens which can be used to associate traces with test cases when
running snapshot tests in subprocesses (like in #2640).

I debated having the flag be read at run-time instead of
initialization-time but opted against it as reading from the env is not
cheap (on the order of 1µs):

    >>> %timeit os.environ.get("")
    1.23 µs ± 49.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: Kyle-Verhoog

ddtrace/internal/writer.py

@@ -2,6 +2,7 @@
 from collections import defaultdict
 from json import loads
 import logging
+import os
 import sys
 from typing import List
 from typing import Optional
@@ -22,6 +23,7 @@
 from ..sampler import BasePrioritySampler
 from ..sampler import BaseSampler
 from ..utils.formats import get_env
+from ..utils.formats import parse_tags_str
 from ..utils.time import StopWatch
 from ._encoding import BufferFull
 from ._encoding import BufferItemTooLarge
@@ -266,6 +268,9 @@ def __init__(
             max_item_size=self._max_payload_size,
         )
         self._headers.update({"Content-Type": self._encoder.content_type})
+        additional_header_str = os.environ.get("_DD_TRACE_WRITER_ADDITIONAL_HEADERS")
+        if additional_header_str is not None:
+            self._headers.update(parse_tags_str(additional_header_str))
         self.dogstatsd = dogstatsd
         self._report_metrics = report_metrics
         self._metrics_reset()

tests/tracer/test_writer.py

@@ -22,6 +22,7 @@
 from ddtrace.span import Span
 from tests.utils import AnyInt
 from tests.utils import BaseTestCase
+from tests.utils import override_env
 
 
 class DummyOutput:
@@ -538,3 +539,10 @@ def do_write(i):
         t.join()
 
     assert len(writer._encoder) == 100
+
+
+def test_additional_headers():
+    with override_env(dict(_DD_TRACE_WRITER_ADDITIONAL_HEADERS="additional-header:additional-value,header2:value2")):
+        writer = AgentWriter(agent_url="http://localhost:9126")
+        assert writer._headers["additional-header"] == "additional-value"
+        assert writer._headers["header2"] == "value2"