#!/bin/bash

# XML escape in string metrics. Gmond handles this for its own metrics, we must do it ourself
xml_escape() {
    local input="$1"
    input="${input//&/\&amp;}"
    input="${input//</\&lt;}"
    input="${input//>/\&gt;}"
    input="${input//\"/\&quot;}"
    input="${input//\'/\&apos;}"
    echo "$input"
}

hn=$(hostname)

# Only run Slurm checks on nodes with slurm installed
if [ $(which sinfo) ]; then
	# Slurm node state
	slurm_node_state=$(sinfo --nodes=$hn --Node --noheader --format='%t' | head -n1)
	if [ "$slurm_node_state" ]; then
		gmetric --name="slurm_node_state" --title="Slurm node state" --value="$(xml_escape "$slurm_node_state")" --type="string" --group="slurm"
	fi

	# Slurm reason for being down
	slurm_down_reason=$(sinfo --nodes=$hn --Node --noheader --format='%E' | head -n1)
	if [ "$slurm_down_reason" ]; then
		gmetric --name="slurm_down_reason" --title="Slurm node down reason" --value="$(xml_escape "$slurm_down_reason")" --type="string" --group="slurm"
	fi

	# Slurm partition
	while read line; do
		slurm_partition=$slurm_partition$line,
	done <<< $(sinfo --nodes=$hn --Node --noheader --format='%R')
	slurm_partition=${slurm_partition%,}
	if [ $slurm_partition ]; then
		gmetric --name="slurm_partition" --title="Slurm partition" --value="$(xml_escape "$slurm_partition")" --type="string" --group="slurm"
	fi

	# Slurm jobs on this node
	node_job_output=$(squeue --nodes=$hn --format='%i|%j|%u|%M|' --noheader --sort='i' 2>/dev/null)
	if [ $? -eq 0 ]; then
		while read line; do
			readarray -td '|' linearray <<< ${line}
			slurm_job_id=${slurm_job_id}${linearray[0]},
			slurm_job_name=${slurm_job_name}${linearray[1]},
			slurm_job_user=${slurm_job_user}${linearray[2]},
			slurm_job_time_elapsed=${slurm_job_time_elapsed}${linearray[3]%},
		done <<< ${node_job_output}
		slurm_job_id=${slurm_job_id%,}
		slurm_job_name=${slurm_job_name%,}
		slurm_job_user=${slurm_job_user%,}
		slurm_job_time_elapsed=${slurm_job_time_elapsed%,}

		gmetric --name="slurm_job_id" --title="Slurm running job ID" --value="$(xml_escape "$slurm_job_id")" --type="string" --group="slurm"
		gmetric --name="slurm_job_name" --title="Slurm running job name" --value="$(xml_escape "$slurm_job_name")" --type="string" --group="slurm"
		gmetric --name="slurm_job_user" --title="Slurm running job user" --value="$(xml_escape "$slurm_job_user")" --type="string" --group="slurm"
		gmetric --name="slurm_job_time_elapsed" --title="Slurm running job elapsed time" --value="$(xml_escape "$slurm_job_time_elapsed")" --type="string" --group="slurm"
	fi
	
fi