DEBUG-2334 respect maxFieldCount in probe specification (#4142)

Maximum capture attribute count can be specified in the probe,
implement support for it in DI.

lib/datadog/di/instrumenter.rb

@@ -103,14 +103,17 @@ def hook_method(probe, &block)
           # target method here.
         end
         rate_limiter = probe.rate_limiter
+        settings = self.settings
 
         mod = Module.new do
           define_method(method_name) do |*args, **kwargs| # steep:ignore
             if rate_limiter.nil? || rate_limiter.allow?
               # Arguments may be mutated by the method, therefore
               # they need to be serialized prior to method invocation.
               entry_args = if probe.capture_snapshot?
-                serializer.serialize_args(args, kwargs)
+                serializer.serialize_args(args, kwargs,
+                  depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
+                  attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count)
               end
               rv = nil
               # Under Ruby 2.6 we cannot just call super(*args, **kwargs)

lib/datadog/di/probe.rb

@@ -36,7 +36,9 @@ class Probe
 
       def initialize(id:, type:,
         file: nil, line_no: nil, type_name: nil, method_name: nil,
-        template: nil, capture_snapshot: false, max_capture_depth: nil, rate_limit: nil)
+        template: nil, capture_snapshot: false, max_capture_depth: nil,
+        max_capture_attribute_count: nil,
+        rate_limit: nil)
         # Perform some sanity checks here to detect unexpected attribute
         # combinations, in order to not do them in subsequent code.
         unless KNOWN_TYPES.include?(type)
@@ -64,6 +66,7 @@ def initialize(id:, type:,
         @template = template
         @capture_snapshot = !!capture_snapshot
         @max_capture_depth = max_capture_depth
+        @max_capture_attribute_count = max_capture_attribute_count
 
         # These checks use instance methods that have more complex logic
         # than checking a single argument value. To avoid duplicating
@@ -91,6 +94,10 @@ def initialize(id:, type:,
       # the global default will be used.
       attr_reader :max_capture_depth
 
+      # Configured maximum capture attribute count. Can be nil in which case
+      # the global default will be used.
+      attr_reader :max_capture_attribute_count
+
       # Rate limit in effect, in invocations per second. Always present.
       attr_reader :rate_limit
 

lib/datadog/di/probe_builder.rb

@@ -37,6 +37,7 @@ module ProbeBuilder
           template: config["template"],
           capture_snapshot: !!config["captureSnapshot"],
           max_capture_depth: config["capture"]&.[]("maxReferenceDepth"),
+          max_capture_attribute_count: config["capture"]&.[]("maxFieldCount"),
           rate_limit: config["sampling"]&.[]("snapshotsPerSecond"),
         )
       rescue KeyError => exc

lib/datadog/di/probe_notification_builder.rb

@@ -65,22 +65,28 @@ def build_snapshot(probe, rv: nil, snapshot: nil, path: nil,
                 arguments: if serialized_entry_args
                   serialized_entry_args
                 else
-                  (args || kwargs) && serializer.serialize_args(args, kwargs)
+                  (args || kwargs) && serializer.serialize_args(args, kwargs,
+                    depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
+                    attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count)
                 end,
                 throwable: nil,
                 # standard:enable all
               },
               return: {
                 arguments: {
-                  "@return": serializer.serialize_value(rv),
+                  "@return": serializer.serialize_value(rv,
+                    depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
+                    attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count),
                 },
                 throwable: nil,
               },
             }
           elsif probe.line?
             {
               lines: snapshot && {
-                probe.line_no => {locals: serializer.serialize_vars(snapshot)},
+                probe.line_no => {locals: serializer.serialize_vars(snapshot,
+                  depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
+                  attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count,)},
               },
             }
           end

lib/datadog/di/serializer.rb

@@ -82,24 +82,28 @@ def initialize(settings, redactor, telemetry: nil)
       # between positional and keyword arguments. We convert positional
       # arguments to keyword arguments ("arg1", "arg2", ...) and ensure
       # the positional arguments are listed first.
-      def serialize_args(args, kwargs)
+      def serialize_args(args, kwargs,
+        depth: settings.dynamic_instrumentation.max_capture_depth,
+        attribute_count: settings.dynamic_instrumentation.max_capture_attribute_count)
         counter = 0
         combined = args.each_with_object({}) do |value, c|
           counter += 1
           # Conversion to symbol is needed here to put args ahead of
           # kwargs when they are merged below.
           c[:"arg#{counter}"] = value
         end.update(kwargs)
-        serialize_vars(combined)
+        serialize_vars(combined, depth: depth, attribute_count: attribute_count)
       end
 
       # Serializes variables captured by a line probe.
       #
       # These are normally local variables that exist on a particular line
       # of executed code.
-      def serialize_vars(vars)
+      def serialize_vars(vars,
+        depth: settings.dynamic_instrumentation.max_capture_depth,
+        attribute_count: settings.dynamic_instrumentation.max_capture_attribute_count)
         vars.each_with_object({}) do |(k, v), agg|
-          agg[k] = serialize_value(v, name: k)
+          agg[k] = serialize_value(v, name: k, depth: depth, attribute_count: attribute_count)
         end
       end
 
@@ -115,7 +119,11 @@ def serialize_vars(vars)
       # (integers, strings, arrays, hashes).
       #
       # Respects string length, collection size and traversal depth limits.
-      def serialize_value(value, name: nil, depth: settings.dynamic_instrumentation.max_capture_depth, type: nil)
+      def serialize_value(value, name: nil,
+        depth: settings.dynamic_instrumentation.max_capture_depth,
+        attribute_count: nil,
+        type: nil)
+        attribute_count ||= settings.dynamic_instrumentation.max_capture_attribute_count
         cls = type || value.class
         begin
           if redactor.redact_type?(value)
@@ -203,7 +211,6 @@ def serialize_value(value, name: nil, depth: settings.dynamic_instrumentation.ma
               serialized.update(notCapturedReason: "depth")
             else
               fields = {}
-              max = settings.dynamic_instrumentation.max_capture_attribute_count
               cur = 0
 
               # MRI and JRuby 9.4.5+ preserve instance variable definition
@@ -229,7 +236,7 @@ def serialize_value(value, name: nil, depth: settings.dynamic_instrumentation.ma
               ivars = value.instance_variables
 
               ivars.each do |ivar|
-                if cur >= max
+                if cur >= attribute_count
                   serialized.update(notCapturedReason: "fieldCount", fields: fields)
                   break
                 end

sig/datadog/di/probe.rbs

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

sig/datadog/di/serializer.rbs

@@ -15,12 +15,12 @@ module Datadog
 
       attr_reader telemetry: Core::Telemetry::Component
 
-      def serialize_args: (untyped args, untyped kwargs) -> untyped
-      def serialize_vars: (untyped vars) -> untyped
-      def serialize_value: (untyped value, ?name: String, ?depth: Integer) -> untyped
+      def serialize_args: (untyped args, untyped kwargs, ?depth: Integer, ?attribute_count: Integer?) -> untyped
+      def serialize_vars: (untyped vars, ?depth: Integer, ?attribute_count: Integer?) -> untyped
+      def serialize_value: (untyped value, ?name: String, ?depth: Integer, ?attribute_count: Integer?) -> untyped
 
       def self.register: (?condition: Proc) { 
-        (serializer: Serializer, value: untyped, name: Symbol, depth: Integer) -> untyped } -> void
+        (serializer: Serializer, value: untyped, name: Symbol, depth: Integer, ?attribute_count: Integer?) -> untyped } -> void
 
       private
       def class_name: (untyped cls) -> untyped

spec/datadog/di/instrumenter_spec.rb

@@ -12,6 +12,7 @@
     allow(settings.dynamic_instrumentation).to receive(:enabled).and_return(true)
     allow(settings.dynamic_instrumentation.internal).to receive(:untargeted_trace_points).and_return(false)
     allow(settings.dynamic_instrumentation).to receive(:max_capture_depth).and_return(2)
+    allow(settings.dynamic_instrumentation).to receive(:max_capture_attribute_count).and_return(2)
     allow(settings.dynamic_instrumentation).to receive(:redacted_type_names).and_return([])
     allow(settings.dynamic_instrumentation).to receive(:redacted_identifiers).and_return([])
   end

spec/datadog/di/integration/probe_notification_builder_spec.rb

@@ -28,6 +28,7 @@
         allow(settings).to receive(:enabled).and_return(true)
         allow(settings).to receive(:untargeted_trace_points).and_return(false)
         allow(settings).to receive(:max_capture_depth).and_return(2)
+        allow(settings).to receive(:max_capture_attribute_count).and_return(2)
         allow(settings).to receive(:max_capture_string_length).and_return(20)
         allow(settings).to receive(:max_capture_collection_size).and_return(20)
         allow(settings).to receive(:redacted_type_names).and_return([])

spec/datadog/di/probe_builder_spec.rb

@@ -22,7 +22,7 @@
          "captureSnapshot" => false,
          # Use a value different from our library default to ensure that
          # it is correctly processed.
-         "capture" => {"maxReferenceDepth" => 33},
+         "capture" => {"maxReferenceDepth" => 33, 'maxFieldCount' => 34},
          # Use a value different from our library default to ensure that
          # it is correctly processed.
          "sampling" => {"snapshotsPerSecond" => 4500},
@@ -37,6 +37,7 @@
         expect(probe.type_name).to be nil
         expect(probe.method_name).to be nil
         expect(probe.max_capture_depth).to eq 33
+        expect(probe.max_capture_attribute_count).to eq 34
         expect(probe.rate_limit).to eq 4500
 
         expect(probe.line?).to be true