#!/bin/bash
# azkaban-io-guard: call BEFORE (and during) heavy disk work on azkaban (rsync/VM images/backups/builds).
# Exit 0 = IO pressure low, go ahead.  Exit 1 = busy, back off.  Exit 2 = critical, stop now.
# Usage: azkaban-io-guard [--wait SECS] [--quiet] [--max-some N] [--max-full N]
#   --wait SECS   poll every 15 s until clear or SECS elapsed (exit 1 on timeout)
# Source: /proc/pressure/io avg10 (host-wide PSI). Thresholds: busy if full>=5 or some>=15; critical if full>=25 or some>=50.
# Typical agent loop (chunked work): while more; do azkaban-io-guard --wait 300 || sleep 60; do_next_chunk; done
# Also run heavy work as: ionice -c3 nice -n19 <cmd> (rsync: --bwlimit=15000, never faster)
MAXSOME=15; MAXFULL=5; WAIT=0; Q=0
while [ $# -gt 0 ]; do case $1 in --wait) WAIT=$2; shift;; --quiet) Q=1;; --max-some) MAXSOME=$2; shift;; --max-full) MAXFULL=$2; shift;; esac; shift; done
read_psi() { awk '/^some/{split($2,a,"=");s=a[2]} /^full/{split($2,a,"=");f=a[2]} END{print s, f}' /proc/pressure/io; }
end=$(( $(date +%s) + WAIT ))
while :; do
  read S F <<<"$(read_psi)"
  state=ok; rc=0
  awk -v s=$S -v f=$F -v ms=$MAXSOME -v mf=$MAXFULL 'BEGIN{exit !(s>=ms||f>=mf)}' && { state=busy; rc=1; }
  awk -v s=$S -v f=$F 'BEGIN{exit !(s>=50||f>=25)}' && { state=critical; rc=2; }
  [ $Q -eq 0 ] && echo "{\"state\":\"$state\",\"io_some_avg10\":$S,\"io_full_avg10\":$F,\"host\":\"azkaban\"}"
  [ $rc -eq 0 ] && exit 0
  [ "$(date +%s)" -ge "$end" ] && exit $rc
  sleep 15
done
