From 3dfdb9f24c8920d9ad2aaa469e9098bdbbe6ca04 Mon Sep 17 00:00:00 2001
From: David Crosby <dcrosby@meta.com>
Date: Tue, 16 Apr 2024 08:06:43 -0700
Subject: [PATCH] Add profiling hooks

Summary: This adds the `--profiler` option for running ruby-prof.

Differential Revision: D56162222

fbshipit-source-id: ec7ccdc1641abfef63a3ff980ca6700bf2467ada
---
 .../files/default/bookworm/bookworm.rb        | 36 ++++++++++---------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/cookbooks/fb_bookworm/files/default/bookworm/bookworm.rb b/cookbooks/fb_bookworm/files/default/bookworm/bookworm.rb
index 56e3b726..d3deacc3 100755
--- a/cookbooks/fb_bookworm/files/default/bookworm/bookworm.rb
+++ b/cookbooks/fb_bookworm/files/default/bookworm/bookworm.rb
@@ -52,11 +52,10 @@ def initialize
       #   'Enable verbose mode',
       # )
 
-      # TODO(dcrosby) get ruby-prof working
-      # parser.on(
-      #   '--profiler',
-      #   '(WIP) Enable profiler for performance debugging',
-      # )
+      parser.on(
+        '--profiler',
+        'Enable profiler for performance debugging (requires ruby-prof)',
+      )
 
       parser.on(
         '--irb-config-step',
@@ -94,11 +93,12 @@ def parse
 end
 parser = Bookworm::CLIParser.new
 options = parser.parse
-# TODO(dcrosby) get ruby-prof working
-# if options[:profiler]
-#   require 'ruby-prof'
-#   RubyProf.start
-# end
+
+if options[:profiler]
+  require 'ruby-prof'
+  Bookworm::Profile = RubyProf::Profile.new
+  Bookworm::Profile.start
+end
 
 # We require the libraries *after* the profiler has a chance to start,
 # also means faster `bookworm -h` response
@@ -299,9 +299,13 @@ def build_report
   run.do_action
 end
 
-# TODO(dcrosby) get ruby-prof working
-# if options[:profiler]
-#   result = RubyProf.stop
-#   printer = RubyProf::FlatPrinter.new(result)
-#   printer.print($stdout)
-# end
+if options[:profiler]
+  result = Bookworm::Profile.stop
+  printer = RubyProf::GraphPrinter.new(result)
+  path = "#{Dir.tmpdir}/bookworm_profile-#{DateTime.now.iso8601(4)}.out"
+  printer = ::RubyProf::GraphPrinter.new(result)
+  File.open(path, 'w+') do |file|
+    printer.print(file)
+  end
+  puts "Wrote profiler output to #{path}"
+end