Sync Riak Partitions over rsync

riak-admin-rsync.sh

#!/usr/bin/env bash

# Prints the usage details of this script
usage() {
        local prog=${0##*/}

        cat <<-EOF
                Transfers Riak partitions between nodes
                Usage: $prog -N <destination>
                Example: $prog prod-2290
                Options:
                  -h    Displays this help screen
                  -l    Lists the partitions to be transferred from the current node
                  -N    The name of the destination Riak node; required
        EOF
        exit 1
}


# Prints an error message and exits
fatal() {
        echo "$@" >&2
        exit 1
}


# Retrieves the partitions to be transferred from the current Riak node
get_partitions() {
sudo riak-admin ring-status |\
awk '
/^Owner: / {
        _owner = $2;
        sub(/^.*@/, "", _owner);
        if (owner != _owner) {
                owner = _owner;
        }
}
/^Next Owner: / {
        nextowner = $3;
        sub(/^.*@/, "", nextowner);
}
/Index: / {
        printf("%s %s %s\n", owner, $2, nextowner);
}'
}


# Lists the partitions to be transferred from the source riak node
list_partitions() {
        local partitions=$(get_partitions | grep $SRC_RIAK_NODE)
        local partitions_count=$(echo "$partitions" | wc -l)

        if [ -z "$partitions" ]; then
                echo "No partitions to be transferred."
                exit 0
        fi

        echo "$partitions"
        echo "$partitions_count partitions to be transferred!"
        exit 0
}


# Returns the name of the Riak node
get_riak_node_name() {
        sudo riak-admin member-status | grep "$1" | awk '{ print $4 }' | tr -d "'"
}


# Starts Riak on the specified node
start_riak() {
        local riak_node=$1

        printf "Attempting to start Riak on $riak_node..."
        ssh $riak_node "sudo svcadm enable riak-epmd" || fatal "Couldn't start Riak-epmd!"
        sleep 5
        ssh $riak_node "sudo svcadm enable riak" || fatal "Couldn't start Riak!"
        echo "done!"
}


# Stops Riak on the specified node
stop_riak() {
        local riak_node=$1

        printf "Attempting to stop Riak on $riak_node..."
        ssh $riak_node "sudo svcadm disable riak" || fatal "Couldn't stop Riak!"
        sleep 5
        ssh $riak_node "sudo svcadm disable riak-epmd" || fatal "Couldn't stop Riak-epmd!"
        echo "done!"

        local psresult=`ssh $riak_node "ps ax | grep 'epmd\|beam.smp' | grep -v 'grep'"`
        until [ $? -eq 1 ]; do
                echo "Riak still running!"
                echo "$psresult"
                echo

                sleep 1
                psresult=`ssh $riak_node "ps ax | grep 'epmd\|beam.smp' | grep -v 'grep'"`
        done
        ssh $riak_node "svcs -a | grep riak"
        echo "Riak stopped!"
}


# Transfers the Riak partitions to the destination
transfer_partitions() {
        local destination_path=$1

        for p in $PARTITIONS; do
                echo "Transferring partition $p from $SRC_RIAK_NODE to $DEST_RIAK_NODE..."
                cd $DIR_LEVELDB/$p
                rsync -azr -f"+ */" -f"- *" --rsync-path="sudo rsync" --progress $DIR_LEVELDB/$p $DEST_RIAK_NODE:$destination_path >> $FILE_LOG || fatal "Error transferring directory structure for $p"
                find . ! -type d -print0 | xargs -0 -n1 -P9 -I% rsync -avz --progress --rsync-path="sudo rsync" % $DEST_RIAK_NODE:$destination_path/$p/% >> $FILE_LOG || fatal "Error transferring files for $p"
        done
}


# Moves the partitions from the tmp location to the riak directory on the destination riak node
move_partitions_in() {
        printf "Moving $PARTITIONS_COUNT partitions into $DIR_LEVELDB on $DEST_RIAK_NODE..."
        for p in $PARTITIONS; do
                ssh $DEST_RIAK_NODE "sudo rm -rf $DIR_LEVELDB/$p" || fatal "Error deleting existing files for $p"
                ssh $DEST_RIAK_NODE "sudo mv $DIR_TMP/$p $DIR_LEVELDB/" || fatal "Error moving $p"
        done
        echo "done!"
}


# Moves the partitions from the riak directory to a tmp location on the source riak node
move_partitions_out() {
        printf "Moving $PARTITIONS_COUNT partitions out of $DIR_LEVELDB on $SRC_RIAK_NODE..."
        for p in $PARTITIONS; do
                sudo mv $DIR_LEVELDB/$p $DIR_ARCHIVE/ || fatal "Error moving $DIR_LEVELDB/$p to $DIR_ARCHIVE on $SRC_RIAK_NODE"
        done
        echo "done!"
}


# Declare the necessary variables
DIR_LEVELDB="/var/db/riak/leveldb"
DIR_ARCHIVE="/var/db/riak/voxer_moved_partitions"
DIR_TMP="/var/tmp/riak/leveldb"

FILE_LOG="/var/tmp/riak_transfer.log"

SRC_RIAK_NODE=$(hostname | cut -d '.' -f 1)
DEST_RIAK_NODE=
SRC_RIAK_NAME=
DEST_RIAK_NAME=


# Process the command line options
while getopts 'hlN:' option; do
        case "$option" in
                h) usage;;
                l) list_partitions;;
                N) DEST_RIAK_NODE=$OPTARG;;
                *) usage;;
        esac
done

# Ensure that all required args are present
if [ -z $DEST_RIAK_NODE ]; then
        fatal "Must specify a destination node using -N <destination>"
fi


# Ensure there is data to be transferred
PARTITIONS=$(get_partitions | grep $SRC_RIAK_NODE | grep $DEST_RIAK_NODE | cut -d ' ' -f 2)
PARTITIONS_COUNT=$(echo "$PARTITIONS" | wc -l)

if [ -z "$PARTITIONS" ]; then
        echo "No partitions to be transferred to $DEST_RIAK_NODE"
        exit 2
fi

# Ensure the destination is accessible
if [ $DEST_RIAK_NODE != $(ssh $DEST_RIAK_NODE "hostname | cut -d '.' -f 1") ]; then
        echo "Destination node $DEST_RIAK_NODE is not accessible!"
        exit 3
fi


# Get the names of the Riak nodes
SRC_RIAK_NAME=$(get_riak_node_name $SRC_RIAK_NODE)
DEST_RIAK_NAME=$(get_riak_node_name $DEST_RIAK_NODE)

# Display stats for the transfer process
echo "Transfers from $SRC_RIAK_NODE to $DEST_RIAK_NODE:"
echo "$PARTITIONS"
echo "$PARTITIONS_COUNT partitions to be transferred."
echo

# Prepare to start the transfers
echo "Starting Riak transfers from $SRC_RIAK_NODE to $DEST_RIAK_NODE" > $FILE_LOG
echo "Moving $PARTITIONS_COUNT partitions:" >> $FILE_LOG
echo "$PARTITIONS" >> $FILE_LOG

ssh $DEST_RIAK_NODE "sudo mkdir -p $DIR_TMP; sudo chown riak:riak $DIR_TMP"
sudo mkdir -p $DIR_ARCHIVE; sudo chown riak:riak $DIR_ARCHIVE

sudo riak-admin transfer-limit 0
sleep 5

# Transfer the data
transfer_partitions $DIR_TMP
stop_riak $DEST_RIAK_NODE
move_partitions_in
stop_riak $SRC_RIAK_NODE
transfer_partitions $DIR_LEVELDB
move_partitions_out

start_riak $DEST_RIAK_NODE
start_riak $SRC_RIAK_NODE
echo "$PARTITIONS_COUNT partitions successfully moved from $SRC_RIAK_NODE to $DEST_RIAK_NODE."

sleep 60
sudo riak-admin transfer-limit $SRC_RIAK_NAME $partition_count
sudo riak-admin transfer-limit $DEST_RIAK_NAME $partition_count