#!/bin/sh -e
#
#    disk_io: calculate the disk io rate
#
#    Copyright (C) 2010 Canonical Ltd.
#    Copyright (C) 2011-2014 Dustin Kirkland
#
#    Authors: Dustin Kirkland <kirkland@byobu.org>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, version 3 of the License.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

__disk_io_detail() {
	if eval $BYOBU_TEST iostat >/dev/null; then
		iostat -d -m -h
	else
		printf "%s\n" "Please install iostat if you want detailed information on your disk throughput"
	fi
}


getdisk() {
	local t=""
	if [ -L "${1}" ]; then
		t=$($BYOBU_READLINK -f "$1")
	else
		t="$1"
	fi
	t="${t##*/}"
	# If it's already a block device, use it as-is
	if [ -h "/sys/block/$t" ]; then
		_RET="$t"
	else
		# Strip partition number (trailing digits)
		rtrim "$t" "0-9"
		t="$_RET"
		# For NVMe (nvme0n1p2), MMC (mmcblk0p2), etc., also strip the 'p'
		# that precedes the partition number. If it ends with 'p' after
		# removing digits, it's likely this naming scheme.
		case "$t" in
			*p) _RET="${t%p}" ;;
			*) _RET="$t" ;;
		esac
	fi
}

__disk_io_openbsd() {
	local mount_point disk part kb xfr t2 t1 cache x1 x2 rate symbol unit
	[ -z "$MONITORED_DISK" ] && mount_point="/" || mount_point="$MONITORED_DISK"
	[ -n "$DISK_IO_THRESHOLD" ] || DISK_IO_THRESHOLD=50
	# Find the disk device for the mount point (e.g., /dev/sd0a -> sd0)
	part=$(mount | awk -v mp="$mount_point" '$3 == mp { print $1; exit }')
	[ -n "$part" ] || return
	disk="${part##*/}"
	# Strip partition letter (sd0a -> sd0)
	disk="${disk%[a-p]}"
	# Parse iostat -DI for cumulative KB transferred for this disk
	# Output format: disk_name  KB  xfr  time
	local iostat_line
	iostat_line=$(iostat -DI 2>/dev/null | awk -v d="$disk" '$1 == d { print $2; exit }')
	[ -n "$iostat_line" ] || return
	# iostat -DI gives total KB transferred (read+write combined)
	x2="$iostat_line"
	[ -n "$x2" ] || return
	cache="$BYOBU_RUN_DIR/cache.$BYOBU_BACKEND/disk.io"
	t2=$(date +%s)
	t1=$(stat -f %m "$cache" 2>/dev/null) || t1=0
	x1=0
	[ -r "$cache" ] && read x1 < "$cache"
	printf "%s" "$x2" > "$cache"
	if [ "$t2" -le "$t1" ] || [ "$t1" -eq 0 ]; then
		return
	fi
	rate=$(( (x2 - x1) / (t2 - t1) ))
	[ "$rate" -lt 0 ] && rate=0
	if [ "$rate" -lt "$DISK_IO_THRESHOLD" ]; then
		rm -f "$BYOBU_RUN_DIR/status.$BYOBU_BACKEND/disk_io"*
		return
	fi
	# Combined read+write from iostat; show both direction symbols
	symbol="${ICON_RD}${ICON_WR}"
	if [ "$rate" -gt 1048576 ]; then
		unit="GB/s"
		fpdiv "$rate" 1048576 0
		rate=${_RET}
	elif [ "$rate" -gt 1024 ]; then
		unit="MB/s"
		fpdiv "$rate" 1024 0
		rate=${_RET}
	else
		unit="kB/s"
	fi
	if [ -z "$rate" ] || [ "$rate" = "0" ]; then
		rm -f "$BYOBU_RUN_DIR/status.$BYOBU_BACKEND/disk_io"*
	else
		color b M W; printf "%s%s" "$symbol" "$rate"; color -; color M W; printf "%s" "$unit"; color --
	fi
}

__disk_io() {
	local part= i=
	# OpenBSD: use iostat -DI instead of /sys/block
	if [ "$BYOBU_OSTYPE" = "OpenBSD" ]; then
		__disk_io_openbsd
		return
	fi
	# Default to disk providing /, but let users override with MONITORED_DISK
	[ -z "$MONITORED_DISK" ] && mount_point="/" ||  mount_point="$MONITORED_DISK"
	# By default, we won't bug the user with the display of disk throughput
	# below DISK_IO_THRESHOLD in kB/s; override in $BYOBU_CONFIG_DIR/status
	[ -n "$DISK_IO_THRESHOLD" ] || DISK_IO_THRESHOLD=50
	case "$mount_point" in
		/dev/*) part="${mount_point}";;
		*) part=$(awk '$2 == mp { print $1 ; exit(0); }' "mp=$mount_point" /etc/mtab);;
	esac
	[ -e "$part" ] || return
	getdisk "$part"
	local disk=${_RET}
	local t2=$(date +%s) t1=
	local over_threshold=0
	for i in "read" "write"; do
		local cache="$BYOBU_RUN_DIR/cache.$BYOBU_BACKEND/disk.$i"
		t1=$(stat -c %Y "$cache") 2>/dev/null || t1=0
		local unit="kB/s"
		local rate=0 x1=0 x2=0 symbol= unit=
		if [ $t2 -le $t1 ]; then
			rate=0
		else
			[ -r "$cache" ] && read x1 < "$cache"
			local a1= a2= a3= a4= a5= a6= a7= a8= a9= a10=
			read a1 a2 a3 a4 a5 a6 a7 a8 a9 a10 other < "/sys/block/$disk/stat"
			if [ "$i" = "read" ]; then
				symbol="$ICON_RD"
				[ -n "$a3" ] && x2="$a3" || x2=0
			else
				symbol="$ICON_WR"
				[ -n "$a7" ] && x2="$a7" || x2=0
			fi
			printf "%s\n" "$x2" > "$cache"
			rate=$((($x2 - $x1) / ($t2 - $t1) * 512 / 1024))
			if [ "$rate" -lt 0 ]; then
				rate=0
			elif [ "$rate" -gt 1048576 ]; then
				unit="GB/s"
				fpdiv "$rate" 1048576 0
				rate=${_RET}
			elif [ "$rate" -gt 1024 ]; then
				unit="MB/s"
				fpdiv "$rate" 1024 0
				rate=${_RET}
			else
				unit="kB/s"
			fi
		fi
		if [ -n "$rate" ] && [ "$rate" != "0" ] && [ $rate -ge $DISK_IO_THRESHOLD ]; then
			over_threshold=1
			color b M W; printf "%s%s" "$symbol" "$rate"; color -; color M W; printf "%s" "$unit"; color --
		fi
	done
	if [ "$over_threshold" = "0" ]; then
		rm -f "$BYOBU_RUN_DIR/status.$BYOBU_BACKEND/disk_io"*
	fi
}

# vi: syntax=sh ts=4 noexpandtab
