cti

#!/bin/bash
# Copyright (c) The Libra Core Contributors
# SPDX-License-Identifier: Apache-2.0
set -e
set -o pipefail

TAG=""
CLUSTER_TEST_TAG=""
PR=""
WORKSPACE=""
ENV=""
REPORT=""
DEPLOY=yes
EXIT_CODE=0

K8S=""
K8S_CONTEXT_PATTERN='arn:aws:eks:us-west-2:853397791086:cluster/CLUSTERNAME-k8s-testnet'

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"

source "${DIR}/ct.vars"

# Colorize Output
RESTORE=$(echo -en '\001\033[0m\002')
BLUE=$(echo -en '\001\033[01;34m\002')

join_args() {
  retval_join_args=""
  for var in $*
  do
    retval_join_args="${retval_join_args}, \"${var}\""
  done
}

join_env_vars() {
  retval_join_env_vars=""
  for var in $*
  do
    IFS='=' read -ra env_var <<< "$var"
    retval_join_env_vars="{\"name\":\"${env_var[0]}\", \"value\":\"${env_var[1]}\"}, ${retval_join_env_vars}"
  done
}

kube_init_context () {
  aws eks --region us-west-2 describe-cluster --name ct-0-k8s-testnet &>/dev/null || (echo "Failed to access EKS, try awsmfa?"; exit 1)
  local highest_pool_index=$(($K8S_POOL_SIZE - 1))
  local context=${K8S_CONTEXT_PATTERN/CLUSTERNAME/ct-${highest_pool_index}}
  if kubectl config get-contexts ${context} &> /dev/null; then
    return
  fi
  for ((i = 0; i < ${K8S_POOL_SIZE}; i++)); do
    aws eks --region us-west-2 update-kubeconfig --name ct-${i}-k8s-testnet
  done
}

kube_select_cluster () {
  retval_kube_select_cluster=""
  for attempt in {1..360} ; do
    for ((i = 0; i < ${K8S_POOL_SIZE}; i++)); do
      local context=${K8S_CONTEXT_PATTERN/CLUSTERNAME/ct-${i}}
      local running_pods=$(kubectl --context="${context}" get pods -l app=cluster-test --field-selector=status.phase==Running 2> /dev/null  | grep -v ^NAME | wc -l)
      local pending_pods=$(kubectl --context="${context}" get pods -l app=cluster-test --field-selector=status.phase==Pending 2> /dev/null | grep -v ^NAME | wc -l)
      local prometheus_healthy_containers=$(kubectl --context="${context}" get pod/libra-testnet-prometheus-server-0 | grep 'libra-testnet-prometheus-server-0' | awk '{print $2}')
      if [[ "${pending_pods}" -gt 0 ]]; then
        echo "ct-${i} has ${pending_pods} pending pods. Skipping."
      elif [[ ${prometheus_healthy_containers} != "2/2" ]]; then
        echo "prometheus is not healthy for ct-${i}. Skipping."
      elif [[ ${running_pods} -gt 0 ]]; then
        echo "ct-${i} has ${running_pods} running pods. Skipping."
      else
        retval_kube_select_cluster="ct-${i}"
        return
      fi
    done
    echo "All clusters have jobs running on them. Retrying in 10 secs."
    sleep 10
  done
  echo "Failed to schedule job on a cluster as all are busy"
  exit 1
}

kube_wait_pod () {
  local pod_name="${1}"
  local context="${2}"
  for i in {1..360} ; do
    local phase=$(kubectl --context="${context}" get pod "${pod_name}" -o jsonpath="{.status.phase}" || echo -n "kubectlfailed")
    if [[ "${phase}" == "kubectlfailed" ]]; then
        echo "kubectl get pod ${pod_name} failed. Retrying."
        sleep 10
        continue
    fi
    if [[ $phase != "Pending" &&  $phase != "Unknown" ]]; then
      echo "${pod_name} reached phase : ${phase}"
      return
    fi
    if kubectl --context="${context}" get pod "${pod_name}" | grep -i -e ImagePullBackOff -e ErrImagePull &>/dev/null; then
      image_name=$(kubectl --context="${context}" get pod "${pod_name}" -o jsonpath="{.spec.containers[0].image}")
      echo "${pod_name} name failed to be scheduled because there was an error pulling the image : ${image_name}"
      # Delete the pod so that it doesn't block other pods from being scheduled on this
      kubectl  --context="${context}" delete pod "${pod_name}"
      exit 1
    fi
    echo "Waiting for ${pod_name} to be scheduled. Current phase : ${phase}"
    sleep 10
  done
  echo "Pod ${pod_name} failed to be scheduled"
  exit 1
}

while (( "$#" )); do
  case "$1" in
    --k8s)
      K8S=1
      shift 1
      ;;
    -R|--report)
      REPORT=$2
      shift 2
      ;;
    -p|--pr)
      PR=$2
      shift 2
      ;;
    -M|--master)
      TAG=master
      shift 1
      ;;
    -T|--tag)
      TAG=$2
      shift 2
      ;;
    --cluster-test-tag)
      CLUSTER_TEST_TAG=$2
      shift 2
      ;;
    -W|--workspace)
      WORKSPACE=$2
      shift 2
      ;;
    -E|--env)
      ENV="$ENV $2"
      shift 2
      ;;
    -c|--container|-i|--image|--deploy)
      echo "$1 command is not allowed in cti"
      exit 1
      ;;
    *) # end argument parsing
      break
      ;;
  esac
done

if [ -z "$PR" ] && [ -z "$TAG" ]
then
      echo "No PR or tag specified"
      echo "Usage:"
      echo "cti [--pr <PR>|--master|--tag <TAG>] [--workspace <WORKSPACE>] [-E <env>] <cluster test arguments>"
      echo
      echo "--pr <PR>: Build image from pull request #<PR>"
      echo "-M|--master: Use latest image available in CI"
      echo "-T|--tag <TAG>: Use image with tag TAG"
      echo "--cluster-test-tag <TAG>: Use this tag for cluster test"
      echo "-W|--workspace <WORKSPACE>: Use custom workplace <WORKSPACE>"
      echo "-E|--env <ENV>: Set environment variable, ex. -E RUST_LOG=debug. Can be repeated, e.g. -E A=B -E C=D"
      echo "-R|--report file.json: Generate json report into file.json"
      echo
      echo "To see cluster test runner arguments run cti --master"
      echo
      echo "Examples:"
      echo "cti --pr <PR> --run bench # Run benchmark"
      echo "cti --master --emit-tx # Submit transactions until Ctrl+C pressed"
      exit 1
fi

if [ -z "$TAG" ]; then
    aws codebuild list-projects >/dev/null || (echo "Failed to access codebuild, try awsmfa?"; exit 1)
    ./docker/build-aws.sh --build-all --version pull/$PR
    TAG=dev_${USER}_pull_${PR}
    echo "**TIP Use cti -T $TAG <...> to restart this run with same tag without rebuilding it"
fi

CLUSTER_TEST_TAG=${CLUSTER_TEST_TAG:-${TAG}}

OUTPUT_TEE=${CTI_OUTPUT_LOG:-$(mktemp)}

if ! which kubectl &>/dev/null; then
  echo "kubectl is not installed. Please install kubectl. On mac, you can use : brew install kubectl"
  exit 1
fi
echo "Running cluster-test on Kubernetes"
kube_init_context
pod_name="cluster-test-$(whoami)-$(date +%s)"
pod_name=${pod_name/_/-} #underscore not allowed in pod name
specfile=$(mktemp)
echo "Pod Spec : ${specfile}"
join_args $*
ENV="$ENV AWS_ROLE_SESSION_NAME=AWS_ROLE_SESSION_NAME"
join_env_vars $ENV
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
sed -e "s/{pod_name}/${pod_name}/g" \
    -e "s/{image_tag}/${TAG}/g" \
    -e "s/{cluster_test_image_tag}/${CLUSTER_TEST_TAG}/g" \
    -e "s^{env_variables}^${retval_join_env_vars}^g" \
    -e "s/{extra_args}/${retval_join_args}/g" \
    ${DIR}/cluster_test_pod_template.yaml > ${specfile}
if [[ -z "${WORKSPACE}" ]]; then
  kube_select_cluster
  WORKSPACE=${retval_kube_select_cluster}
fi
echo "Using cluster : ${WORKSPACE}"
context=${K8S_CONTEXT_PATTERN/CLUSTERNAME/${WORKSPACE}}
kubectl --context=${context} apply -f ${specfile} || (echo "Failed to create cluster-test pod"; exit 1)
kube_wait_pod ${pod_name} ${context}
START_UTC=$(TZ=UTC date +"%Y-%m-%dT%H:%M:%SZ")
START_TS_MS=$(date +%s)000
echo "**********"
echo "${BLUE}Auto refresh Dashboard:${RESTORE} http://grafana.${WORKSPACE}-k8s-testnet.aws.hlw3truzy4ls.com/d/2XqUIhnWz/performance?from=${START_TS_MS}&to=now&refresh=5s"
echo "${BLUE}Tail logs:${RESTORE} http://kibana.${WORKSPACE}-k8s-testnet.aws.hlw3truzy4ls.com/app/kibana#/discover?_g=(refreshInterval:(pause:!f,value:10000),time:(from:'$START_UTC',to:now))"
echo "**********"
kubectl --context=${context} logs -f "${pod_name}" | tee $OUTPUT_TEE
pod_status=$(kubectl --context=${context} get pods "${pod_name}" -o jsonpath="{.status.phase}")
END_UTC=$(TZ=UTC date +"%Y-%m-%dT%H:%M:%SZ")
END_TS_MS=$(date +%s)000
echo "**********"
echo "${BLUE}Logs snapshot:${RESTORE} http://kibana.${WORKSPACE}-k8s-testnet.aws.hlw3truzy4ls.com/app/kibana#/discover?_g=(time:(from:'$START_UTC',to:'$END_UTC'))"
echo "${BLUE}Dashboard snapshot:${RESTORE} http://grafana.${WORKSPACE}-k8s-testnet.aws.hlw3truzy4ls.com/d/2XqUIhnWz/performance?from=${START_TS_MS}&to=${END_TS_MS}"
echo "**********"
if [[ "${pod_status}" != "Succeeded" ]]; then
  echo "${pod_name} status: ${pod_status}"
  EXIT_CODE=1
fi

if [ ! -z "$REPORT" ]; then
    cat $OUTPUT_TEE | awk '/====json-report-begin===/{f=1;next} /====json-report-end===/{f=0} f' > $REPORT
fi

exit ${EXIT_CODE}