wenshu/k8s-orchestrator/test/e2e/full_api_e2e.sh

#!/usr/bin/env bash
set -euo pipefail

# Full API e2e for k8s-orchestrator based on README interfaces.
# Tested against real kubeadm cluster and real worker hosts.

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"

API_HOST="${E2E_API_HOST:-127.0.0.1}"
API_PORT="${E2E_API_PORT:-18080}"
API="http://${API_HOST}:${API_PORT}"

SECRET_NAMESPACE="${E2E_SECRET_NAMESPACE:-orchestrator-system}"
SSH_USER="${E2E_SSH_USER:-root}"
SSH_PASSWORD="${E2E_SSH_PASSWORD:-Wenshu_test}"
SPARK_OPERATOR_NAMESPACE="${E2E_SPARK_OPERATOR_NAMESPACE:-spark-operator}"
SPARK_OPERATOR_RELEASE="${E2E_SPARK_OPERATOR_RELEASE:-spark-operator}"
SPARK_OPERATOR_CHART="${E2E_SPARK_OPERATOR_CHART:-spark-operator/spark-operator}"
STARROCKS_OPERATOR_NAMESPACE="${E2E_STARROCKS_OPERATOR_NAMESPACE:-starrocks}"
STARROCKS_OPERATOR_DEPLOYMENT="${E2E_STARROCKS_OPERATOR_DEPLOYMENT:-kube-starrocks-operator}"
STARROCKS_CRD_URL="${E2E_STARROCKS_CRD_URL:-https://raw.githubusercontent.com/StarRocks/starrocks-kubernetes-operator/main/deploy/starrocks.com_starrocksclusters.yaml}"
STARROCKS_OPERATOR_URL="${E2E_STARROCKS_OPERATOR_URL:-https://raw.githubusercontent.com/StarRocks/starrocks-kubernetes-operator/main/deploy/operator.yaml}"

# Comma-separated nodeId=ip
WORKER_MAP="${E2E_WORKER_MAP:-worker-1=47.237.255.159,worker-2=47.237.255.162,worker-3=47.237.254.197,worker-4=47.84.84.212}"

# true: reset workers to no kubeadm/kubelet/containerd before AddNode
RESET_WORKERS="${E2E_RESET_WORKERS:-false}"
# true: build/start local api/controller automatically
START_LOCAL="${E2E_START_LOCAL:-true}"

# default inventory capacity used for AddNode request
NODE_CPU="${E2E_NODE_CPU:-4}"
NODE_RAM="${E2E_NODE_RAM:-7}"
NODE_SSD="${E2E_NODE_SSD:-100}"

SPARK_OP_MAX="${E2E_SPARK_OP_MAX:-420}"
STARROCKS_OP_MAX="${E2E_STARROCKS_OP_MAX:-420}"
DEFAULT_OP_MAX="${E2E_DEFAULT_OP_MAX:-360}"

LOG_DIR="${E2E_LOG_DIR:-/tmp}"
API_BIN="${E2E_API_BIN:-/tmp/orchestrator-api}"
CONTROLLER_BIN="${E2E_CONTROLLER_BIN:-/tmp/orchestrator-controller}"
API_LOG="${LOG_DIR}/orchestrator-api.log"
CONTROLLER_LOG="${LOG_DIR}/orchestrator-controller.log"
API_PID_FILE="${LOG_DIR}/orchestrator-api.pid"
CONTROLLER_PID_FILE="${LOG_DIR}/orchestrator-controller.pid"
JOIN_CMD_FILE="${LOG_DIR}/kubeadm_join_cmd"

SCRIPT_TAG="$(date +%s)"
SPARK_CLUSTER_ID="spark-e2e-${SCRIPT_TAG}"
SPARK_NS="spark-e2e-${SCRIPT_TAG}"
SR_CLUSTER_ID="sr-e2e-${SCRIPT_TAG}"
SR_NS="starrocks-sr-e2e-${SCRIPT_TAG}"

IFS=',' read -r -a WORKERS <<<"${WORKER_MAP}"

need_cmd() {
  command -v "$1" >/dev/null 2>&1 || {
    echo "missing required command: $1" >&2
    exit 1
  }
}

log() {
  echo "[$(date '+%F %T')] $*"
}

cleanup_local_processes() {
  if [[ -f "${API_PID_FILE}" ]]; then
    kill "$(cat "${API_PID_FILE}")" >/dev/null 2>&1 || true
  fi
  if [[ -f "${CONTROLLER_PID_FILE}" ]]; then
    kill "$(cat "${CONTROLLER_PID_FILE}")" >/dev/null 2>&1 || true
  fi
}

on_exit() {
  if [[ "${START_LOCAL}" == "true" ]]; then
    cleanup_local_processes
  fi
}
trap on_exit EXIT

wait_api_ready() {
  for _ in $(seq 1 60); do
    if curl -sf "${API}/healthz" >/dev/null; then
      return 0
    fi
    sleep 1
  done
  echo "API not ready: ${API}" >&2
  tail -n 200 "${API_LOG}" || true
  exit 1
}

ensure_k8s_prereq() {
  log "ensure namespace/crd/rbac"
  kubectl get ns "${SECRET_NAMESPACE}" >/dev/null 2>&1 || kubectl create ns "${SECRET_NAMESPACE}"
  kubectl apply -k "${ROOT_DIR}/config/crd" >/dev/null
  kubectl apply -k "${ROOT_DIR}/config/rbac" >/dev/null

  log "refresh kubeadm-join secret"
  local join_cmd
  join_cmd="$(kubeadm token create --print-join-command)"
  kubectl -n "${SECRET_NAMESPACE}" create secret generic kubeadm-join \
    --from-literal=joinCommand="${join_cmd}" \
    --dry-run=client -o yaml | kubectl apply -f - >/dev/null
  echo "${join_cmd}" >"${JOIN_CMD_FILE}"
}

ensure_external_operators() {
  log "ensure Spark operator"
  helm repo add spark-operator https://kubeflow.github.io/spark-operator >/dev/null 2>&1 || true
  helm repo update spark-operator >/dev/null
  local spark_operator_job_namespaces_json
  spark_operator_job_namespaces_json="${E2E_SPARK_OPERATOR_JOB_NAMESPACES_JSON:-[\"\"]}"
  local spark_operator_ivy_cache_path
  spark_operator_ivy_cache_path="${E2E_SPARK_OPERATOR_IVY_CACHE_PATH:-/var/lib/wenshu/spark-ivy}"

  helm upgrade --install "${SPARK_OPERATOR_RELEASE}" "${SPARK_OPERATOR_CHART}" \
    --namespace "${SPARK_OPERATOR_NAMESPACE}" \
    --create-namespace \
    --set-json "spark.jobNamespaces=${spark_operator_job_namespaces_json}" \
    --set-json "controller.volumes=[{\"name\":\"tmp\",\"emptyDir\":{\"sizeLimit\":\"1Gi\"}},{\"name\":\"wenshu-spark-ivy-cache\",\"emptyDir\":{}}]" \
    --set-json "controller.volumeMounts=[{\"name\":\"tmp\",\"mountPath\":\"/tmp\",\"readOnly\":false},{\"name\":\"wenshu-spark-ivy-cache\",\"mountPath\":\"${spark_operator_ivy_cache_path}\",\"readOnly\":false}]" >/dev/null
  kubectl -n "${SPARK_OPERATOR_NAMESPACE}" rollout status deploy/spark-operator-controller --timeout=300s >/dev/null
  if kubectl -n "${SPARK_OPERATOR_NAMESPACE}" get deploy spark-operator-webhook >/dev/null 2>&1; then
    kubectl -n "${SPARK_OPERATOR_NAMESPACE}" rollout status deploy/spark-operator-webhook --timeout=300s >/dev/null
  fi

  log "ensure StarRocks operator"
  kubectl get ns "${STARROCKS_OPERATOR_NAMESPACE}" >/dev/null 2>&1 || kubectl create ns "${STARROCKS_OPERATOR_NAMESPACE}" >/dev/null
  if ! kubectl get crd starrocksclusters.starrocks.com >/dev/null 2>&1; then
    kubectl apply -f "${STARROCKS_CRD_URL}" >/dev/null
  fi
  if ! kubectl -n "${STARROCKS_OPERATOR_NAMESPACE}" get deploy "${STARROCKS_OPERATOR_DEPLOYMENT}" >/dev/null 2>&1; then
    kubectl apply -f "${STARROCKS_OPERATOR_URL}" >/dev/null
  fi
  kubectl -n "${STARROCKS_OPERATOR_NAMESPACE}" rollout status deploy/"${STARROCKS_OPERATOR_DEPLOYMENT}" --timeout=300s >/dev/null
}

start_local_binaries() {
  log "build local api/controller"
  (cd "${ROOT_DIR}" && go build -o "${API_BIN}" ./cmd/orchestrator-api)
  (cd "${ROOT_DIR}" && go build -o "${CONTROLLER_BIN}" ./cmd/orchestrator-controller)

  cleanup_local_processes

  log "start local api/controller"
  setsid env API_BIND_ADDRESS="${API_HOST}:${API_PORT}" SECRET_NAMESPACE="${SECRET_NAMESPACE}" \
    "${API_BIN}" >"${API_LOG}" 2>&1 < /dev/null &
  echo $! >"${API_PID_FILE}"

  setsid env KUBEADM_JOIN_COMMAND="$(cat "${JOIN_CMD_FILE}")" \
    "${CONTROLLER_BIN}" --leader-elect=false >"${CONTROLLER_LOG}" 2>&1 < /dev/null &
  echo $! >"${CONTROLLER_PID_FILE}"

  wait_api_ready
}

cleanup_orchestrator_resources() {
  log "cleanup old orchestrator resources"
  kubectl delete clusteroperations.orchestrator.spark.cyrus.io --all --ignore-not-found >/dev/null 2>&1 || true
  kubectl delete sparkvirtualclusters.orchestrator.spark.cyrus.io --all --ignore-not-found >/dev/null 2>&1 || true
  kubectl delete starrocksvirtualclusters.orchestrator.spark.cyrus.io --all --ignore-not-found >/dev/null 2>&1 || true
  kubectl delete nodeinventories.orchestrator.spark.cyrus.io --all --ignore-not-found >/dev/null 2>&1 || true
}

clear_node_isolation() {
  for pair in "${WORKERS[@]}"; do
    local node_id
    node_id="${pair%%=*}"
    kubectl taint nodes "${node_id}" spark-cluster- --overwrite >/dev/null 2>&1 || true
    kubectl taint nodes "${node_id}" starrocks-cluster- --overwrite >/dev/null 2>&1 || true
    kubectl label nodes "${node_id}" spark.cluster.id- starrocks.cluster.id- orchestrator.spark.cyrus.io/managed-by- --overwrite >/dev/null 2>&1 || true
  done
}

reset_workers_if_needed() {
  if [[ "${RESET_WORKERS}" != "true" ]]; then
    return 0
  fi

  log "reset workers to no kubeadm/kubelet/containerd"
  for pair in "${WORKERS[@]}"; do
    local node_id ip
    node_id="${pair%%=*}"
    ip="${pair##*=}"

    kubectl delete node "${node_id}" --ignore-not-found >/dev/null 2>&1 || true

    sshpass -p "${SSH_PASSWORD}" ssh -o StrictHostKeyChecking=no "${SSH_USER}@${ip}" 'bash -s' <<'EOS'
set -euo pipefail
apt-mark unhold kubeadm kubelet kubectl || true
apt-get purge -y kubeadm kubelet kubectl kubernetes-cni containerd conntrack socat apt-transport-https || true
apt-get autoremove -y || true
kubeadm reset -f || true
systemctl stop kubelet || true
# kubelet may leave projected volume mounts; lazy-unmount before cleanup.
findmnt -R /var/lib/kubelet -o TARGET -n 2>/dev/null | sort -r | xargs -r umount -l || true
rm -rf /etc/kubernetes /var/lib/kubelet /var/lib/cni /etc/cni/net.d /etc/containerd || true
rm -f /etc/apt/sources.list.d/kubernetes.list /etc/apt/keyrings/kubernetes-apt-keyring.gpg
EOS

    sshpass -p "${SSH_PASSWORD}" ssh -o StrictHostKeyChecking=no "${SSH_USER}@${ip}" '
      command -v kubeadm >/dev/null && { echo "kubeadm should be missing"; exit 1; } || true
      command -v kubelet >/dev/null && { echo "kubelet should be missing"; exit 1; } || true
      command -v containerd >/dev/null && { echo "containerd should be missing"; exit 1; } || true
    '
  done
}

wait_op() {
  local op_id=$1
  local max=${2:-${DEFAULT_OP_MAX}}

  for i in $(seq 1 "${max}"); do
    local opjson phase message
    opjson="$(curl -sS "${API}/v1/operations/${op_id}")"
    phase="$(echo "${opjson}" | jq -r '.status.phase // empty')"
    message="$(echo "${opjson}" | jq -r '.status.message // empty')"
    log "op=${op_id} step=${i}/${max} phase=${phase} msg=${message}"

    if [[ "${phase}" == "SUCCEEDED" ]]; then
      return 0
    fi
    if [[ "${phase}" == "FAILED" ]]; then
      echo "${opjson}" >&2
      return 1
    fi
    sleep 5
  done

  echo "operation timeout: ${op_id}" >&2
  return 1
}

call_mutation() {
  local method=$1
  local path=$2
  local key=$3
  local body=${4:-}
  local resp

  if [[ -n "${body}" ]]; then
    resp="$(curl -sS -X "${method}" "${API}${path}" \
      -H 'Content-Type: application/json' \
      -H "Idempotency-Key: ${key}" \
      -d "${body}")"
  else
    resp="$(curl -sS -X "${method}" "${API}${path}" \
      -H "Idempotency-Key: ${key}")"
  fi

  echo "${resp}" | jq . >&2
  echo "${resp}" | jq -r '.operationId'
}

add_node() {
  local node_id=$1
  local ip=$2
  local key=$3
  local body op

  body="$(jq -nc \
    --arg id "${node_id}" \
    --arg ip "${ip}" \
    --argjson cpu "${NODE_CPU}" \
    --argjson ram "${NODE_RAM}" \
    --argjson ssd "${NODE_SSD}" \
    --arg user "${SSH_USER}" \
    --arg pw "${SSH_PASSWORD}" \
    '{nodeId:$id,cpu:$cpu,ram:$ram,ssd:$ssd,ipAddress:$ip,sshUser:$user,sshPassword:$pw,sshPort:22}')"

  op="$(call_mutation POST /v1/nodes "${key}" "${body}")"
  wait_op "${op}"
}

assert_operator_crds() {
  kubectl get crd sparkapplications.sparkoperator.k8s.io >/dev/null
  kubectl get crd starrocksclusters.starrocks.com >/dev/null
  kubectl -n "${SPARK_OPERATOR_NAMESPACE}" get deploy spark-operator-controller >/dev/null
  kubectl -n "${STARROCKS_OPERATOR_NAMESPACE}" get deploy "${STARROCKS_OPERATOR_DEPLOYMENT}" >/dev/null
}

run_flow() {
  local op
  local spark_create_body spark_job_body
  local sr_create_body sr_job_body

  log "GET /healthz"
  curl -sS "${API}/healthz" | jq .

  log "register nodes via POST /v1/nodes"
  for pair in "${WORKERS[@]}"; do
    local node_id ip
    node_id="${pair%%=*}"
    ip="${pair##*=}"
    add_node "${node_id}" "${ip}" "${SCRIPT_TAG}-add-${node_id}"
  done

  log "GET /v1/nodes"
  curl -sS "${API}/v1/nodes" | jq '.items | length'

  spark_create_body="$(jq -nc \
    --arg cid "${SPARK_CLUSTER_ID}" \
    --arg ns "${SPARK_NS}" \
    '{clusterId:$cid,namespace:$ns,nodeIds:["worker-1","worker-2"],sparkConfig:{"spark.dynamicAllocation.enabled":"false","spark.image":"apache/spark:3.5.1","spark.imagePullPolicy":"IfNotPresent","spark.prepull.enabled":"true"}}')"

  log "POST /v1/spark-clusters"
  op="$(call_mutation POST /v1/spark-clusters "${SCRIPT_TAG}-spark-create" "${spark_create_body}")"
  wait_op "${op}" "${SPARK_OP_MAX}"

  log "GET /v1/spark-clusters/{id}"
  curl -sS "${API}/v1/spark-clusters/${SPARK_CLUSTER_ID}" | jq '.metadata.name,.status.phase'

  log "POST /v1/spark-clusters/{id}/scale-out"
  op="$(call_mutation POST "/v1/spark-clusters/${SPARK_CLUSTER_ID}/scale-out" "${SCRIPT_TAG}-spark-scaleout" '{"nodeIds":["worker-3"]}')"
  wait_op "${op}" "${SPARK_OP_MAX}"

  log "POST /v1/spark-clusters/{id}/scale-in"
  op="$(call_mutation POST "/v1/spark-clusters/${SPARK_CLUSTER_ID}/scale-in" "${SCRIPT_TAG}-spark-scalein" '{"nodeIds":["worker-2"]}')"
  wait_op "${op}" "${SPARK_OP_MAX}"

  spark_job_body="$(jq -nc \
    --arg cid "${SPARK_CLUSTER_ID}" \
    '{clusterId:$cid,sparkApplicationSpec:{type:"Scala",mode:"cluster",image:"apache/spark:3.5.1",mainClass:"org.apache.spark.examples.SparkPi",mainApplicationFile:"local:///opt/spark/examples/jars/spark-examples_2.12-3.5.1.jar",driver:{cores:1,memory:"1g"},executor:{instances:1,cores:1,memory:"1g"}}}')"

  log "POST /v1/spark-jobs"
  op="$(call_mutation POST /v1/spark-jobs "${SCRIPT_TAG}-spark-job" "${spark_job_body}")"
  wait_op "${op}" "${SPARK_OP_MAX}"

  for i in $(seq 1 120); do
    local running
    running="$(kubectl -n "${SPARK_NS}" get pods --no-headers 2>/dev/null | awk '$3=="Running" || $3=="Pending" {c++} END {print c+0}')"
    log "spark pods running_or_pending=${running}"
    [[ "${running}" == "0" ]] && break
    sleep 5
  done

  log "POST /v1/spark-clusters/{id}/release"
  op="$(call_mutation POST "/v1/spark-clusters/${SPARK_CLUSTER_ID}/release" "${SCRIPT_TAG}-spark-release")"
  wait_op "${op}" "${SPARK_OP_MAX}"

  sr_create_body="$(jq -nc \
    --arg cid "${SR_CLUSTER_ID}" \
    --arg ns "${SR_NS}" \
    '{clusterId:$cid,namespace:$ns,feNodeIds:["worker-1"],beNodeIds:["worker-2"],starRocksConfig:{feResources:{requests:{cpu:"1",memory:"1Gi"},limits:{cpu:"2",memory:"2Gi"}},beResources:{requests:{cpu:"1",memory:"1Gi"},limits:{cpu:"2",memory:"2Gi"}}}}')"

  log "POST /v1/starrocks-clusters"
  op="$(call_mutation POST /v1/starrocks-clusters "${SCRIPT_TAG}-sr-create" "${sr_create_body}")"
  wait_op "${op}" "${STARROCKS_OP_MAX}"

  log "GET /v1/starrocks-clusters/{id}"
  curl -sS "${API}/v1/starrocks-clusters/${SR_CLUSTER_ID}" | jq '.metadata.name,.status.phase'

  log "POST /v1/starrocks-clusters/{id}/scale-out"
  op="$(call_mutation POST "/v1/starrocks-clusters/${SR_CLUSTER_ID}/scale-out" "${SCRIPT_TAG}-sr-scaleout" '{"nodeIds":["worker-3"]}')"
  wait_op "${op}" "${STARROCKS_OP_MAX}"

  log "POST /v1/starrocks-clusters/{id}/scale-in"
  op="$(call_mutation POST "/v1/starrocks-clusters/${SR_CLUSTER_ID}/scale-in" "${SCRIPT_TAG}-sr-scalein" '{"nodeIds":["worker-2"]}')"
  wait_op "${op}" "${STARROCKS_OP_MAX}"

  sr_job_body="$(jq -nc \
    --arg cid "${SR_CLUSTER_ID}" \
    '{clusterId:$cid,jobSpec:{backoffLimit:1,template:{spec:{restartPolicy:"Never",containers:[{name:"task",image:"busybox:1.36",command:["sh","-c","echo starrocks-job-ok"]}]}}}}')"

  log "POST /v1/starrocks-jobs"
  op="$(call_mutation POST /v1/starrocks-jobs "${SCRIPT_TAG}-sr-job" "${sr_job_body}")"
  wait_op "${op}" "${STARROCKS_OP_MAX}"

  log "POST /v1/starrocks-clusters/{id}/release"
  op="$(call_mutation POST "/v1/starrocks-clusters/${SR_CLUSTER_ID}/release" "${SCRIPT_TAG}-sr-release")"
  wait_op "${op}" "${STARROCKS_OP_MAX}"

  log "DELETE /v1/nodes/{id}"
  for pair in "${WORKERS[@]}"; do
    local node_id
    node_id="${pair%%=*}"
    op="$(call_mutation DELETE "/v1/nodes/${node_id}" "${SCRIPT_TAG}-del-${node_id}")"
    wait_op "${op}" "${DEFAULT_OP_MAX}"
  done

  log "final GET /v1/nodes"
  curl -sS "${API}/v1/nodes" | jq '.items | length'
}

main() {
  need_cmd kubectl
  need_cmd curl
  need_cmd jq
  need_cmd sshpass
  need_cmd kubeadm
  need_cmd helm

  ensure_k8s_prereq
  ensure_external_operators
  if [[ "${START_LOCAL}" == "true" ]]; then
    start_local_binaries
  else
    wait_api_ready
  fi

  assert_operator_crds
  cleanup_orchestrator_resources
  clear_node_isolation
  reset_workers_if_needed
  run_flow

  log "E2E PASSED"
}

main "$@"