apiVersion: v1
kind: ConfigMap
metadata:
  name: scale-ci-workload-script
data:
  run.sh: |
    #!/bin/sh
    set -eo pipefail
    workload_log() { echo "$(date -u) $@" >&2; }
    export -f workload_log
    workload_log "Configuring pbench for running Scale workload"
    mkdir -p /var/lib/pbench-agent/tools-default/
    echo "${USER_NAME:-default}:x:$(id -u):0:${USER_NAME:-default} user:${HOME}:/sbin/nologin" >> /etc/passwd
    if [ "${ENABLE_PBENCH_AGENTS}" = true ]; then
      echo "" > /var/lib/pbench-agent/tools-default/disk
      echo "" > /var/lib/pbench-agent/tools-default/iostat
      echo "workload" > /var/lib/pbench-agent/tools-default/label
      echo "" > /var/lib/pbench-agent/tools-default/mpstat
      echo "" > /var/lib/pbench-agent/tools-default/oc
      echo "" > /var/lib/pbench-agent/tools-default/perf
      echo "" > /var/lib/pbench-agent/tools-default/pidstat
      echo "" > /var/lib/pbench-agent/tools-default/sar
      master_nodes=`oc get nodes -l pbench_agent=true,node-role.kubernetes.io/master= --no-headers | awk '{print $1}'`
      for node in $master_nodes; do
        echo "master" > /var/lib/pbench-agent/tools-default/remote@$node
      done
      infra_nodes=`oc get nodes -l pbench_agent=true,node-role.kubernetes.io/infra= --no-headers | awk '{print $1}'`
      for node in $infra_nodes; do
        echo "infra" > /var/lib/pbench-agent/tools-default/remote@$node
      done
      worker_nodes=`oc get nodes -l pbench_agent=true,node-role.kubernetes.io/worker= --no-headers | awk '{print $1}'`
      for node in $worker_nodes; do
        echo "worker" > /var/lib/pbench-agent/tools-default/remote@$node
      done
    fi
    source /opt/pbench-agent/profile
    workload_log "Done configuring pbench for Scale workload run"

    workload_log "Running scale workload"
    if [ "${PBENCH_INSTRUMENTATION}" = "true" ]; then
      pbench-user-benchmark -- sh /root/workload/workload.sh
      result_dir="/var/lib/pbench-agent/$(ls -t /var/lib/pbench-agent/ | grep "pbench-user" | head -1)"/1/sample1
      if [ "${ENABLE_PBENCH_COPY}" = "true" ]; then
        pbench-copy-results --prefix ${SCALE_TEST_PREFIX}
      fi
    else
      sh /root/workload/workload.sh
      result_dir=/tmp
    fi
    workload_log "Completed Scale workload run"

    workload_log "Checking Test Results"
    workload_log "Checking Test Exit Code"
    if [ $(jq '.exit_code==0' ${result_dir}/exit.json) == "false" ]; then
      workload_log "Test Failure"
      workload_log "Test Analysis: Failed"
      exit 1
    fi
    workload_log "Comparing scale duration to expected duration"
    workload_log "Scaling Duration: $(jq '.duration' ${result_dir}/exit.json)"
    if [ $(jq '.duration>'${EXPECTED_SCALE_DURATION}'' ${result_dir}/exit.json) == "true" ]; then
      workload_log "EXPECTED_SCALE_DURATION (${EXPECTED_SCALE_DURATION}) exceeded ($(jq '.duration' ${result_dir}/exit.json))"
      workload_log "Test Analysis: Failed"
      exit 1
    fi
    # TODO: Check pbench-agent collected metrics for Pass/Fail
    # TODO: Check prometheus collected metrics for Pass/Fail
    workload_log "Test Analysis: Passed"
  workload.sh: |
    #!/bin/sh

    result_dir=/tmp
    if [ "${PBENCH_INSTRUMENTATION}" = "true" ]; then
      result_dir=${benchmark_results_dir}
    fi
    machinesets=$(oc get machineset -n openshift-machine-api | egrep "\-worker\-|\-w\-" | awk '{print $1}')
    IFS=$'\n' read -rd '' -a ms_arr <<<"$machinesets"

    start_time=$(date +%s)
    if [ "${#ms_arr[@]}" == "1" ]; then
      workload_log "Scaling ${ms_arr[0]} to ${SCALE_WORKER_COUNT}"
      oc patch machineset ${ms_arr[0]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${SCALE_WORKER_COUNT}' }}'
    elif [ "${#ms_arr[@]}" == "2" ]; then
      az_a_count=$(((${SCALE_WORKER_COUNT}+1)/2))
      az_b_count=$(((${SCALE_WORKER_COUNT})/2))

      workload_log "Scaling ${ms_arr[0]} to ${az_a_count}"
      oc patch machineset ${ms_arr[0]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_a_count}' }}'
      workload_log "Scaling ${ms_arr[1]} to ${az_b_count}"
      oc patch machineset ${ms_arr[1]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_b_count}' }}'
    elif [ "${#ms_arr[@]}" == "3" ]; then
      az_a_count=$(((${SCALE_WORKER_COUNT}+2)/3))
      az_b_count=$(((${SCALE_WORKER_COUNT}+1)/3))
      az_c_count=$((${SCALE_WORKER_COUNT}/3))

      workload_log "Scaling ${ms_arr[0]} to ${az_a_count}"
      oc patch machineset ${ms_arr[0]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_a_count}' }}'
      workload_log "Scaling ${ms_arr[1]} to ${az_b_count}"
      oc patch machineset ${ms_arr[1]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_b_count}' }}'
      workload_log "Scaling ${ms_arr[2]} to ${az_c_count}"
      oc patch machineset ${ms_arr[2]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_c_count}' }}'
    elif [ "${#ms_arr[@]}" == "4" ]; then
      az_a_count=$(((${SCALE_WORKER_COUNT}+3)/4))
      az_b_count=$(((${SCALE_WORKER_COUNT}+2)/4))
      az_c_count=$(((${SCALE_WORKER_COUNT}+1)/4))
      az_d_count=$((${SCALE_WORKER_COUNT}/4))

      workload_log "Scaling ${ms_arr[0]} to ${az_a_count}"
      oc patch machineset ${ms_arr[0]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_a_count}' }}'
      workload_log "Scaling ${ms_arr[1]} to ${az_b_count}"
      oc patch machineset ${ms_arr[1]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_b_count}' }}'
      workload_log "Scaling ${ms_arr[2]} to ${az_c_count}"
      oc patch machineset ${ms_arr[2]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_c_count}' }}'
      workload_log "Scaling ${ms_arr[3]} to ${az_d_count}"
      oc patch machineset ${ms_arr[3]} --type=merge -n openshift-machine-api -p '{"spec": {"replicas": '${az_d_count}' }}'
    else
      workload_log "Unhandled number of machinesets: ${#ms_arr[@]}"
      exit 1
    fi

    retries=0
    while [ ${retries} -le ${SCALE_POLL_ATTEMPTS} ] ; do
      worker_count=`oc get nodes -l node-role.kubernetes.io/worker= --no-headers | grep " Ready" -c`
      if [ "${worker_count}" == "${SCALE_WORKER_COUNT}" ]; then
        workload_log "Cluster Scaled to ${SCALE_WORKER_COUNT}"
        break
      else
        workload_log "Cluster Scaling from ${worker_count} to ${SCALE_WORKER_COUNT}, Poll attempts: ${retries}/${SCALE_POLL_ATTEMPTS}"
        sleep 2
      fi
      retries=$[${retries} + 1]
    done
    end_time=$(date +%s)
    duration=$((end_time-start_time))
    exit_code=0
    if [ "${worker_count}" != "${SCALE_WORKER_COUNT}" ]; then
      workload_log "Cluster failed to scale to ${SCALE_WORKER_COUNT} in (${SCALE_POLL_ATTEMPTS} * 2s)"
      exit_code=1
    fi
    workload_log "Writing Exit Code and Duration"
    jq -n '. | ."exit_code"='${exit_code}' | ."duration"='${duration}'' > "${result_dir}/exit.json"
