#!/bin/sh
# vzubc - a tool for displaying OpenVZ user beancounters.
# Copyright (C) 2011, Parallels, Inc. All rights reserved.
# Licensed under GNU GPL v2 or later version.

umask 0077

# Ratio of held/limit to bypass quiet
Q_HELD=0.8
# Ratio of maxheld/limit to bypass quiet
Q_MAXHELD=0.5
# Directory to store previous failcnt values for relative mode
STOREDIR=/tmp/vzubc.store

usage() {
	cat << EOF
Usage: vzubc [option ...] [<filename>] [<ctid> ...]

 -w, --watch:
     Run itself under watch(1) (a la top mode, Ctrl-C to exit).
     -wd: show differences between runs
     -wn <time>: refresh every <time> seconds
 -q, --quiet:
     Quiet mode (only show beancounters with fails and close to limits).
     -qh <ratio>: quiet watermark for held/limit ratio (default: $Q_HELD)
     -qm <ratio>: quiet watermark for maxheld/limit ratio (default: $Q_MAXHELD)
 -r, --relative:
     Relative mode (show fail counters delta to previous run).
     -rc: clear all saved fail counter data for relative mode
     -rd <dir>: directory for relative data (default: $STOREDIR)

 <filename>: file to read UBC values from (defaults are
             /proc/bc/resources or /proc/user_beancounters)
 <ctid>: container name/ID to show info about (can be used multiple times)
EOF
	exit $1
}

CTIDS=""
FILE=""
WATCH=""
ARGV=$*

# For vzlist to work, make sure @SBINDIR@ is in $PATH
if ! echo ":${PATH}:" | fgrep -q ':@SBINDIR@:'; then
	PATH="@SBINDIR@:$PATH"
fi

chk_zeroarg() {
	if test -z "$2"; then
		echo "Error: option $1 requires an argument" 1>&2
		usage 1
	fi
}

while test $# -gt 0; do
	case $1 in
	-W)	# Internal use only
		watch=
		;;
	-w|--watch)
		if ! which watch >/dev/null; then
			echo "Error: watch binary not found" 1>&2
			exit 1
		fi
		watch=yes
		;;
	-wd)
		watch=yes
		WATCH_ARGS="$WATCH_ARGS -d"
		;;
	-wn)
		chk_zeroarg $1 $2
		watch=yes
		WATCH_ARGS="$WATCH_ARGS -n $2"
		shift
		;;
	-q|--quiet)
		quiet=1
		;;
	-qh)
		chk_zeroarg $1 $2
		quiet=1
		Q_HELD=$2
		shift
		;;
	-qm)
		chk_zeroarg $1 $2
		quiet=1
		Q_MAXHELD=$2
		shift
		;;
	-r|--relative)
		relative=1
		;;
	-rc)
		relative=1
		rm -fr $STOREDIR/
		;;
	-rd)
		chk_zeroarg $1 $2
		relative=1
		STOREDIR=$2
		;;
	-h|--help)
		usage 0
		;;
	-)
		# Use stdin for input
		FILE="-"
		;;
	-*)
		echo "Unrecognized option: $1" 1>&2
		usage 1
		;;
	0)
		# CTID 0, i.e. host system, vzlist doesn't know it
		CTIDS=$(echo $CTIDS 0)
		;;
	*)
		# A file name or CT id/name? Check CT first
		ID=$(vzlist -H -o ctid "$1" 2>/dev/null)
		if test $? -eq 0; then
			CTIDS=$(echo $CTIDS $ID)
		else
			FILE="$1"
		fi
		;;
	esac
	shift
done

if test -z "$FILE"; then
	# No file name given, substitute sane default
	FILE=/proc/bc/resources
	test -f $FILE || file=/proc/user_beancounters
fi

# Test that input file is readable
if test "$FILE" != "-"; then
	cat < "$FILE" > /dev/null || exit 1
fi

# Cumulative mode preparations and sanity checks
if ! test -z "$relative"; then
	# Create dir if needed
	if ! test -d $STOREDIR; then
		mkdir $STOREDIR || exit 1
	fi
	# Check we can write to it
	touch $STOREDIR/test || exit 1
	rm -f $STOREDIR/test || exit 1
fi

# Re-exec ourselves under watch
test -z "$watch" || exec watch $WATCH_ARGS -- $0 $ARGV -W

cat $FILE | LANG=C awk -v CTIDS=" ${CTIDS} " -v quiet="$quiet"  \
	-v qheld="$Q_HELD" -v qmaxheld="$Q_MAXHELD" \
	-v rel="$relative" -v storedir="$STOREDIR" \
'
BEGIN {
	if (qmaxheld > qheld)
		qmaxheld=qheld
}
function hr(res, v) {
	if ((v == 9223372036854775807) || (v == 2147483647) || (v == 0))
		return "- ";
	i=1
	if ((res ~ /pages$/) && (v != 0)) {
		v = v*4; i++
	}
	while (v >= 1024) {
		v=v/1024
		i++
	}
	fmt="%d%c"
	if (v < 100)
		fmt="%.3g%c"
	return sprintf(fmt, v, substr(" KMGTPEZY", i, 1))
}

function dp(p, d) {
	if ((d == 0) || (d == 9223372036854775807) || (d == 2147483647))
		return "- "
	r = sprintf("%.2f", p / d * 100);
	fmt="%d"
	if (r < 10)
		fmt="%.2g"
	r = sprintf(fmt, r)
	if (r == 0)
		return "- "
	return r "%"
}

function important(id, held, maxheld, barrier, limit, failcnt) {
	if (failcnt > 0)
		return 2;
	if (barrier == 0)
		barrier = limit;
	if ((barrier == 9223372036854775807) || (barrier == 2147483647))
		return 0;
	if (held > barrier)
		return 2;
	if (held > barrier * qheld)
		return 1;
	if (maxheld > barrier * qmaxheld)
		return 1;
	return 0;
}

/^Version: / {
	if ($2 != "2.5") {
		print "Error: unknown version:",
			$2 > "/dev/stderr"
		exit 1
	}
	next
}
/^[[:space:]]*uid / {
	next
}
/^[[:space:]]*dummy/ {
	id=""
	next
}
/^[[:space:]]*[0-9]+:/ {
	header=""
	bcid=strtonum($1)
	if ((CTIDS !~ /^[ ]*$/) && (CTIDS !~ " " bcid " ")) {
		bcid=""
		next
	}
	header=sprintf( \
"-------------------------------------------------------------------\n" \
"CT %-10s| HELD  Bar%%  Lim%%| MAXH  Bar%%  Lim%%| BAR | LIM |%cFAIL\n" \
"-------------+-----------------+-----------------+-----+-----+-----",
bcid, rel ? "+" : " ")
	id=$2
	held=$3
	maxheld=$4
	barrier=$5
	limit=$6
	failcnt=$7
}
/^[[:space:]]*[a-z]+/ {
	id=$1
	held=$2
	maxheld=$3
	barrier=$4
	limit=$5
	failcnt=$6
}
((id!="")  && (bcid!="")) {
	newfc=failcnt
	if ( (rel) && (failcnt > 0) ) {
		storefile=storedir "/" bcid "." id
		getline oldfc < storefile
		if (oldfc > 0)
			failcnt=failcnt-oldfc
		if (failcnt < 0)
			failcnt=newfc
	}
	imp=important(id, held, maxheld, barrier, limit, failcnt)
	if ((quiet) && (!imp)) {
		id=""
		next
	}
	if (header != "") {
		if (printed)
			print ""
		print header
		printed=1
		header=""
	}
	printf "%13s|%5s %5s %5s|%5s %5s %5s|%5s|%5s| %5s\n",
		id,
		hr(id, held), dp(held, barrier), dp(held, limit),
		hr(id, maxheld), dp(maxheld, barrier), dp(maxheld, limit),
		hr(id, barrier), hr(id, limit), hr(id, failcnt)
	# Save failcnt to store file
	if ( (rel) && (newfc > 0) )
		print newfc > storefile
	id=""
}
END {
	if (printed)
		printf "-------------------------------------------------------------------\n"
}
'
