Implements connection pool for SSH connections.

ssh-pool.sh

#!/bin/sh

#
# ssh-pool.sh
#
# Implements connection pool for SSH connections.
#

##### start script name

if [ -z "${script}" ]
then
  script=ssh-pool.sh
fi

##### end script name

# get script start time
start_dt=$( date "+%s" )

# make temporary folder
tmp=/tmp/tmp.${script}.$$

mkdir -m 777 ${tmp} >> /dev/null 2>&1

# log messages to terminal and file
log_file=/tmp/${script}.log

log() {
  now=$( date "+[%Y-%m-%d %H:%M:%S]" )
  pid=[$$]
  message="${1}"

  echo ${now} ${pid} "${message}" | tee -a ${log_file}
}

passthru() {
  tee -a ${log_file}
}

# logger pipe and pid
logger_pipe=${tmp}/logger
logger_pid_file=${tmp}/logger.pid

# make named pipe (FIFO)
mkfifo ${logger_pipe}

# keep pipe open using sleep
sleep 1d > ${logger_pipe} &

# save sleep pid
echo $! > ${logger_pid_file}

# run logger
while read message
do
  # indent message
  log "    ${message}"
done < ${logger_pipe} &

# cleanup
finished="no"

cleanup() {
  if [ "${finished}" = "no" ]
  then
    normal_exit=${1}

    # kill logger
    xargs -a ${logger_pid_file} kill -KILL >> /dev/null 2>&1

    # remove logger pid and pipe
    rm -f ${logger_pid_file} ${logger_pipe}

    ##### start script cleanup

    teardown_connection_pool

    log "wait for connections to finish ..."

    wait

    cleanup_connection_pool

    ##### end script cleanup

    # remove temporary folder
    rm -rf ${tmp}

    # get script end time
    end_dt=$( date "+%s" )

    seconds=$(( ${end_dt} - ${start_dt} ))

    h=$(( ${seconds} / 3600 ))
    m=$(( ( ${seconds} % 3600 ) / 60 ))
    s=$(( ${seconds} % 60 ))

    if [ ${h} -gt 0 ]
    then
      if [ ${h} -eq 1 ]
      then
        hour="1 hour"
      else
        hour="${h} hours"
      fi
    fi

    if [ ${m} -gt 0 ]
    then
      if [ ${m} -eq 1 ]
      then
        minute="1 minute"
      else
        minute="${m} minutes"
      fi
    fi

    if [ ${s} -gt 0 ]
    then
      if [ ${s} -eq 1 ]
      then
        second="1 second"
      else
        second="${s} seconds"
      fi
    fi

    # use xargs to trim leading, trailing, and internal whitespace
    runtime=$( echo ${hour} ${minute} ${second} | xargs )

    if [ -z "${runtime}" ]
    then
      runtime="0 seconds"
    fi

    # display script runtime
    log "script runtime: ${script} ${runtime}"

    finished="yes"

    # handle abnormal exit
    if [ "${normal_exit}" = "no" ]
    then
      # indicate failure
      exit 1
    fi
  fi
}

# abnormal cleanup

# normal exit = no
trap "cleanup no" INT
trap "cleanup no" TERM

# normal cleanup

# normal exit = yes
trap "cleanup yes" EXIT

# display running script
log "running script: ${script}"

##### start script action

# SSH identity file
ssh_identity_file="THE_SSH_KEY"

# SSH user
ssh_user="THE_SSH_USER"

# SSH hostname
ssh_hostname="THE_SSH_HOSTNAME"

# connections
connections=10

ssh="${ssh_user}@${ssh_hostname}"

setup_connection_pool() {
  log "setup connection pool ..."

  # connections = 1, 2, 3, ..., connections
  all_list=$( seq 1 ${connections} )

  # timeout, 300 seconds = 5 minutes
  timeout=300

  # timeout = 1, 2, 3, ..., timeout
  timeout_list=$( seq 1 ${timeout} )

  # internal connection id
  connection_id=0

  # setup all connections
  set -- ${all_list}
  for index
  do
    config_file=${tmp}/ssh_config.${index}
    master_config_file=${tmp}/ssh_config.master.${index}
    socket_file=${tmp}/socket.${index}
    ready_file=${tmp}/ready.${index}

    # just in case
    rm -f ${config_file} ${master_config_file} ${socket_file} ${ready_file}

    # config file
    cat > ${config_file} <<config_file
User ${ssh_user}
HostName ${ssh_hostname}
IdentityFile ${ssh_identity_file}
BatchMode yes
ControlPath ${socket_file}
LogLevel QUIET
config_file

    # master config file
    cat > ${master_config_file} <<master_config_file
User ${ssh_user}
HostName ${ssh_hostname}
IdentityFile ${ssh_identity_file}
BatchMode yes
ControlPath ${socket_file}
LogLevel QUIET
ControlMaster yes
ControlPersist yes
master_config_file

    # open SSH connection, creates socket file
    ssh -f -F ${master_config_file} ${ssh} "sleep 1d" >> /dev/null 2>&1

    # connection ready
    touch ${ready_file}
  done
}

teardown_connection_pool() {
  log "teardown connection pool ..."

  # teardown all connections
  set -- ${all_list}
  for index
  do
    socket_file=${tmp}/socket.${index}
    config_file=${tmp}/ssh_config.${index}
    ready_file=${tmp}/ready.${index}

    if [ -S ${socket_file} ]
    then
      # "stop" = request master to stop accepting further multiplexing requests
      ssh -F ${config_file} -O "stop" ${ssh} >> /dev/null 2>&1
    fi

    if [ -f ${ready_file} ]
    then
      rm -f ${ready_file}
    fi
  done
}

cleanup_connection_pool() {
  log "cleanup connection pool ..."

  # cleanup all connections
  set -- ${all_list}
  for index
  do
    socket_file=${tmp}/socket.${index}
    config_file=${tmp}/ssh_config.${index}

    if [ -S ${socket_file} ]
    then
      # "exit" = request master to exit
      ssh -F ${config_file} -O "exit" ${ssh} >> /dev/null 2>&1

      # remove socket file
      rm -f ${socket_file}
    fi

    # remove config file
    rm -f ${config_file}
  done
}

get_connection() {
  ready=-1

  # timeout, 300 seconds = 5 minutes
  set -- ${timeout_list}
  for time
  do
    set -- ${all_list}
    for index
    do
      ready_file=${tmp}/ready.${index}

      if [ -f ${ready_file} ]
      then
        rm -f ${ready_file}

        ready=${index}

        break
      fi
    done

    if [ ${ready} -lt 0 ]
    then
      sleep 1
    else
      break
    fi
  done

  connection=${ready}

  # increment internal connection id
  connection_id=$(( connection_id + 1 ))
}

ssh_command() {
  command=${1}
  output=${2}
  background=${3}
  before_message="${4}"
  after_message="${5}"

  # get connection
  get_connection

  if [ -n "${before_message}" ]
  then
    # log before message
    log "${before_message}"
  fi

  config_file=${tmp}/ssh_config.${connection}
  ready_file=${tmp}/ready.${connection}

  if [ "${background}" = "yes" ]
  then
    # run in background
    {
      ssh -F ${config_file} ${ssh_hostname} ${command} 2>> /dev/null > ${output} ;

      if [ -n "${after_message}" ]
      then
        # log after message
        log "${after_message}" ;
      fi

      # touch ready file
      touch ${ready_file} ;
    } &
  else
    # run in foreground
    ssh -F ${config_file} ${ssh_hostname} ${command} 2>> /dev/null > ${output}

    if [ -n "${after_message}" ]
    then
      # log after message
      log "${after_message}"
    fi

    # touch ready file
    touch ${ready_file}
  fi
}

##### end script action