#!/usr/bin/env bash
set -e;

jobname="${1}";
taskname="${2}";

help() {
	echo "nomad_smartrestart, v0.1
	By Starbeamrainbowlabs

USAGE
	path/to/nomad_smartrestart.sh {JOB_NAME} {TASK_NAME}

ENVIRONMENT VARIABLES
	URL_NOMAD	The URL of the nomad server to talk to (using the Consul domain instead of a direct IP is recommended)"
}

if [[ -z "${jobname}" ]]; then
	echo "Error: No jobname specified." >&2;
	exit 1;
fi
if [[ -z "${taskname}" ]]; then
	echo "Error: No taskname specified." >&2;
	exit 1;
fi

URL_NOMAD="${URL_NOMAD:-http://nomad.service.mooncarrot.space:4646}";

###############################################################################


# Finds allocations for a given job / task name combo.
# $1	The job name
# $2	The task name
nomad_allocs_find() {
	local jobname="${1}"
	local taskname="${2}"
	
	if [[ -z "${taskname}" ]]; then
		curl -sS "${URL_NOMAD}/v1/allocations" | JOB_NAME="${jobname}" jq --raw-output '.[] | select(.JobID == env.JOB_NAME) | .ID'
	else
		curl -sS "${URL_NOMAD}/v1/allocations" | JOB_NAME="${jobname}" TASK_NAME="${taskname}" jq --raw-output '.[] | select(.JobID == env.JOB_NAME and (.TaskStates|keys|.[]|. == env.TASK_NAME)) | .ID'
	fi
}

# Restarts a Nomad allocation in-place.
# $1	The allocation id to restart.
nomad_alloc_restart() {
	local alloc_id="${1}"
	
	curl -sS "${URL_NOMAD}/v1/client/allocation/${alloc_id}/restart";
}


###############################################################################

count_restarted=0;
while read -r alloc_id jobname taskname; do
	alloc_id_short="$(echo "${alloc_id}" | head -c7)";
	
	log_msg "[nomad_smartrestart] Restarting allocation ${alloc_id_short} in-place for ${jobname}:${taskname}";
	
	nomad_alloc_restart "${alloc_id}";
	count_restarted="$((count_restarted+1))"
done < <(nomad_allocs_find "${jobname}" "${taskname}");

echo "[nomad_smartrestart] Restarted ${count_restarted} allocations in-place";