cronlock

#!/usr/bin/env bash
# cronlock - https://github.com/kvz/cronlock
# Version v0.1.0
# Use a central redis server to globally lock cronjobs across a distributed system
#
# Licensed under MIT
# Copyright 2012 by Kevin van Zonneveld (http://kevin.vanzonneveld.net)
# Contains code from https://github.com/reinh/redis.bash
#
# License:
# - web: https://github.com/kvz/cronlock/blob/master/LICENSE.txt
# - cli: curl https://raw.githubusercontent.com/kvz/cronlock/master/LICENSE.txt
#
# Readme:
# - web: https://github.com/kvz/cronlock/blob/master/README.md
# - cli: curl https://raw.githubusercontent.com/kvz/cronlock/master/README.md

function info {
  if [ "${CRONLOCK_VERBOSE}" = "yes" ]; then
    echo "${1}" >&2
  fi
}
function fatal {
  echo "${1}" >&2
  exit 201
}

function get_master {
  exec 3<>/dev/tcp/${CRONLOCK_SENTINEL_HOST}/${CRONLOCK_SENTINEL_PORT} || fatal "Sentinel down"
  { printf "SENTINEL get-master-addr-by-name $CRONLOCK_SENTINEL_MASTER\r\n"; sleep 0.01; } >&3 # sleep to avoid a socket close race condition.

  # Parse the response according to redis protocol
  read -r response <&3
  response="$(echo "${response#\$}"| tr -d '\n\r\0\t')"

  case $response in
    \*-1)
      echo "ERR - Unknown master $CRONLOCK_SENTINEL_MASTER" >&2
      exit 201
      ;;
    \*2) # Array
      read response <&3
      
      nchars="$(echo "${response#\$}"| tr -d '\n\r\0\t')"      
      nchars="$(echo "${nchars%\r}"| tr -d '\n\r\0\t')"
      read -n $nchars CRONLOCK_HOST <&3

      read response <&3 # Read in the newline
      read response <&3
      
      nchars="$(echo "${response#\$}"| tr -d '\n\r\0\t')"
      nchars="$(echo "${nchars%\r}"| tr -d '\n\r\0\t')"
      read -n $nchars CRONLOCK_PORT <&3
      ;;
    *) # net yet handled
      echo "ERR - Unknown response\n" >&2
      exit 201
      ;;
    esac  
}

function connect {
  exec 3<>/dev/tcp/${CRONLOCK_HOST}/${CRONLOCK_PORT}
}

function cmd {
  local  __resultvar=$2
  attempts=0
  until raw_cmd "${@}" || [[ $? != 1 && $? != 141 && $? != 142 ]]; do
    # Connection either closed, errored or timed out
    if [ $attempts -ge ${CRONLOCK_RECONNECT_ATTEMPTS} ]; then
      fatal "Unable to reconnect to redis"
    fi
    sleep $(($attempts * ${CRONLOCK_RECONNECT_BACKOFF}))
    info "Redis connection lost, reconnecting..."

    ## close socket, as we maybe reconnecting due to a cluster MOVE command
    exec 3>&-  

    connect > /dev/null 2>&1
    attempts=$(($attempts + 1))
  done
  
  if [[ "$__resultvar" ]]; then
    eval $__resultvar="'${redis_response}'"
  else
    echo -n ${redis_response}
  fi
}

function raw_cmd {
  if [ ! -S /dev/fd/3 ]; then
    return 1 # No redis connection so exit as if connection was closed so reconnect logic kicks in
  fi

  { printf "${1}\r\n"; sleep 0.01; } >&3 # sleep to avoid a socket close race condition.

  # Parse the response according to redis protocol
  read -t ${CRONLOCK_REDIS_TIMEOUT} -r response <&3 || exit $?
  case $response in
    +*) # Status
      redis_response="${response#+}"
      ;;
    -*) # Error
      redis_response="${response#-}"
      if [[ ${response#-} =~ ([[:upper:]]+)[[:space:]]+[[:digit:]]+[[:space:]]([[:digit:]]+\.[[:digit:]]+\.[[:digit:]]+\.[[:digit:]]+)\:([[:digit:]]+) ]]; then
        if [[ ${BASH_REMATCH[1]} = "MOVED" || ${BASH_REMATCH[1]} = "ASK" ]]; then
          ### We got a REDIS Cluster MOVED or ASK response. Reset HOST and PORT and return 1 to allow for reconnection and reattempt
          CRONLOCK_HOST=${BASH_REMATCH[2]}
          CRONLOCK_PORT=${BASH_REMATCH[3]}

          info "${BASH_REMATCH[1]} response, resetting host to: ${BASH_REMATCH[2]} and port to: ${BASH_REMATCH[3]}"
          return 1
        fi
      fi
      echo "${response#-}" >&2
      exit 201
      ;;
    :*) # Integer
      [[ ${response#:} =~ ([[:digit:]]+) ]]
      redis_response=${BASH_REMATCH[1]}
      ;;
    \$*) # Bulk reply
      nchars="$(echo "${response#\$}"| tr -d '\n\r\0\t')"
      nchars="$(echo "${nchars%\r}"| tr -d '\n\r\0\t')"
      if [ $nchars -ge 0 ]; then # If chars is less than 0, redis returned null
        read -t ${CRONLOCK_REDIS_TIMEOUT} -r -n $nchars -d '' response <&3 || exit $?
        redis_response=${response}
        read -t ${CRONLOCK_REDIS_TIMEOUT} -r response <&3 || exit $? # Clear remaining newline
      fi
      ;;
    *) # net yet handled
      echo "ERR - Unknown response" >&2
      info "${response}"
      exit 201
      ;;
    esac
}

# Set magic variables for current FILE & DIR
__FILE__="$(test -L "$0" && readlink "$0" || echo "$0")"
__DIR__="$(dirname "${__FILE__}")"

# Try different locations for config file unless CRONLOCK_CONFIG was set
TRY="${__DIR__}/cronlock.conf"
[ -z "${CRONLOCK_CONFIG}" ] && [ -f "${TRY}" ] && CRONLOCK_CONFIG="${TRY}"
TRY="/etc/cronlock.conf"
[ -z "${CRONLOCK_CONFIG}" ] && [ -f "${TRY}" ] && CRONLOCK_CONFIG="${TRY}"

# Include config file
if [ -n "${CRONLOCK_CONFIG}" ]; then
  info "Including config from ${CRONLOCK_CONFIG}"
  source "${CRONLOCK_CONFIG}"
fi

# Defaults
[ -z "${CRONLOCK_HOST}" ]               && CRONLOCK_HOST="localhost"
[ -z "${CRONLOCK_PORT}" ]               && CRONLOCK_PORT=6379
[ -z "${CRONLOCK_DB}" ]                 && CRONLOCK_DB=0
[ -z "${CRONLOCK_REDIS_TIMEOUT}" ]      && CRONLOCK_REDIS_TIMEOUT=30
[ -z "${CRONLOCK_GRACE}" ]              && CRONLOCK_GRACE=40
[ -z "${CRONLOCK_RELEASE}" ]            && CRONLOCK_RELEASE=86400
[ -z "${CRONLOCK_RECONNECT_ATTEMPTS}" ] && CRONLOCK_RECONNECT_ATTEMPTS=5
[ -z "${CRONLOCK_RECONNECT_BACKOFF}" ]  && CRONLOCK_RECONNECT_BACKOFF=5
[ -z "${CRONLOCK_PREFIX}" ]             && CRONLOCK_PREFIX="cronlock."
[ -z "${CRONLOCK_VERBOSE}" ]            && CRONLOCK_VERBOSE="no"
[ -z "${CRONLOCK_RESET}" ]              && CRONLOCK_RESET="no"
[ -z "${CRONLOCK_NTPDATE}" ]            && CRONLOCK_NTPDATE="no"
[ -z "${CRONLOCK_TIMEOUT}" ]            && CRONLOCK_TIMEOUT=0

[ -z "${CRONLOCK_USE_SENTINEL}" ]       && CRONLOCK_USE_SENTINEL="no"
[ -z "${CRONLOCK_SENTINEL_MASTER}" ]    && CRONLOCK_SENTINEL_MASTER="mymaster"
[ -z "${CRONLOCK_SENTINEL_HOST}" ]      && CRONLOCK_SENTINEL_HOST="localhost"
[ -z "${CRONLOCK_SENTINEL_PORT}" ]      && CRONLOCK_SENTINEL_PORT=26379

# Mac OSX md5 != ubuntu md5
md5="$(which md5sum)"
if [ ! -x "${md5}" ]; then
  # osx
  md5="$(which md5)"
fi

# Fail on no md5
if [ ! -x "${md5}" ]; then
  fatal "Unable to find md5"
fi

# Mac OSX timelimit != ubuntu timeout
if [ "${CRONLOCK_TIMEOUT}" -gt 0 ]; then
  # ubuntu: aptitude install timeout
  timeout="$(which timeout)"
  # if we are on alpine version between 2.6 and 3.9, use different timeout command structure
  if [[ $(readlink $timeout) == "/bin/busybox" ]] && /bin/busybox | head -n 1 | egrep -q 'v1\.2[0-9]\.[0-9]+'; then
    timeoutcmd="timeout -t ${CRONLOCK_TIMEOUT} -s 9"
  else
    timeoutcmd="timeout -s 9 ${CRONLOCK_TIMEOUT}"
  fi

  if [ ! -x "${timeout}" ]; then
    # osx: brew install timelimit
    timeout="$(which timelimit)"
    timeoutcmd="timelimit -t ${CRONLOCK_TIMEOUT} -s 9"
  fi

  # Fail on no timeout/timelimit
  if [ ! -x "${timeout}" ]; then
    fatal "Unable to find timeout or timelimit. Please install or use CRONLOCK_TIMEOUT=0"
  fi
fi

# Sync time
if [ "${CRONLOCK_NTPDATE}" = "yes" ]; then
  if [ -x "$(which ntpd)" ]; then
    info "Ignoring request to synchronize time as ntpd is installed which does this continously. "
  else
    sudo ntpdate pool.ntp.org
  fi
fi

# Setup the key
[ -z "${CRONLOCK_KEY}" ] && CRONLOCK_KEY="$(echo "${@}" | ${md5} | awk '{print $1}')"
CRONLOCK_KEY="${CRONLOCK_PREFIX}${CRONLOCK_KEY}"

# Lookup the redis instance using Sentinel
if [ "${CRONLOCK_USE_SENTINEL}" = "yes" ]; then
  get_master
fi

# Connect to redis
info "Connecting to redis at '${CRONLOCK_HOST}' port '${CRONLOCK_PORT}'"
connect 2>/dev/null

if [ ! -z "${CRONLOCK_AUTH}" ]; then
  cmd "AUTH ${CRONLOCK_AUTH}" | grep 'OK' > /dev/null || fatal "Invalid auth for Redis server '${CRONLOCK_HOST}' port '${CRONLOCK_PORT}'."
fi

# Test redis connection
cmd "PING" | grep 'PONG' > /dev/null || fatal "Redis was not found at '${CRONLOCK_HOST}' port '${CRONLOCK_PORT}'. Please install, or see README for targeting a different host"

# Select Db
cmd "SELECT ${CRONLOCK_DB}" | grep 'OK' > /dev/null || fatal "Redis db '${CRONLOCK_DB}' does not exist. Please check number of redis db's, or see README for targeting a different db"

# Reset mode
if [ "${CRONLOCK_RESET}" = "yes" ]; then
  removed="$(cmd "DEL ${CRONLOCK_KEY}")"
  if [ "${removed}" = "1" ] || [ "${removed}" = "0" ]; then
    exit 200
  else
    fatal "Unable to remove ${CRONLOCK_KEY}"
  fi
fi

# Grace periods
let "expire_at_min = $(date -u "+%s") + ${CRONLOCK_GRACE} + 1"
let "expire_at_max = $(date -u "+%s") + ${CRONLOCK_RELEASE} + 1"

# Acquire
cmd "SETNX ${CRONLOCK_KEY} ${expire_at_max}" acquired

# Here come the edge cases
if [ "${acquired}" = "1" ]; then
  info "Cronlock ${CRONLOCK_KEY} acquired by me"
elif [ "${acquired}" = "0" ]; then
  cmd "GET ${CRONLOCK_KEY}" expires_at
  let "expires_in = ${expires_at} - $(date -u "+%s")"

  if [ "${expires_in}" -gt "0" ]; then
    info "Cronlock ${CRONLOCK_KEY} was already acquired (expires in ${expires_in}s)"
    exit 200
  fi

  if [ "${expires_in}" -eq "0" ]; then
    info "Cronlock ${CRONLOCK_KEY} was already acquired but expiring now. "
    exit 200
  else
    info "Cronlock ${CRONLOCK_KEY} was already acquired but expired ${expires_in#-}s ago"
  fi

  cmd "GETSET ${CRONLOCK_KEY} ${expire_at_max}" respire
  let "expires_in = ${respire} - $(date -u "+%s")"
  if [ "${expires_in}" -gt "0" ]; then
    info "Cronlock ${CRONLOCK_KEY} was just now acquired by a different box (expires in ${expires_in}s)"
    exit 200
  fi
else
  fatal "Unexpected return value while acquiring Cronlock: ${CRONLOCK_KEY} at: ${CRONLOCK_HOST}"
fi

# Handle timeout
if [ "${CRONLOCK_TIMEOUT}" -le 0 ]; then
  timeoutcmd=""
fi

# Execute original command
${timeoutcmd} "${@}"
exitcode=${?}

if [ "${CRONLOCK_TIMEOUT}" -gt 0 ]; then
  if [ "${exitcode}" = 137 ]; then
    echo "emergency: had to kill '${@}' after ${CRONLOCK_TIMEOUT}s timeout"
    exitcode=202
  fi
fi

# Setup min-grace period to compensate for small jobs on inconstistent system clocks
cmd "GETSET ${CRONLOCK_KEY} ${expire_at_min}" grace

# Issue Expire so we clean up the keys (more likely to end up actually obtaining the aquired lock first time - as the key wont remain indefinitely)
cmd "EXPIREAT ${CRONLOCK_KEY} ${expire_at_min}" expire

# Pass original command's exit code. Cronlock codes all exit above 200
exit ${exitcode}