rubyprof.rb in fluent-rubyprof-0.2.0

- old
+ new
@@ -1,89 +1,121 @@
 require 'optparse'
 require 'drb/drb'
 
-def parse_options
-  op = OptionParser.new
-  op.banner += ' <start/stop> [output_file]'
+module Fluent
+  class Rubyprof
 
-  (class<<self;self;end).module_eval do
-    define_method(:usage) do |msg|
-      puts op.to_s
-      puts "error: #{msg}" if msg
-      exit 1
-    end
-  end
+    PRINTERS = {
+      'flat' => 'FlatPrinter',
+      'flat_with_line_numbers' => 'FlatPrinterWithLineNumbers',
+      'graph' => 'GraphPrinter',
+      'graph_html' => 'GraphHtmlPrinter',
+      'call_tree' => 'CallTreePrinter',
+      'call_stack' => 'CallStackPrinter',
+      'dot' => 'DotPrinter',
+      'multi' => 'MultiPrinter',
+    }
 
-  opts = {
-    host: '127.0.0.1',
-    port: 24230,
-    unix: nil,
-    command: nil, # start or stop
-    output: '/tmp/fluent-rubyprof.txt'
-  }
+    def parse_options(argv = ARGV)
+      op = OptionParser.new
+      op.banner += ' <start/stop> [output_file]'
 
-  op.on('-h', '--host HOST', "fluent host (default: #{opts[:host]})") {|v|
-    opts[:host] = v
-  }
+      (class<<self;self;end).module_eval do
+        define_method(:usage) do |msg|
+          puts op.to_s
+          puts "error: #{msg}" if msg
+          exit 1
+        end
+      end
 
-  op.on('-p', '--port PORT', "debug_agent tcp port (default: #{opts[:host]})", Integer) {|v|
-    opts[:port] = v
-  }
+      opts = {
+        host: '127.0.0.1',
+        port: 24230,
+        unix: nil,
+        command: nil, # start or stop
+        output: '/tmp/fluent-rubyprof.txt',
+        measure_mode: 'PROCESS_TIME',
+        printer: 'flat',
+      }
 
-  op.on('-u', '--unix PATH', "use unix socket instead of tcp") {|v|
-    opts[:unix] = v
-  }
+      op.on('-h', '--host HOST', "fluent host (default: #{opts[:host]})") {|v|
+        opts[:host] = v
+      }
 
-  op.on('-o', '--output PATH', "output path (default: #{opts[:output]})") {|v|
-    opts[:output] = v
-  }
+      op.on('-p', '--port PORT', "debug_agent tcp port (default: #{opts[:host]})", Integer) {|v|
+        opts[:port] = v
+      }
 
-  begin
-    op.parse!(ARGV)
-    opts[:command] = ARGV.shift
-    unless %w[start stop].include?(opts[:command])
-      usage "`start` or `stop` must be specified as the 1st argument"
-    end
-  rescue
-    usage $!.to_s
-  end
+      op.on('-u', '--unix PATH', "use unix socket instead of tcp") {|v|
+        opts[:unix] = v
+      }
 
-  opts
-end
+      op.on('-o', '--output PATH', "output path (default: #{opts[:output]})") {|v|
+        opts[:output] = v
+      }
 
-def main
-  opts = parse_options
+      op.on('-m', '--measure_mode MEASURE_MODE', "ruby-prof measure mode (default: #{opts[:measure_mode]})") {|v|
+        opts[:measure_mode] = v
+      }
 
-  unless opts[:unix].nil?
-    uri = "drbunix:#{opts[:unix]}"
-  else
-    uri = "druby://#{opts[:host]}:#{opts[:port]}"
-  end
+      op.on('-P', '--printer PRINTER', PRINTERS.keys,
+                "ruby-prof print format (default: #{opts[:printer]})",
+                "currently one of: #{PRINTERS.keys.join(', ')}") {|v|
+        opts[:printer] = v
+      }
+      op.parse!(argv)
 
-  $remote_engine = DRb::DRbObject.new_with_uri(uri)
+      opts[:command] = argv.shift
+      unless %w[start stop].include?(opts[:command])
+        raise OptionParser::InvalidOption.new("`start` or `stop` must be specified as the 1st argument")
+      end
 
-  case opts[:command]
-  when 'start'
-    remote_code = <<-CODE
-      require 'ruby-prof'
-      RubyProf.start
-    CODE
-  when 'stop'
-    remote_code = <<-"CODE"
-      result = RubyProf.stop
-      File.open('#{opts[:output]}', 'w') {|f|
-        RubyProf::FlatPrinter.new(result).print(f)
-      }
-    CODE
-  end
+      measure_modes = %w[PROCESS_TIME WALL_TIME CPU_TIME ALLOCATIONS MEMORY GC_RUNS GC_TIME] 
+      unless measure_modes.include?(opts[:measure_mode])
+        raise OptionParser::InvalidOption.new("-m allows one of #{measure_modes.join(', ')}")
+      end
 
-  $remote_engine.method_missing(:instance_eval, remote_code)
+      opts
+    end
 
-  case opts[:command]
-  when 'start'
-    $stdout.puts 'fluent-rubyprof: started'
-  when 'stop'
-    $stdout.puts "fluent-rubyprof: outputs to #{opts[:output]}"
+    def run
+      begin
+        opts = parse_options
+      rescue OptionParser::InvalidArgument, OptionParser::InvalidOption => e
+        usage e.message
+      end
+
+      unless opts[:unix].nil?
+        uri = "drbunix:#{opts[:unix]}"
+      else
+        uri = "druby://#{opts[:host]}:#{opts[:port]}"
+      end
+
+      $remote_engine = DRb::DRbObject.new_with_uri(uri)
+
+      case opts[:command]
+      when 'start'
+        remote_code = <<-CODE
+        require 'ruby-prof'
+        RubyProf.measure_mode = eval("RubyProf::#{opts[:measure_mode]}")
+        RubyProf.start
+        CODE
+      when 'stop'
+        remote_code = <<-"CODE"
+        result = RubyProf.stop
+        File.open('#{opts[:output]}', 'w') {|f|
+          RubyProf::#{PRINTERS[opts[:printer]]}.new(result).print(f)
+        }
+        CODE
+      end
+
+      $remote_engine.method_missing(:instance_eval, remote_code)
+
+      case opts[:command]
+      when 'start'
+        $stdout.puts 'fluent-rubyprof: started'
+      when 'stop'
+        $stdout.puts "fluent-rubyprof: outputs to #{opts[:output]}"
+      end
+    end
   end
 end
-
-main