require 'find'

$subsys_array = Array.new
$allsubsys = Hash.new
$allmounts = Hash.new

class Sub_system
  def initialize(name, mount, option)
    @name = name
    @mount= mount
    @hierarchy = Array.new
    if (option =~ /.*noprefix.*/) then
      @prefix =""
    else
      @prefix = name +"."
    end
   @option = option
   @writable_files = Array.new
  end

  def mount_point
    @mount
  end

  def type
    @name
  end

  def myfile(name, attr)
    name + "/" + @prefix + attr
  end

  def option
    @option
  end
  #
  # walk directroy tree and add Cgroups to hash.
  #
  def reload
    @hierarchy.clear
    len = @mount.size
    Find.find(@mount) do |file|
      if File.directory?(file) then
	@hierarchy.push(file);
      end
    end
  end

  def each_cgroup(&block)
    @hierarchy.each(&block)
  end

  def ent(id)
    if (id < 0) then return nil
    end
    return @hierarchy.at(id)
  end

  def size
    @hierarchy.size
  end

  def stat (name)
    [["Not implemented", ""]]
  end

  def writable_files(name)
    list=Array.new
    @writable_files.each {|x| list.push(myfile(name,x))}
    list.push(name+"/tasks")
    return list
  end

  def tasks(name)
    list=Array.new
    begin
      File.open(name+"/tasks", "r") do |file|
        file.each_line do |x|
        x.chomp!
        list.push(x)
        end
      end
    rescue
      return nil
    end
    return list
  end
end

def read_oneline_file(filename)
  val=nil
  begin
    f = File.open(filename, "r")
    line = f.readline
    val = line.to_i
  rescue
    throw :readfailure,false
  ensure
    f.close if f != nil
  end
  return val
end

#
#for CPU subsystem
#
class Cpu_Subsys < Sub_system
  def initialize(mount, option)
    super("cpu", mount, option)
    @writable_files += ["shares"]
  end

  def read_share (name)
    ret = nil
    catch :readfailure do
      val = read_oneline_file(myfile(name,"shares"))
      return [val.to_s, val.to_s+" (100%)"] if (name == @mount)
      all=0

      dirname = File.dirname(name)
      Dir.foreach(dirname) do |x|
        next if ((x == ".") || (x == ".."))
        x = "#{dirname}/#{x}"
        next unless File.directory?(x)
        next unless File.exist?(myfile(name,"shares"))
        got = read_oneline_file(myfile(name,"shares"))
        all+=got
      end
      share=sprintf("%d (%.1f%%)", all, val*100.0/all)
      ret = [val.to_s, share]
    end
    return ret
  end

  def stat(name)
    level=0
    data = Array.new
    pos = @mount
    name_array = Array.new
    loop do
      name_array.push(name)
      break if name == @mount
      name = File.dirname(name)
    end
    name_array.reverse!
    name_array.each do |x|
      val = read_share(x)
      if val == nil then
        data = nil
        break
      end
      str = sprintf("%5s / %s", val[0], val[1])
      data.push([x, str])
    end
    return data if (data != nil && data.size > 0)
    return nil
  end
end

#
#for CPUacct subsystem
#
class Cpuacct_Subsys < Sub_system
  def initialize(mount, option)
    super("cpuacct", mount, option)
  end
  def stat(name)
    data = Array.new
    catch :read_failure do
      val = read_oneline_file(myfile(name, "usage"))
      data.push(["All", val.to_s])
      begin
        f = File.open(myfile(name,"usage_percpu"), "r")
        id=0
        line = f.readline
        while (line =~/\d+/) do
          line =$'
          data.push(["cpu"+id.to_s, $&])
          id += 1
        end 
      rescue
        data.clear
      ensure
        f.close if f != nil
      end
    return data if data.size > 0
    return nil
    end
  end
end

#
# For cpuset
#
class Cpuset_Subsys < Sub_system
  def initialize(mount, option)
    super("cpuset", mount, option)
    @elements =["cpu_exclusive","cpus","mems", "mem_exclusive","mem_hardwall",
                "memory_migrate", "memory_pressure", "memory_pressure_enabled",
                "memory_spread_page","memory_spread_slab",
                "sched_load_balance","sched_relax_domain_level"]
    @writable_files += @elements
  end
  def stat(name)
    data = Array.new
    for x in @elements
      begin
        filename = myfile(name, x)
        next unless (File.file?(filename))
        File.open(filename, "r") do | file |
          str = file.readline
          str.chomp!
          case x
          when "cpus", "mems"
            str = "empty" if (str == "")
          end
          data.push([x,str])
        end
      rescue
        #data = nil
        break
      end
    end
    return data
  end
end
#
#for Memory Subsys
#
def convert_bytes(bytes, precise)
  case
  when (precise == 0) && (bytes > 64 * 1024*1024*1024*1024)
    sprintf("Unlimited")
  when (precise == 0) && (bytes > 1024*1024*1024*1024)
    sprintf("%dT",bytes/1024/1024/1024/1024)
  when (precise == 0) && (bytes > 1024*1024*1024)
    sprintf("%dG", bytes/1024/1024/1024)
  when (bytes > 1024*1024)
    sprintf("%dM", bytes/1024/1024)
  when (bytes > 1024)
    sprintf("%dk", bytes/1024)
  else
    sprintf("%d", bytes)
  end
end

#
#for Memory Subsystem
#
class Memory_Subsys < Sub_system
  def initialize(mount, option)
    super("memory", mount, option)
    if (File.exist?("#{mount}/memory.memsw.usage_in_bytes")) then
      @memsw=true
    else
      @memsw=false
    end
    @writable_files += ["limit_in_bytes", "use_hierarchy","swappiness"]
    if (@memsw) then
      @writable_files += ["memsw.limit_in_bytes"]
    end
  end
  #
  # Find a root directroy of hierarchy.
  #
  def find_hierarchy_root(name)
    cur=[name, File.dirname(name)]
    ret=@mount
    while (cur[0] != @mount)
      ret="hoge"
      under = read_oneline_file("#{cur[1]}/memory.use_hierarchy")
      if (under == 0) then
        return cur[0]
      end
      cur[0] = cur[1]
      cur[1] = File.dirname(cur[1])
    end
    return ret
  end
  #
  # Generate an array for reporintg status
  #
  def stat(name)
    data = Array.new

    success = catch(:readfailure) do
      under =read_oneline_file(myfile(name,"use_hierarchy"))
      if (under == 1) then
        str=find_hierarchy_root(name)
        if (str != name) then
          str="under #{str}"
          under=2
        else
          str="hierarchy ROOT"
        end
      else #Not under hierarchy
        str=""
      end
      ent = ["Memory Subsys", str]
      data.push(ent)
      
      # Limit and Usage
      x=Array.new
      x.push("Usage/Limit")

      bytes = read_oneline_file(myfile(name,"usage_in_bytes"))
      usage = convert_bytes(bytes, 1)

      if (@memsw) then
        bytes = read_oneline_file(myfile(name,"memsw.usage_in_bytes"))
        usage2 = convert_bytes(bytes, 1)
        usage = "#{usage} (#{usage2})"
      end
      bytes = read_oneline_file(myfile(name,"limit_in_bytes"))
      limit = convert_bytes(bytes, 0)
      usage = "#{usage} / #{limit}"
      if (@memsw) then
        bytes = read_oneline_file(myfile(name,"memsw.limit_in_bytes"))
        limit2 = convert_bytes(bytes, 0)
        usage = "#{usage} (#{limit2})"
      end
      x.push(usage)

      data.push(x)

      # MAX USAGE
      x = Array.new
      x.push("Max Usage")
      bytes = read_oneline_file(myfile(name, "max_usage_in_bytes"))
      usage = convert_bytes(bytes, 1)
      if (@memsw) then
        bytes = read_oneline_file(myfile(name,"memsw.max_usage_in_bytes"))
        usage2 = convert_bytes(bytes, 1)
        usage = "#{usage} (#{usage2})"
      end
      x.push(usage)
      data.push(x)

      # failcnt
      x = Array.new
      x.push("Fail Count")
      cnt = read_oneline_file(myfile(name,"failcnt"))
      failcnt = cnt.to_s
      if (@memsw) then
        cnt = read_oneline_file(myfile(name, "memsw.failcnt"))
        failcnt="#{failcnt} (#{cnt.to_s})"
      end
      x.push(failcnt)
      data.push(x)

      begin
        f = File.open(myfile(name,"stat"), "r")
        for x in ["Cache","Rss","Pagein","Pageout",nil, nil, nil, nil, nil,
                  "HierarchyLimit","SubtreeCache","SubtreeRss"]
          line =f.readline
          next if x == nil
          line =~ /^\S+\s+(.+)/
          val=$1
          case x
          when "Cache","Rss","SubtreeCache","SubtreeRss"
            bytes = convert_bytes(val.to_i, 1)
            data.push([x, bytes])
          when "Pagein","Pageout"
            data.push([x, val])
          when "HierarchyLimit"
            memlimit = convert_bytes(val.to_i, 0)
            if (@memsw) then
              line =f.readline
              line =~ /^\S+\s+(.+)/
              memswlimit = convert_bytes(val.to_i, 0)
              memlimit += " (" + memswlimit + ")"
            end
            data.push([x, memlimit])
          end
        end
      ensure
        f.close if f != nil
      end
    true
    end
    return data if success==true
    return nil
  end
end

#
# Read /proc/mounts and parse each lines.
# When cgroup mount point is found, each subsystem's cgroups are added
# to subsystem's Hash.
#

def register_subsys(name, mount, option)
  if $allsubsys[name] == nil then
    subsys = nil
    case name
    when "cpu" then subsys = Cpu_Subsys.new(mount, option)
    when "cpuacct" then subsys = Cpuacct_Subsys.new(mount, option)
    when "memory" then  subsys = Memory_Subsys.new(mount, option)
    when "cpuset" then  subsys = Cpuset_Subsys.new(mount, option)
    end
    if subsys != nil then
      $subsys_array.push(name)
      $allsubsys[name] = subsys
    end
  end
end

#
# Read /proc/mounts and prepare subsys array
#
def parse_mount(line)
  parsed = line.split(/\s+/)
  if parsed[2] == "cgroup" then
    mount=parsed[1]
    opts=parsed[3].split(/\,/)
    opts.each do |name|
      case name
      when "rw" then next
      else
        register_subsys(name, mount, parsed[3])
        $allmounts[mount]=name
      end
    end
  end
end


def read_mount
  File.open("/proc/mounts", "r") do |file|
    file.each_line {|line| parse_mount(line) }
  end
  $subsys_array.sort!
end

#
# Read all /proc/mounts and scan directory under mount point.
#
def refresh_all
   $allmounts.clear
   $subsys_array.clear
   $allsubsys.clear
   read_mount
end

def check_and_refresh_mount_info
  
  mysubsys=Array.new
  File.open("/proc/mounts", "r") do |file|
    file.each_line do |line|
      parsed = line.split(/\s+/)
      if (parsed[2] == "cgroup")  then
        mysubsys.push(parsed[1])
      end
    end
  end

  if (mysubsys.size != $allmounts.size) then
    refresh_all
    return true 
  end
  
  mysubsys.each do |x|
    if ($allmounts[x] == nil) then
      refresh_all
      return true
    end
  end
  return false
end