#!/usr/bin/ruby
#
# NAME
#  custodian-queue - Work with the queue.
#
# SYNOPSIS
#  custodian-queue  [ -h | --help ]
#                   [ -m | --manual ]
#                   [ -M | --monitor ]
#                   [ -f | --flush ]
#                   [ -s | --stats ]
#                   [ -S | --server 1.2.3.4:123 ]
#
# OPTIONS
#
#  -h, --help          Show a help message, and exit.
#
#  -m, --manual        Show this manual, and exit.
#
#  -M, --monitor       If the queue size exceeds the threshold report that via an exit code.
#
#  -f, --flush         Flush the queue, removing all jobs.
#
#  -s, --stats         Show the count of pending jobs.
#
#  -S, --server        Specify the host:port for the beanstalkd queue.
#
#
# ABOUT
#
#  This tool is designed to inspect the beanstalkd queue, or flush the pending
# jobs from it.
#
#
# AUTHOR
#
#  Steve Kemp  <steve@bytemark.co.uk>
#



#
# Standard modules
#
require 'beanstalk-client'
require 'getoptlong'


#
# Our code
#
require 'custodian/settings'



#
#  Entry-point to our code.
#
if __FILE__ == $0 then

  $FLUSH   = false
  $STATS   = false
  $help    = false
  $manual  = false
  $MONITOR = nil

  #
  #  The beanstalkd server address.
  #
  settings = Custodian::Settings.instance()
  $SERVER  = settings.queue_server
  $QUEUE   = settings.queue_name

  begin
    opts = GetoptLong.new(
                          [ "--flush",   "-f", GetoptLong::NO_ARGUMENT ],
                          [ "--help",    "-h", GetoptLong::NO_ARGUMENT ],
                          [ "--manual",  "-m", GetoptLong::NO_ARGUMENT ],
                          [ "--monitor", "-M", GetoptLong::OPTIONAL_ARGUMENT ],
                          [ "--server",  "-S", GetoptLong::REQUIRED_ARGUMENT ],
                          [ "--queue",   "-q", GetoptLong::REQUIRED_ARGUMENT ],
                          [ "--stats",   "-s", GetoptLong::NO_ARGUMENT ]
                          )
    opts.each do |opt, arg|
      case opt
      when "--monitor" then
          if ( arg )
              $MONITOR = arg.to_i
          else
              $MONITOR = 5000
          end
      when "--stats" then
          $STATS = true
      when "--flush" then
          $FLUSH = true
      when "--server" then
          $SERVER = arg
      when "--queue" then
          $QUEUE = arg
      when "--help" then
          $help = true
      when "--manual" then
          $manual = true
      end
    end
  rescue StandardError => ex
    puts "Option parsing failed: #{ex.to_s}"
    exit
  end

  #
  # CAUTION! Here be quality kode.
  #
  if $manual or $help

    # Open the file, stripping the shebang line
    lines = File.open(__FILE__){|fh| fh.readlines}[1..-1]

    found_synopsis = false

    lines.each do |line|

      line.chomp!
      break if line.empty?

      if $help and !found_synopsis
        found_synopsis = (line =~ /^#\s+SYNOPSIS\s*$/)
        next
      end

      puts line[2..-1].to_s

      break if $help and found_synopsis and line =~ /^#\s*$/

    end

    exit 0
  end

  #
  #  Create the queue object.
  #
  queue = Beanstalk::Pool.new([$SERVER], $QUEUE )

  #
  # Alerting on a queue that is too-full?
  #
  if ( !$MONITOR.nil? )

    #
    # Find the number of jobs
    #
    stats = queue.stats()
    jobs  = stats['current-jobs-ready'] || 0

    if ( jobs > $MONITOR )
        exit 1
    else
        exit 0
    end
  end


  #
  # Showing stats?
  #
  if ( $STATS )
     stats = queue.stats()
     puts "There are #{stats['current-jobs-ready'] || 0} jobs pending."
     exit( 0 )
  end


  #
  #  Are we flushing the queue?
  #
  if ( $FLUSH )
    count = 0
    while( true )
      begin
        job = queue.reserve(1)
        job.delete
        count += 1
      rescue Beanstalk::TimedOut => ex
        puts "Flushed #{count} pending jobs."
        exit( 0 )
      end
    end
  end


end