From fd4074798e8fbff31244d2ab00d629eafc6aa2cd Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Sat, 5 Sep 2020 19:24:18 +0100
Subject: [PATCH] nomad_smartrestart: add pretty colours

---
 nomad_smartrestart.sh | 76 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 57 insertions(+), 19 deletions(-)

diff --git a/nomad_smartrestart.sh b/nomad_smartrestart.sh
index 71da1bc..f3e5d7d 100755
--- a/nomad_smartrestart.sh
+++ b/nomad_smartrestart.sh
@@ -1,29 +1,52 @@
 #!/usr/bin/env bash
 set -e;
 
+cd "$(dirname "$(readlink -f "$0")")" || { echo "Error: Failed to cd to script directory" >&2; exit 1; };
+
+imagename="${1}"
+
 jobname="${1}";
 taskname="${2}";
 
+if [[ -z "${NO_COLOR}" ]]; then
+	RS="\033[0m"    # reset
+	HC="\033[1m"    # hicolor
+	
+	FRED="\033[31m" # foreground red
+	FGRN="\033[32m" # foreground green
+	FYEL="\033[33m" # foreground yellow
+	FBLE="\033[34m" # foreground blue
+fi
+
 help() {
 	echo "nomad_smartrestart, v0.1
 	By Starbeamrainbowlabs
 
-USAGE
-	path/to/nomad_smartrestart.sh {JOB_NAME} {TASK_NAME}
+${FBLE}${HC}USAGE${RS}
+	path/to/nomad_smartrestart.sh {IMAGE_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)"
+${FBLE}${HC}ENVIRONMENT VARIABLES${RS}
+	${FYEL}URL_NOMAD${RS}	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;
+log_msg() {
+	echo "${HC}>>> nomad_smartrestart: ${RS}$*" >&2;
+}
+log_error() {
+	echo "${HC}>>> nomad_smartrestart: ${FRED}$*${RS}";
+}
+
+if [[ -z "${imagename}" ]]; then
+	log_error "Error: No image name specified." >&2;
 	exit 1;
 fi
-if [[ -z "${taskname}" ]]; then
-	echo "Error: No taskname specified." >&2;
-	exit 1;
+
+if [[ ! -f "./images/${imagename}/nomad.txt" ]]; then
+	log_msg "No nomad.txt present for ${HC}${FGRN}${imagename}${RS}, nothing to do";
+	exit 0;
 fi
 
+
 URL_NOMAD="${URL_NOMAD:-http://nomad.service.mooncarrot.space:4646}";
 
 ###############################################################################
@@ -54,14 +77,29 @@ nomad_alloc_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";
+count_tasks=0;
+count_restarted=0;
+while read -r jobname taskname; do
+	if [[ -z "${jobname}" ]]; then
+		log_error "Error: No jobname specified." >&2;
+		continue;
+	fi
+	if [[ -z "${taskname}" ]]; then
+		log_error "Error: No taskname specified." >&2;
+		continue;
+	fi
+	
+	log_msg "Checking ${FYEL}${HC}${jobname}:${taskname}${RS}";
+	
+	while read -r alloc_id jobname taskname; do
+		alloc_id_short="$(echo "${alloc_id}" | head -c7)";
+		log_msg "Restarting allocation ${HC}${FBLE}${alloc_id_short}${RC} in-place";
+		
+		nomad_alloc_restart "${alloc_id}";
+		count_restarted="$((count_restarted+1))";
+	done < <(nomad_allocs_find "${jobname}" "${taskname}");
+	count_tasks="$((count_tasks+1))";
+done
+
+log_msg "Restarted ${HC}${count_restarted}${RS} allocations for ${HC}${count_tasks}${RS} tasks in-place";