Fix symbol extraction for instrumented methods

lib/datadog/symbol_database/extractor.rb

@@ -43,7 +43,7 @@ module SymbolDatabase
     #    These are expected failures — no logging needed.
     #
     # 2. **Method-level rescues** (`rescue => e` with logging):
-    #    Catch failures in extract_method_scope, find_source_file, etc. Log at debug
+    #    Catch failures in build_class_method_scopes, find_source_file, etc. Log at debug
     #    for post-hoc diagnosis, return nil or empty array. One bad method/module
     #    doesn't kill the entire class extraction.
     #
@@ -133,16 +133,16 @@ def extract(mod)
         return nil unless Module === mod
         mod_name = safe_mod_name(mod)
         return nil unless mod_name
-
-        return nil unless user_code_module?(mod)
+        return nil unless user_code_module_name?(mod)
 
         source_file = find_source_file(mod)
         return nil unless source_file
+        return nil unless user_code_path?(source_file)
 
         inner_scope = if Class === mod
-          extract_class_scope(mod)
+          extract_class_scope(mod, source_file)
         else
-          extract_module_scope(mod)
+          extract_module_scope(mod, source_file)
         end
 
         wrap_in_file_scope(source_file, [inner_scope])
@@ -275,10 +275,12 @@ def resolves_to_same_module?(mod_name, mod)
         false
       end
 
-      # Check if module is from user code (not gems or stdlib)
-      # @param mod [Module] The module to check
-      # @return [Boolean] true if user code
-      def user_code_module?(mod)
+      # Return whether +mod+ passes the name-based user-code exclusions
+      # (named, outside Datadog, not a Ruby root class). Does not resolve
+      # a source file.
+      # @param mod [Module]
+      # @return [Boolean]
+      def user_code_module_name?(mod)
         mod_name = safe_mod_name(mod)
         return false unless mod_name
 
@@ -297,6 +299,15 @@ def user_code_module?(mod)
         return false if mod.equal?(Object) || mod.equal?(BasicObject) ||
           mod.equal?(Kernel) || mod.equal?(Module) || mod.equal?(Class)
 
+        true
+      end
+
+      # Check if module is from user code (not gems or stdlib)
+      # @param mod [Module] The module to check
+      # @return [Boolean] true if user code
+      def user_code_module?(mod)
+        return false unless user_code_module_name?(mod)
+
         source_file = find_source_file(mod)
         return false unless source_file
 
@@ -331,6 +342,25 @@ def user_code_path?(path)
         true
       end
 
+      # Resolve the UnboundMethod +mod+ itself declares for +method_name+,
+      # skipping prepended wrapper modules (e.g. a DI method-logpoint wrapper)
+      # that shadow it.
+      #
+      # @param mod [Module] Module or class expected to declare the method
+      # @param method_name [Symbol] Name of the method to resolve
+      # @return [UnboundMethod] Method whose owner is +mod+, or the
+      #   prepend-resolved method when +mod+ owns no method in the super chain
+      def declared_instance_method(mod, method_name)
+        original_method = mod.instance_method(method_name)
+        method = original_method
+        until method.owner.equal?(mod)
+          super_method = method.super_method
+          return original_method unless super_method
+          method = super_method
+        end
+        method
+      end
+
       # Find source file for a module.
       # Prefers user code paths over gem/stdlib paths. ActiveRecord models have
       # generated methods (autosave callbacks) whose source is in the gem, but
@@ -349,18 +379,23 @@ def find_source_file(mod)
 
         # Try instance methods first
         mod.instance_methods(false).each do |method_name|
-          method = mod.instance_method(method_name)
+          method = declared_instance_method(mod, method_name)
           location = method.source_location
           next unless location
 
           path = location[0]
           return path if user_code_path?(path)
 
           fallback ||= path # steep:ignore
+        rescue Exception => e # standard:disable Lint/RescueException
+          Datadog::DI.reraise_if_fatal(e)
+          @logger.debug { "symdb: error resolving #{safe_mod_name(mod)}##{method_name}: #{e.class}: #{e.message}" }
         end
 
         # Try singleton methods
         mod.singleton_methods(false).each do |method_name|
+          # Method-logpoint probes prepend instance methods only, so singleton
+          # methods are never shadowed by a DI wrapper.
           method = mod.method(method_name)
           location = method.source_location
           next unless location
@@ -369,6 +404,9 @@ def find_source_file(mod)
           return path if user_code_path?(path)
 
           fallback ||= path # steep:ignore
+        rescue Exception => e # standard:disable Lint/RescueException
+          Datadog::DI.reraise_if_fatal(e)
+          @logger.debug { "symdb: error resolving #{safe_mod_name(mod)}.#{method_name}: #{e.class}: #{e.message}" }
         end
 
         # Use const_source_location to find where this class/module is declared.
@@ -460,10 +498,9 @@ def wrap_in_file_scope(file_path, inner_scopes)
       # Extract MODULE scope (without file_hash — that belongs on the FILE root scope).
       # Does not include nested classes — nesting is handled by extract_all via FQN splitting.
       # @param mod [Module] The module
+      # @param source_file [String] Source file resolved by the caller
       # @return [Scope] The module scope
-      def extract_module_scope(mod)
-        source_file = find_source_file(mod)
-
+      def extract_module_scope(mod, source_file)
         Scope.new(
           scope_type: "MODULE",
           name: safe_mod_name(mod),
@@ -476,11 +513,11 @@ def extract_module_scope(mod)
 
       # Extract CLASS scope
       # @param klass [Class] The class
+      # @param source_file [String] Source file resolved by the caller
       # @return [Scope] The class scope
-      def extract_class_scope(klass)
-        methods = klass.instance_methods(false)
-        start_line, end_line = calculate_class_line_range(klass, methods)
-        source_file = find_source_file(klass)
+      def extract_class_scope(klass, source_file)
+        method_scopes = build_class_method_scopes(klass)
+        start_line, end_line = line_range_from_scopes(method_scopes)
 
         Scope.new(
           scope_type: "CLASS",
@@ -489,36 +526,20 @@ def extract_class_scope(klass)
           start_line: start_line,
           end_line: end_line,
           language_specifics: build_class_language_specifics(klass),
-          scopes: extract_method_scopes(klass),
+          scopes: method_scopes,
           symbols: extract_scope_symbols(klass)
         )
       end
 
-      # Calculate class line range from method locations.
-      # Start from the earliest method start, end at the latest method end (derived
-      # from iseq trace_points so methods spanning multiple lines aren't truncated).
-      # @param klass [Class] The class
-      # @param methods [Array<Symbol>] Method names
+      # Derive a [start_line, end_line] range from method scopes, ignoring
+      # UNKNOWN sentinel values. The end is the max of method end_lines so a
+      # multi-line method's last line is not truncated to its start_line.
+      # @param method_scopes [Array<Scope>]
       # @return [Array<Integer, Integer>] [start_line, end_line]
-      def calculate_class_line_range(klass, methods)
-        starts = []
-        ends = []
-        methods.each do |method_name|
-          method = klass.instance_method(method_name)
-          location = method.source_location
-          next unless location && location[0]
-          starts << location[1]
-          _ranges, method_end = extract_targetable_lines(method, location[1])
-          ends << method_end
-        end
-
-        return [UNKNOWN_MIN_LINE, UNKNOWN_MAX_LINE] if starts.empty?
-
-        [starts.min, ends.max]
-      rescue Exception => e # standard:disable Lint/RescueException
-        Datadog::DI.reraise_if_fatal(e)
-        @logger.debug { "symdb: error calculating line range for #{safe_mod_name(klass)}: #{e.class}: #{e.message}" }
-        [UNKNOWN_MIN_LINE, UNKNOWN_MAX_LINE]
+      def line_range_from_scopes(method_scopes)
+        starts = method_scopes.map(&:start_line).reject { |line| line == UNKNOWN_MIN_LINE }
+        ends = method_scopes.map(&:end_line).reject { |line| line == UNKNOWN_MAX_LINE }
+        [starts.min || UNKNOWN_MIN_LINE, ends.max || UNKNOWN_MAX_LINE]
       end
 
       # Build language specifics for CLASS
@@ -566,66 +587,33 @@ def build_class_language_specifics(klass)
         {}
       end
 
-      # Extract method scopes from a class
-      # @param klass [Class] The class
-      # @return [Array<Scope>] Method scopes
-      def extract_method_scopes(klass)
-        scopes = []
-
-        # Get all instance methods (public, protected, private)
-        all_instance_methods = klass.instance_methods(false) +
+      # Build a METHOD scope for each user-code instance method +klass+
+      # declares.
+      # @param klass [Class]
+      # @return [Array<Scope>]
+      def build_class_method_scopes(klass)
+        method_names = (klass.instance_methods(false) +
           klass.protected_instance_methods(false) +
-          klass.private_instance_methods(false)
-        all_instance_methods.uniq!
+          klass.private_instance_methods(false)).uniq
 
-        all_instance_methods.each do |method_name|
-          method_scope = extract_method_scope(klass, method_name, :instance)
-          scopes << method_scope if method_scope
-        end
+        Core::Utils::EnumerableCompat.filter_map(method_names) do |method_name|
+          method = declared_instance_method(klass, method_name)
+          location = method.source_location
+          next unless location
+          next unless user_code_path?(location[0])
 
-        scopes
+          build_instance_method_scope(klass, method_name, method)
+        rescue Exception => e # standard:disable Lint/RescueException
+          Datadog::DI.reraise_if_fatal(e)
+          @logger.debug { "symdb: error resolving #{safe_mod_name(klass)}##{method_name}: #{e.class}: #{e.message}" }
+          nil
+        end
       rescue Exception => e # standard:disable Lint/RescueException
         Datadog::DI.reraise_if_fatal(e)
         @logger.debug { "symdb: failed to extract methods from #{safe_mod_name(klass)}: #{e.class}: #{e.message}" }
         []
       end
 
-      # Extract a single method scope
-      # @param klass [Class] The class
-      # @param method_name [Symbol] Method name
-      # @param method_type [Symbol] :instance or :class
-      # @return [Scope, nil] Method scope or nil
-      def extract_method_scope(klass, method_name, method_type)
-        method = klass.instance_method(method_name)
-        location = method.source_location
-
-        return nil unless location  # Skip methods without source location
-
-        source_file, line = location
-        return nil unless user_code_path?(source_file)  # Skip gem/stdlib methods
-
-        targetable_lines, end_line = extract_targetable_lines(method, line)
-
-        Scope.new(
-          scope_type: "METHOD",
-          name: method_name.to_s,
-          source_file: source_file,
-          start_line: line,
-          end_line: end_line,
-          targetable_lines: targetable_lines,
-          language_specifics: {
-            visibility: method_visibility(klass, method_name),
-            method_type: method_type.to_s,
-            arity: method.arity,
-          },
-          symbols: extract_method_parameters(method)
-        )
-      rescue Exception => e # standard:disable Lint/RescueException
-        Datadog::DI.reraise_if_fatal(e)
-        @logger.debug { "symdb: failed to extract method #{safe_mod_name(klass)}##{method_name}: #{e.class}: #{e.message}" }
-        nil
-      end
-
       # Get method visibility
       # @param klass [Class] The class
       # @param method_name [Symbol] Method name
@@ -808,7 +796,7 @@ def collect_method_names_by_file(mod)
         # visibilities without an intermediate merged array.
         [mod.instance_methods(false), mod.private_instance_methods(false)].each do |method_names|
           method_names.each do |method_name|
-            method = mod.instance_method(method_name)
+            method = declared_instance_method(mod, method_name)
             loc = method.source_location
             next unless loc
             next unless user_code_path?(loc[0])
@@ -852,7 +840,7 @@ def build_file_scope(file_path, entries)
           # live only as long as the tree node holds them; they are released when
           # convert_tree_to_scope finishes building the file's Scope.
           method_infos = Core::Utils::EnumerableCompat.filter_map(method_names) do |name|
-            method = mod.instance_method(name)
+            method = declared_instance_method(mod, name)
             # Pass 1 (build_per_file_index) recorded this method under file_path.
             # If the method has been redefined in another file between the two
             # passes (e.g. a class reopened during a Rails reload while extract_all
@@ -989,13 +977,7 @@ def convert_node_to_scope(node)
         # Recurse into child scopes (nested modules/classes)
         child_scopes = node[:children].values.map { |child| convert_node_to_scope(child) }
 
-        # Compute line range: start from the earliest method start, end at the latest
-        # method end. Using max(start_line) would underreport the class's end_line for
-        # classes whose last method spans multiple lines.
-        starts = method_scopes.map(&:start_line).reject { |l| l == UNKNOWN_MIN_LINE } # steep:ignore
-        ends = method_scopes.map(&:end_line).reject { |l| l == UNKNOWN_MAX_LINE } # steep:ignore
-        start_line = starts.empty? ? UNKNOWN_MIN_LINE : starts.min
-        end_line = ends.empty? ? UNKNOWN_MAX_LINE : ends.max
+        start_line, end_line = line_range_from_scopes(method_scopes)
 
         # Extract symbols (constants, class variables) if we have the actual module object
         symbols = node[:mod] ? extract_scope_symbols(node[:mod]) : []

sig/datadog/symbol_database/extractor.rbs

@@ -43,25 +43,27 @@ module Datadog
 
       def resolves_to_same_module?: (String mod_name, Module mod) -> bool
 
+      def user_code_module_name?: (Module mod) -> bool
+
       def user_code_module?: (Module mod) -> bool
 
       def user_code_path?: (String path) -> bool
 
+      def declared_instance_method: (Module mod, ::Symbol method_name) -> UnboundMethod
+
       def find_source_file: (Module mod) -> String?
 
       def wrap_in_file_scope: (String file_path, Array[Scope] inner_scopes) -> Scope
 
-      def extract_module_scope: (Module mod) -> Scope
+      def extract_module_scope: (Module mod, String source_file) -> Scope
 
-      def extract_class_scope: (Class klass) -> Scope
+      def extract_class_scope: (Class klass, String source_file) -> Scope
 
-      def calculate_class_line_range: (Class klass, Array[::Symbol] methods) -> [Integer, Integer]
+      def line_range_from_scopes: (Array[Scope] method_scopes) -> [Integer, Integer]
 
       def build_class_language_specifics: (Class klass) -> language_specifics
 
-      def extract_method_scopes: (Class klass) -> Array[Scope]
-
-      def extract_method_scope: (Class klass, ::Symbol method_name, ::Symbol method_type) -> Scope?
+      def build_class_method_scopes: (Class klass) -> Array[Scope]
 
       def method_visibility: (Class klass, ::Symbol method_name) -> String