# Benchmark Class for SystemTap
# Copyright (C) 2006, 2007 Red Hat Inc.
#
# This file is part of systemtap, and is free software.  You can
# redistribute it and/or modify it under the terms of the GNU General
# Public License (GPL); either version 2, or (at your option) any
# later version.

# Where to find laptop frequency files
MAXFILE = "/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq"
MINFILE = "/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq"

# more constants
STREAM = 1
BULK= 2

at_exit {Bench.done}

class Bench
  def initialize(desc)
    @desc = desc
    @code = nil
    @file = nil
    @trans = STREAM
    @outfile= "xxx"
    @failures = []
    @results = []
    if @@printed_header == 0
      print_header
    end
    if @@staprun.nil?
      for path in ['/usr/bin/staprun', '/usr/local/bin/staprun'] 
	if File.exist?(path)
	  @@staprun = path
	  break
	end
      end
      # now do make
      `make`
      if $? != 0 || !File.exist?('itest')
	puts "ERROR: Compiling itest failed\n"
	exit
      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
    Signal.trap("INT") do
      cleanup
      exit
    end
  end

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

  def run
    compile
    @results = []
    @failures = []
    @@num_threads.each do |threads|
      load
      sum=0
      `./itest #{threads} > #{@dir}/bench`
      `sudo killall -HUP stapio`
      Process.wait	# wait for staprun to exit
      @results[threads] = `cat #{@dir}/bench`.split[0].to_i - @@ftime[threads]
      File.open("#{@dir}/xxx.out") do |file|
	file.each_line do |line|
	  m = line.match(/WARNING: There were ([\d]*)/)
	  if (m)
	    @failures[threads] = m[1]
	    break
	  end
	end
      end
    end
    cleanup
  end

  def print
    if @results
      if self.kind_of? Stapbench
	printf("S")
      else
	printf("R")
      end
      if @trans == BULK
	printf("B")
      else
	printf(" ")
      end
      if @outfile == "/dev/null"
	printf("N")
      else
	printf(" ")
      end
      printf(": %-20s", @desc)
      @@num_threads.each {|thread| printf("\t%5d",@results[thread])}
      printf("\n")
      @@num_threads.each do |thread|
	printf("WARNING: %d transport failures with %d threads\n", \
	       @failures[thread], thread) unless @failures[thread].nil?
      end
    end
  end

  def Bench.done
    if @@minfreq != 0
      `sudo /bin/sh -c \"echo #{@@minfreq} > #{MINFILE}\"`
    end
  end

  private
  @dir = nil

  protected

  @@ftime = []
  @@printed_header = 0
  @@staprun = nil
  @@runtime = nil
  @@num_threads = []
  @@minfreq = 0

  def cleanup
    system('sudo killall -HUP stapio &> /dev/null')
    system('sudo /sbin/rmmod bench &> /dev/null')
    `/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?
    `/bin/rm -f stpd_cpu* probe.out`
  end

  def load
    fork do exec "sudo #{@@staprun} -o #{@outfile} #{@dir}/bench.ko &> #{@dir}/xxx.out" 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
	getuid="0x"+`grep sys_getuid$ /proc/kallsyms`.split[0]
	f = File.new("#{@dir}/bench.c","w")
	f << "#include \"runtime.h\"
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}\n"
	f << " static struct kprobe kp = {\n"
	f << "  .addr = (kprobe_opcode_t *)#{getuid}, \n"
	f << ".pre_handler = inst_sys_getuid\n  };\n
int probe_start(void)\n{\n  return register_kprobe(&kp);\n}\n
void probe_exit (void)\n{\n  unregister_kprobe (&kp); \n}\n"
      rescue
	puts "Error writing source file"
	exit
      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
"
      if @trans == BULK then makefile << "CFLAGS += -DSTP_BULKMODE" end
      makefile.close
    else
      puts "NO CODE!"
    end
  end

  def print_header
    @@printed_header = 1
    nproc=`grep ^processor /proc/cpuinfo`.count("\n")
    if nproc >= 16
      @@num_threads = [1,2,4,16]
    elsif nproc >= 8
      @@num_threads = [1,2,4,8]
    elsif nproc >= 4
      @@num_threads = [1,2,4]
    elsif nproc >= 2
      @@num_threads = [1,2]
    else
      @@num_threads = [1]
    end
    arch=`uname -m`.strip
    if (arch.match(/ppc64/))
	cpu=`grep "cpu" /proc/cpuinfo`.match(/(cpu\t\t: )([^\n]*)/)[2]
	clock=`grep "clock" /proc/cpuinfo`.match(/(clock\t\t: )([^\n]*)/)[2]
	revision=`grep "revision" /proc/cpuinfo`.match(/(revision\t: )([^\n]*)/)[2]
	cpuinfo=cpu + " " + clock + " revision: " + revision
    else
       	physical_cpus=`grep "physical id" /proc/cpuinfo`.split("\n").uniq.length
       	model=`grep "model name" /proc/cpuinfo`.match(/(model name\t: )([^\n]*)/)[2]
	cpuinfo="(#{physical_cpus} physical) #{model}"
    end	
    puts "SystemTap BENCH2 \t" + `date`
    puts "kernel: " + `uname -r`.strip + " " + `uname -m`.strip
    puts IO.read("/etc/redhat-release") if File.exists?("/etc/redhat-release")
    puts `uname -n`.strip + ": " + `uptime`
    puts "processors: #{nproc} #{cpuinfo}"

    begin
      mem=IO.read("/proc/meminfo").split("\n")
      puts mem[0] + "\t" + mem[1]
    rescue
    end
    puts "-"*64
    check_cpuspeed
    @@num_threads.each do |threads|
      @@ftime[threads] = `./itest #{threads} 10000000`.to_i
    end
    puts "Times below are nanoseconds per probe and include kprobe overhead."
    puts "-"*64
    puts "+---- S = Script, R = Runtime"
    puts "|+--- B = Bulk"
    puts "||+-- N = No output\t\t  Threads"
    printf "||| NAME                 "
    @@num_threads.each {|n| printf("\t    %d",n)}
    printf "\n"
  end

  def check_cpuspeed
    if @@minfreq == 0 && File.exist?(MAXFILE)
      maxfreq = `cat #{MAXFILE}`.to_i
      @@minfreq = `cat #{MINFILE}`.to_i
      if @@minfreq != maxfreq
	`sudo /bin/sh -c \"echo #{maxfreq} > #{MINFILE}\"`
	sleep 1
      end
    end
  end
end

#### STAPBENCH ######

class Stapbench < Bench

  def run
    compile
    @results = []
    @failures = []
    @@num_threads.each do |threads|
      load
      sum=0
      `./itest #{threads} > bench`
      `sudo killall -HUP stapio`
      if ($? != 0)
	puts "status=#{$?}"
	system("tail xxx.out")
      end
      Process.wait	# wait for stapio to exit
      @results[threads] = `cat bench`.split[0].to_i - @@ftime[threads]
      File.open("xxx.out") do |file|
	file.each_line do |line|
	  m = line.match(/WARNING: There were ([\d]*)/)
	  if (m)
	    @failures[threads] = m[1]
	    break
	  end
	end
      end
    end
    cleanup
  end

  protected

  def load
    fork do exec "sudo #{@@staprun} bench.ko &> xxx.out" end
    sleep 5
  end
  
  def compile
    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
    args = "-p4 -vv -DSTP_NO_OVERLOAD"
    if @trans == BULK then args = "-p4 -bvv -DSTP_NO_OVERLOAD" end
    `stap #{args} -m bench.ko bench.stp &> xxx.out`
    if ($? != 0)
      puts "compile failed. status=#{$?}"
      system("tail xxx.out")
      exit
    end
  end
  
end