From 9c09dead5add7f500d18e968b12e15c5f413d3da Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 29 Sep 2023 17:11:47 +0000
Subject: [PATCH 1/3] feat(scaletest): add service banner status for runner

---
 .../templates/scaletest-runner/startup.sh     | 37 ++++++++++++++++++-
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/scaletest/templates/scaletest-runner/startup.sh b/scaletest/templates/scaletest-runner/startup.sh
index df0f128b1ccd9..8d8a6dad91d2e 100755
--- a/scaletest/templates/scaletest-runner/startup.sh
+++ b/scaletest/templates/scaletest-runner/startup.sh
@@ -21,6 +21,26 @@ fi
 # shellcheck disable=SC2153 source=scaletest/templates/scaletest-runner/scripts/lib.sh
 . "${SCRIPTS_DIR}/lib.sh"
 
+appearance_json="$(
+	curl -sSL \
+		-H "Coder-Session-Token: ${CODER_USER_TOKEN}" \
+		"${CODER_URL}/api/v2/appearance"
+)"
+service_banner_message=$(jq -r '.service_banner.message' <<<"${appearance_json}")
+service_banner_message="${service_banner_message/% | */}"
+service_banner_color="#D65D0F" # Orange.
+
+set_appearance() {
+	local color=$1 message=$2
+	jq --arg color "${color}" --arg message "${message}" '. | .service_banner.message |= $message | .service_banner.background_color |= $color' <<<"${appearance_json}" \
+		| curl -sSL \
+			-X 'PUT' \
+			-H 'Content-Type: application/json' \
+			-H "Coder-Session-Token: ${CODER_USER_TOKEN}" \
+			--data @- \
+			"${CODER_URL}/api/v2/appearance"
+}
+
 annotate_grafana "workspace" "Agent running" # Ended in shutdown.sh.
 
 {
@@ -55,8 +75,9 @@ annotate_grafana "workspace" "Agent running" # Ended in shutdown.sh.
 } &
 pprof_pid=$!
 
+set_appearance "${service_banner_color}" "${service_banner_message} | Scaletest running: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE})!"
+
 # Show failure in the UI if script exits with error.
-failed_status=Failed
 on_exit() {
 	code=${?}
 	trap - ERR EXIT
@@ -64,25 +85,37 @@ on_exit() {
 
 	kill -INT "${pprof_pid}"
 
+	message_color="#4CD473" # Green.
+	message_status=COMPLETE
+	if ((code > 0)); then
+		message_color="#D94A5D" # Red.
+		message_status=FAILED
+	fi
+
 	case "${SCALETEST_PARAM_CLEANUP_STRATEGY}" in
 	on_stop)
 		# Handled by shutdown script.
 		;;
 	on_success)
 		if ((code == 0)); then
+			set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
 			"${SCRIPTS_DIR}/cleanup.sh" "${SCALETEST_PARAM_CLEANUP_STRATEGY}"
 		fi
 		;;
 	on_error)
 		if ((code > 0)); then
+			set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
 			"${SCRIPTS_DIR}/cleanup.sh" "${SCALETEST_PARAM_CLEANUP_STRATEGY}"
 		fi
 		;;
 	*)
+		set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
 		"${SCRIPTS_DIR}/cleanup.sh" "${SCALETEST_PARAM_CLEANUP_STRATEGY}"
 		;;
 	esac
 
+	set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE})!"
+
 	annotate_grafana_end "" "Start scaletest"
 }
 trap on_exit EXIT
@@ -93,7 +126,7 @@ on_err() {
 	set +e
 
 	log "Scaletest failed!"
-	GRAFANA_EXTRA_TAGS=error set_status "${failed_status} (exit=${code})"
+	GRAFANA_EXTRA_TAGS=error set_status "Failed (exit=${code})"
 	"${SCRIPTS_DIR}/report.sh" failed
 	lock_status # Ensure we never rewrite the status after a failure.
 

From b21dfb4066eec639f76d4460f3df8dc00c844605 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 29 Sep 2023 17:24:54 +0000
Subject: [PATCH 2/3] refactor, reset on shutdown

---
 .../templates/scaletest-runner/scripts/lib.sh | 30 +++++++++++++++++++
 .../templates/scaletest-runner/shutdown.sh    |  7 +++++
 .../templates/scaletest-runner/startup.sh     | 27 ++++-------------
 3 files changed, 43 insertions(+), 21 deletions(-)

diff --git a/scaletest/templates/scaletest-runner/scripts/lib.sh b/scaletest/templates/scaletest-runner/scripts/lib.sh
index a9df0a10eba0a..ec94e952fda3c 100644
--- a/scaletest/templates/scaletest-runner/scripts/lib.sh
+++ b/scaletest/templates/scaletest-runner/scripts/lib.sh
@@ -210,3 +210,33 @@ wait_baseline() {
 	maybedryrun "$DRY_RUN" sleep $((s * 60))
 	PHASE_TYPE="phase-wait" end_phase
 }
+
+get_appearance() {
+	session_token=$CODER_USER_TOKEN
+	if [[ -f "${CODER_CONFIG_DIR}/session" ]]; then
+		session_token="$(<"${CODER_CONFIG_DIR}/session")"
+	fi
+	curl -sSL \
+		-H "Coder-Session-Token: ${session_token}" \
+		"${CODER_URL}/api/v2/appearance"
+}
+set_appearance() {
+	local json=$1 color=$2 message=$3
+
+	session_token=$CODER_USER_TOKEN
+	if [[ -f "${CODER_CONFIG_DIR}/session" ]]; then
+		session_token="$(<"${CODER_CONFIG_DIR}/session")"
+	fi
+	newjson="$(
+		jq \
+			--arg color "${color}" \
+			--arg message "${message}" \
+			'. | .service_banner.message |= $message | .service_banner.background_color |= $color' <<<"${json}"
+	)"
+	maybedryrun "${DRY_RUN}" curl -sSL \
+		-X PUT \
+		-H 'Content-Type: application/json' \
+		-H "Coder-Session-Token: ${session_token}" \
+		--data "${newjson}" \
+		"${CODER_URL}/api/v2/appearance"
+}
diff --git a/scaletest/templates/scaletest-runner/shutdown.sh b/scaletest/templates/scaletest-runner/shutdown.sh
index 14d6023aaaa62..972fe07944753 100755
--- a/scaletest/templates/scaletest-runner/shutdown.sh
+++ b/scaletest/templates/scaletest-runner/shutdown.sh
@@ -16,3 +16,10 @@ annotate_grafana "workspace" "Agent stopping..."
 "${SCRIPTS_DIR}/cleanup.sh" shutdown
 
 annotate_grafana_end "workspace" "Agent running"
+
+appearance_json="$(get_appearance)"
+service_banner_message=$(jq -r '.service_banner.message' <<<"${appearance_json}")
+service_banner_message="${service_banner_message/% | */}"
+service_banner_color="#4CD473" # Green.
+
+set_appearance "${appearance_json}" "${service_banner_color}" "${service_banner_message}"
diff --git a/scaletest/templates/scaletest-runner/startup.sh b/scaletest/templates/scaletest-runner/startup.sh
index 8d8a6dad91d2e..300ff40466b6f 100755
--- a/scaletest/templates/scaletest-runner/startup.sh
+++ b/scaletest/templates/scaletest-runner/startup.sh
@@ -21,26 +21,11 @@ fi
 # shellcheck disable=SC2153 source=scaletest/templates/scaletest-runner/scripts/lib.sh
 . "${SCRIPTS_DIR}/lib.sh"
 
-appearance_json="$(
-	curl -sSL \
-		-H "Coder-Session-Token: ${CODER_USER_TOKEN}" \
-		"${CODER_URL}/api/v2/appearance"
-)"
+appearance_json="$(get_appearance)"
 service_banner_message=$(jq -r '.service_banner.message' <<<"${appearance_json}")
 service_banner_message="${service_banner_message/% | */}"
 service_banner_color="#D65D0F" # Orange.
 
-set_appearance() {
-	local color=$1 message=$2
-	jq --arg color "${color}" --arg message "${message}" '. | .service_banner.message |= $message | .service_banner.background_color |= $color' <<<"${appearance_json}" \
-		| curl -sSL \
-			-X 'PUT' \
-			-H 'Content-Type: application/json' \
-			-H "Coder-Session-Token: ${CODER_USER_TOKEN}" \
-			--data @- \
-			"${CODER_URL}/api/v2/appearance"
-}
-
 annotate_grafana "workspace" "Agent running" # Ended in shutdown.sh.
 
 {
@@ -75,7 +60,7 @@ annotate_grafana "workspace" "Agent running" # Ended in shutdown.sh.
 } &
 pprof_pid=$!
 
-set_appearance "${service_banner_color}" "${service_banner_message} | Scaletest running: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE})!"
+set_appearance "${appearance_json}" "${service_banner_color}" "${service_banner_message} | Scaletest running: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE})!"
 
 # Show failure in the UI if script exits with error.
 on_exit() {
@@ -98,23 +83,23 @@ on_exit() {
 		;;
 	on_success)
 		if ((code == 0)); then
-			set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
+			set_appearance "${appearance_json}" "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
 			"${SCRIPTS_DIR}/cleanup.sh" "${SCALETEST_PARAM_CLEANUP_STRATEGY}"
 		fi
 		;;
 	on_error)
 		if ((code > 0)); then
-			set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
+			set_appearance "${appearance_json}" "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
 			"${SCRIPTS_DIR}/cleanup.sh" "${SCALETEST_PARAM_CLEANUP_STRATEGY}"
 		fi
 		;;
 	*)
-		set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
+		set_appearance "${appearance_json}" "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE}), cleaning up..."
 		"${SCRIPTS_DIR}/cleanup.sh" "${SCALETEST_PARAM_CLEANUP_STRATEGY}"
 		;;
 	esac
 
-	set_appearance "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE})!"
+	set_appearance "${appearance_json}" "${message_color}" "${service_banner_message} | Scaletest ${message_status}: [${CODER_USER}/${CODER_WORKSPACE}](${CODER_URL}/@${CODER_USER}/${CODER_WORKSPACE})!"
 
 	annotate_grafana_end "" "Start scaletest"
 }

From aba6d272e782dde2675784ac61277893491db203 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 29 Sep 2023 17:31:56 +0000
Subject: [PATCH 3/3] remove session on shutdown

---
 scaletest/templates/scaletest-runner/shutdown.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scaletest/templates/scaletest-runner/shutdown.sh b/scaletest/templates/scaletest-runner/shutdown.sh
index 972fe07944753..d5c81366b1217 100755
--- a/scaletest/templates/scaletest-runner/shutdown.sh
+++ b/scaletest/templates/scaletest-runner/shutdown.sh
@@ -8,6 +8,7 @@ set -e
 
 cleanup() {
 	coder tokens remove scaletest_runner >/dev/null 2>&1 || true
+	rm -f "${CODER_CONFIG_DIR}/session"
 }
 trap cleanup EXIT