From 9bbdb3132cc0f2af4d17902f87e4585916d3ff0f Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Sun, 16 Jun 2024 23:03:04 +0300
Subject: [PATCH] Add ability to profile commands via CLI

---
 CHANGELOG.md                   |  2 ++
 README.md                      | 15 ++++++++--
 lib/memory_profiler/autorun.rb | 17 ++++++++++++
 lib/memory_profiler/cli.rb     | 50 ++++++++++++++++------------------
 test/test_cli.rb               |  5 ++++
 test/test_reporter.rb          |  2 +-
 6 files changed, 61 insertions(+), 30 deletions(-)
 create mode 100644 lib/memory_profiler/autorun.rb

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a876d9a..977d754 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,7 @@
 # Changelog
 
+- Add ability to profile commands via CLI @fatkodima
+
 ## 1.0.1 - 23-10-2022
 
 - Adapts tests to Ruby 3.0 / 3.1
diff --git a/README.md b/README.md
index 7da53a0..4d55f02 100644
--- a/README.md
+++ b/README.md
@@ -32,14 +32,23 @@ There are two ways to use `memory_profiler`:
 ### Command Line
 
 The easiest way to use memory_profiler is via the command line, which requires no modifications to your program. The basic usage is:
+
+```
+$ ruby-memory-profiler [options] run [--] command [command-options]
 ```
-$ ruby-memory-profiler [options] <script.rb> [--] [script-options]
+
+Example:
+
+```
+$ ruby-memory-profiler --pretty run -- rubocop --cache false
+
+$ ruby-memory-profiler --max=10 --pretty run -- ruby notify_users.rb 1 2 3 --quiet
 ```
-Where `script.rb` is the program you want to profile.
 
 For a full list of options, execute the following command:
+
 ```
-ruby-memory-profiler -h
+$ ruby-memory-profiler -h
 ```
 
 ### Convenience API
diff --git a/lib/memory_profiler/autorun.rb b/lib/memory_profiler/autorun.rb
new file mode 100644
index 0000000..745dfd3
--- /dev/null
+++ b/lib/memory_profiler/autorun.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+require "memory_profiler"
+require "base64"
+
+def deserialize_hash(data)
+  Marshal.load(Base64.urlsafe_decode64(data)) if data
+end
+
+options = deserialize_hash(ENV["MEMORY_PROFILER_OPTIONS"]) || {}
+
+at_exit do
+  report = MemoryProfiler.stop
+  report.pretty_print(**options)
+end
+
+MemoryProfiler.start(options)
diff --git a/lib/memory_profiler/cli.rb b/lib/memory_profiler/cli.rb
index 71bd052..ee8a613 100644
--- a/lib/memory_profiler/cli.rb
+++ b/lib/memory_profiler/cli.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "optparse"
+require "base64"
 
 module MemoryProfiler
   class CLI
@@ -14,20 +15,6 @@ class CLI
       ignore_files: "memory_profiler/lib"
     }.freeze
 
-    REPORTER_KEYS = [
-      :top, :trace, :ignore_files, :allow_files
-    ].freeze
-
-    RESULTS_KEYS = [
-      :to_file, :color_output, :retained_strings, :allocated_strings,
-      :detailed_report, :scale_bytes, :normalize_paths
-    ].freeze
-
-    private_constant :BIN_NAME, :VERSION_INFO,:STATUS_SUCCESS, :STATUS_ERROR,
-                     :DEFAULTS, :REPORTER_KEYS, :RESULTS_KEYS
-
-    #
-
     def run(argv)
       options = {}
       parser = option_parser(options)
@@ -43,17 +30,24 @@ def run(argv)
         return STATUS_ERROR
       end
 
-      MemoryProfiler.start(reporter_options(options))
-      load script
-
-      STATUS_SUCCESS
+      if script == "run"
+        # We are profiling a command.
+        profile_command(options, argv)
+      else
+        # We are profiling a ruby file.
+        begin
+          MemoryProfiler.start(options)
+          load(script)
+        ensure
+          report = MemoryProfiler.stop
+          report.pretty_print(**options)
+        end
+        STATUS_SUCCESS
+      end
     rescue OptionParser::InvalidOption, OptionParser::InvalidArgument, OptionParser::MissingArgument => e
       puts parser
       puts e.message
       STATUS_ERROR
-    ensure
-      report = MemoryProfiler.stop
-      report&.pretty_print(**results_options(options))
     end
 
     private
@@ -66,7 +60,7 @@ def option_parser(options)
               A Memory Profiler for Ruby
 
           Usage:
-              #{BIN_NAME} [options] <script.rb> [--] [script-options]
+              #{BIN_NAME} [options] run [--] command [command-options]
         BANNER
 
         opts.separator ""
@@ -138,12 +132,16 @@ def option_parser(options)
       end
     end
 
-    def reporter_options(options)
-      options.select { |k, _v| REPORTER_KEYS.include?(k) }
+    def profile_command(options, argv)
+      env = {}
+      env["MEMORY_PROFILER_OPTIONS"] = serialize_hash(options) if options.any?
+      gem_path = File.expand_path('../', __dir__)
+      env["RUBYOPT"] = "-I #{gem_path} -r memory_profiler/autorun #{ENV['RUBYOPT']}"
+      exec(env, *argv)
     end
 
-    def results_options(options)
-      options.select { |k, _v| RESULTS_KEYS.include?(k) }
+    def serialize_hash(hash)
+      Base64.urlsafe_encode64(Marshal.dump(hash))
     end
   end
 end
diff --git a/test/test_cli.rb b/test/test_cli.rb
index db0330f..bcd9109 100644
--- a/test/test_cli.rb
+++ b/test/test_cli.rb
@@ -111,4 +111,9 @@ def test_returns_error_when_script_not_specified
       assert_equal 1, result
     end
   end
+
+  def test_profiles_commands
+    out = `bin/ruby-memory-profiler --pretty run -- ruby -e '[]'`
+    assert_includes(out, "Total allocated:")
+  end
 end
diff --git a/test/test_reporter.rb b/test/test_reporter.rb
index e33c24a..257c64d 100644
--- a/test/test_reporter.rb
+++ b/test/test_reporter.rb
@@ -12,7 +12,7 @@ def setup
   def default_block
     # Create an object from a gem outside memory_profiler which allocates
     # its own objects internally
-    MiniTest::Reporter.new
+    Minitest::Reporter.new
 
     # Create 10 strings
     10.times { |i| i.to_s }