nutcracker.sh

#!/usr/bin/env bash

STAGE_STARTED=${STAGE_STARTED:-0}
VERSION="1.1.0"

DEV_MODE=0
FAKE_UPS_FILE="/tmp/nutcracker_ups"

### CONFIG ###

STATE_FILE="/run/nut/nutcracker.state"
CONFIG_FILE="/etc/nut/nutcracker.conf"

HEALTH_FILE="/run/nut/nutcracker.health"
HEALTH_INTERVAL=30   # loops (≈60s)

GRACE_DELAY=60
GRACE_RUNTIME=420
FORCE_RUNTIME=240
HOST_RUNTIME=120

LOOP_DELAY=2

GRACE_DEP_FORCE_DELAY=5
GRACE_NON_DEP_START=8
NON_DEP_INTERVAL=5

UPS_WARN_FAIL=6        # already have
UPS_CRIT_FAIL=30       # ~60 seconds

OL_STABLE_COUNT=0
OL_DEBOUNCE_CYCLES=15

### GLOBAL STATE ###

STAGE=0
OUTAGE_START=0
FORCE_DONE=0

UPS_STATUS=""
UPS_RUNTIME=0

UPS_FAIL_COUNT=0

LOOP_COUNT=0

GRACEFUL_START_LOOP=0
NON_DEP_INDEX=0
CURRENT_VM=""

RUNNING_VMS=""

LAST_STATUS=""


### LOGGING ###

log() {
    msg="$1"
    logger -t nutcracker "$msg"
    echo "$(date '+%F %T') $msg" >> /var/log/nutcracker.log
}


write_health() {

    ELAPSED=$(elapsed)

    cat > "$HEALTH_FILE" <<EOF
version=$VERSION
stage=$STAGE
stage_name=$(stage_name "$STAGE")
status=$UPS_STATUS
runtime=$UPS_RUNTIME
elapsed=$ELAPSED
running_vms="$RUNNING_VMS"
last_update=$(date -Iseconds)
EOF

}


### STATE IO ###

load_state() {

    if [[ -f $STATE_FILE ]]; then
        source "$STATE_FILE" 2>/dev/null || true
    fi

}


init_state_defaults() {

    STAGE=${STAGE:-0}
    OUTAGE_START=${OUTAGE_START:-0}
    STAGE_STARTED=${STAGE_STARTED:-0}
    FORCE_DONE=${FORCE_DONE:-0}
    RUNNING_VMS=${RUNNING_VMS:-""}
    LOOP_COUNT=${LOOP_COUNT:-0}
    GRACEFUL_START_LOOP=${GRACEFUL_START_LOOP:-0}
    NON_DEP_INDEX=${NON_DEP_INDEX:-0}
    CURRENT_VM=${CURRENT_VM:-""}
    DEP_VMS=${DEP_VMS:-""}
    NON_DEP_VMS=${NON_DEP_VMS:-""}
    OL_STABLE_COUNT=${OL_STABLE_COUNT:-0}

}


save_state() {

    tmp=$(mktemp)

    cat > "$tmp" <<EOF
STAGE=$STAGE
OUTAGE_START=$OUTAGE_START
STAGE_STARTED=$STAGE_STARTED
FORCE_DONE=$FORCE_DONE
RUNNING_VMS="$RUNNING_VMS"
GRACEFUL_START_LOOP=$GRACEFUL_START_LOOP
NON_DEP_INDEX=$NON_DEP_INDEX
CURRENT_VM="$CURRENT_VM"
OL_STABLE_COUNT=$OL_STABLE_COUNT
EOF

    mv "$tmp" "$STATE_FILE"

}

snapshot_running_vms() {

    RUNNING_VMS=$(qm list | awk 'NR>1 && $3=="running" {print $1}' | tr '\n' ' ')

    if [[ -z $RUNNING_VMS ]]; then
        log "No running VMs detected"
    else
        COUNT=$(wc -w <<< "$RUNNING_VMS")
        log "Running VMs: $RUNNING_VMS ($COUNT total)"
    fi

}


### UPS READ ###

read_ups() {

    if (( DEV_MODE == 1 )) && [[ -f $FAKE_UPS_FILE ]]; then
        UPS_DATA=$(cat "$FAKE_UPS_FILE")
    else
        UPS_DATA=$(upsc ups@localhost 2>/dev/null)
    fi

    #UPS_DATA=$(upsc ups@localhost 2>/dev/null)

    UPS_STATUS=$(grep "^ups.status:" <<< "$UPS_DATA" | cut -d' ' -f2-)
    UPS_RUNTIME=$(grep "^battery.runtime:" <<< "$UPS_DATA" | awk '{print $2}')

    if [[ -z $UPS_STATUS ]]; then
        log "UPS read failed"
        UPS_STATUS="UNKNOWN"
        return
    fi

    UPS_RUNTIME=${UPS_RUNTIME:-0}

}


### HELPERS ###

now() {
    date +%s
}


vm_now() {
    date '+%F %H:%M:%S'
}


elapsed() {

    if (( OUTAGE_START == 0 )); then
        echo 0
        return
    fi

    echo $(( $(now) - OUTAGE_START ))

}


is_in_running() {
    
    [[ -z "$1" ]] && return 1

    [[ " $RUNNING_VMS " =~ [[:space:]]$1[[:space:]] ]]

}


grace_elapsed() {

    echo $((LOOP_COUNT - GRACEFUL_START_LOOP))

}


stage_name() {

    local s="$1"

    case "$s" in
        0) echo NORMAL ;;
        1) echo OUTAGE ;;
        2) echo GRACEFUL ;;
        3) echo FORCE ;;
        4) echo HOST ;;
        5) echo TERMINAL ;;
    esac

}


set_stage() {

    NEW_STAGE=$1
    REASON="${2:-unspecified}"

    # never downgrade stage (unless reset elsewhere)
    if (( NEW_STAGE < STAGE )); then
        return
    fi

    if (( STAGE == NEW_STAGE )); then
        return
    fi

    OLD_STAGE=$STAGE
    NOW=$(now)

    # duration log (only if we already entered a stage)
    if (( STAGE_STARTED > 0 )); then
        STAGE_DURATION=$((NOW - STAGE_STARTED))
        log "Stage $OLD_STAGE duration: ${STAGE_DURATION}s"
    fi

    ELAPSED=$(elapsed)

    log "Stage $(stage_name $OLD_STAGE) ($OLD_STAGE) → $(stage_name $NEW_STAGE) ($NEW_STAGE) | status=$UPS_STATUS | runtime=${UPS_RUNTIME}s | elapsed=${ELAPSED}s | reason=$REASON"

    STAGE="$NEW_STAGE"
    STAGE_STARTED=$NOW

    stage_enter "$STAGE"

    save_state

    write_health

}


reset_state() {

    #log "Power restored, resetting state"
    OUTAGE_TOTAL=$(elapsed)
    log "Total outage duration: ${OUTAGE_TOTAL}s"

    OLD_STAGE=$STAGE

    STAGE=0
    OUTAGE_START=0
    STAGE_STARTED=0
    FORCE_DONE=0
    GRACEFUL_START_LOOP=0
    NON_DEP_INDEX=0
    CURRENT_VM=""
    OL_STABLE_COUNT=0

    if [[ $OLD_STAGE -ne 0 ]]; then
        save_state
    fi

}


is_vm_running() {

    qm status "$1" | grep -q running

}


shutdown_vm() {

    local vm="$1"

    log "Graceful shutdown requested for VM $vm | t=$(vm_now)"

    qm shutdown "$vm" --timeout 60 >/dev/null 2>&1 &

}


force_vm() {

    local vm="$1"

    log "Force stop requested for VM $vm | t=$(vm_now)"

    qm stop "$vm" --skiplock >/dev/null 2>&1 &

}


### STATE MACHINE ###


handle_normal() {

    if [[ $UPS_STATUS == *OB* ]] || [[ $UPS_STATUS == *LB* ]]; then

        log "Power loss detected"

        OUTAGE_START=$(now)

        set_stage 1 "power_loss"
    fi

}



check_power_restore() {

    if [[ $UPS_STATUS == *OL* ]] && (( STAGE > 0 )); then

        ((OL_STABLE_COUNT++))

        if (( OL_STABLE_COUNT == 1 )); then
            log "Power restored, waiting for stability"
        fi

        if (( OL_STABLE_COUNT >= OL_DEBOUNCE_CYCLES )); then
            log "Power stable for $((OL_DEBOUNCE_CYCLES * LOOP_DELAY))s"
            reset_state
            save_state
            return 0
        fi

    else

        if (( OL_STABLE_COUNT > 0 )); then
            log "Power unstable again, cancelling reset"
            OL_STABLE_COUNT=0
            save_state
        fi

    fi

    return 1

}


handle_outage_timer() {

    if check_power_restore; then
        return
    fi

    if [[ $UPS_RUNTIME -lt $HOST_RUNTIME ]]; then
        set_stage 4 "critical_runtime"
        return
    fi

    if [[ $UPS_RUNTIME -lt $FORCE_RUNTIME ]]; then
        set_stage 3 "low_runtime"
        return
    fi

    if [[ $UPS_STATUS == *LB* ]]; then
        set_stage 3 "low_battery"
        return
    fi

    if [[ $(elapsed) -gt $GRACE_DELAY ]]; then
        set_stage 2 "grace_timer"
        return
    fi

    if [[ $UPS_RUNTIME -lt $GRACE_RUNTIME ]]; then
        set_stage 2 "runtime_threshold"
    fi

}



handle_graceful() {

    graceful_progress

    if check_power_restore; then
        return
    fi

    if [[ $UPS_RUNTIME -lt $HOST_RUNTIME ]]; then
        set_stage 4 "critical_runtime"
        return
    fi

    if [[ $UPS_RUNTIME -lt $FORCE_RUNTIME ]]; then
        set_stage 3 "low_runtime"
        return
    fi

    if [[ $UPS_STATUS == *LB* ]]; then
        set_stage 3 "low_battery"
    fi

}



handle_force() {

    force_progress

    if check_power_restore; then
        return
    fi

    if [[ $UPS_RUNTIME -lt $HOST_RUNTIME ]]; then
        set_stage 4 "critical_runtime"
    fi

}



handle_host() {

    if [[ $STAGE -ne 4 ]]; then
        return
    fi

    log "HOST shutdown stage reached"

    snapshot_running_vms

    if [[ -z $RUNNING_VMS ]]; then
        log "No VMs remaining"
    else

        log "Host shutdown imminent - final VM sweep"

        for vm in $RUNNING_VMS; do
            force_vm "$vm"
        done

        sleep 1

    fi

    STAGE=5
    save_state

    sleep 2

    #log "Shutting down host"

    OUTAGE_TOTAL=$(elapsed)
    log "Final outage duration before shutdown: ${OUTAGE_TOTAL}s"

    sync
    sleep 2

    shutdown -h now

    exit 0
}


graceful_progress() {

    if (( GRACEFUL_START_LOOP == 0 )); then
        return
    fi

    local elapsed
    elapsed=$(grace_elapsed)


    DEP_COUNT=0

    for vm in $DEP_VMS; do
        is_vm_running "$vm" && ((DEP_COUNT++))
    done

    (( DEP_COUNT > 0 )) && log "Graceful phase: $DEP_COUNT DEP VMs scheduled"


    # DEP shutdown
    if (( elapsed == 1 )); then
        for vm in $DEP_VMS; do

            if is_in_running "$vm"; then
                shutdown_vm "$vm"
            fi

        done
    fi

    # DEP force
    if (( elapsed == GRACE_DEP_FORCE_DELAY + 1 )); then
        for vm in $DEP_VMS; do
            if is_in_running "$vm" && is_vm_running "$vm"; then
                log "VM $vm still running after graceful delay → forcing"
                force_vm "$vm"
            fi
        done
    fi

    # NON_DEP queue
    if (( elapsed < GRACE_NON_DEP_START )); then
        return
    fi

    if (( elapsed % NON_DEP_INTERVAL != 0 )); then
        return
    fi

    if (( elapsed == GRACE_NON_DEP_START )); then
        log "Graceful phase: starting NON_DEP shutdown queue"
    fi

    NONDEP_COUNT=0

    for vm in $NON_DEP_VMS; do
        is_vm_running "$vm" && ((NONDEP_COUNT++))
    done

    (( NONDEP_COUNT > 0 )) && log "Graceful phase: $NONDEP_COUNT NON_DEP VMs scheduled"


    # verify previous VM
    if [[ -n $CURRENT_VM ]]; then

        if is_vm_running "$CURRENT_VM"; then

            log "VM $CURRENT_VM still running → forcing"

            force_vm "$CURRENT_VM"

            return
        fi

    fi

    CURRENT_VM=""

    # next NON_DEP VM
    VMID=$(echo $NON_DEP_VMS | awk "{print \$$((NON_DEP_INDEX+1))}")

    if [[ -z $VMID ]]; then
        return
    fi

    if ! is_in_running "$VMID"; then
        ((NON_DEP_INDEX++))
        save_state
        return
    fi

    log "Graceful shutdown requested for NON_DEP VM $VMID"

    shutdown_vm "$VMID"

    CURRENT_VM="$VMID"

    ((NON_DEP_INDEX++))

    save_state

}


force_progress() {

    if (( FORCE_DONE == 2 )); then
        return
    fi

    snapshot_running_vms

    if [[ -z $RUNNING_VMS ]]; then
        log "No running VMs detected"
        FORCE_DONE=2
        save_state
        return
    fi

    # Phase 1: DEP only
    if (( FORCE_DONE == 0 )); then

        log "Force stopping DEP VMs"

        for vm in $DEP_VMS; do
            if is_vm_running "$vm"; then
                force_vm "$vm"
            fi
        done

        FORCE_DONE=1
        save_state
        return
    fi


    # Phase 2: everything else (run ONCE)
    if (( FORCE_DONE == 1 )); then

        COUNT=$(wc -w <<< "$RUNNING_VMS")
        log "Force phase: $COUNT VMs still running"

        for vm in $RUNNING_VMS; do
            if is_vm_running "$vm"; then
                force_vm "$vm"
            fi
        done

        FORCE_DONE=2
        save_state
        return
    fi

}


stage_enter() {

    case "$1" in

        1)
            [[ $STAGE_STARTED -eq 0 ]] && STAGE_STARTED=$(now)
        ;;

        2)

            log "Entering graceful shutdown stage"

            GRACEFUL_START_LOOP=$LOOP_COUNT

            snapshot_running_vms
            COUNT=$(wc -w <<< "$RUNNING_VMS")
            log "Initial running VMs: $RUNNING_VMS ($COUNT total)"

            # load config AFTER snapshot
            if [[ -f $CONFIG_FILE ]]; then
                source "$CONFIG_FILE"
            else
                DEP_VMS=""
                NON_DEP_VMS="$RUNNING_VMS"
            fi

            save_state
        ;;

        3)

	    log "Entering force shutdown stage"

            snapshot_running_vms

            FORCE_DONE=0

            CURRENT_VM=""
            NON_DEP_INDEX=0

            save_state

        ;;

        4)
            log "Entering host shutdown stage"
        ;;

    esac
}


maybe_write_health() {

    mkdir -p /run/nut

    if (( LOOP_COUNT % HEALTH_INTERVAL == 0 )); then
        write_health
        chmod 644 "$HEALTH_FILE"
    fi

}


### MAIN LOOP ###


log "Nutcracker $VERSION started"

mkdir -p /run/nut

while true; do

    load_state
    init_state_defaults

    read_ups

    if [[ $UPS_STATUS == "UNKNOWN" ]]; then

        ((UPS_FAIL_COUNT++))

        if (( UPS_FAIL_COUNT == UPS_WARN_FAIL )); then
            log "UPS communication lost"
        fi

        if (( UPS_FAIL_COUNT == UPS_CRIT_FAIL )); then
            log "UPS communication critical"
        fi

        #sleep "$LOOP_DELAY"
        #continue

    else

        UPS_FAIL_COUNT=0


        if [[ $STAGE -gt 0 ]] && [[ $OUTAGE_START -eq 0 ]]; then
            OUTAGE_START=$(now)
            save_state
        fi

        # handle startup during outage
        if [[ $STAGE -eq 0 ]] && [[ -n $UPS_STATUS ]] && [[ $UPS_STATUS != *OL* ]]; then

            log "Startup on battery detected"

            OUTAGE_START=$(now)

            set_stage 1 "startup_on_battery"

        fi
        

        if [[ "$UPS_STATUS" != "$LAST_STATUS" ]]; then
            log "UPS status changed: $LAST_STATUS → $UPS_STATUS"
            LAST_STATUS="$UPS_STATUS"
        fi


        case "$STAGE" in

            0)
                handle_normal
            ;;

            1)
                handle_outage_timer
            ;;

            2)
                handle_graceful
            ;;

            3)
                handle_force
            ;;

            4)
                handle_host
            ;;

            5)
                exit 0
            ;;

        esac

    fi

    ((LOOP_COUNT++))

    maybe_write_health

    sleep "$LOOP_DELAY"

done