flamegraph.rb

#
# Step 1:
#
#   FlameGraph.generate(output: 'asdf') { ... },
#
# Step 2:
#
#   git clone https://github.com/brendangregg/FlameGraph
#   cd FlameGraph
#   ./flamegraph.pl < ../asdf > out.svg

module FlameGraph
  def self.generate(max_depth: 10, output:)
    stack_n = ''
    stack_t = []
    res = Hash.new(0.0)
    first = true
    ignore_depth = 0
    tpoint = TracePoint.new(:call, :return, :c_call, :c_return) do |tp|
      if first # reject the Tracepoint.enable :c_return
        first = false
        next
      end
      next if stack_t.empty? && tp.defined_class == 'TracePoint'

      if tp.event == :call || tp.event == :c_call
        if stack_t.size == max_depth
          ignore_depth += 1
        else
          stack_t << Process.clock_gettime(Process::CLOCK_MONOTONIC)
          stack_n << ";#{tp.defined_class}##{tp.method_id}"
        end
      else
        if ignore_depth > 0
          ignore_depth -= 1
        else
          top = stack_t.pop
          diff = Process.clock_gettime(Process::CLOCK_MONOTONIC) - top
          res[stack_n.dup] += diff
          stack_n.sub!(/;[^;]+$/, '')
        end
      end
    end
    tpoint.enable
    ret = yield
    tpoint.disable
    File.open(output, 'w') do |f|
      res.each { |k, v| f.puts "#{k[1..-1]} #{(v * 1_000_000_000).to_i}" }
    end
    ret
  end
end