lib/heap_profiler/reporter.rb in heap-profiler-0.2.1 vs lib/heap_profiler/reporter.rb in heap-profiler-0.3.0

@@ -7,12 +7,10 @@
     # and for anonymous modules and classes this mean naming them.
     #
     # So we name them at the start of the profile to avoid that.
     #
     # See: https://github.com/ruby/ruby/pull/3349
-    #
-    # TODO: Could we actually do the dump ourselves? objspace is a extension already.
     if RUBY_VERSION < '2.8'
       def name_anonymous_modules!
         ObjectSpace.each_object(Module) do |mod|
           next if mod.singleton_class?
           next if real_mod_name(mod)
@@ -54,14 +52,18 @@
     end
 
     def stop
       HeapProfiler.name_anonymous_modules!
       ObjectSpace.trace_object_allocations_stop if @enable_tracing
+
+      # we can't use partial dump for allocated.heap, because we need old generations
+      # as well to build the classes and strings indexes.
       dump_heap(@allocated_heap)
+
       GC.enable
       4.times { GC.start }
-      dump_heap(@retained_heap)
+      dump_heap(@retained_heap, partial: true)
       @allocated_heap.close
       @retained_heap.close
       write_info("generation", @generation.to_s)
     end
 
@@ -82,20 +84,27 @@
 
     def write_info(key, value)
       File.write(File.join(@dir_path, "#{key}.info"), value)
     end
 
-    # ObjectSpace.dump_all does allocate a few objects in itself (https://bugs.ruby-lang.org/issues/17045)
-    # because of this even en empty block of code will report a handful of allocations.
-    # To filter them more easily we attribute call `dump_all` from a method with a very specific `file`
-    # property.
-    class_eval <<~RUBY, '__hprof', __LINE__
-      # frozen_string_literal: true
-      def dump_heap(file)
-        ObjectSpace.dump_all(output: file)
+    if RUBY_VERSION >= '3.0'
+      def dump_heap(file, partial: false)
+        ObjectSpace.dump_all(output: file, since: partial ? @generation : nil)
         file.close
       end
-    RUBY
+    else
+      # ObjectSpace.dump_all does allocate a few objects in itself (https://bugs.ruby-lang.org/issues/17045)
+      # because of this even en empty block of code will report a handful of allocations.
+      # To filter them more easily we attribute call `dump_all` from a method with a very specific `file`
+      # property.
+      class_eval <<~RUBY, '__hprof', __LINE__
+        # frozen_string_literal: true
+        def dump_heap(file, partial: false)
+          ObjectSpace.dump_all(output: file)
+          file.close
+        end
+      RUBY
+    end
 
     def open_heap(name)
       File.open(File.join(@dir_path, "#{name}.heap"), 'w+')
     end