feat(di): add capture expressions

We add support for the new capture expression feature that allows
capturing specific values instead of the whole scope snapshot.

Author: P403n1x87

ddtrace/debugging/_exception/replay.py

@@ -159,6 +159,7 @@ def build(cls, exc_id: uuid.UUID, frame: FrameType) -> "SpanExceptionProbe":
             template=message,
             segments=[LiteralTemplateSegment(message)],
             take_snapshot=True,
+            capture_expressions=[],
             limits=DEFAULT_CAPTURE_LIMITS,
             condition=None,
             condition_error_rate=0.0,

ddtrace/debugging/_origin/span.py

@@ -65,6 +65,7 @@ def build(cls, name: str, module: str, function: str) -> "EntrySpanProbe":
             template=message,
             segments=[LiteralTemplateSegment(message)],
             take_snapshot=True,
+            capture_expressions=[],
             limits=DEFAULT_CAPTURE_LIMITS,
             condition=None,
             condition_error_rate=0.0,
@@ -89,6 +90,7 @@ def build(cls, name: str, filename: str, line: int) -> "ExitSpanProbe":
             template=message,
             segments=[LiteralTemplateSegment(message)],
             take_snapshot=True,
+            capture_expressions=[],
             limits=DEFAULT_CAPTURE_LIMITS,
             condition=None,
             condition_error_rate=0.0,

ddtrace/debugging/_probe/model.py

@@ -245,28 +245,35 @@ def _to_str(value):
         return "".join([_to_str(s.eval(scope)) for s in self.segments])
 
 
+@dataclass
+class CaptureExpression:
+    name: str
+    expr: DDExpression
+
+
 @dataclass
 class LogProbeMixin(AbstractProbeMixIn):
     template: str
     segments: List[TemplateSegment]
     take_snapshot: bool
+    capture_expressions: List[CaptureExpression]
     limits: CaptureLimits = field(compare=False)
 
     @property
     def __budget__(self) -> int:
-        return 10 if self.take_snapshot else 100
+        return 10 if self.take_snapshot or self.capture_expressions else 100
 
 
 @dataclass(eq=False)
 class LogLineProbe(Probe, LineLocationMixin, LogProbeMixin, ProbeConditionMixin, RateLimitMixin):
     def is_global_rate_limited(self) -> bool:
-        return self.take_snapshot
+        return self.take_snapshot or bool(self.capture_expressions)
 
 
 @dataclass(eq=False)
 class LogFunctionProbe(Probe, FunctionLocationMixin, TimingMixin, LogProbeMixin, ProbeConditionMixin, RateLimitMixin):
     def is_global_rate_limited(self) -> bool:
-        return self.take_snapshot
+        return self.take_snapshot or bool(self.capture_expressions)
 
 
 @dataclass

ddtrace/debugging/_probe/remoteconfig.py

@@ -17,6 +17,7 @@
 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 CaptureExpression
 from ddtrace.debugging._probe.model import CaptureLimits
 from ddtrace.debugging._probe.model import ExpressionTemplateSegment
 from ddtrace.debugging._probe.model import FunctionProbe
@@ -146,6 +147,7 @@ def update_args(cls, args, attribs):
             else DEFAULT_CAPTURE_LIMITS,
             condition_error_rate=DEFAULT_PROBE_CONDITION_ERROR_RATE,  # TODO: should we take rate limit out of Probe?
             take_snapshot=take_snapshot,
+            capture_expressions=[CaptureExpression(**_) for _ in attribs.get("captureExpressions", [])],
             template=attribs.get("template"),
             segments=[_compile_segment(segment) for segment in attribs.get("segments", [])],
         )

ddtrace/debugging/_signal/snapshot.py

@@ -7,12 +7,14 @@
 from types import ModuleType
 from typing import Any
 from typing import Dict
+from typing import List
 from typing import Mapping
 from typing import Optional
 from typing import cast
 
 from ddtrace.debugging._expressions import DDExpressionEvaluationError
 from ddtrace.debugging._probe.model import DEFAULT_CAPTURE_LIMITS
+from ddtrace.debugging._probe.model import CaptureExpression
 from ddtrace.debugging._probe.model import CaptureLimits
 from ddtrace.debugging._probe.model import FunctionLocationMixin
 from ddtrace.debugging._probe.model import LineLocationMixin
@@ -85,6 +87,26 @@ def timeout(_):
         }
 
 
+def _capture_expressions(
+    exprs: List[CaptureExpression],
+    scope: Mapping[str, Any],
+    limits: CaptureLimits = DEFAULT_CAPTURE_LIMITS,
+) -> Dict[str, Any]:
+    with HourGlass(duration=CAPTURE_TIME_BUDGET) as hg:
+
+        def timeout(_):
+            return not hg.trickling()
+
+        return {
+            "captureExpressions": {
+                e.name: utils.capture_value(
+                    e.expr.eval(scope), limits.max_level, limits.max_len, limits.max_size, limits.max_fields, timeout
+                )
+                for e in exprs
+            }
+        }
+
+
 _EMPTY_CAPTURED_CONTEXT: Dict[str, Any] = {"arguments": {}, "locals": {}, "staticFields": {}, "throwable": None}
 
 
@@ -131,7 +153,13 @@ def _do(self, retval, exc_info, scope):
 
         self._stack = utils.capture_stack(self.frame)
 
-        return _capture_context(frame, exc_info, retval=retval, limits=probe.limits) if probe.take_snapshot else None
+        if probe.take_snapshot:
+            return _capture_context(frame, exc_info, retval=retval, limits=probe.limits)
+
+        if probe.capture_expressions:
+            return _capture_expressions(probe.capture_expressions, scope, probe.limits)
+
+        return None
 
     def enter(self, scope: Mapping[str, Any]) -> None:
         self.entry_capture = self._do(_NOTSET, (None, None, None), scope)
@@ -166,7 +194,7 @@ def data(self):
         probe = self.probe
 
         captures = {}
-        if isinstance(probe, LogProbeMixin) and probe.take_snapshot:
+        if isinstance(probe, LogProbeMixin) and (probe.take_snapshot or probe.capture_expressions):
             if isinstance(probe, LineLocationMixin):
                 captures = {"lines": {str(probe.line): self.line_capture or _EMPTY_CAPTURED_CONTEXT}}
             elif isinstance(probe, FunctionLocationMixin):

releasenotes/notes/feat-di-capture-expressions-cc9d1973a6119888.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    dynamic instrumentation: added support for capture expressions in log
+    probes.

tests/debugging/test_debugger.py

@@ -32,7 +32,10 @@
 from ddtrace.internal.utils.formats import format_trace_id
 from ddtrace.internal.utils.inspection import linenos
 from tests.debugging.mocking import debugger
+from tests.debugging.utils import compile_capture_expressions
 from tests.debugging.utils import compile_template
+from tests.debugging.utils import create_capture_expressions_function_probe
+from tests.debugging.utils import create_capture_expressions_line_probe
 from tests.debugging.utils import create_log_function_probe
 from tests.debugging.utils import create_log_line_probe
 from tests.debugging.utils import create_metric_line_probe
@@ -916,6 +919,83 @@ def test_debugger_log_line_probe_generate_messages(stuff):
         assert not msg1["debugger"]["snapshot"]["captures"]
 
 
+def test_debugger_capture_expressions_line_probe_(stuff):
+    with debugger(upload_flush_interval=float("inf")) as d:
+        d.add_probes(
+            create_capture_expressions_line_probe(
+                probe_id="foo",
+                source_file="tests/submod/stuff.py",
+                line=36,
+                rate=float("inf"),
+                **compile_capture_expressions([{"name": "foo", "expr": {"dsl": "bar", "json": {"ref": "bar"}}}]),
+            ),
+        )
+
+        sentinel = {"foo": 42}
+        # recursive reference to itself for infinite depth
+        sentinel["self"] = sentinel  # type: ignore
+
+        stuff.Stuff().instancestuff(sentinel)
+
+        (msg,) = d.uploader.wait_for_payloads(1)
+
+        assert msg["debugger"]["snapshot"]["captures"]["lines"]["36"] == {
+            "captureExpressions": {
+                "foo": {
+                    "type": "dict",
+                    "entries": [
+                        [{"type": "str", "value": "'foo'"}, {"type": "int", "value": "42"}],
+                        [
+                            {"type": "str", "value": "'self'"},
+                            {
+                                "type": "dict",
+                                "entries": [
+                                    [{"type": "str", "value": "'foo'"}, {"type": "int", "value": "42"}],
+                                    [
+                                        {"type": "str", "value": "'self'"},
+                                        {
+                                            "type": "dict",
+                                            "entries": [
+                                                [{"type": "str", "value": "'foo'"}, {"type": "int", "value": "42"}],
+                                                [
+                                                    {"type": "str", "value": "'self'"},
+                                                    {"type": "dict", "notCapturedReason": "depth", "size": 2},
+                                                ],
+                                            ],
+                                            "size": 2,
+                                        },
+                                    ],
+                                ],
+                                "size": 2,
+                            },
+                        ],
+                    ],
+                    "size": 2,
+                }
+            }
+        }
+
+
+def test_debugger_capture_expressions_function_probe(stuff):
+    with debugger() as d:
+        d.add_probes(
+            create_capture_expressions_function_probe(
+                probe_id="exit-probe",
+                module="tests.submod.stuff",
+                func_qname="mutator",
+                evaluate_at=ProbeEvalTiming.EXIT,
+                **compile_capture_expressions(
+                    [{"name": "retval", "expr": {"dsl": "@return", "json": {"ref": "@return"}}}]
+                ),
+            )
+        )
+
+        stuff.mutator(arg=[])
+
+        with d.assert_single_snapshot() as snapshot:
+            assert snapshot.return_capture["captureExpressions"]["retval"] == {"isNull": True, "type": "NoneType"}
+
+
 class SpanProbeTestCase(TracerTestCase):
     def setUp(self):
         super(SpanProbeTestCase, self).setUp()

tests/debugging/utils.py

@@ -6,6 +6,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 CaptureExpression
 from ddtrace.debugging._probe.model import ExpressionTemplateSegment
 from ddtrace.debugging._probe.model import LiteralTemplateSegment
 from ddtrace.debugging._probe.model import LogFunctionProbe
@@ -37,6 +38,14 @@ def compile_template(*args):
     return {"template": template, "segments": segments}
 
 
+def compile_capture_expressions(exprs):
+    return {
+        "capture_expressions": [
+            CaptureExpression(name=expr["name"], expr=DDRedactedExpression.compile(expr["expr"])) for expr in exprs
+        ],
+    }
+
+
 def ddexpr(json, dsl="test"):
     return DDExpression(dsl=dsl, callable=dd_compile(json))
 
@@ -76,6 +85,7 @@ def _wrapper(*args, **kwargs):
         kwargs.setdefault("take_snapshot", False)
         kwargs.setdefault("rate", DEFAULT_PROBE_RATE)
         kwargs.setdefault("limits", DEFAULT_CAPTURE_LIMITS)
+        kwargs.setdefault("capture_expressions", [])
         return f(*args, **kwargs)
 
     return _wrapper
@@ -88,6 +98,19 @@ def _wrapper(*args, **kwargs):
         kwargs.setdefault("limits", DEFAULT_CAPTURE_LIMITS)
         kwargs.setdefault("template", "")
         kwargs.setdefault("segments", [])
+        kwargs.setdefault("capture_expressions", [])
+        return f(*args, **kwargs)
+
+    return _wrapper
+
+
+def capture_expressions_probe_defaults(f):
+    def _wrapper(*args, **kwargs):
+        kwargs.setdefault("take_snapshot", False)
+        kwargs.setdefault("rate", DEFAULT_SNAPSHOT_PROBE_RATE)
+        kwargs.setdefault("limits", DEFAULT_CAPTURE_LIMITS)
+        kwargs.setdefault("template", "")
+        kwargs.setdefault("segments", [])
         return f(*args, **kwargs)
 
     return _wrapper
@@ -142,6 +165,21 @@ def create_snapshot_function_probe(**kwargs):
     return LogFunctionProbe(**kwargs)
 
 
+@create_probe_defaults
+@probe_conditional_defaults
+@capture_expressions_probe_defaults
+def create_capture_expressions_line_probe(**kwargs):
+    return LogLineProbe(**kwargs)
+
+
+@create_probe_defaults
+@probe_conditional_defaults
+@timing_defaults
+@capture_expressions_probe_defaults
+def create_capture_expressions_function_probe(**kwargs):
+    return LogFunctionProbe(**kwargs)
+
+
 @create_probe_defaults
 @probe_conditional_defaults
 @log_probe_defaults