#!/usr/bin/env ruby

require 'getoptlong'
require 'ftools'


def kmod_suffix
  suff = 'ko'
  v = `uname -r`
  suff = 'o' if v.slice(0..2) == '2.4'
  return suff
end

$mod_list = [
  TRACE = "kmemprof-trace",
  CGPROF = "kmemprof-cgprof",
  OBJFILT = "kmemprof-objfilter",
  FILT = "kmemprof-filter",
  FREQCTRL = "kmemprof-freqctrl",
  PLUG = "kmemprof-plug"
  ]

$load_status = {}

IO.popen("lsmod").readlines.slice(1..-1).collect{|l| l.split[0]}.each do |m|
  # little ugry
  m.sub!(/[_]/, '-')
  $load_status[m] = 1
end
  
def loaded? (mod)
  return true if $load_status[mod] == 1
  false
end

def insmod (mod, *args)
  mod = ENV['KMEMPROF_MODULES'] + '/' + mod + '.' + kmod_suffix
  comm = ['insmod', mod, *args]
  system(comm.join(' '))
end
  
def rmmod (mod, *args)
  mod = mod
  comm = ['rmmod', mod, *args]
  system(comm.join(' '))
end

def insmod_if_not_loaded(mod)
  insmod mod if ! loaded? mod
end

def rmmod_if_loaded(mod)
  rmmod mod if loaded? mod
end

### blocks

def save_ksyms
  proc {|file| File.copy '/proc/ksyms', file.to_s }
end

def set_freq
  proc do |arg|
    rmmod_if_loaded (FREQCTRL)
    if arg.to_i != 1
      insmod FREQCTRL, 'freq='+([arg]*32).join(',')
    end
  end
end

def set_filt
  proc do |arg|
    rmmod_if_loaded (FILT)
    if arg != ''
      insmod FILT, arg.to_s
    end
  end
end

def set_objfilt
  proc do |arg|
    rmmod_if_loaded(OBJFILT)
    if arg != ''
      insmod OBJFILT, 'objfilter=' + arg.to_s
    end
  end
end

def load_cgprof
  proc do |arg|
    # if TRACE has already been loaded...

    insmod_if_not_loaded (PLUG)
    IO.popen("cgprof-config " + arg) do |f|
      insmod CGPROF, f.readlines
    end
  end
end

def unload_cgprof
  proc { rmmod_if_loaded(CGPROF) }
end

def load_trace
  proc do |arg|
    if arg == ""
      nr_cpus=1
    else
      nr_cpus=arg.to_i
    end
    # if TRACE has already been loaded...

    insmod_if_not_loaded (PLUG)
    insmod TRACE, 'num_cpus=' + nr_cpus.to_s
  end
end

def unload_trace
  proc { rmmod_if_loaded(TRACE) }
end

def start
  proc { rmmod_if_loaded(PLUG) }
end

def stop
  proc { insmod_if_not_loaded(PLUG) }
end

def shutdown
  proc { $mod_list.each {|x| rmmod_if_loaded(x)} }
end

show_help = Proc.new do
  $options.each do |x|
    printf "\t"
    puts x.opts[0] + "\t\t" + x.description
  end
end

show_status = Proc.new do
  $load_status.each_key { |k| puts k}
end

class Option
  attr_reader :opts, :description, :handler
  
  def initialize(opts, handler, desc)
    @opts = opts
    @handler = handler
    @description= desc
  end

  def accept(opt, arg)
    handler.call(arg) if (opt == opts[0])
  end
end

options_args = [
  [ [ "--save-ksyms", "-l", GetoptLong::REQUIRED_ARGUMENT],
    save_ksyms,
    "writes out /proc/ksyms"],

  [ [ "--set-frequency", "-f", GetoptLong::REQUIRED_ARGUMENT],
    set_freq,
    "Sets the sampling frequency"],

  [ [ "--set-filter", "-r", GetoptLong::OPTIONAL_ARGUMENT], set_filt,
  "sets miscellaneous filtering"],

  [ [ "--set-objfilter", "-o", GetoptLong::OPTIONAL_ARGUMENT], set_objfilt,
  "limits sampling to paticular objects"],

  [ [ "--load-cgprof", "-p", GetoptLong::OPTIONAL_ARGUMENT], load_cgprof,
  "load the allocation call-graph profiling module"],

  [ [ "--unload-cgprof", "-q", GetoptLong::NO_ARGUMENT], unload_cgprof,
  "unload the allocation call-graph profiling module"],

  [ [ "--load-trace", "-t", GetoptLong::OPTIONAL_ARGUMENT], load_trace,
  "load the allocation tracing module"],
  
  [ [ "--unload-trace", "-u", GetoptLong::OPTIONAL_ARGUMENT], unload_trace,
  "unload the allocation tracing module"],

  [ [ "--start", "-b", GetoptLong::NO_ARGUMENT], start,
  "start profiling and tracing"],

  [ [ "--stop", "-e", GetoptLong::NO_ARGUMENT], stop,
  "stop profiling and tracing"],

  [ [ "--shutdown", "-d", GetoptLong::NO_ARGUMENT], shutdown,
  "unload all modules"],

  [ [ "--help", "-h", GetoptLong::NO_ARGUMENT], show_help,
  "display help message"],

  [ [ "--status", "-s", GetoptLong::NO_ARGUMENT], show_status,
  "display the current status of kmemprof"]
  ]

$options = options_args.collect{|args| Option.new(*args)}
optlong = (GetoptLong.new).set_options(*($options.collect{|x| x.opts}))

optlong.each do |opt, arg|
  $options.each do |k|
    k.accept(opt.to_s, arg)
  end
end

