#!/bin/bash

WARNING_QUEUE_SIZE=10000
CRITICAL_QUEUE_SIZE=20000

function usage() {
    cat <<EOF
Usage $(basename $0) [--help|-h] [command]

Commands:
  help         - Print this help message
  master       - Check master process exists
  queue-size   - Check if task queue size exceeds preconfigured limits
  has-workers  - Check there are workers to execute tasks

Options:
  --help, -h   - Print this help message
EOF
    exit 1
}

# Extracts field value from 'field: value' pair.
function get_stat_field() {
    stat=$1
    field=$2
    echo "$stat" | grep "$field" | cut -d: -f2 | tr --delete '[[:space:]]'
}

function check_master() {
    stat=$1
    master=$(get_stat_field "$stat" 'masterId')
    if [ -z "$master" ]; then
        echo '2;No grinder master found'
    else
        echo "0;Ok: master is [$master]"
    fi
}

function check_queue_size() {
    stat=$1
    task_count=$(get_stat_field "$stat" 'taskCount')
    if [ $task_count -gt $CRITICAL_QUEUE_SIZE ]; then
        echo "2;Error: task queue size $task_count exceeds critical value $CRITICAL_QUEUE_SIZE"
    elif [ $task_count -gt $WARNING_QUEUE_SIZE ]; then
        echo "1;Warning: task queue size $task_count exceeds $WARNING_QUEUE_SIZE"
    else
        echo "0;Ok: task queue size is $task_count"
    fi
}

function check_have_workers() {
    stat=$1
    task_count=$(get_stat_field "$stat" 'taskCount')
    workers_count=$(get_stat_field "$stat" 'workerCount')
    if [ $task_count -gt 0 -a $workers_count -lt 1 ]; then
        echo "2;Error: have $task_count tasks but have no workers to execute them"
    else
        echo "0;Ok: have $workers_count workers and $task_count tasks"
    fi
}

stat=''

function init_stat() {
    if [ -z "$stat" ]; then
        stat=$(/usr/bin/grinder-stat)
        if [ $? -ne 0 ]; then
            echo "2;Error: $stat"
            exit 1
        fi
    fi
}

for cmd in "$@"; do
    case $cmd in
        "-h" | "--help" | "help")
            usage
            ;;
       "master")
            init_stat
            check_master "$stat"
            ;;
       "queue-size")
           init_stat
           check_queue_size "$stat"
           ;;
        "has-workers")
            init_stat
            check_have_workers "$stat"
            ;;
        *)
            echo "Bad command: $cmd"
            usage
            ;;
    esac
done
