#!/usr/bin/env ruby

# Reports current CPU, disk, load average, and memory use to riemann.

require File.expand_path('../../lib/riemann/tools', __FILE__)

class Riemann::Tools::Health
  include Riemann::Tools
  
  opt :cpu_warning, "CPU warning threshold (fraction of total jiffies)", :default => 0.9
  opt :cpu_critical, "CPU critical threshold (fraction of total jiffies)", :default => 0.95
  opt :disk_warning, "Disk warning threshold (fraction of space used)", :default => 0.9
  opt :disk_critical, "Disk critical threshold (fraction of space used)", :default => 0.95
  opt :load_warning, "Load warning threshold (load average / core)", :default => 3
  opt :load_critical, "Load critical threshold (load average / core)", :default => 8
  opt :memory_warning, "Memory warning threshold (fraction of RAM)", :default => 0.85
  opt :memory_critical, "Memory critical threshold (fraction of RAM)", :default => 0.95

  def initialize
    @limits = {
      :cpu => {:critical => opts[:cpu_critical], :warning => opts[:cpu_warning]},
      :disk => {:critical => opts[:disk_critical], :warning => opts[:disk_warning]},
      :load => {:critical => opts[:load_critical], :warning => opts[:load_warning]},
      :memory => {:critical => opts[:memory_critical], :warning => opts[:memory_warning]}
    }
  end

  def alert(service, state, metric, description)
    report(
      :service => service,
      :state => state.to_s,
      :metric => metric.to_f,
      :description => description
    )
  end

  def cores
    i = 0;
    File.read("/proc/cpuinfo").split(/\n\n/).inject({}) do |cores, p|
      physical_id = p[/physical id\s+:\s+(\d+)/, 1]
      core_id = p[/core id\s+:\s+(\d+)/, 1]
      if physical_id and core_id
        cores["#{physical_id}:#{core_id}"] = true
      elsif physical_id
        cores["#{physical_id}:"] = true
      else
        cores[i += 1] = true;
      end

      cores
    end.size
  end

  def cpu
    new = File.read('/proc/stat')
    unless new[/cpu\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)/]
      alert 'cpu', :unknown, nil, "/proc/stat doesn't include a CPU line"
      return false
    end
    u2, n2, s2, i2 = [$1, $2, $3, $4].map { |e| e.to_i }
    
    if @old_cpu
      u1, n1, s1, i1 = @old_cpu

      used = (u2+n2+s2) - (u1+n1+s1)
      total = used + i2-i1
      fraction = used.to_f / total

      if fraction > @limits[:cpu][:critical]
        alert "cpu", :critical, fraction, "#{sprintf("%.2f", fraction * 100)}% user+nice+sytem\n\n#{cpu_report}"
      elsif fraction > @limits[:cpu][:warning]
        alert "cpu", :warning, fraction, "#{sprintf("%.2f", fraction * 100)}% user+nice+sytem\n\n#{cpu_report}"
      else
        alert "cpu", :ok, fraction, "#{sprintf("%.2f", fraction * 100)}% user+nice+sytem\n\n#{cpu_report}"
      end
    end
    
    @old_cpu = [u2, n2, s2, i2]
  end

  def cpu_report
    `ps -eo pcpu,pid,args | sort -nrb -k1 | head -10`.chomp
  end

  def disk
    `df -P`.split(/\n/).each do |r|
      f = r.split(/\s+/)
      next unless f[0] =~ /^\//
      next if f[0] == 'Filesystem'
      x = f[4].to_f/100

      if x > @limits[:disk][:critical]
        alert "disk #{f[5]}", :critical, x, "#{f[4]} used"
      elsif x > @limits[:disk][:warning]
        alert "disk #{f[5]}", :warning, x, "#{f[4]} used"
      else
        alert "disk #{f[5]}", :ok, x, "#{f[4]} used"
      end
    end
  end
  
  def load
    load = File.read('/proc/loadavg').split(/\s+/)[2].to_f / cores
    if load > @limits[:load][:critical]
      alert "load", :critical, load, "15-minute load average/core is #{load}"
    elsif load > @limits[:load][:warning]
      alert "load", :warning, load, "15-minute load average/core is #{load}"
    else
      alert "load", :ok, load, "15-minute load average/core is #{load}"
    end
  end

  def memory
    m = File.read('/proc/meminfo').split(/\n/).inject({}) { |info, line|
      x = line.split(/:?\s+/)
      # Assume kB...
      info[x[0]] = x[1].to_i
      info
    }

    free = m['MemFree'] + m['Buffers'] + m['Cached']
    total = m['MemTotal']
    fraction = 1 - (free.to_f / total)

    if fraction > @limits[:memory][:critical]
      alert "memory", :critical, fraction, "#{sprintf("%.2f", fraction * 100)}% used\n\n#{memory_report}"
    elsif fraction > @limits[:memory][:warning]
      alert "memory", :warning, fraction, "#{sprintf("%.2f", fraction * 100)}% used\n\n#{memory_report}"
    else
      alert "memory", :ok, fraction, "#{sprintf("%.2f", fraction * 100)}% used\n\n#{memory_report}"
    end
  end

  def memory_report
    `ps -eo pmem,pid,args | sort -nrb -k1 | head -10`.chomp
  end

  def tick
    cpu
    memory
    load
    disk
  end
end
  
Riemann::Tools::Health.run