#!/usr/bin/env ruby
#
#	unicorn_status - A munin plugin for Linux to monitor unicorn processes
#
#  Copyright (C) 2010 Shinji Furuya - shinji.furuya@gmail.com
#  Licensed under the MIT license:
#  http://www.opensource.org/licenses/mit-license.php
#

# set path to your rails app
RAILS_ROOT = '/path/to/rails/app'.freeze

# set name to your unicorn.pid
PID_NAME = 'unicorn.pid'.freeze

module Munin
  class UnicornStatus
    attr_reader :pid_file

    def initialize(rails_root, pid_name)
      @pid_file = "#{rails_root}/tmp/pids/#{pid_name}"
    end

    def master_pid
      File.read(pid_file).to_i
    end

    def worker_pids
      result = []
      ps_output = `ps w --ppid #{master_pid}`
      ps_output.split("\n").each do |line|
        chunks = line.strip.split(/\s+/, 5)
        pid = chunks[0]
        pcmd = chunks[4]
        next if pid !~ /\A\d+\z/ || pcmd !~ /worker/

        result << pid.to_i
      end
      result
    end

    def worker_count
      worker_pids.size
    end

    def idle_worker_count
      result = 0
      before_cpu = {}
      worker_pids.each do |pid|
        before_cpu[pid] = cpu_time(pid)
      end
      sleep 1
      after_cpu = {}
      worker_pids.each do |pid|
        after_cpu[pid] = cpu_time(pid)
      end
      worker_pids.each do |pid|
        result += 1 if after_cpu[pid] - before_cpu[pid] == 0
      end
      result
    end

    def cpu_time(pid)
      usr, sys = `cat /proc/#{pid}/stat | awk '{print $14,$15 }'`.strip.split(/\s+/).collect { |i| i.to_i }
      usr + sys
    end
  end
end

case ARGV[0]
when 'autoconf'
  puts 'yes'
when 'config'
  puts "graph_title Unicorn [#{File.basename(__FILE__).gsub(/^unicorn_status_/, '')}] - Status"
  puts 'graph_args -l 0'
  puts 'graph_vlabel number of workers'
  puts 'graph_category webserver'
  puts 'total_worker.label total_workers'
  puts 'idle_worker.label idle_workers'
else
  m = Munin::UnicornStatus.new(ENV['rails_root'] || RAILS_ROOT, ENV['pid_name'] || PID_NAME)
  puts "total_worker.value #{m.worker_count}"
  puts "idle_worker.value #{m.idle_worker_count}"
end