#!/usr/bin/env ruby

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

require 'trollop'
require 'riemann/client'

class Riemann::Health
  def initialize(opts)
    @host = opts[:host]
    @port = opts[:port]
    @interval = opts[:interval]
    @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]}
    }

    @client = Riemann::Client.new(:host => @host, :port => @port)
  end

  def alert(service, state, metric, description)
    @client << {
      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`.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
    begin
      cpu
      memory
      load
      disk
    rescue => e
      $stderr.puts "#{e.class} #{e}\n#{e.backtrace.join "\n"}"
      sleep 10
    end
  end

  def run
    loop do
      tick
      sleep @interval
    end
  end
end
  
Riemann::Health.new(Trollop.options do
  opt :host, "Host", :default => '127.0.0.1'
  opt :port, "Port", :default => 5555
  opt :interval, "Seconds between updates", :default => 5
  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
end).run