Update nomad-scale

nomad-scale/script.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+NOMAD_UNIX_ADDR="${NOMAD_UNIX_ADDR:-}"
+NOMAD_TOKEN="${NOMAD_TOKEN:-}"
+NOMAD_ADDR="${NOMAD_ADDR:-localhost}"
+
+TARGET_NAMESPACE="${TARGET_NAMESPACE:-default}"
+TARGET_JOB="${TARGET_JOB:-}"
+TARGET_GROUP="${TARGET_GROUP:-}"
+
+TARGET_COUNT="${TARGET_COUNT:-0}"
+
+CURL_ARGS=(
+  --silent
+  --show-error
+  --fail-with-body
+)
+
+if [[ -n "${NOMAD_TOKEN}" ]]; then
+  CURL_ARGS+=(-H "Authorization: Bearer ${NOMAD_TOKEN}")
+fi
+if [[ -n "${NOMAD_UNIX_ADDR}" ]]; then
+  CURL_ARGS+=(--unix-socket "${NOMAD_UNIX_ADDR}")
+fi
+
+if ! [[ "$TARGET_COUNT" =~ ^[+-]?[0-9]+$ ]]; then
+  echo "$TARGET_COUNT is not a number"
+  exit 1
+fi
+
+curl "${CURL_ARGS[@]}" \
+  -X POST \
+  "${NOMAD_ADDR}/v1/job/${TARGET_JOB}/scale?namespace=${TARGET_NAMESPACE}" \
+  -H "Content-Type: application/json" \
+  -d "$(jq -n \
+    --arg group "$TARGET_GROUP" \
+    --arg count "$TARGET_COUNT" \
+    '{
+      Count: ($count|tonumber),
+      Target: {
+        "Group": ($group)
+      },
+      Message: "Scaling up to perform backup"
+    }'
+  )"
+
+while true; do
+  allocations=$(
+    curl "${CURL_ARGS[@]}" \
+      "${NOMAD_ADDR}/v1/job/${TARGET_JOB}/allocations?namespace=${TARGET_NAMESPACE}"
+  )
+
+  count=$(
+    jq --arg group "$TARGET_GROUP" '
+      [
+        .[]
+        | select(.TaskGroup == $group)
+        | select(
+            .ClientStatus != "complete"
+            and .ClientStatus != "dead"
+          )
+      ] | length
+    ' <<< "$allocations"
+  )
+
+  echo "Remaining allocations: ${count}"
+
+  if [[ "$count" -eq $TARGET_COUNT ]]; then
+    break
+  fi
+
+  sleep 5
+done