mongodb-replica.sh

#!/usr/bin/env bash

# MongoDB port
PORT=27017
# Consul key to setup the lock
MONGO_KEY="mongodb/replicaset"

#
# Retrieve MongoDB instance IPs from Consul
#
getMongoIps() {
  dig @127.0.0.1 -p8600 mongo.service.consul +short
}

#
# Deal with the "lock" key on Consul
#
setLock() {
  curl -X PUT -d @- localhost:8500/v1/kv/${MONGO_KEY} <<< "true"
}
isLock() {
  curl -s localhost:8500/v1/kv/${MONGO_KEY}?raw
}
consulStatus() {
  curl -s localhost:8500/v1/status/leader
}

#
# Initialize the replicaset with the private ip
#
initReplica() {
  mongo --quiet <<EOF
rs.initiate({
  _id: "rsa",
  version: 1,
  members: [{_id: 0, host: "${1}"}]
})
EOF
}

#
# Check if we aleady have a RS configuration
#
checkReplica() {
  mongo --quiet --eval "rs.status().set"
}

#
# Check if $1 knows who is the master and set $2 variable with the master address
#
replicaSetMaster() {
  local output
  output=$(mongo --host $1 --quiet --eval "db.isMaster()['primary']")
  eval "$2=${output}"
}

#
# List all nodes of the replica set from $1 node
#
replicaSetNodes() {
  mongo --host $1 --quiet --eval "db.isMaster()['hosts']" \
  | jq '.[]?' \
  | awk -F\" '{print $2}'
}

#
# Add $2 to the RS config
#
addMember() {
  echo "Adding $2 to the replicaSet"
  mongo --host $1 --quiet --eval "rs.add('${2}')"
}

#
# Reconfigure the replica set by removing the non-healthy nodes
#
reconfig() {
  mongo --quiet <<EOF
cfg=rs.conf();
healthy=rs.status().members.filter(function(d){ return d.health > 0; }).map(function(d){ return d.name });
cfg.members=cfg.members.filter(function(d){ return healthy.indexOf(d.host) > = 0 });
rs.reconfig(cfg, { force:true });
EOF
}

#
# Find who is the master node
#
findMaster() {
  local mongo_ips
  mongo_ips=$(getMongoIps)
  for mongo_ip in $mongo_ips;
  do
    replicaSetMaster $mongo_ip MASTER
    [[ -n "$MASTER" ]] && break
  done
}

#
# Main
#
main() {
  [[ "$(which mongo)" = "" ]] && \
    echo "Can't find mongo binary. Is mongo-org-shell installed ?" && \
    exit 1

  # Waiting for consult to be ready
  WAIT_CONSUL=true
  while $WAIT_CONSUL; do
    [[ -n "$(consulStatus)" ]] && WAIT_CONSUL=false
    sleep 5
  done

  CURRENT_NODE="$(hostname -i):${PORT}"
  echo "Current node : ${CURRENT_NODE}"

  [[ -n "$(checkReplica)" ]] && \
    echo "Replica set already configured. Exiting." && \
    exit 0

  # Get IPs from Consul
  MONGO_IPS=$(getMongoIps)
  echo -e "Nodes in the mongodb cluster : \n$MONGO_IPS\n"

  # Did another node start the replica set ?
  #TODO: the lock should be set with the IP of the instance which creates it
  #      then when checking for the master, if none is found we should check that the value in the lock is valid
  IS_LOCK=$(isLock)
  if [ "${IS_LOCK}" = "true" ]; then
    echo "The replica set is already present. Waiting for the master..."
    MASTER=""
    RETRIES=3
    while [ "${MASTER}" = "" -a $RETRIES -gt 0 ]; do
      findMaster
      if [ -n "${MASTER}" ]; then
        if [[ ! "${MASTER}" = "${CURRENT_NODE}" ]];then
          addMember $MASTER $CURRENT_NODE
        else
          echo "Current node (${CURRENT_NODE}) is the master. Nothing to do."
        fi
      else
        echo "No master has been found. Trying again in 5sec..."
        sleep 5
        (( RETRIES -= 1 ))
      fi
    done

    if [ "${MASTER}" = "" ]; then
      echo "No master found after ${RETRIES}. Reconfiguring the replica set..."
      reconfig
    fi
  else
    echo "Creating lock with value : ${CURRENT_NODE}"
    setLock $CURRENT_NODE

    echo "Initialize replica set"
    initReplica $CURRENT_NODE
  fi

  echo "Current member in the replica set :"
  replicaSetNodes $MASTER
}

main