add rate limiter and propagation

ddtrace/debugging/_live.py

@@ -1,22 +1,10 @@
-import typing as t
-
 from ddtrace.debugging._session import Session
 from ddtrace.internal import core
 
 
-def handle_distributed_context(context: t.Any) -> None:
-    debug_tag = context._meta.get("_dd.p.debug")
-    if debug_tag is None:
-        return
-
-    for session in debug_tag.split(","):
-        ident, _, level = session.partition(":")
-        Session(ident=ident, level=int(level or 0)).link_to_trace(context)
-
-
 def enable() -> None:
-    core.on("distributed_context.activated", handle_distributed_context, "live_debugger")
+    core.on("distributed_context.activated", Session.from_trace_context, "live_debugger")
 
 
 def disable() -> None:
-    core.reset_listeners("distributed_context.activated", handle_distributed_context)
+    core.reset_listeners("distributed_context.activated", Session.from_trace_context)

ddtrace/debugging/_probe/model.py

@@ -27,6 +27,7 @@
 
 DEFAULT_PROBE_RATE = 5000.0
 DEFAULT_SNAPSHOT_PROBE_RATE = 1.0
+DEFAULT_TRIGGER_PROBE_RATE = 1.0
 DEFAULT_PROBE_CONDITION_ERROR_RATE = 1.0 / 60 / 5
 
 

ddtrace/debugging/_probe/remoteconfig.py

@@ -14,6 +14,7 @@
 from ddtrace.debugging._probe.model import DEFAULT_PROBE_CONDITION_ERROR_RATE
 from ddtrace.debugging._probe.model import DEFAULT_PROBE_RATE
 from ddtrace.debugging._probe.model import DEFAULT_SNAPSHOT_PROBE_RATE
+from ddtrace.debugging._probe.model import DEFAULT_TRIGGER_PROBE_RATE
 from ddtrace.debugging._probe.model import CaptureLimits
 from ddtrace.debugging._probe.model import ExpressionTemplateSegment
 from ddtrace.debugging._probe.model import FunctionProbe
@@ -211,6 +212,7 @@ class TriggerProbeFactory(ProbeFactory):
     @classmethod
     def update_args(cls, args, attribs):
         args.update(
+            rate=attribs.get("sampling", {}).get("cooldownInSeconds", DEFAULT_TRIGGER_PROBE_RATE),
             session_id=attribs["sessionId"],
             level=attribs["level"],
         )

ddtrace/debugging/_session.py

@@ -8,16 +8,51 @@
 SessionId = str
 
 
+def _sessions_from_debug_tag(debug_tag: str) -> t.Generator["Session", None, None]:
+    for session in debug_tag.split(","):
+        ident, _, level = session.partition(":")
+        yield Session(ident=ident, level=int(level or 0))
+
+
+def _sessions_to_debug_tag(sessions: t.Iterable["Session"]) -> str:
+    # TODO: Validate tag length
+    return ",".join(f"{session.ident}:{session.level}" for session in sessions)
+
+
 @dataclass
 class Session:
     ident: SessionId
     level: int
 
+    @classmethod
+    def from_trace_context(cls, context: t.Any) -> None:
+        debug_tag = context._meta.get("_dd.p.debug")
+        if debug_tag is None:
+            return
+
+        for session in _sessions_from_debug_tag(debug_tag):
+            session.link_to_trace(context)
+
+    def to_trace_context(self, context: t.Any) -> None:
+        sessions = list(_sessions_from_debug_tag(context))
+        for session in sessions:
+            if self.ident == session.ident:
+                # The session is already in the tags so we don't need to add it
+                if self.level > session.level:
+                    # We only need to update the level if it's higher
+                    session.level = self.level
+                break
+        else:
+            # The session is not in the tags so we need to add it
+            sessions.append(self)
+
+        context._meta["_dd.p.debug"] = _sessions_to_debug_tag(sessions)
+
     def link_to_trace(self, trace_context: t.Optional[t.Any] = None):
         SessionManager.link_session_to_trace(self, trace_context)
 
     def unlink_from_trace(self, trace_context: t.Optional[t.Any] = None):
-        SessionManager.unlink_session_to_trace(self, trace_context)
+        SessionManager.unlink_session_from_trace(self, trace_context)
 
     @classmethod
     def from_trace(cls) -> t.Iterable["Session"]:
@@ -43,7 +78,7 @@ def link_session_to_trace(cls, session, trace_context: t.Optional[t.Any] = None)
         cls._sessions_trace_map.setdefault(context, {})[session.ident] = session
 
     @classmethod
-    def unlink_session_to_trace(cls, session, trace_context: t.Optional[t.Any] = None) -> None:
+    def unlink_session_from_trace(cls, session, trace_context: t.Optional[t.Any] = None) -> None:
         context = trace_context or tracer.current_trace_context()
         if context is None:
             # Nothing to unlink from

ddtrace/debugging/_signal/trigger.py

@@ -20,7 +20,14 @@ class Trigger(LogSignal):
 
     def _link_session(self) -> None:
         probe = t.cast(SessionMixin, self.probe)
-        Session(probe.session_id, probe.level).link_to_trace(self.trace_context)
+        session = Session(probe.session_id, probe.level)
+
+        # Link the session to the running trace
+        session.link_to_trace(self.trace_context)
+
+        # Ensure that the new session information is included in the debug
+        # propagation tag for distributed debugging
+        session.to_trace_context(self.trace_context)
 
     def enter(self, scope: t.Mapping[str, t.Any]) -> None:
         self._link_session()

tests/debugging/test_session.py

@@ -0,0 +1,11 @@
+from ddtrace.debugging._session import _sessions_from_debug_tag
+from ddtrace.debugging._session import _sessions_to_debug_tag
+
+
+def test_tag_identity():
+    """
+    Test that the combination of _sessions_from_debug_tag and
+    _sessions_to_debug_tag is the identity function.
+    """
+    debug_tag = "session1:1,session2:2"
+    assert _sessions_to_debug_tag(_sessions_from_debug_tag(debug_tag)) == debug_tag