spinup

#!/bin/bash

## Manage a local cockroach cluster, for testing or writing docs
##
## Supports the following commands:
##
## - start:  start a new cockroach cluster and init
## - stop:   stop all running cockroach instances started by this script
## - status: show all running cockroach instances started by this script
## - clean:  deletes all local directories from cockroach instances started
##              by this script
## - web:    launch the DB Console web UI to the first cockroach instance
##              started by this script
## - sql:    launch an interactive SQL shell to the first cockroach instance
##              started by this script
## - log:    open the logfile of the first cockroach instance started by this
##              script
##
## Supports the following flags:
##
## - n: (start only) Number of nodes to spinup and join together. If 1, uses
##         cockroach start-single-node and skips init (default: 3)
## - v: (all except clean) Specific version of cockroach to use (see NOTE ON
##         VERSION below)
## - i: (start only) Import movr database
##
## NOTE ON VERSION:
## This script uses your environment's `cockroach` command, whether alias,
## symlink, or through $PATH. Be sure you have this set to the version of
## CRDB that you expect. The following is recommended (with ~/crl/ in $PATH):
##
## $ ls -ld ~/crl/cockroach*
##   ~/crl/cockroach -> cockroach-22.1
##   ~/crl/cockroach-21.1 -> cockroach-v21.1.19.darwin-10.9-amd64/cockroach
##   ~/crl/cockroach-21.2 -> cockroach-v21.2.12.darwin-10.9-amd64/cockroach
##   ~/crl/cockroach-22.1 -> cockroach-v22.1.2.darwin-10.9-amd64/cockroach
##   ~/crl/cockroach-v21.1.19.darwin-10.9-amd64
##   ~/crl/cockroach-v21.2.12.darwin-10.9-amd64
##   ~/crl/cockroach-v22.1.2.darwin-10.9-amd64
##
## With a setup like the above, we can then freely run the following easily:
##    spinup start           :: starts up with default `cockroach`, 22.1 currently
##    spinup -v 22.1 start   :: same as above, currently
##    spinup -v 21.2 start   :: starts up with `which cockroach-21.2`
##
## Similarly, we can fine-tune the behavior of the other commands as well:
##    spinup stat            :: shows all cockroach procs started by us
##    spinup -v 22.1 stat    :: shows only 22.1 cockroach procs started by us
##    spinup -v 21.2 stat    :: shows only 21.2 cockroach procs started by us
##
##    spinup stop            :: stops all cockroach procs started by us
##    spinup -v 22.1 stop    :: stops only 22.1 cockroach procs started by us
##    spinup -v 21.2 stop    :: stops only 21.2 cockroach procs started by us
##
##    spinup web            :: opens DB Console of first cockroach proc started by us
##    spinup -v 22.1 web    :: opens DB Console of first cockroach 22.1 proc started by us
##    spinup -v 21.2 web    :: opens DB Console of first cockroach 21.2 proc started by us
##
##    spinup sql            :: opens SQL shell to first cockroach proc started by us
##    spinup -v 22.1 sql    :: opens SQL shell to first cockroach 22.1 proc started by us
##    spinup -v 21.2 sql    :: opens SQL shell to first cockroach 21.2 proc started by us
##
##    spinup log            :: opens logfile of first cockroach proc started by us
##    spinup -v 22.1 log    :: opens logfile of first cockroach 22.1 proc started by us
##    spinup -v 21.2 log    :: opens logfile of first cockroach 21.2 proc started by us

## TODO: 
##   - Add --pid-file support per Raph.
##   - Improve resilience working with staging dir, $ID clobbering, etc.
##         Likely with --pid-file.
##   - Improve version handling for WEB, SQL, LOG, where no handling is
##         present for requesting a version that is not running
##   - Improve log management using --log or --log-config-file
##   - Add -b break flag to randomly break an active cluster, based on
##         firedrill docs
##   - Add cert support, built around the following:
##     mkdir -p /Users/andf/certs
##     cockroach-22.2 cert create-ca --certs-dir /Users/andf/certs/ --ca-key="/Users/andf/certs/crl-ca-key"
##     cockroach-22.2 cert create-node localhost --certs-dir /Users/andf/certs/ --ca-key="/Users/andf/certs/crl-ca-key"
##     cockroach-22.2 cert create-client andf --certs-dir /Users/andf/certs/ --ca-key="/Users/andf/certs/crl-ca-key"

############################################################
## YOUR VALUES ##
#################

## The staging directory: where to store the following:
##    DATA: --store=$STAGING/$ID/data/node$NODE
##    LOG:  >> $STAGING/$ID/log/cockroach_node1.log
## If you use a ~ in the path, don't quote the value:
STAGING=~/staging

############################################################
## BEGIN SCRIPT ##
##################

NODE_COUNT=3
VERSION=""
IMPORT_MOVR=0
SINGLE_NODE=0
PORT_OFFSET=0
JOIN_HOSTS=""

## Get this script's current running PID. This is how we'll uniqueify
## our local folder structure later:
ID=`echo $$`

## Check for arguments passed to this script:
while getopts ":n:v:i" opt; do
  case ${opt} in
    n )
      NODE_COUNT="$OPTARG"
      ;;
    v)
      VERSION=`echo $OPTARG | sed 's%^v%%g'`
      ;;
    i)
      IMPORT_MOVR=1
      ;;
    \? )
      echo -e "\nERROR: Invalid option: $OPTARG. The following are the supported flags:" 1>&2
      echo "        -n: (start only) Number of nodes to spinup and join together. If 1, uses"
      echo "            cockroach start-single-node and skips init (default: 3)"
      echo "        -v: (all except clean) Specific version of cockroach to use"
      echo -e "        -i: (start only) Import movr database\n"
      exit;
      ;;
    : )
      echo -e "\nERROR: If you specify -$OPTARG, you must provide a value." 1>&2
      ;;
  esac
done

shift $((OPTIND-1))

## Sanity check on node count:
if [ $NODE_COUNT -lt 1 ]; then
   echo "ERROR: Impossible to spin up a 0-node cluster, and you know it!"
   exit;
fi

## Sanity check on version selection:
if [[ ! -z $VERSION && `echo $VERSION | egrep -c '^[0-9]{1,2}\.[0-9]$'` -lt 1 ]]; then
   echo "ERROR: Malformed version specified to -v"
   echo "       It must be something like \`21.2\` or \`21.1\`."
   echo "       Leave off the -v flag to use your environment's /`cockroach/`."
   exit;
elif [[ -z $VERSION && ! `which cockroach` ]]; then
   echo "ERROR: No matching cockroach binary found in \$PATH."
   echo "       This script expects either \`cockroach\` or \`cockroach-21.2\`"
   echo "       (or similar other version) to be findable in your \$PATH."
   exit;
elif [[ ! -z $VERSION && ! `which cockroach-$VERSION` ]]; then
   echo "ERROR: No matching cockroach binary found in \$PATH for requested version $VERSION."
   echo "       This script expects either \`cockroach\` or \`cockroach-21.2\`"
   echo "       (or similar other version) to be findable in your \$PATH."
   exit;
elif [[ ! -z $VERSION && $VERSION != `cockroach --version | grep 'Build Tag:' | egrep -o '[0-9]{1,2}\.[0-9]'` ]]; then
   COCKROACH="cockroach-$VERSION"
else
   COCKROACH="cockroach"
fi

## Collect search string from provided parameter:
PARAMETER=$1

# Parse user-provided parameter and determine what to do:
if [ -z $PARAMETER ]; then
   echo -e "\nERROR: No command specified. You must specify one of:"
   echo "            start  : start a new cockroach cluster and init"
   echo "            stop   : stop all running cockroach instances started by this script"
   echo "            status : show all running cockroach instances started by this script"
   echo "            clean  : deletes all local directories from cockroach instances started"
   echo "                        by this script"
   echo "            web    : launch the DB Console web UI to the first cockroach instance"
   echo "                        started by this script"
   echo "            sql    : launch an interactive SQL shell to the first cockroach instance"
   echo "                        started by this script"
   echo "            log    : open the logfile of the first cockroach instance started by this"
   echo -e "                        script\n"
   exit;
elif [ ! -z $2 ]; then
   echo -e "\nERROR: Too many commands provided. You must provide only one of:"
   echo "            start  : start a new cockroach cluster and init"
   echo "            stop   : stop all running cockroach instances started by this script"
   echo "            status : show all running cockroach instances started by this script"
   echo "            clean  : deletes all local directories from cockroach instances started"
   echo "                        by this script"
   echo "            web    : launch the DB Console web UI to the first cockroach instance"
   echo "                        started by this script"
   echo "            sql    : launch an interactive SQL shell to the first cockroach instance"
   echo "                        started by this script"
   echo "            log    : open the logfile of the first cockroach instance started by this"
   echo -e "                        script\n"
   exit;
elif [[ "$PARAMETER" =~ ^(start|START|Start)$ ]]; then
   START=1
   echo "Starting new $NODE_COUNT node cluster ..."
elif [[ "$PARAMETER" =~ ^(stop|STOP|Stop)$ ]]; then
   STOP=1
   echo "Stopping all $COCKROACH processes started earlier by this script ..."
elif [[ "$PARAMETER" =~ ^(status|STATUS|Status|stat|STAT|Stat|show|SHOW|Show)$ ]]; then
   STATUS=1
   echo "Getting status ..."
elif [[ "$PARAMETER" =~ ^(clean|CLEAN|Clean)$ ]]; then
   CLEAN=1
   echo "Cleaning up any previous staging dirs left over from previous runs of this script ..."
elif [[ "$PARAMETER" =~ ^(web|WEB|Web|http|HTTP|Http)$ ]]; then
   OPEN_WEB=1
elif [[ "$PARAMETER" =~ ^(sql|SQL|Sql)$ ]]; then
   OPEN_SQL=1
elif [[ "$PARAMETER" =~ ^(log|LOG|Log)$ ]]; then
   OPEN_LOG=1
else
   echo -e "\nERROR: You specified an unsupported command. Valid commands are:"
   echo "            start  : start a new cockroach cluster and init"
   echo "            stop   : stop all running cockroach instances started by this script"
   echo "            status : show all running cockroach instances started by this script"
   echo "            clean  : deletes all local directories from cockroach instances started"
   echo "                        by this script"
   echo "            web    : launch the DB Console web UI to the first cockroach instance"
   echo "                        started by this script"
   echo "            sql    : launch an interactive SQL shell to the first cockroach instance"
   echo "                        started by this script"
   echo "            log    : open the logfile of the first cockroach instance started by this"
   echo -e "                        script\n"
   exit;
fi

## Good ol' Bash field-separator
IFS="
"

## Build local folder structure in $STAGING. Only run if we have been asked
## to $START:
if [ $START ]; then
   mkdir -p $STAGING
   mkdir -m 750 -p $STAGING/$ID/data
   mkdir -m 750 -p $STAGING/$ID/log
fi

## Function to build an egrep-compatible list of ports we intend to use
## and then check to see if they are already in use:
check_for_conflict()
{
   CHECK_PORTS=""
   for NODE in `seq 1 $NODE_COUNT`; do
      if [ -z $CHECK_PORTS ]; then
         CHECK_PORTS=".$((26256+NODE+PORT_OFFSET))"
         CHECK_PORTS="${CHECK_PORTS}|.$((8079+NODE+PORT_OFFSET))"
      else
         CHECK_PORTS="${CHECK_PORTS}|.$((26256+NODE+PORT_OFFSET))"
         CHECK_PORTS="${CHECK_PORTS}|.$((8079+NODE+PORT_OFFSET))"
      fi
   done

   PORT_CONFLICT=`netstat -an | egrep '(ESTABLISHED|LISTEN)' | egrep -c "($CHECK_PORTS)"`
   if [[ $PORT_CONFLICT -gt 0 ]]; then
      return 0
   else
      return 1
   fi
}

## Uses above function to check for potential port conflicts, continually
## bumping our intended ports by 100 each time we find a conflict: for both
## SQL port and HTTP port. Only run if we have been asked to $START:
if [ $START ]; then
   while check_for_conflict; do
      PORT_OFFSET="$((PORT_OFFSET+100))"
   done
fi

## Notify user if $STAGING is over 50GB, recommending a spinup clean at some point:
USED_RAW=`du -sk $STAGING/* 2>/dev/null | cut -f 1 | grep -v ^0$ | paste -sd+ -| bc`
if [[ $CLEAN -ne 1 && $USED_RAW -gt 52428800 ]]; then
   USED=`echo "scale=2;$USED_RAW/1024/1024" | bc`
   echo -e "\nWARNING: Staging dir ($STAGING) is over 50 GB (currently: $USED GB)!"
   echo -e "         Run \`spinup clean\` at some point to clear it out!\n"
fi

## Spinup a cluster, if -n 1 was passed (using start-single-node);
## note that start-single-node does not require --join or running init:
if [[ $START && $NODE_COUNT -eq 1 ]]; then
   $COCKROACH start-single-node \
   --insecure \
   --store=$STAGING/$ID/data/node1 \
   --listen-addr=localhost:$((26257+PORT_OFFSET)) \
   --http-addr=localhost:$((8080+PORT_OFFSET)) \
   --background >> $STAGING/$ID/log/cockroach_node1.log 2>&1
   HTTP_HOST=localhost:$((8080+PORT_OFFSET))
   HOST=localhost:$((26257+PORT_OFFSET))

   ## If -p was passed, populate the cluster with sample data from movr:
   if [ $IMPORT_MOVR -eq 1 ]; then
      $COCKROACH workload init movr "postgresql://root@$HOST?sslmode=disable"
   fi

   ## Report our success!
   echo -e "\n##############################  CLUSTER READY!  ##############################\n"
   echo -e "   Single node cluster is live!\n"

   ps aux | grep "$COCKROACH start" | grep $STAGING | grep $ID | awk '{for (i=1;i<=10;i++){$i=""};print}'

   echo -e "\n   SQL Shell: $COCKROACH sql --insecure --host=$HOST"
   echo "   DB Console: http://$HTTP_HOST"
   echo -e "\n##############################################################################\n"

## Spinup a cluster, with $NODE_COUNT nodes with -n, or 3 nodes otherwise (default):
elif [ $START ]; then
   ## Assemble JOIN_HOSTS:
   for NODE in `seq 1 $NODE_COUNT`; do
      if [ -z $JOIN_HOSTS ]; then
         JOIN_HOSTS=localhost:$((26256+NODE+PORT_OFFSET))
      else
         JOIN_HOSTS=${JOIN_HOSTS},localhost:$((26256+NODE+PORT_OFFSET))
      fi
   done

   for NODE in `seq 1 $NODE_COUNT`; do
      $COCKROACH start \
      --insecure \
      --store=$STAGING/$ID/data/node$NODE \
      --listen-addr=localhost:$((26256+NODE+PORT_OFFSET)) \
      --http-addr=localhost:$((8079+NODE+PORT_OFFSET)) \
      --join=$JOIN_HOSTS \
      --background >> $STAGING/$ID/log/cockroach_node$NODE.log 2>&1
      if [[ $NODE -eq 1 ]]; then
         HOST=localhost:$((26256+NODE+PORT_OFFSET))
         HTTP_HOST=localhost:$((8079+NODE+PORT_OFFSET))
      fi
   done

   sleep 1

   $COCKROACH init --insecure --host=$HOST 2>&1 > /dev/null

   sleep 1

   ## if the -i flag was passed, import sample data from movr into the cluster:
   if [ $IMPORT_MOVR -eq 1 ]; then
      $COCKROACH workload init movr "postgresql://root@$HOST?sslmode=disable"
   fi

   ## Report our success!
   echo -e "\n##############################  CLUSTER READY!  ##############################\n"
   echo -e "   $NODE_COUNT node cluster is live!\n"

   ps aux | grep "$COCKROACH start" | grep $STAGING | grep $ID | awk '{for (i=1;i<=10;i++){$i=""};print}'

   echo -e "\n   SQL Shell: $COCKROACH sql --insecure --host=$HOST"
   echo "   DB Console: http://$HTTP_HOST"
   echo -e "\n##############################################################################\n"

elif [ $STOP ]; then
   if [[ -z $VERSION  && `ps aux | grep "cockroach.* start" | grep $STAGING | wc -l` -gt 0 ]]; then
      for PID in `ps aux | grep "cockroach.* start" | grep $STAGING | awk '{print $2}'`; do
         # This is what roachprod does, soooo ....
         kill -9 $PID
      done
   elif [[ `ps aux | grep "$COCKROACH start" | grep $STAGING | wc -l` -gt 0 ]]; then
      for PID in `ps aux | grep "$COCKROACH start" | grep $STAGING | awk '{print $2}'`; do
         # This is what roachprod does, soooo ....
         kill -9 $PID
      done
   elif [[ `ps aux | grep "cockroach start" | grep $STAGING | wc -l` -gt 0 ]]; then
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script."
      echo -e "       But there are other versions of cockroach running on this system presently."
      echo -e "       Run \`spinup stat\` to see them.\n"
   else
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script.\n"
   fi
elif [ $STATUS ]; then
   if [[ -z $VERSION  && `ps aux | grep "cockroach.* start" | grep $STAGING | wc -l` -gt 0 ]]; then
      echo -e "\nThe following $COCKROACH processes are running that were started by this script:\n"
      ps aux | grep "cockroach.* start" | grep $STAGING
      echo ""
   elif [[ `ps aux | grep "$COCKROACH start" | grep $STAGING | wc -l` -gt 0 ]]; then
      echo -e "\nThe following $COCKROACH processes are running that were started by this script:\n"
      ps aux | grep "$COCKROACH start" | grep $STAGING
      echo ""
   elif [[ `ps aux | grep "cockroach start" | grep $STAGING | wc -l` -gt 0 ]]; then
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script."
      echo -e "       But there are other versions of cockroach running on this system presently."
      echo -e "       Run \`spinup stat\` to see them.\n"
   else
      echo -e "\nNo $COCKROACH processes are running that were started by this script.\n"
   fi
elif [ $CLEAN ]; then
   USED_RAW=`du -sk $STAGING/* 2>/dev/null | cut -f 1 | grep -v ^0$ | paste -sd+ -| bc`
   if [[ `lsof +D $STAGING | grep -v ' cwd ' | grep -c $STAGING` -gt 0 ]]; then
      echo -e "\nERROR: Files in $STAGING are in use! Please run \`spinup stop\` first"
      echo -e "       before attempting to clean the staging directory.\n"
      exit;
   elif [[ ! -z $USED_RAW ]]; then
      USED=`echo "scale=2;$USED_RAW/1024/1024" | bc`
      while true; do
         echo ""
         read -p "Are you sure you wish to delete $USED GB of data in $STAGING? [y/n] " yn
         case $yn in
            [Yy]* ) rm -rf $STAGING/*; break;;
            [Nn]* ) exit;;
            * ) echo "Please answer yes or no.";;
         esac
      done
      echo -e "\nFreed up $USED GB!\n"
   else
      echo -e "\nNothing to clean up! Staging dir $STAGING uses 0KB\n"
   fi
elif [ $OPEN_WEB ]; then
   if [[ -z $VERSION  && `ps aux | grep "cockroach.* start" | grep $STAGING | wc -l` -gt 0 ]]; then
      HTTP_HOST=`ps aux | grep "cockroach.* start" | grep $STAGING | sort -k 16 | head -1 | awk '{print $16}' | sed 's%--http-addr=%%g'`
      open http://$HTTP_HOST
   elif [[ `ps aux | grep "$COCKROACH start" | grep $STAGING | wc -l` -gt 0 ]]; then
      HTTP_HOST=`ps aux | grep "$COCKROACH start" | grep $STAGING | sort -k 16 | head -1 | awk '{print $16}' | sed 's%--http-addr=%%g'`
      open http://$HTTP_HOST
   elif [[ `ps aux | grep "cockroach start" | grep $STAGING | wc -l` -gt 0 ]]; then
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script."
      echo -e "       But there are other versions of cockroach running on this system presently."
      echo -e "       Run \`spinup stat\` to see them.\n"
   else
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script.\n"
   fi
elif [ $OPEN_SQL ]; then
   if [[ -z $VERSION  && `ps aux | grep "cockroach.* start" | grep $STAGING | wc -l` -gt 0 ]]; then
      SQL_PROC=`ps aux | grep "cockroach.* start" | grep $STAGING | sort -k 15 | head -1`
      SQL_VERSION=`echo $SQL_PROC | awk '{print $11}'`
      SQL_HOST=`echo $SQL_PROC | awk '{print $15}' | sed 's%--listen-addr=%%g'`
      $SQL_VERSION sql --insecure --host=$SQL_HOST
   elif [[ `ps aux | grep "$COCKROACH start" | grep $STAGING | wc -l` -gt 0 ]]; then
      SQL_HOST=`ps aux | grep "$COCKROACH start" | grep $STAGING | sort -k 15 | head -1 | awk '{print $15}' | sed 's%--listen-addr=%%g'`
      $COCKROACH sql --insecure --host=$SQL_HOST
   elif [[ `ps aux | grep "cockroach start" | grep $STAGING | wc -l` -gt 0 ]]; then
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script."
      echo -e "       But there are other versions of cockroach running on this system presently."
      echo -e "       Run \`spinup stat\` to see them.\n"
   else
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script.\n"
   fi
elif [ $OPEN_LOG ]; then
   if [[ `ps aux | grep "cockroach.* start" | grep $STAGING | wc -l` -gt 0 ]]; then
      ## We still want to sort by a port, to be sure we are finding the first started
      ## Ports are sorted by first-started, per this script's design, but $IDs
      ## (i.e. PIDs of _this script_) as reported via column $14 cannot be assumed to
      ## be in any order. So we sort by $15 but return $14:
      LOG_HOST_DATA_PATH=`ps aux | grep "$COCKROACH start" | grep $STAGING | sort -k 15 | head -1 | awk '{print $14}' | sed 's%--store=%%g'`

      ## This will always be 1, but eventually this script will support passing the desired node number
      ## as a param, to select which node's logfile we open. That will require this line here:
      LOG_HOST_NODE_NUM=`echo $LOG_HOST_DATA_PATH | awk -F '/' '{print $(NF)}' | sed 's%node%%g'`
      LOG_HOST_ID=`echo $LOG_HOST_DATA_PATH | sed "s%/data/node$LOG_HOST_NODE_NUM%%g" | awk -F '/' '{print $(NF)}'`
      LOGFILE="$STAGING/$LOG_HOST_ID/log/cockroach_node$LOG_HOST_NODE_NUM.log"

      ## Just a plain cat here, so we can pipe to our preferred text processing commands
      ## like tail or head. TODO: Support replacing cat with tail so we can tail -f:
      cat $LOGFILE
   else
      echo -e "\nERROR: No $COCKROACH processes are running that were started by this script.\n"
   fi
else
   ## It's not actually possible to get here, but:
   echo "ERROR: Unexpected command.\n"
   exit;
fi

exit;