Add benchmarks for memory profiler

2 files changed, 94 insertions, 0 deletions

diff --git a/Rakefile b/Rakefile
index a3c3fde..3c69c17 100644
--- a/Rakefile
+++ b/Rakefile
@@ -92,6 +92,36 @@ namespace :benchmark do
     )
   end
 
+  desc 'Memory profiler'
+  task :memory, [:top_x, :mime_types_only] => 'benchmark:support' do |_, args|
+    require 'benchmarks/memory_profiler'
+    Benchmarks::ProfileMemory.report(
+      mime_types_only: args.mime_types_only,
+      top_x: args.top_x
+    )
+  end
+
+  desc 'Columnar memory profiler'
+  task 'memory:columnar', [:top_x, :mime_types_only] => 'benchmark:support' do |_, args|
+    require 'benchmarks/memory_profiler'
+    Benchmarks::ProfileMemory.report(
+      columnar: true,
+      mime_types_only: args.mime_types_only,
+      top_x: args.top_x
+    )
+  end
+
+  desc 'Columnar allocation counts (full load)'
+  task 'memory:columnar:full', [:top_x, :mime_types_only] => 'benchmark:support' do |_, args|
+    require 'benchmarks/memory_profiler'
+    Benchmarks::ProfileMemory.report(
+      columnar: true,
+      full: true,
+      top_x: args.top_x,
+      mime_types_only: args.mime_types_only
+    )
+  end
+
   desc 'Object counts'
   task objects: 'benchmark:support' do
     require 'benchmarks/object_counts'
diff --git a/support/benchmarks/memory_profiler.rb b/support/benchmarks/memory_profiler.rb
new file mode 100644
index 0000000..c5ae405
--- /dev/null
+++ b/support/benchmarks/memory_profiler.rb
@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+if RUBY_VERSION < '2.3'
+  warn 'Cannot use memory_profiler on less than 2.3.8.'
+  exit 1
+end
+
+begin
+  require 'memory_profiler'
+rescue Exception # rubocop:disable Lint/RescueException
+  warn "Memory profiling requires the gem 'memory_profiler'."
+  exit 1
+end
+
+module Benchmarks
+  # Use Memory Profiler to profile memory
+  class ProfileMemory
+    def self.report(columnar: false, full: false, mime_types_only: false, top_x: nil)
+      new(
+        columnar: columnar,
+        full: full,
+        mime_types_only: mime_types_only,
+        top_x: top_x
+      ).report
+    end
+
+    def initialize(columnar:, full:, mime_types_only:, top_x:)
+      @columnar = !!columnar
+      @full = !!full
+      @mime_types_only = !!mime_types_only
+
+      @top_x = top_x
+
+      return unless @top_x
+
+      @top_x = top_x.to_i
+      @top_x = 10 if @top_x <= 0
+    end
+
+    def report
+      collect.pretty_print
+    end
+
+    private
+
+    def collect
+      report_params = {
+        top: @top_x,
+        allow_files: @mime_types_only ? %r{mime-types/lib} : nil
+      }.delete_if { |_k, v| v.nil? }
+
+      if @columnar
+        MemoryProfiler.report(**report_params) do
+          require 'mime/types'
+          MIME::Types.first.to_h if @full
+        end
+      else
+        MemoryProfiler.report(**report_params) do
+          require 'mime/types/full'
+        end
+      end
+    end
+  end
+end