azworker.sh

#!/usr/bin/env bash

# Copyright 2017 The Cockroach Authors.
#
# Use of this software is governed by the CockroachDB Software License
# included in the /LICENSE file.

#
# Prerequisites:
# - Cockroach Labs employees: ask an admin to create an Azure account for you.
# - Install the Azure XPlat CLI 2.0:
# https://docs.microsoft.com/en-us/cli/azure/install-azure-cli?view=azure-cli-latest
# - Run "az login" and sign in to your Azure account.

set -euo pipefail

cd"$(dirname "${0}")/.."&&source build/shlib.sh

LOCATION=${LOCATION-eastus}
MACHINE_SIZE=${MACHINE_SIZE-Standard_F16}
USER=${USER-$(id -un)}
CLUSTER=azworker-${USER}
NAME=${AZWORKER_NAME-${CLUSTER}}

# Names for various resources just reuse cluster/vm name depending on scope.
RG=${CLUSTER}
NET=${CLUSTER}
SUBNET=${CLUSTER}
NIC=${NAME}
IP=${NAME}
DOMAIN=cockroach-${NAME}
FQDN=${DOMAIN}.${LOCATION}.cloudapp.azure.com

# TODO(radu): switch to the newer "az" CLI.

case${1-}in
 create)
if [ $(az group show -n "${RG}"| wc -l)-gt 0 ];then
echo"Resource group ${RG} already exists; adding worker VM to it"
else
echo"Creating resource group ${RG}"
 az group create -n "${RG}" -l "${LOCATION}"
 az network vnet create -g "${RG}" -n "${NET}" --address-prefixes 192.168.0.0/16 -l "${LOCATION}"
 az network vnet subnet create -g "${RG}" --vnet-name "${NET}" -n "${SUBNET}" --address-prefix 192.168.1.0/24
fi

 az network public-ip create -g "${RG}" -n "${IP}" -l "${LOCATION}" --dns-name "${DOMAIN}"
 az network nic create -g "${RG}" -n "${NIC}" -l "${LOCATION}" --public-ip-address "${IP}" --subnet "${SUBNET}" --vnet-name "${NET}"

 az vm create \
 --resource-group "${RG}" \
 --name "${NAME}" \
 --location "${LOCATION}" \
 --image canonical:UbuntuServer:16.04-LTS:latest \
 --ssh-key-value ~/.ssh/id_rsa.pub \
 --admin-username "${USER}" \
 --size "${MACHINE_SIZE}" \
 --nics "${NIC}"

# Clear any cached host keys for this hostname and accept the new one.
 ssh-keygen -R "${FQDN}"
# Retry while vm and sshd to start up.
 retry ssh -o StrictHostKeyChecking=no "${USER}@${FQDN}"true

 rsync -az "build/bootstrap/""${USER}@${FQDN}:bootstrap/"
 rsync -az "build/disable-hyperv-timesync.sh""${USER}@${FQDN}:bootstrap/"
 ssh -A "${USER}@${FQDN}" ./bootstrap/bootstrap-debian.sh
 ssh -A "${USER}@${FQDN}" ./bootstrap/disable-hyperv-timesync.sh
 ssh -A "${USER}@${FQDN}" ./bootstrap/install-azure-cli.sh

# Copy the azure config (including credentials!).
 ssh -A "${USER}@${FQDN}""mkdir .azure"
 rsync -az ~/.azure/*.json "${USER}@${FQDN}:.azure/"

# Set up tmux configuration (for persistent SSH).
if [ -e~/.tmux.conf ];then
 rsync -azL ~/.tmux.conf "${USER}@${FQDN}:./"
else
# Present a color terminal and disable tmux scrollback.
 ssh -A "${USER}@${FQDN}""echo 'set -g default-terminal screen-256color' > .tmux.conf"
 ssh -A "${USER}@${FQDN}""echo 'set -g terminal-overrides \"xterm*:XT:smcup@:rmcup@\"' >> .tmux.conf"
fi

# shellcheck disable=SC2029
 ssh "${USER}@${FQDN}""echo '* * * * * /home/${USER}/bootstrap/autoshutdown.cron.sh 10 az vm deallocate --resource-group \"${RG}\" --name \"${NAME}\"' | crontab -"

echo"VM now running at ${FQDN}"
 ;;
 start)
 az vm start -g "${RG}" -n "${NAME}"
 ;;
 stop)
 az vm deallocate -g "${RG}" -n "${NAME}"
 ;;
 delete|destroy)
# The straightforward thing to do would be to first delete the VM, then
# check if there are any virtual machines left in the group. However, the
# deleted VM doesn't disappear right away from the listing. So we instead
# count the initial number of VMs.
 NUM_VMS=$(az vm list -g "${RG}"| (grep -c "\"type\":.*virtualMachines"|| true))
 az vm delete -g "${RG}" -n "${NAME}"
if [ "$NUM_VMS"-gt 1 ];then
 az network nic delete -g "${RG}" -n "${NIC}"
 az network public-ip delete -g "${RG}" -n "${IP}"
echo"Resource group ${RG} still contains VMs; leaving in place"
else
echo"Deleting resource group ${RG}"
 az group delete -n "${RG}" --yes
fi
 ;;
 ssh)
shift
# shellcheck disable=SC2029
 ssh -A "${USER}@${FQDN}""$@"
 ;;
 sshmux)
shift
if [ -z"${1:-}" ];then
 ssh -A "${USER}@${FQDN}""tmux list-sessions"
else
# shellcheck disable=SC2029
 ssh -A -t "${USER}@${FQDN}""tmux attach -t $1 || tmux new-session -s $1"
fi
 ;;
*)
if [ -n"${1:-}" ];then
echo"$0: unknown command: ${1-}"
fi
 cat <<EOF
Usage:
$0 create
 Creates a new azure worker VM.

$0 start
 Powers on an azure worker VM.

$0 stop
 Powers off an azure worker VM.

$0 delete
 Deletes an azure worker VM.

$0 ssh
 SSH into an azure worker VM.

$0 sshmux <session-name>
 Creates or reconnects to a persistent SSH session with the given name.

$0 sshmux
 List persistent SSH sessions.

For all commands, worker VM name can be customized via AZWORKER_NAME.
EOF
exit 1
 ;;
esac