class Bench
  def initialize(desc)
    @desc = desc
    @code = nil
    @file = nil
    @failures = ""
    @results = 0
    if @@printed_header == 0
      print_header
    end
    if @@ftime == 0
      @@ftime = `./itest 5`.to_i
    end
    if @@stpd.nil?
      for path in ['/usr/libexec/systemtap/stpd', '/usr/local/libexec/systemtap/stpd'] 
	if File.exist?(path)
	  @@stpd = path
	  break
	end
      end
    end
    if @@runtime.nil?
      for path in ['/usr/share/systemtap/runtime', '/usr/local/share/systemtap/runtime'] 
	if File.exist?(path)
	  @@runtime = path
	  break
	end
      end
    end
  end

  attr_writer :code, :file
  attr_reader :failures, :results

  def run
    Signal.trap("INT") do
      cleanup
      exit
    end
    compile
    load
    @results = `./itest 1`.to_i - @@ftime
    sleep 5
    `sudo killall -HUP stpd`
    `tail #{@dir}/xxx > #{@dir}/xxx.out`
    File.open("#{@dir}/xxx.out") do |file|
      file.each_line {|line| @failures += line if line =~ /WARNING/}
    end
    Process.wait
    cleanup
  end

  def print
    if @results
      printf("runtime: %-20s\t%d\n", @desc, @results)
    end
  end

  private
  @dir = nil

  protected

  # function call overhead
  @@ftime = 0
  @@printed_header = 0
  @@stpd = nil
  @@runtime = nil

  def cleanup
    `/bin/rm -f stap.out` if File.exists?("stap.out")
    `/bin/rm -f bench.stp` if File.exists?("bench.stp")
    `/bin/rm -rf #{@dir}` unless @dir.nil?
  end

  def load
    fork do exec "sudo #{@@stpd} -rmq #{@dir}/bench.ko > #{@dir}/xxx" end
    sleep 5
  end

  def compile
    @dir = `mktemp -dt benchXXXXX`.strip
    emit_code
    modules = "/lib/modules/" + `uname -r`.strip + "/build"
    res=`make -C \"#{modules}\" M=\"#{@dir}\" modules V=1`
    if !File.exist?("#{@dir}/bench.ko")
      puts res
      cleanup
      exit
    end
  end

  def emit_code
    if @code
      begin
	f = File.new("#{@dir}/bench.c","w")
	f << "#include \"runtime.h\"
#include \"probes.c\"\n
MODULE_DESCRIPTION(\"SystemTap probe: bench\");
MODULE_AUTHOR(\"automatically generated by bench2/run_bench\");\n\n"
	f << "static int inst_sys_getuid (struct kprobe *p, struct pt_regs *regs) {\n"
	f << @code
	f << "\n  return 0;\n}
static struct kprobe kp[] = {\n  {\n    .addr = \"sys_getuid\",
    .pre_handler = inst_sys_getuid\n  }\n};\n
#define NUM_KPROBES 1\n
int probe_start(void)\n{\n  return _stp_register_kprobes (kp, NUM_KPROBES);\n}\n
void probe_exit (void)\n{\n  _stp_unregister_kprobes (kp, NUM_KPROBES); \n}\n"
      rescue
	puts "Error writing source file"
      ensure
	f.close unless f.nil?
      end
      makefile = File.new("#{@dir}/Makefile","w")
      makefile << "CFLAGS += -Wno-unused -Werror
CFLAGS += -I \"#{@@runtime}\"
obj-m := bench.o
"
      makefile.close
    else
      puts "NO CODE!"
    end
  end

  def print_header
    @@printed_header = 1
    nproc=`grep ^processor /proc/cpuinfo`.count("\n")
    physical_cpus=`grep "physical id" /proc/cpuinfo`.split("\n").uniq.length
    model=`grep "model name" /proc/cpuinfo`.match(/(model name\t: )([^\n]*)/)[2]
    puts "SystemTap BENCH2 \t" + `date`
    puts "kernel: " + `uname -r`.strip + " " + `uname -m`.strip
    begin
      puts IO.read("/etc/redhat-release")
    rescue
    end

    puts `uname -n`.strip + ": " + `uptime`
    if nproc > 1
      puts "processors: #{nproc} (#{physical_cpus} physical) #{model}"
    else
      puts "processors: #{nproc} #{model}"
    end

    begin
      mem=IO.read("/proc/meminfo").split("\n")
      puts mem[0] + "\t" + mem[1]
    rescue
    end
    puts "-"*64
    check_cpuspeed
    @@ftime = `./itest 5`.to_i
    puts "For comparison, function call overhead is #@@ftime nsecs."
    puts "Times below are in nanoseconds and include kprobe overhead."
    puts "-"*64
  end

  def check_cpuspeed
    maxfile = "/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq"
    minfile = "/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq"
    if File.exist?(maxfile)
      maxfreq = `cat #{maxfile}`.to_i
      minfreq = `cat #{minfile}`.to_i
      if minfreq != maxfreq
	`sudo /bin/sh -c \"echo #{maxfreq} > #{minfile}\"`
	puts "CPU frequency scaling detected and disabled."
	puts "After testing, enable it again by typing (as root)"
	puts "echo #{minfreq} > #{minfile}"
	puts "-"*64
	end
    end
  end
end

class Stapbench < Bench
  def print
    if @results
      printf("script : %-20s\t%d\n", @desc, @results)
    end
  end

  protected

  def load
    # we do this in several steps because the compilation phase can take a long time
    `stap -kvvp4 -m bench bench.stp &> stap.out`
    IO.foreach("stap.out") {|line| @dir = line if line =~ /Created temporary directory/} 
    @dir = @dir.match(/"([^"]*)/)[1]
    if !File.exist?("#{@dir}/bench.ko")
      puts `cat stap.out`
      cleanup
      exit
    end
    fork do exec "sudo #{@@stpd} -rmq #{@dir}/bench.ko > #{@dir}/xxx" end
    sleep 5
  end
  
  def compile
    emit_code
  end
  
  def emit_code
    if @file
      File.open("bench.stp","w") do |b|
	File.open(@file,"r") do |f|
	  f.each_line do |line|
	    b.puts(line.sub(/TEST/,'kernel.function("sys_getuid")'))
	  end
	end
      end
    elsif @code
      begin
	f = File.new("bench.stp","w")
	f << "probe kernel.function(\"sys_getuid\") {\n"
	f << @code
	f << "\n}\n"
      rescue
	puts "Error writing source file"
      ensure
	f.close unless f.nil?
      end
    else
      puts "NO CODE!"
    end
  end
end