From 0f023399ccfac3a01848328cd1c4f7269d07f1ab Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <156273877+p-datadog@users.noreply.github.com>
Date: Tue, 17 Sep 2024 11:42:42 -0400
Subject: [PATCH] DEBUG-2334 dynamic instrumentation Probe class (#3912)

---
 lib/datadog/di/error.rb       |  31 ++++++
 lib/datadog/di/probe.rb       | 128 +++++++++++++++++++++++
 sig/datadog/di/error.rbs      |  12 +++
 sig/datadog/di/probe.rbs      |  46 ++++++++
 spec/datadog/di/probe_spec.rb | 190 ++++++++++++++++++++++++++++++++++
 5 files changed, 407 insertions(+)
 create mode 100644 lib/datadog/di/error.rb
 create mode 100644 lib/datadog/di/probe.rb
 create mode 100644 sig/datadog/di/error.rbs
 create mode 100644 sig/datadog/di/probe.rbs
 create mode 100644 spec/datadog/di/probe_spec.rb

diff --git a/lib/datadog/di/error.rb b/lib/datadog/di/error.rb
new file mode 100644
index 00000000000..5a90d259959
--- /dev/null
+++ b/lib/datadog/di/error.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Datadog
+  module DI
+    # Base class for Dynamic Instrumentation exceptions.
+    #
+    # None of these exceptions should be propagated out of DI to user
+    # applications, therefore these exceptions are not considered to be
+    # part of the public API of the library.
+    #
+    # @api private
+    class Error < StandardError
+      # Probe does not contain a line number (i.e., is not a line probe).
+      class MissingLineNumber < Error
+      end
+
+      # Failed to communicate to the local Datadog agent (e.g. to send
+      # probe status or a snapshot).
+      class AgentCommunicationError < Error
+      end
+
+      # Attempting to instrument a method or file which does not exist.
+      #
+      # This could be due to the code that is referenced in the probe
+      # having not been loaded yet, or due to the probe referencing code
+      # that does not in fact exist anywhere (e.g. due to a misspelling).
+      class DITargetNotDefined < Error
+      end
+    end
+  end
+end
diff --git a/lib/datadog/di/probe.rb b/lib/datadog/di/probe.rb
new file mode 100644
index 00000000000..48941812413
--- /dev/null
+++ b/lib/datadog/di/probe.rb
@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+require_relative "error"
+require_relative "../core/rate_limiter"
+
+module Datadog
+  module DI
+    # Encapsulates probe information (as received via remote config)
+    # and state (e.g. whether the probe was installed, or executed).
+    #
+    # It is possible that remote configuration will specify an unsupported
+    # probe type or attribute, due to new DI functionality being added
+    # over time. We want to have predictable behavior in such cases, and
+    # since we can't guarantee that there will be enough information in
+    # a remote config payload to construct a functional probe, ProbeBuilder
+    # and remote config code must be prepared to deal with exceptions
+    # raised by Probe constructor in particular. Therefore, Probe constructor
+    # will raise an exception if it determines that there is not enough
+    # information (or confilcting information) in the arguments to create a
+    # functional probe, and upstream code is tasked with not spamming logs
+    # with notifications of such errors (and potentially limiting the
+    # attempts to construct probe from a given payload).
+    #
+    # Note that, while remote configuration provides line numbers as an
+    # array, the only supported line number configuration is a single line
+    # (this is the case for all languages currently). Therefore Probe
+    # only supports one line number, and ProbeBuilder is responsible for
+    # extracting that one line number out of the array received from RC.
+    #
+    # Note: only some of the parameter/attribute values are currently validated.
+    #
+    # @api private
+    class Probe
+      def initialize(id:, type:,
+        file: nil, line_no: nil, type_name: nil, method_name: nil,
+        template: nil, capture_snapshot: false)
+        # Perform some sanity checks here to detect unexpected attribute
+        # combinations, in order to not do them in subsequent code.
+        if line_no && method_name
+          raise ArgumentError, "Probe contains both line number and method name: #{id}"
+        end
+
+        if type_name && !method_name || method_name && !type_name
+          raise ArgumentError, "Partial method probe definition: #{id}"
+        end
+
+        @id = id
+        @type = type
+        @file = file
+        @line_no = line_no
+        @type_name = type_name
+        @method_name = method_name
+        @template = template
+        @capture_snapshot = !!capture_snapshot
+
+        # These checks use instance methods that have more complex logic
+        # than checking a single argument value. To avoid duplicating
+        # the logic here, use the methods and perform these checks after
+        # instance variable assignment.
+        unless method? || line?
+          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
+      end
+
+      attr_reader :id
+      attr_reader :type
+      attr_reader :file
+      attr_reader :line_no
+      attr_reader :type_name
+      attr_reader :method_name
+      attr_reader :template
+
+      # For internal DI use only
+      attr_reader :rate_limiter
+
+      def capture_snapshot?
+        @capture_snapshot
+      end
+
+      # Returns whether the probe is a line probe.
+      #
+      # Method probes may still specify a file name (to aid in locating the
+      # method or for stack traversal purposes?), therefore we do not check
+      # for file name/path presence here and just consider the line number.
+      def line?
+        !line_no.nil?
+      end
+
+      # Returns whether the probe is a method probe.
+      def method?
+        !!(type_name && method_name)
+      end
+
+      # Returns the line number associated with the probe, raising
+      # Error::MissingLineNumber if the probe does not have a line number
+      # associated with it.
+      #
+      # This method is used by instrumentation driver to ensure a line number
+      # that is passed into the instrumentation logic is actually a line number
+      # and not nil.
+      def line_no!
+        if line_no.nil?
+          raise Error::MissingLineNumber, "Probe #{id} does not have a line number associated with it"
+        end
+        line_no
+      end
+
+      # Source code location of the probe, for diagnostic reporting.
+      def location
+        if method?
+          "#{type_name}.#{method_name}"
+        elsif line?
+          "#{file}:#{line_no}"
+        else
+          # This case should not be possible because constructor verifies that
+          # the probe is a method or a line probe.
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end
diff --git a/sig/datadog/di/error.rbs b/sig/datadog/di/error.rbs
new file mode 100644
index 00000000000..3f3dddc7780
--- /dev/null
+++ b/sig/datadog/di/error.rbs
@@ -0,0 +1,12 @@
+module Datadog
+  module DI
+    class Error < StandardError
+      class MissingLineNumber < Error
+      end
+      class AgentCommunicationError < Error
+      end
+      class DITargetNotDefined < Error
+      end
+    end
+  end
+end
diff --git a/sig/datadog/di/probe.rbs b/sig/datadog/di/probe.rbs
new file mode 100644
index 00000000000..c98df5119ac
--- /dev/null
+++ b/sig/datadog/di/probe.rbs
@@ -0,0 +1,46 @@
+module Datadog
+  module DI
+    class Probe
+      @id: String
+
+      @type: String
+
+      @file: String?
+
+      @line_no: Integer?
+
+      @type_name: String?
+
+      @method_name: String?
+
+      @template: String
+
+      @capture_snapshot: bool
+
+      @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
+
+      attr_reader id: String
+
+      attr_reader type: String
+
+      attr_reader file: String?
+
+      attr_reader line_no: Integer?
+
+      attr_reader type_name: String?
+
+      attr_reader method_name: String?
+
+      attr_reader template: String
+      attr_reader rate_limiter: Datadog::Core::RateLimiter
+
+      def capture_snapshot?: () -> bool
+      def line?: () -> bool
+      def method?: () -> bool
+      def line_no!: () -> Integer
+      def location: () -> ::String
+    end
+  end
+end
diff --git a/spec/datadog/di/probe_spec.rb b/spec/datadog/di/probe_spec.rb
new file mode 100644
index 00000000000..e37832ac46e
--- /dev/null
+++ b/spec/datadog/di/probe_spec.rb
@@ -0,0 +1,190 @@
+require "datadog/di/probe"
+
+RSpec.describe Datadog::DI::Probe do
+  shared_context "method probe" do
+    let(:probe) do
+      described_class.new(id: "42", type: "foo", type_name: "Foo", method_name: "bar")
+    end
+  end
+
+  shared_context "line probe" do
+    let(:probe) do
+      described_class.new(id: "42", type: "foo", file: "foo.rb", line_no: 4)
+    end
+  end
+
+  describe ".new" do
+    context "method probe" do
+      include_context "method probe"
+
+      it "creates an instance" do
+        expect(probe).to be_a(described_class)
+        expect(probe.id).to eq "42"
+        expect(probe.type).to eq "foo"
+        expect(probe.type_name).to eq "Foo"
+        expect(probe.method_name).to eq "bar"
+        expect(probe.file).to be nil
+        expect(probe.line_no).to be nil
+      end
+    end
+
+    context "line probe" do
+      include_context "line probe"
+
+      it "creates an instance" do
+        expect(probe).to be_a(described_class)
+        expect(probe.id).to eq "42"
+        expect(probe.type).to eq "foo"
+        expect(probe.type_name).to be nil
+        expect(probe.method_name).to be nil
+        expect(probe.file).to eq "foo.rb"
+        expect(probe.line_no).to eq 4
+      end
+    end
+
+    context "neither method nor line" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo")
+      end
+
+      it "raises ArgumentError" do
+        expect do
+          probe
+        end.to raise_error(ArgumentError, /neither method nor line/)
+      end
+    end
+
+    context "both method and line" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo",
+          type_name: "foo", method_name: "bar", file: "baz", line_no: 4)
+      end
+
+      it "raises ArgumentError" do
+        expect do
+          probe
+        end.to raise_error(ArgumentError, /both line number and method name/)
+      end
+    end
+  end
+
+  describe "#line?" do
+    context "line probe" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo", file: "bar.rb", line_no: 5)
+      end
+
+      it "is true" do
+        expect(probe.line?).to be true
+      end
+    end
+
+    context "method probe" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar")
+      end
+
+      it "is false" do
+        expect(probe.line?).to be false
+      end
+    end
+
+    context "method probe with file name" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar", file: "quux.rb")
+      end
+
+      it "is false" do
+        expect(probe.line?).to be false
+      end
+    end
+  end
+
+  describe "#method?" do
+    context "line probe" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo", file: "bar.rb", line_no: 5)
+      end
+
+      it "is false" do
+        expect(probe.method?).to be false
+      end
+    end
+
+    context "method probe" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar")
+      end
+
+      it "is true" do
+        expect(probe.method?).to be true
+      end
+    end
+
+    context "method probe with file name" do
+      let(:probe) do
+        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar", file: "quux.rb")
+      end
+
+      it "is true" do
+        expect(probe.method?).to be true
+      end
+    end
+  end
+
+  describe "#line_no" do
+    context "one line number" do
+      let(:probe) { described_class.new(id: "x", type: "log", line_no: 5) }
+
+      it "returns the line number" do
+        expect(probe.line_no).to eq 5
+      end
+    end
+
+    context "nil line number" do
+      let(:probe) { described_class.new(id: "id", type: "LOG", type_name: "x", method_name: "y", line_no: nil) }
+
+      it "returns nil" do
+        expect(probe.line_no).to be nil
+      end
+    end
+  end
+
+  describe "#line_no!" do
+    context "one line number" do
+      let(:probe) { described_class.new(id: "x", type: "log", line_no: 5) }
+
+      it "returns the line number" do
+        expect(probe.line_no!).to eq 5
+      end
+    end
+
+    context "nil line number" do
+      let(:probe) { described_class.new(id: "id", type: "LOG", type_name: "x", method_name: "y", line_no: nil) }
+
+      it "raises MissingLineNumber" do
+        expect do
+          probe.line_no!
+        end.to raise_error(Datadog::DI::Error::MissingLineNumber, /does not have a line number/)
+      end
+    end
+  end
+
+  describe "#location" do
+    context "method probe" do
+      include_context "method probe"
+
+      it "returns method location" do
+        expect(probe.location).to eq "Foo.bar"
+      end
+    end
+
+    context "line probe" do
+      include_context "line probe"
+
+      it "returns line location" do
+        expect(probe.location).to eq "foo.rb:4"
+      end
+    end
+  end
+end