Add possibility to convert rbtrace output to flamegraph format

.github/workflows/ci.yml

@@ -23,4 +23,4 @@ jobs:
           bundler-cache: true
 
       - name: Tests
-        run: ./test.sh
+        run: ./test.sh && rake test

README.md

@@ -62,6 +62,23 @@ for popular ruby libraries and functions.
 
     % rbtrace -p <PID> -c activerecord io
 
+## flamegraph
+
+You could generate flamegraph from traces
+
+```
+rbtrace -p <PID> > trace.log
+rbtrace --convert=flamegraph trace.log > out.stacks
+```
+
+Then run [flamegraph](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl)
+```
+flamegraph.pl --color=io --width=800 --title="Rbtrace" out.stacks > rbtrace.svg
+firefox rbtrace.svg
+```
+
+Or upload `out.stacks` to e.g. https://speedscope.app
+
 ## detailed example
 
 ### require rbtrace into a process

Rakefile

@@ -1,4 +1,5 @@
 require "bundler/gem_tasks"
+require "rake/testtask"
 
 desc "Compile the c extension"
 task :compile do
@@ -11,4 +12,12 @@ end
 
 task :build => :compile
 
+task :default => :test
 
+desc 'Run the rbtracer test suite'
+Rake::TestTask.new do |t|
+  t.libs += %w(lib ext test)
+  t.test_files = Dir['test/**_test.rb']
+  t.verbose = true
+  t.warning = true
+end

lib/rbtrace/cli.rb

@@ -1,6 +1,7 @@
 require 'optimist'
 require 'rbtrace/rbtracer'
 require 'rbtrace/version'
+require 'rbtrace/converter'
 
 class RBTraceCLI
   singleton_class.send(:attr_accessor, :tracer)
@@ -228,6 +229,10 @@ def self.run
       opt :shapesdump,
         "generate a shapes dump for the process in FILENAME",
         :default => "AUTO"
+
+      opt :convert,
+        "convert a trace file to another format",
+        :default => "flamegraph"
     end
 
     opts = Optimist.with_standard_exception_handling(parser) do
@@ -243,8 +248,8 @@ def self.run
       ARGV.clear
     end
 
-    unless %w[ fork eval interactive backtrace backtraces slow slowcpu firehose methods config gc memory heapdump].find{ |n| opts[:"#{n}_given"] }
-      $stderr.puts "Error: --slow, --slowcpu, --gc, --firehose, --methods, --interactive, --backtraces, --backtrace, --memory, --heapdump, --shapesdump or --config required."
+    unless %w[ fork eval interactive backtrace backtraces slow slowcpu firehose methods config gc memory heapdump convert].find{ |n| opts[:"#{n}_given"] }
+      $stderr.puts "Error: --slow, --slowcpu, --gc, --firehose, --methods, --interactive, --backtraces, --backtrace, --memory, --heapdump, --shapesdump, --convert or --config required."
       $stderr.puts "Try --help for help."
       exit(-1)
     end
@@ -360,6 +365,14 @@ def self.run
       end
     end
 
+    if opts[:convert_given]
+      require 'rbtrace/converter'
+      converter = RBTrace::Converter.build(opts[:convert], ARGF)
+      converter.convert
+
+      return
+    end
+
     if opts[:exec_given]
       tracee = fork{
         Process.setsid

lib/rbtrace/converter.rb

@@ -0,0 +1,77 @@
+# Converts the output of rbtrace to a different format
+class RBTrace::Converter
+  # Public: The IO where converter input is read.
+  attr_accessor :input
+
+  # Public: The IO where converter output is written (default: STDOUT).
+  attr_accessor :output
+
+  # Build a converter of the given type and input
+  #
+  # type - The Symbol type of the converter to build.
+  # input - The IO where converter input is read.
+  #
+  # Returns a converter
+  def self.build(type, input)
+    case type.to_sym
+    when :flamegraph
+      Flamegraph.new(input)
+    else
+      raise "Unknown type: #{type}"
+    end
+  end
+
+  # Creates a new converter
+  #
+  # input - The IO where converter input is read.
+  #
+  # Returns a converter
+  def initialize(input)
+    @input = input
+    @output = $stdout
+  end
+
+  # Converts the input to a format suitable for the flamegraph tool,
+  # e.g. https://github.com/brendangregg/FlameGraph or https://www.speedscope.app/
+  class Flamegraph < RBTrace::Converter
+
+    # Converts and writes the output to the output IO.
+    # Returns nothing.
+    def convert
+      stack = []
+
+      input.each_line do |line|
+        match = line =~ /\S/
+        next if match.nil?
+
+        # Calculate depth by the number of leading spaces (each indent is 2 spaces in the example)
+        depth = match / 2
+        method_line = line.strip
+
+        matches = method_line.match(/(.*) <\s*(.*)>$/)
+        method =
+          if matches
+            method_name, timing = matches[1], matches[2]
+            timing = (timing.to_f * 1000000).to_i # int microseconds
+            [method_name.strip, timing]
+          else
+            [method_line]
+          end
+
+        if method.size == 2 && stack.length == depth # on the same depth
+          output.puts (stack + ["#{method[0]} #{method[1]}"]).join(';')
+          next
+        end
+
+        if method.size == 1 # start of new block
+          stack.push(method)
+          next
+        end
+
+        if stack.length > depth # go to lesser level
+          stack.pop
+        end
+      end
+    end
+  end
+end

rbtrace.gemspec

@@ -24,6 +24,7 @@ Gem::Specification.new do |s|
   s.add_dependency 'msgpack',  '>= 0.4.3'
 
   s.add_development_dependency "rake"
+  s.add_development_dependency "minitest"
 
   s.license = "MIT"
   s.summary = 'rbtrace: like strace but for ruby code'

test/converter_flamegraph_test.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require File.expand_path('../test_helper', __FILE__)
+require 'stringio'
+
+class ConverterFlamegraphTest < TestCase
+  def convert(input)
+    output = StringIO.new
+    converter = RBTrace::Converter::Flamegraph.new(StringIO.new(input))
+    converter.output = output
+    converter.convert
+
+    output.rewind
+
+    output.read
+  end
+
+  def test_convert_flamegraph
+    input = <<~INPUT
+IO.select <0.000044>
+IO.select <0.000055>
+
+Puma::Client#eagerly_finish
+  IO.select <0.000010>
+  Puma::Client#try_to_finish
+    BasicSocket#read_nonblock
+      BasicSocket#__read_nonblock <0.000018>
+    BasicSocket#read_nonblock <0.000029>
+  Puma::Client#try_to_finish <0.000153>
+Puma::Client#eagerly_finish <0.000179>
+Puma::ThreadPool#<<
+  Thread::Mutex#synchronize
+    Thread::ConditionVariable#signal <0.000008>
+  Thread::Mutex#synchronize <0.000027>
+Puma::ThreadPool#<< <0.000035>
+    INPUT
+
+    expected_output = <<~OUTPUT
+IO.select 44
+IO.select 55
+Puma::Client#eagerly_finish;IO.select 10
+Puma::Client#eagerly_finish;Puma::Client#try_to_finish;BasicSocket#read_nonblock;BasicSocket#__read_nonblock 18
+Puma::ThreadPool#<<;Thread::Mutex#synchronize;Thread::ConditionVariable#signal 8
+    OUTPUT
+
+    assert_equal expected_output, convert(input)
+  end
+end

test/test_helper.rb

@@ -0,0 +1,11 @@
+# encoding: UTF-8
+
+ext_path = File.expand_path(File.join(__dir__, '..', 'ext'))
+$LOAD_PATH.unshift(File.expand_path(ext_path))
+
+require 'rbtrace'
+require 'rbtrace/cli'
+require 'minitest/autorun'
+
+class TestCase < Minitest::Test
+end