IRB’s Custom Measure Procedures

If you came here from my post about IRB’s built-in measure command, thanks for reading part two too! If not, that is a good place to start learning about IRB’s measure command!

Custom measure procedures

Measure can also output the results of custom measurement procedures. The rest of this blog post will work through an example of creating a custom measurement procedure, and in doing so, demonstrate how you could write your own ones for your needs.

I’m a big fan of using Speedscope for flamegraph visualizations. In order to produce a flamegraph, Speedscope needs a profile of our code. Speedscope can take json formatted Stackprof output. So, let’s continue with our Stackprof theme, and work through creating a custom measure procedure which will open a flamegraph in Speedscope.

How will we call the procedure we write? All of the measure procedures are defined in IRB.conf[:MEASURE_PROC]. We can take a look at the defaults:

irb(main):002:0> IRB.conf[:MEASURE_PROC]
=> {:TIME=> #<Proc:0x00007fbca19038e0 /path/to/lib/ruby/3.0.0/irb/init.rb:116>,
 :STACKPROF=> #<Proc:0x00007fbca19038b8 /path/to/lib/ruby/3.0.0/irb/init.rb:123>}

We see our two built-in procedures! And if we look closely, can see exactly where they’re defined in the IRB source. To write our own, we can add to this IRB.conf[:MEASURE_PROC] hash.

Next, where will we put this procedure we add to the IRB.conf[:MEASURE_PROC] hash? A good place to set this would be in your ~/.irbrc file. If you’ve never used this before, it’s where you can define anything to be loaded into each IRB console you open. Many people define methods or require gems in their ~/.irbrc to save themselves from rewriting the same snippets in each IRB console they open.

Enough of the ~/.irbrc ramble, let’s get back to talking about measure. It turns out, measure (or measure :on) will actually default first to running any proc we’ve defined in IRB.conf[:MEASURE_PROC][:CUSTOM] and fall back to the proc defined by IRB.conf[:MEASURE_PROC][:TIME]. So if we’d like a non-time default proc, we can define it in IRB.conf[:MEASURE_PROC][:CUSTOM], like this:

irb > IRB.conf[:MEASURE_PROC][:CUSTOM] = Proc.new do
irb >   puts "New default measure proc!"
irb > end
=> #<Proc:0x00007fe76b0a41c8 (irb):1>

irb > measure
CUSTOM is added.
=> nil

irb > 1
New default measure proc!
=> nil

However, for our case, we want a named procedure so that we can call measure :that_name (like we do measure :stackprof). The lowercased key for our procedure in the IRB.conf[:MEASURE_PROC] hash will be the argument we pass to measure in order to call that procedure. For instance, we can add :PRINTING to IRB.conf[:MEASURE_PROC] and then call measure :printing:

irb > IRB.conf[:MEASURE_PROC][:PRINTING] = Proc.new do
irb >   puts "Wahoo! We're in a custom measure proc!"
irb > end
=> #<Proc:0x00007fc2f500f838 (irb):1>

irb > measure :printing
PRINTING is added.
=> nil

irb > 1
Wahoo! We're in a custom measure proc!
=> nil

Arguments for custom procedures

In the examples above, you might have noticed that we didn’t actually return any values when adding the custom measure procedures. Instead, we puts‘d some string, and then returned nil. In most cases though, we’ll want to actually execute the block of code that a user passes, in addition to printing out some measurement output.

Measure procedures have five parameters. They are not keyword parameters, so order is relevant here:

context (IRB::Context) Describes the state of the current IRB session
code (String): The snippet of code entered into the IRB console
line_number (Integer): The line number in the IRB console. This is a counter which increments with each new line of code typed into the console
arg (Any!): Any arguments passed when calling the measure procedure. For example, measure :stackprof, :wall would have an arg value of :wall
block (Proc): A proc surrounding the code entered. (Calling block.() will execute the code)

Jumping back to the speedscope custom measure proc we’re going to write, we’ll definitely need to use block so we can execute the code. We’ll also need arg so we can allow the user to set the sampling mode (like in the Stackprof example).

Lastly, it could actually be counterproductive to open a speedscope window each time a user inputs some value into IRB. Instead, let’s have a workaround where we’ll open the flamegraph in speedscope for the first time after a user calls measure :speedscope, and then any subsequent time they wish to open a speedscope window, they can type :speedscope into the console, and we’ll open it on the next command. For this to work, we’ll need to use code from above, so we know if the user has entered :speedscope.

Before trying this snippet on your own, be sure to install speedscope (npm install -g speedscope), the JSON gem (gem install json) and the stackprof gem (gem install stackprof). Enough words, I’ll let the code do the rest of the talking:

IRB.conf[:MEASURE_PROC][:SPEEDSCOPE] = Proc.new do |_, code, _, arg, &block|
  begin
    # IRB.conf[:SPEEDSCOPED] is the flag we'll use to determine if we're going
    # to speedscope the command or not
    if IRB.conf[:SPEEDSCOPED]
      if code == ":speedscope\n"
        IRB.conf[:SPEEDSCOPED] = false
        puts "Will speedscope the next command!"
      end

      # Run the code the user inputs
      block.()
    else
      require "stackprof"
      require "json"

      result = nil
      json_profile = JSON.generate(
        StackProf.run(mode: arg || :cpu, raw: true) { result = block.() }
      )

      puts "Opening a web console with a flamegraph visualized in speedscope!"

      Tempfile.create do |f|
        f.write(json_profile)
        f.rewind
        `speedscope #{f.path}`
      end

      IRB.conf[:SPEEDSCOPED] = true
      result
    end
  rescue LoadError
    puts "In order to see speedscope results, " \
      "both `json` and `stackprof` gems must be installed"
    block.()
  rescue Errno::ENOENT
    puts "In order to see speedscope results, " \
      "install speedscope with `npm install -g speedscope`"
    block.()
  end
end

We can put this block of code in our ~/.irbrc. IRB will then load this custom measure proc each time we open a new IRB console. We can test this by running measure :speedscope from a new IRB console:

irb > measure :speedscope
SPEEDSCOPE is added.
=> nil

irb > snippet
Opening a web console with a flamegraph visualized in speedscope!
=> 10000

Tada! This opened up a web console with a Speedscope flamegraph visualization. I took a screenshot of a section of the output both to show you what Speedscope looks like, and to funnel my excitement somewhere:

Speedscope

In order to turn on Speedscope again, we can simply enter :speedscope into the console like this:

irb > snippet
=> 10000

irb > :speedscope
Will speedscope the next command!
=> :speedscope

irb > snippet
Opening a web console with a flamegraph visualized in speedscope!
=> 10000

We’ve now successfully written our own custom measure procedure!

TL;DR

In summary, here is what we learned about IRB’s new measure functionality

We can define custom measure procs in the IRB.conf[:MEASURE_PROC] hash
We can call these procs using measure :lowercase_key_in_the_hash
Custom proc parameters are context, code, line_number, arg, &block