Script for dynamically adding dockerized Selenium nodes to a running SeLion-Grid enhanced Selenium hub

elastic-node.sh

#!/bin/bash
# ---------------------------------------------------------------------------------------------------------------------\
#   Copyright (C) 2016 Doug Simmons                                                                                    |
#                                                                                                                      |
#   Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance     |
#   with the License.                                                                                                  |
#                                                                                                                      |
#   You may obtain a copy of the License at                                                                            |
#                                                                                                                      |
#        http://www.apache.org/licenses/LICENSE-2.0                                                                    |
#                                                                                                                      |
#   Unless required by applicable law or agreed to in writing, software distributed under the License is distributed   |
#   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for  |
#   the specific language governing permissions and limitations under the License.                                     |
# ---------------------------------------------------------------------------------------------------------------------/


# Requirements:
# DOCKER_MACHINE defined or available at default
# SELION_HUB_HOST defined or available at 127.0.0.1
# SELION_HUB_PORT defined or available at 4444
# MAX_CONTAINERS (optional)
# LOOP_SLEEP_INTERVAL (optional)
# QUEUE_WAITING_THRESHOLD (optional)
# CONTAINER_UPPER_PORT_RANGE (optional)
# CONTAINER_LOWER_PORT_RANGE (optional)
# SELION_CHROME_CONTAINER_ID (optional)
# SELION_FIREFOX_CONTAINER_ID (optional)
# docker command available on PATH
# curl available on PATH

BLUE='\033[1;34m'
GREEN='\033[0;32m'
CYAN='\033[0;36m'
YELLOW='\033[0;33m'
NC='\033[0m'

DOCKER_MACHINE=${DOCKER_MACHINE-default}
trap '[ "$?" -eq 0 ] || shutdown' EXIT SIGTERM SIGINT

function check_deps () {
  STEP="checking PATH for dependencies"
  docker-machine -h &> /dev/null
  if [[ $? -ne 0 ]]; then
    echo -e "${GREEN}docker-machine${NC} is not in the PATH.\n"
    exit 1
  fi

  docker -h &> /dev/null
  if [[ $? -ne 0 ]]; then
    echo -e "${GREEN}docker${NC} is not in the PATH.\n"
    exit 1
  fi

  curl -h &> /dev/null
  if [[ $? -ne 0 ]]; then
    echo -e "${GREEN}curl${NC} is not in the PATH.\n"
    exit 1
  fi

  STEP="checking for presence of docker machine"
  DOCKER_MACHINE_IP=`docker-machine ip ${DOCKER_MACHINE}`
  if [[ $? -ne 0 ]]; then
    exit 1
  fi
}

function check_hub_status () {
  STEP="Selenium hub status check"
  curl -s http://$SELION_HUB_HOST:$SELION_HUB_PORT/wd/hub/status &> /dev/null
  if [[ $? -ne 0 ]]; then
    echo -e "The hub ${GREEN}${SELION_HUB_HOST}:${SELION_HUB_PORT}${NC} can not be reached"
    exit 1
  fi
}

function get_env () {
  SELION_HUB_HOST=${SELION_HUB_HOST-127.0.0.1}
  SELION_HUB_PORT=${SELION_HUB_PORT-4444}
  SELION_CHROME_CONTAINER_ID=${SELION_CHROME_CONTAINER_ID-selion/node-chrome:latest}
  SELION_FIREFOX_CONTAINER_ID=${SELION_FIREFOX_CONTAINER_ID-selion/node-firefox:latest}
  MAX_CONTAINERS=${MAX_CONTAINERS-16}
  CONTAINER_LOWER_PORT_RANGE=${CONTAINER_LOWER_PORT_RANGE-5000}
  CONTAINER_UPPER_PORT_RANGE=${CONTAINER_UPPER_PORT_RANGE-6000}
  LOOP_SLEEP_INTERVAL=${LOOP_SLEEP_INTERVAL-60}
  QUEUE_WAITING_THRESHOLD=${QUEUE_WAITING_THRESHOLD-10}

  echo -e "${CYAN}Using Docker Machine:${NC} ${DOCKER_MACHINE}@${DOCKER_MACHINE_IP}"
  echo -e "${CYAN}Target Selenium Hub:${NC} ${SELION_HUB_HOST}:${SELION_HUB_PORT}"
  echo -e "${CYAN}Using Containers:${NC} ${SELION_CHROME_CONTAINER_ID}, ${SELION_FIREFOX_CONTAINER_ID}"
  echo -e "${CYAN}Max Containers Allowed:${NC} ${MAX_CONTAINERS}"
  echo -e "${CYAN}Using Port Range:${NC} ${CONTAINER_LOWER_PORT_RANGE}-${CONTAINER_UPPER_PORT_RANGE}"
  echo -e "${CYAN}Main Loop Sleep Interval:${NC} ${LOOP_SLEEP_INTERVAL}"
  echo -e "${CYAN}Queue Waiting Threshold:${NC} ${QUEUE_WAITING_THRESHOLD}\n"
}

function get_running_count () {
  RUNNING=`docker ps -a | grep 'selion/node' | grep 'Up' | awk '{print $1}' | wc -l  | sed -e 's/^[ \t]*//'`
}

function get_free_port () {
  PORT=$((PORT+1))
  # if we reach the upper limit start looking from lower port range again
  if [[ $PORT -eq $CONTAINER_UPPER_PORT_RANGE ]]; then
    PORT=$CONTAINER_LOWER_PORT_RANGE
  fi

  local containers=`docker ps -a | grep 'Up' | awk '{print $1}'`
  for container in $containers
  do
    docker port $container | sed -e 's/\/tcp.*//g' >> /tmp/elastic-nodes.$DOCKER_MACHINE.ports.tmp
  done

  for (( $PORT; $PORT<=$CONTAINER_UPPER_PORT_RANGE; PORT=$((PORT+1)) ))
  do
    [ -f /tmp/elastic-nodes.$DOCKER_MACHINE.ports.tmp ] && \
    local free=`grep $PORT /tmp/elastic-nodes.$DOCKER_MACHINE.ports.tmp | wc -l | sed -e 's/^[ \t]*//'` || \
    local free=0
    if [[ $free == 0 ]]; then
      rm /tmp/elastic-nodes.$DOCKER_MACHINE.ports.tmp &> /dev/null
      return
    fi
  done

  rm /tmp/elastic-nodes.$DOCKER_MACHINE.ports.tmp &> /dev/null
  echo -e "${YELLOW}Exhausted free port range. Quitting.${NC}"
  exit 1
}

function start_node () {
  local containerId=$SELION_FIREFOX_CONTAINER_ID
  if [[ $1 = "chrome" ]]; then
    containerId=$SELION_CHROME_CONTAINER_ID
  fi

  get_running_count
  if [[ $RUNNING -lt $MAX_CONTAINERS ]]; then
    get_free_port
    echo -e "Starting a ${GREEN}${1}${NC} container on port ${BLUE}${PORT}${NC} ..."

    nohup docker run --rm -e HUB_PORT_4444_TCP_ADDR=$SELION_HUB_HOST \
    -e HUB_PORT_4444_TCP_PORT=$SELION_HUB_PORT \
    -e SELION_OPTS="-host $DOCKER_MACHINE_IP -port $PORT" \
    -p $PORT:$PORT --name=node-$PORT $containerId \
    &> /dev/null &

    STARTED=$((STARTED+1))
  fi
}

function get_queue () {
  if [[ $1 = "chrome" ]]; then
    local output=`curl -s http://$SELION_HUB_HOST:$SELION_HUB_PORT/grid/admin/GridStatistics \
    | sed -e 's/.*chrome//g' \
    | sed -e 's/,"maxBrowser.*//g' \
    | sed 's/.*waitingRequests"://g'`
  else
    local output=`curl -s http://$SELION_HUB_HOST:$SELION_HUB_PORT/grid/admin/GridStatistics \
    | sed -e 's/.*firefox//g' \
    | sed -e 's/,"maxBrowser.*//g' \
    | sed 's/.*waitingRequests"://g'`
  fi

  CREATENODE=false
  if [[ $output = "[]" ]]; then
    CREATENODE=true
    return
  fi

  echo -e "$output requests waiting for ${GREEN}$1${NC}"
  if [[ $output -gt $QUEUE_WAITING_THRESHOLD ]]; then
    CREATENODE=true
  fi
}

function cleanup_exited_nodes () {
  local containers=`docker ps -a | grep 'selion/node' | grep 'Exited' | awk '{print $1}'`
  for container in $containers
  do
    echo -e "Cleaning up exited container ${BLUE}${container}${NC} ..."
    docker rm $container &> /dev/null
  done
}

function elastic-node () {
  while [ true ]
  do
    for flavor in chrome firefox
    do
      get_queue $flavor
      if [[ "$CREATENODE" = true ]]; then
        start_node $flavor
      fi
    done

    get_running_count
    echo -e "\n${CYAN}Total Containers Running:${NC} ${RUNNING}"
    echo -e "${CYAN}Max Containers Allowed:${NC} ${MAX_CONTAINERS}"
    echo -e "${CYAN}Total Containers Started:${NC} ${STARTED}\n"

    sleep $LOOP_SLEEP_INTERVAL
    cleanup_exited_nodes
  done
}

function shutdown () {
  if [[ $STEP != "main loop" ]]; then
    echo -e "\nProcess aborted at ${STEP}."
  else
    echo -e "\n${YELLOW}Shutting down${NC} ..."
    rm /tmp/elastic-nodes.$DOCKER_MACHINE.ports.tmp &> /dev/null
    cleanup_exited_nodes
    echo -e "Shutdown of elastic-node complete."
    exit 0
  fi
}

check_deps
get_env
check_hub_status
STEP="main loop"
PORT=$((CONTAINER_LOWER_PORT_RANGE-1))
CREATENODE=false
STARTED=0
eval $(docker-machine env $DOCKER_MACHINE)
cleanup_exited_nodes
elastic-node 
wait $!