Add support for setting trace sample rates

README.md

@@ -203,6 +203,12 @@ Specify trace ids as comma separated values (eg. `12345,7890,2468`)
 Initiate a _synchronous_ session. All subsequent traces received will be
 associated with the required test token provided.
 
+#### [optional] `?agent_sample_rate_by_service=`
+
+Sample rates to be returned by the agent in response to trace v0.4 and v0.5 requests.
+
+Example: `"{'service:test,env:staging': 0.5, 'service:test2,env:prod': 0.2}"` (note the JSON has to be URL-encoded).
+
 #### [optional] `?test_session_token=`
 #### [optional] `X-Datadog-Test-Session-Token`
 

ddapm_test_agent/agent.py

@@ -3,6 +3,8 @@
 import base64
 from collections import OrderedDict
 from collections import defaultdict
+from dataclasses import dataclass
+from dataclasses import field
 import json
 import logging
 import os
@@ -11,6 +13,7 @@
 import sys
 from typing import Awaitable
 from typing import Callable
+from typing import DefaultDict
 from typing import Dict
 from typing import List
 from typing import Literal
@@ -179,13 +182,19 @@ def default_value_trace_results_summary():
     }
 
 
+@dataclass
+class _AgentSession:
+    """Maintain Agent state across requests."""
+
+    sample_rate_by_service_env: Dict[str, float] = field(default_factory=dict)
+
+
 class Agent:
     def __init__(self) -> None:
-        """Only store the requests sent to the agent. There are many representations
-        of data but typically information is lost while transforming the data.
-
-        Storing exactly what is sent to the agent enables us to transform the data
-        however we desire later on.
+        """
+        Try to only store the requests sent to the agent. There are many representations
+        of data but typically information is lost while transforming the data so it is best
+        to keep the original and compute transformation when needed.
         """
         # Token to be used if running test cases synchronously
         self._requests: List[Request] = []
@@ -204,6 +213,12 @@ def __init__(self) -> None:
             "/tracer_flare/v1",
         ]
 
+        # Note that sessions are not cleared at any point since we don't know
+        # definitively when a session is over.
+        self._sessions: DefaultDict[Optional[str], _AgentSession] = defaultdict(
+            lambda: _AgentSession(sample_rate_by_service_env={})
+        )
+
     async def traces(self) -> TraceMap:
         """Return the traces stored by the agent in the order in which they
         arrived.
@@ -667,11 +682,14 @@ async def _handle_traces(self, request: Request, version: Literal["v0.4", "v0.5"
             except MsgPackExtraDataException as e:
                 log.error(f"Error unpacking trace bytes with Msgpack: {str(e)}, error {e}")
 
-        # TODO: implement sampling logic
-        return web.json_response(data={"rate_by_service": {}})
+        return web.json_response(data={"rate_by_service": self._sessions[token].sample_rate_by_service_env})
 
     async def handle_session_start(self, request: Request) -> web.Response:
+        rates = json.loads(request.url.query.get("agent_sample_rate_by_service", "{}"))
         self._requests.append(request)
+        session = self._sessions[_session_token(request)]
+        session.sample_rate_by_service_env = rates
+        log.info("Starting new session with token %r: %r", _session_token(request), session)
         return web.HTTPOk()
 
     async def handle_snapshot(self, request: Request) -> web.Response:

releasenotes/notes/sample-rate-1e06ae4cdc933b14.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Add the ability to set the trace sample rates returned by the agent.

tests/test_agent.py

@@ -53,6 +53,27 @@ async def test_trace_clear_token(
     assert await resp.text() == "[[]]"
 
 
+async def test_trace_agent_sample_rate(
+    agent,
+    v04_reference_http_trace_payload_headers,
+    v04_reference_http_trace_payload_data,
+):
+    agent_rates = {"service:test,env:staging": 0.5}
+    resp = await agent.get(
+        "/test/session/start",
+        params={"test_session_token": "1", "agent_sample_rate_by_service": json.dumps(agent_rates)},
+    )
+    assert resp.status == 200, await resp.text()
+    resp = await agent.put(
+        "/v0.4/traces",
+        params={"test_session_token": "1"},
+        headers=v04_reference_http_trace_payload_headers,
+        data=v04_reference_http_trace_payload_data,
+    )
+    assert resp.status == 200, await resp.text()
+    assert await resp.json() == {"rate_by_service": {"service:test,env:staging": 0.5}}
+
+
 async def test_info(agent):
     resp = await agent.get("/info")
     assert resp.status == 200