DEBUG-2334 Probe Builder for dynamic instrumentation (#3921)

Co-authored-by: Oleg Pudeyev <code@olegp.name>

lib/datadog/di/probe.rb

@@ -33,7 +33,7 @@ module DI
     class Probe
       def initialize(id:, type:,
         file: nil, line_no: nil, type_name: nil, method_name: nil,
-        template: nil, capture_snapshot: false)
+        template: nil, capture_snapshot: false, max_capture_depth: nil, rate_limit: nil)
         # Perform some sanity checks here to detect unexpected attribute
         # combinations, in order to not do them in subsequent code.
         if line_no && method_name
@@ -52,6 +52,7 @@ def initialize(id:, type:,
         @method_name = method_name
         @template = template
         @capture_snapshot = !!capture_snapshot
+        @max_capture_depth = max_capture_depth
 
         # These checks use instance methods that have more complex logic
         # than checking a single argument value. To avoid duplicating
@@ -61,11 +62,8 @@ def initialize(id:, type:,
           raise ArgumentError, "Unhandled probe type: neither method nor line probe: #{id}"
         end
 
-        @rate_limiter = if capture_snapshot
-          Datadog::Core::TokenBucket.new(1)
-        else
-          Datadog::Core::TokenBucket.new(5000)
-        end
+        @rate_limit = rate_limit || (@capture_snapshot ? 1 : 5000)
+        @rate_limiter = Datadog::Core::TokenBucket.new(@rate_limit)
       end
 
       attr_reader :id
@@ -76,7 +74,14 @@ def initialize(id:, type:,
       attr_reader :method_name
       attr_reader :template
 
-      # For internal DI use only
+      # Configured maximum capture depth. Can be nil in which case
+      # the global default will be used.
+      attr_reader :max_capture_depth
+
+      # Rate limit in effect, in invocations per second. Always present.
+      attr_reader :rate_limit
+
+      # Rate limiter object. For internal DI use only.
       attr_reader :rate_limiter
 
       def capture_snapshot?

lib/datadog/di/probe_builder.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require_relative "probe"
+
+module Datadog
+  module DI
+    # Creates Probe instances from remote configuration payloads.
+    #
+    # Due to the dynamic instrumentation product evolving over time,
+    # it is possible that the payload corresponds to a type of probe that the
+    # current version of the library does not handle.
+    # For now ArgumentError is raised in such cases (by ProbeBuilder or
+    # Probe constructor), since generally DI is meant to rescue all exceptions
+    # internally and not propagate any exceptions to applications.
+    # A dedicated exception could be added in the future if there is a use case
+    # for it.
+    #
+    # @api private
+    module ProbeBuilder
+      module_function def build_from_remote_config(config)
+        # The validations here are not yet comprehensive.
+        Probe.new(
+          id: config.fetch("id"),
+          type: config.fetch("type"),
+          file: config["where"]&.[]("sourceFile"),
+          # Sometimes lines are sometimes received as an array of nil
+          # for some reason.
+          line_no: config["where"]&.[]("lines")&.compact&.map(&:to_i)&.first,
+          type_name: config["where"]&.[]("typeName"),
+          method_name: config["where"]&.[]("methodName"),
+          template: config["template"],
+          capture_snapshot: !!config["captureSnapshot"],
+          max_capture_depth: config["capture"]&.[]("maxReferenceDepth"),
+          rate_limit: config["sampling"]&.[]("snapshotsPerSecond"),
+        )
+      rescue KeyError => exc
+        raise ArgumentError, "Malformed remote configuration entry for probe: #{exc.class}: #{exc}: #{config}"
+      end
+    end
+  end
+end

sig/datadog/di/probe.rbs

@@ -19,7 +19,8 @@ module Datadog
 
       @rate_limiter: Datadog::Core::RateLimiter
 
-      def initialize: (id: String, type: String, ?file: String?, ?line_no: Integer?, ?type_name: String?, ?method_name: String?, ?template: String?, ?capture_snapshot: bool) -> void
+      def initialize: (id: String, type: String, ?file: String?, ?line_no: Integer?, ?type_name: String?, ?method_name: String?, ?template: String?, ?capture_snapshot: bool,
+	?max_capture_depth: Integer, ?rate_limit: Integer) -> void
 
       attr_reader id: String
 

sig/datadog/di/probe_builder.rbs

@@ -0,0 +1,7 @@
+module Datadog
+  module DI
+    module ProbeBuilder
+      def self?.build_from_remote_config: (Hash[untyped,untyped] config) -> Probe
+    end
+  end
+end

spec/datadog/di/probe_builder_spec.rb

@@ -0,0 +1,143 @@
+require "datadog/di/probe_builder"
+
+RSpec.describe Datadog::DI::ProbeBuilder do
+  describe ".build_from_remote_config" do
+    let(:probe) do
+      described_class.build_from_remote_config(rc_probe_spec)
+    end
+
+    context "typical line probe" do
+      let(:rc_probe_spec) do
+        {"id" => "3ecfd456-2d7c-4359-a51f-d4cc44141ffe",
+         "version" => 0,
+         "type" => "LOG_PROBE",
+         "language" => "python",
+         "where" => {"sourceFile" => "aaa.rb", "lines" => [4321]},
+         "tags" => [],
+         "template" => "In aaa, line 1",
+         "segments" => [{"str" => "In aaa, line 1"}],
+         "captureSnapshot" => false,
+         # Use a value different from our library default to ensure that
+         # it is correctly processed.
+         "capture" => {"maxReferenceDepth" => 33},
+         # Use a value different from our library default to ensure that
+         # it is correctly processed.
+         "sampling" => {"snapshotsPerSecond" => 4500},
+         "evaluateAt" => "EXIT"}
+      end
+
+      it "creates line probe with corresponding values" do
+        expect(probe.id).to eq "3ecfd456-2d7c-4359-a51f-d4cc44141ffe"
+        expect(probe.type).to eq "LOG_PROBE"
+        expect(probe.file).to eq "aaa.rb"
+        expect(probe.line_no).to eq 4321
+        expect(probe.type_name).to be nil
+        expect(probe.method_name).to be nil
+        expect(probe.max_capture_depth).to eq 33
+        expect(probe.rate_limit).to eq 4500
+
+        expect(probe.line?).to be true
+        expect(probe.method?).to be false
+      end
+    end
+
+    context "minimum set of fields" do
+      # This is a made up payload to test attribute defaulting.
+      # In practice payloads like this should not be seen.
+      let(:rc_probe_spec) do
+        {"id" => "3ecfd456-2d7c-4359-a51f-d4cc44141ffe",
+         "type" => "LOG_PROBE",
+         "where" => {"sourceFile" => "aaa.rb", "lines" => [4321]},}
+      end
+
+      describe ".max_capture_depth" do
+        it "is nil" do
+          expect(probe.max_capture_depth).to be nil
+        end
+      end
+
+      describe ".rate_limit" do
+        it "is defaulted to 5000" do
+          expect(probe.rate_limit).to eq 5000
+        end
+      end
+    end
+
+    context "when lines is an array of nil" do
+      let(:rc_probe_spec) do
+        {"id" => "3ecfd456-2d7c-4359-a51f-d4cc44141ffe",
+         "version" => 0,
+         "type" => "LOG_PROBE",
+         "language" => "python",
+         "where" => {"sourceFile" => "aaa.rb", "lines" => [nil]},
+         "tags" => [],
+         "template" => "In aaa, line 1",
+         "segments" => [{"str" => "In aaa, line 1"}],
+         "captureSnapshot" => false,
+         "capture" => {"maxReferenceDepth" => 3},
+         "sampling" => {"snapshotsPerSecond" => 5000},
+         "evaluateAt" => "EXIT"}
+      end
+
+      describe "construction" do
+        it "fails with exception" do
+          expect do
+            probe
+          end.to raise_error(ArgumentError, /neither method nor line probe/)
+        end
+      end
+    end
+
+    context "RC payload with capture snapshot" do
+      let(:rc_probe_spec) do
+        {"id" => "3ecfd456-2d7c-4359-a51f-d4cc44141ffe",
+         "version" => 0,
+         "type" => "LOG_PROBE",
+         "language" => "python",
+         "where" => {"sourceFile" => "aaa", "lines" => [2]},
+         "tags" => [],
+         "template" => "In aaa, line 1",
+         "segments" => [{"str" => "In aaa, line 1"}],
+         "captureSnapshot" => true,
+         "capture" => {"maxReferenceDepth" => 3},
+         "sampling" => {"snapshotsPerSecond" => 5000},
+         "evaluateAt" => "EXIT"}
+      end
+
+      it "capture_snapshot? is true" do
+        expect(probe.capture_snapshot?).to be true
+      end
+    end
+
+    context "RC payload without capture snapshot" do
+      let(:rc_probe_spec) do
+        {"id" => "3ecfd456-2d7c-4359-a51f-d4cc44141ffe",
+         "version" => 0,
+         "type" => "LOG_PROBE",
+         "language" => "python",
+         "where" => {"sourceFile" => "aaa", "lines" => [4]},
+         "tags" => [],
+         "template" => "In aaa, line 1",
+         "segments" => [{"str" => "In aaa, line 1"}],
+         "captureSnapshot" => false,
+         "capture" => {"maxReferenceDepth" => 3},
+         "sampling" => {"snapshotsPerSecond" => 5000},
+         "evaluateAt" => "EXIT"}
+      end
+
+      it "capture_snapshot? is false" do
+        expect(probe.capture_snapshot?).to be false
+      end
+    end
+
+    context "empty input" do
+      let(:rc_probe_spec) { {} }
+
+      it "raises ArgumentError" do
+        expect do
+          probe
+        end.to raise_error(ArgumentError, /Malformed remote configuration entry/)
+      end
+    end
+  end
+end