run

#!/bin/bash

err_tag="[error]"
info_tag="[info]"

log_error() {
  >&2 echo "$err_tag $1"
}
log() {
  echo "$info_tag $1"
}

exit_with_error() {
  local msg="$1"
  local exit_code=${2:-1}
  log_error "$msg (exit code $exit_code)"
  exit $2
}

exit_if_error() {
  if [ $? -eq 0 ]; then
    return
  fi
  exit_with_error "${1:-Command failed}" $?
}

check_executables() {
  for exe in "${executables[@]}"; do
    if !which "$exe" &>/dev/null; then
      exit_with_error "Executable is missing: $exe"
    fi
  done
}

check_monitor_runtime_executables() {
  executables=(
    tail inotifywait wc cut
  )
  check_executables
}

check_app_runtime_executables() {
  if [ -e ~/.cargo/env ]; then
    . ~/.cargo/env
  fi
  executables=(
    rustup cargo git bash
  )
  check_executables
}

remote_repo_name="bench-bot"
remote_repo="https://github.com/paritytech/$remote_repo_name"

benchbot_user="benchbot"
benchbot_session="/tmp/bench-bot"
install_location="/home/$benchbot_user/bench-bot"

exec_log_dir_parent="/home/$benchbot_user"
exec_log_dir="$install_location"
exec_log_file_name="log.txt"
exec_log_file="$exec_log_dir/$exec_log_file_name"

monitor_service="benchbot-monitor"
monitor_service_dir="/usr/lib/systemd/system"
monitor_service_file="$monitor_service_dir/$monitor_service.service"

print_help_and_exit() {
echo "
Usage: run [command]

Commands:
  bootstrap:
    Bootstrap the bot to its predefined location ($install_location).

  start, stop, restart:
    Execute the relevant subcommand for the bot's process.

  monitor:
    Use 'monitor install' or 'monitor uninstall' for setting up the bot's
    monitoring service.
    Otherwise, the arguments are forwarded to systemctl.

  update [ref]:
    Pull a ref (branch or pull request) from $remote_repo, install it and
    restart the bot.

    For pull requests:
      update pull/number/head:branch (e.g. pull/1/head:master)

    For branches:
      update branch

  help:
    Print this message and exit
"
  exit $1
}

current_user="${USER:-$(whoami 2>/dev/null)}"
if [ "$current_user" != "$benchbot_user" ]; then
  as_benchbot="sudo -u $benchbot_user"
fi

install_deps() {
  # needed to detect rustup if it's installed
  if [ -e ~/.cargo/env ]; then
    . ~/.cargo/env
  fi

  if [ "${2:-}" == "--force" ] || ! which rustup &>/dev/null; then
    log "Installing Rust"

    curl https://sh.rustup.rs -sSf | sh -s -- -y
    exit_if_error "Failed to install rustup"

    # For ensuring consistency, it's _required_ that the default toolchain is
    # the same that the release team's Substrate benchmark scripts for Substrate
    # uses. It's _good_ to have the toolchain's versions also match.
    rustup default stable

    rustup toolchain install nightly
    exit_if_error "Failed to install nightly toolchain"

    rustup target add wasm32-unknown-unknown --toolchain nightly
    exit_if_error "Failed to add wasm target"
  fi
}

create_bot_user() {
  if id "$benchbot_user" &>/dev/null; then
    return
  fi

  log "Creating $benchbot_user current_user"

  sudo useradd "$benchbot_user"
  exit_if_error "Failed to create current_user $benchbot"

  sudo mkhomedir_helper "$benchbot_user"
  exit_if_error "Failed to create home directory for $benchbot"
}

install_repo() {
  if [ "${2:-}" != "--force" ] && [ -e "$install_location" ]; then
    return
  fi

  mkdir -p "$install_location"
  exit_if_error "Failed to create $install_parent"

  git clone "$remote_repo" "$install_location"
  exit_if_error "Failed to clone $remote_repo to $install_location"

  cd "$install_location" && yarn
  exit_if_error "Failed to install dependencies in $install_location"
}

install_ref() {
  local ref="${1:-}"
  if [ ! "$ref" ]; then
    log_error "Ref needs to be supplied"
    print_help_and_exit 1
  fi

  cd "$install_location"
  exit_if_error "Failed to cd into $install_location"

  local detached_head="$(git rev-parse HEAD)"
  exit_if_error "Failed to get current HEAD sha"

  git checkout "$detached_head" >/dev/null
  exit_if_error "Failed to checkout to current HEAD sha"
  # NOTE: it's NECESSARY to checkout to the detached HEAD since the branch's ref
  # will be deleted in the following step, which might be the current checked-out
  # ref; deleting the currently checked-out ref might put the git tree in a
  # unrecoverable state.

  # Parse pull requests with pull/ID/head:BRANCHNAME as specified by Github
  if [[ "$ref" =~ ^pull/[[:digit:]]+/head:(.*) ]]; then
    local branch="${BASH_REMATCH[1]}"
  else
    local branch="$ref"
    ref="$ref:$ref"
  fi

  local branch_ref="refs/heads/$branch"
  git update-ref -d "$branch_ref"
  exit_if_error "Failed to clean up ref $branch_ref before fetching the branch $branch"
  # NOTE: be sure to check the step above since proceeding without it working
  # might potentially put the git tree in a unrecoverable state.
  # This depends on the git tree being in a detached HEAD state as done by `git
  # checkout "$detached_head"` above.

  git fetch origin "$ref"
  exit_if_error "Failed to fetch $ref from remote"

  while IFS= read -r line; do
    if  [[ "$line" =~ ^[[:space:]]*([^[:space:]]+)[[:space:]]+refs/heads/(.*) ]] &&
        [ "${BASH_REMATCH[2]}" == "$branch" ]; then
      local ref_commit="${BASH_REMATCH[1]}"
      break
    fi
  done < <(git show-ref)

  if [ ! "${ref_commit:-}" ]; then
    exit_with_error "Failed to find commit reference for $ref (branch $branch)"
  fi

  git branch -D "$branch"

  git checkout "$ref_commit" >/dev/null
  exit_if_error "Failed to checkout commit $ref_commit (ref $ref, branch $branch)"

  git switch -c "$branch"
  exit_if_error "Failed to switch from detached head to branch $branch (ref $ref, commit $ref_commit)"

  local head_sha="$(git rev-parse HEAD)"
  exit_if_error "Failed to parse the HEAD commit SHA for $branch (ref $ref, commit $ref_commit)"

  log "Installed branch '$branch' at $head_sha"
}

handle_exec() {
  local cmd="$1"
  shift

  case "$cmd" in
    start)
      if pgrep -u benchbot &>/dev/null; then
        exit_with_error "the $benchbot_user user is already running a process"
      fi

      if [ -e "$exec_log_file" ]; then
        local start_from_line="$(wc -l "$exec_log_file" | cut -d ' ' -f1)"
        exit_if_error "Failed to count the lines in $exec_log_file"
        start_from_line=$(( start_from_line + 1 ))
      else
        echo "" > "$exec_log_file"
        unset start_from_line
      fi

      unset env_vars

      case "${1:-}" in
        debug)
          local env_vars="DEBUG=true"
        ;;
      esac

      sudo ionice -c 1 -n 0 sudo nice -n -19 sudo -u $benchbot_user \
        tmux new-session -d bash -c "
          . ~/.cargo/env &&
          cd \"$install_location\" &&
          git config --local user.name 'Parity Bot' &&
          git config --local user.email admin@parity.io &&
          yarn &&
          ${env_vars:-} yarn start 2>&1 | tee -a \"$exec_log_file\"
        "
      exit_if_error "Failed to create tmux session for user $benchbot_user"

      echo -e "\nNote: the command will still be running after quitting this terminal. Use \"run stop\" for stopping it.\n"

      tail "--lines=+${start_from_line:-0}" -f "$exec_log_file"
    ;;
    stop)
      if pgrep -u benchbot &>/dev/null; then
        sudo pkill -u benchbot
      else
        return 0
      fi
    ;;
    restart)
      handle_exec stop
      handle_exec start "$@"
    ;;
    *)
      exit_with_error "Unknown handle_exec command $cmd"
    ;;
  esac
}

stop_follow_log_file() {
  if [ ! "${follow_log_file_tail_pid:-}" ]; then
    return
  fi

  kill -9 "$follow_log_file_tail_pid"
  exit_if_error "Failed to kill tail process $follow_log_file_tail_pid"

  unset follow_log_file_tail_pid
}

start_follow_log_file() {
  stop_follow_log_file

  local start_from_line="$(wc -l "$exec_log_file" | cut -d ' ' -f1)"
  exit_if_error "Failed to count the lines in $exec_log_file"
  start_from_line=$(( start_from_line + 1 ))
  tail "--lines=+$start_from_line" -f "$exec_log_file" | awk '{ print "bb: " $0 }' &
  follow_log_file_tail_pid=$?
}

parse_log_file_notification_line() {
  if [[ ! "$1" =~ ^([^[:space:]]+)[[:space:]]+(.*)  ]]; then
    exit_with_error "Notification line did not have the expected format"
  fi
}

follow_log_file() {
  while true; do
    # Monitor the log file while it exists
    if [ -e "$exec_log_dir" ]; then
      start_follow_log_file

      while IFS= read line; do
        parse_log_file_notification_line "$line"

        local event="${BASH_REMATCH[1]}"
        case "$event" in
          DELETE_SELF)
            break
          ;;
        esac

        local file="${BASH_REMATCH[2]}"
        if [ "$file" != "$exec_log_file_name" ]; then
          continue
        fi

        case "$event" in
          CREATE)
            start_follow_log_file
          ;;
          DELETE)
            stop_follow_log_file
          ;;
          *)
            exit_with_error "Unhandled event $event for $exec_log_dir"
          ;;
        esac
      done < <(inotifywait -e create,delete,delete_self --format '%e %f' --monitor --quiet "$exec_log_dir")
    # If the log file does not exist, then wait for the log file's directory to
    # be created
    elif [ -e "$exec_log_dir_parent" ]; then
      while IFS= read line; do
        parse_log_file_notification_line "$line"

        local event="${BASH_REMATCH[1]}"
        case "$event" in
          DELETE_SELF)
            break
          ;;
          CREATE)
            if [ "$exec_log_dir_parent/$file" = "$exec_log_dir" ]; then
              break
            fi
          ;;
          *)
            exit_with_error "Unhandled event $event for $exec_log_dir_parent"
          ;;
        esac
      done < <(inotifywait -e create,delete_self --format '%e %f' --monitor --quiet "$exec_log_dir_parent")
    else
      exit_with_error "Unable to watch '$exec_log_dir_parent' for '$exec_log_dir'"
    fi
  done
}

handle_monitor() {
  local cmd="$1"
  shift

  case "$cmd" in
    install)
      if [ "${1:-}" != "--force" ] && [ -e "$monitor_service_file" ]; then
        return
      fi

      &>/dev/null sudo mkdir -p "$monitor_service_dir"

      echo "
      [Unit]
      Description=Bench Bot Monitor
      Documentation=$remote_repo

      [Service]
      ExecStart=sh -c \"'$install_location/run' follow_log_file\"
      Restart=always
      RestartSec=30
      CapabilityBoundingSet=
      LockPersonality=true
      NoNewPrivileges=true
      PrivateDevices=true
      PrivateMounts=true
      PrivateTmp=true
      PrivateUsers=true
      ProtectControlGroups=true
      ProtectHostname=true
      ProtectKernelModules=true
      ProtectKernelTunables=true
      ProtectSystem=strict
      RemoveIPC=true
      RestrictNamespaces=true
      RestrictSUIDSGID=true
      SystemCallArchitectures=native

      [Install]
      WantedBy=default.target
      " | sudo tee "$monitor_service_file" >/dev/null

      exit_if_error "Failed to create service file at $monitor_service_file"

      if [ -e "$exec_log_file" ] || [ -e "$exec_log_dir_parent" ]; then
        sudo systemctl enable --now "$monitor_service"
      else
        log "The service \"$monitor_service\" was not activated (is the bot cloned at \"$install_location\"?). Start it later with with \"run monitor enable --now\"."
      fi
    ;;
    uninstall)
      if systemctl is-active --quiet "$monitor_service"; then
        sudo systemctl disable --now "$monitor_service"
        exit_if_error "Failed to disable service $monitor_service"
      fi
      sudo rm "$monitor_service_file"
    ;;
    *)
      sudo systemctl "$cmd" "$monitor_service" "$@"
    ;;
  esac
}

main() {
  local cmd="$1"
  shift

  # Initial checks before running the actual commands
  case "$cmd" in
    start|stop|restart|update)
      $as_benchbot bash -c "'${BASH_SOURCE[0]}' check_app_runtime_executables"
      exit_if_error
    ;;
    monitor)
      check_monitor_runtime_executables
    ;;
  esac

  case "$cmd" in
    start|stop|restart)
      handle_exec "$cmd" "$@"
      local exit_code=$?
      echo "Exit code: $exit_code"
      exit $exit_code
    ;;
    update)
      local ref="${1:-}"
      if [ ! "$ref" ]; then
        log_error "Ref needs to be supplied"
        print_help_and_exit 1
      fi
      shift

      handle_exec stop

      $as_benchbot bash -c "'${BASH_SOURCE[0]}' install_ref '$ref'"
      exit_if_error "Failed to install ref '$ref'"

      bash -c "'${BASH_SOURCE[0]}' start $@"
      exit_if_error "Failed to start"
    ;;
    monitor)
      "handle_$cmd" "$@"
      local exit_code=$?
      echo "Exit code: $exit_code"
      exit $exit_code
    ;;
    follow_log_file            | \
    install_repo               | \
    install_ref                | \
    install_deps               | \
    check_app_runtime_executables)
      "$cmd" "$@"
    ;;
    bootstrap)
      create_bot_user

      $as_benchbot bash -c "'${BASH_SOURCE[0]}' install_deps"
      exit_if_error "Failed to install dependencies"

      $as_benchbot bash -c "'${BASH_SOURCE[0]}' install_repo"
      exit_if_error "Failed to install repository"

      handle_monitor install
    ;;
    help)
      print_help_and_exit 0
    ;;
    *)
      log_error "Invalid command $cmd"
      print_help_and_exit 1
    ;;
  esac
}

main "$@"