allow more flexibility in cleanup

mafredri · mafredri · commit 96100c0b2497 · 2024-01-23T16:42:46.000+02:00
diff --git a/scaletest/templates/scaletest-runner/main.tf b/scaletest/templates/scaletest-runner/main.tf
@@ -169,6 +169,16 @@ data "coder_parameter" "cleanup_strategy" {
   }
 }
 
+data "coder_parameter" "cleanup_prepare" {
+  order       = 14
+  type        = "bool"
+  name        = "Cleanup before scaletest"
+  default     = true
+  description = "Cleanup existing scaletest users and workspaces before the scaletest starts (prepare phase)."
+  mutable     = true
+  ephemeral   = true
+}
+
 
 data "coder_parameter" "workspace_template" {
   order        = 20
@@ -563,6 +573,7 @@ resource "coder_agent" "main" {
     SCALETEST_PARAM_NUM_WORKSPACES : data.coder_parameter.num_workspaces.value,
     SCALETEST_PARAM_CREATE_CONCURRENCY : "${data.coder_parameter.create_concurrency.value}",
     SCALETEST_PARAM_CLEANUP_STRATEGY : data.coder_parameter.cleanup_strategy.value,
+    SCALETEST_PARAM_CLEANUP_PREPARE : data.coder_parameter.cleanup_prepare.value ? "1" : "0",
     SCALETEST_PARAM_LOAD_SCENARIOS : data.coder_parameter.load_scenarios.value,
     SCALETEST_PARAM_LOAD_SCENARIO_RUN_CONCURRENTLY : data.coder_parameter.load_scenario_run_concurrently.value ? "1" : "0",
     SCALETEST_PARAM_LOAD_SCENARIO_SSH_TRAFFIC_DURATION : "${data.coder_parameter.load_scenario_ssh_traffic_duration.value}",
diff --git a/scaletest/templates/scaletest-runner/scripts/cleanup.sh b/scaletest/templates/scaletest-runner/scripts/cleanup.sh
@@ -21,17 +21,20 @@ if [[ $event = manual ]]; then
 	fi
 fi
 
-start_phase "Cleanup (${event})"
-coder exp scaletest cleanup \
-	--cleanup-job-timeout 2h \
-	--cleanup-timeout 5h |
-	tee "${SCALETEST_RESULTS_DIR}/cleanup-${event}.txt"
-end_phase
+if [[ $event != shutdown_scale_down_only ]]; then
+	start_phase "Cleanup (${event})"
+	coder exp scaletest cleanup \
+		--cleanup-job-timeout 2h \
+		--cleanup-timeout 5h \
+		| tee "${SCALETEST_RESULTS_DIR}/cleanup-${event}.txt"
+	end_phase
+fi
 
 if [[ $event != prepare ]]; then
-	start_phase "Scaling down provisioners..."
+	start_phase "Scale down provisioners"
 	maybedryrun "$DRY_RUN" kubectl scale deployment/coder-provisioner --replicas 1
 	maybedryrun "$DRY_RUN" kubectl rollout status deployment/coder-provisioner
+	end_phase
 fi
 
 if [[ $event = manual ]]; then
diff --git a/scaletest/templates/scaletest-runner/scripts/prepare.sh b/scaletest/templates/scaletest-runner/scripts/prepare.sh
@@ -47,8 +47,10 @@ unset CODER_SESSION_TOKEN
 echo -n "${token}" >"${CODER_CONFIG_DIR}/session"
 [[ $VERBOSE == 1 ]] && set -x # Restore logging (if enabled).
 
-log "Cleaning up from previous runs (if applicable)..."
-"${SCRIPTS_DIR}/cleanup.sh" "prepare"
+if [[ ${SCALETEST_PARAM_CLEANUP_PREPARE} == 1 ]]; then
+	log "Cleaning up from previous runs (if applicable)..."
+	"${SCRIPTS_DIR}/cleanup.sh" prepare
+fi
 
 log "Preparation complete!"
 
diff --git a/scaletest/templates/scaletest-runner/shutdown.sh b/scaletest/templates/scaletest-runner/shutdown.sh
@@ -14,7 +14,11 @@ trap cleanup EXIT
 
 annotate_grafana "workspace" "Agent stopping..."
 
-"${SCRIPTS_DIR}/cleanup.sh" shutdown
+shutdown_event=shutdown_scale_down_only
+if [[ ${SCALETEST_PARAM_CLEANUP_STRATEGY} == on_stop ]]; then
+	shutdown_event=shutdown
+fi
+"${SCRIPTS_DIR}/cleanup.sh" "${shutdown_event}"
 
 annotate_grafana_end "workspace" "Agent running"
 
