#!/bin/bash

# Author: Carl Baldwin & Alan Pippin
# Description: This script replicates a remote zfs filesystem to a local zfs pool.
#              This script will keep all snapshots in sync, removing the ones
#              that have been deleted since the last replicate was performed.
#              This script will only send the new, or missing, snapshots since
#              the last replicate was performed.
# Usage: replicate <hostname> <zfs filesystem>

# source our configuration
config="${0%/*}/zfs-scripts.conf"
[ -e "${config}.dist" ] && . ${config}.dist
[ -e "${config}" ] && . ${config}

# command line arg parsing
remote=$1
remote_fs=$2
remote_pool=${2%%/*}

# return non-zero exit code if any command in the pipe fails
set -o pipefail

# get the backup pool from the command line or the config file if not specified
if [[ -n $3 ]]; then
    backup_pool=$3
else
    backup_pool=${backup_pool%% *} # use the first backup pool if none specified
fi

# Setup our cleanup and exit trap
cleanup() {
  if [[ -e "$local_list" ]]; then
    rm -f $local_list
  fi
  if [[ -e "$remote_list" ]]; then
    rm -f $remote_list
  fi
  $ssh $remote ls -d "$lockdir" > /dev/null 2>&1
  if [[ $? == 0 ]]; then
    $ssh $remote rm -rf "$lockdir"
  fi
}
fatal_and_exit() {
  echo -e 2>&1 "$1"
  # Destroy the current backup markers from the local backup_pool and remote_pool if they exist
  if [[ -n "$current_backup_marker" ]]; then
    # Local backup pool current backup marker
    $zfs list -t snapshot ${backup_pool}/${current_backup_marker} > /dev/null 2>&1
    if [ $? == 0 ]; then
      $zfs destroy ${backup_pool}/${current_backup_marker}
    fi
    # Remote pool current backup marker
    $ssh $remote zfs list -t snapshot ${current_backup_marker} > /dev/null 2>&1
    if [ $? == 0 ]; then
      $ssh $remote $zfs destroy ${current_backup_marker}
    fi
  fi
  # send email notification
  if [[ -n "$2" ]]; then
    echo -e "$1" | $mailx -s "zfs replicate on $hostname failed" "$2"
  fi
  # exit
  exit 1
}
trap fatal_and_exit INT
trap cleanup EXIT

# Declare a function to handle the replicate operation
replicate() {
  zfs_send="$1"
  zfs_recv="$zfs receive -vF -d ${backup_pool}/${remote_pool}"
  glue="$throttle $throttle_opt"
  if [[ $throttle_enable == 1 && -e $throttle ]]; then
    # handle using the glue in the local and remote host case properly
    if [[ -z "$ssh" ]]; then
      # local host glue case
      $zfs_send | $glue | $zfs_recv 
    else
      # remote host glue case
      $ssh $remote "$zfs_send | $glue" | $zfs_recv
    fi
  else 
    # no glue case - works for both the local and remote host cases
    $ssh $remote $zfs_send | $zfs_recv
  fi
  # The return code of the zfs_send | zfs_recv operation will be returned to the caller
}

# Make sure we have valid arguments
if [[ -z "$remote" ]] || [[ -z "$remote_fs" ]]; then
  fatal_and_exit "Usage: $0 <hostname> <zfs filesystem>"
fi

# check for localhost
if [[ $remote = "localhost" ]]; then
  remote=""
  ssh=""
fi

# Make sure the local backup pool and local receiving filesystem exist, or print some errors
zpool list -H "$backup_pool" >/dev/null 2>&1
if [ $? != 0 ]; then
  fatal_and_exit "-E- The local backup pool on $hostname, '$backup_pool' doesn't seem to exist." $mailto
fi
zfs list "$backup_pool/$remote_pool" >/dev/null 2>&1
if [ $? != 0 ]; then
  echo >&2 "-I- The local filesystem for the remote pool, '$backup_pool/$remote_pool' doesn't seem to exist."
  echo >&2 "    Creating the local filesystem to receive the remote pool into: $backup_pool/$remote_pool"
  $zfs create $backup_pool/$remote_pool
  if [ $? != 0 ]; then
    fatal_and_exit "-E- remote $zfs on $hostname create command failed" $mailto
  fi
fi

# Obtain the zpool guid for the local backup pool
backup_pool_guid=`zpool get guid $backup_pool 2>&1 | grep $backup_pool | awk '{ print $3 }'`
zpool get guid $backup_pool > /dev/null 2>&1
if [ $? != 0 ]; then
  fatal_and_exit "-E- Unable to extract the guid for the local backup pool on $hostname: $backup_pool" $mailto
fi

# Turn on shell verbosity
set -x

# Create the remote lockdir before continuing with the replicate
# Spinlock on creating the lock
maxsleeptime=60
maxattempts=500
attempts=0
while true; do
  $ssh $remote mkdir "$lockdir" >/dev/null 2>&1
  if [ $? != 0 ]; then
    # Another zfs admin tool is running.
    # Wait a random amount of time and try again
    ransleep=$(($RANDOM % $maxsleeptime))
    sleep $ransleep
    ((attempts=attempts+1))
  else 
    # No other zfs admin tool is running, we can now.
    break
  fi
  if [[ $attempts -gt $maxattempts ]]; then
    # We've exceeded our maximum while loop count
    echo "-E- The zfs filesystem has been locked down. Skipping replicate operation."
    fail_msg=`$ssh $remote ls -ld $lockdir 2>&1`
    fatal_and_exit "zfs-replicate-all on $hostname unable to obtain zfs admin lock:\n$fail_msg" $mailto
  fi
done

# Setup our backup marker names
current_backup_marker=${remote_fs}@current-backup-${backup_pool_guid}
previous_backup_marker=${remote_fs}@previous-backup-${backup_pool_guid}

# List the snapshots on the remote machine.
remote_list=$(mktemp /tmp/replicate.XXXXXX)
$ssh $remote \
    $zfs list -H -t snapshot |
    grep ^${remote_fs}@ |
    awk '{print$1}' > $remote_list
if [ $? != 0 ]; then
  fatal_and_exit "-E- remote $zfs list on $hostname command failed" $mailto
fi

# List the snapshots on the local machine.
# Don't list the current backup marker if it exists on the local side.
# If you do, it can mess up the common finding algorithm below.
local_list=$(mktemp /tmp/replicate.XXXXXX)
$zfs list -H -t snapshot |
    grep ^${backup_pool}/${remote_fs}@ |
    grep -v ^${backup_pool}/${current_backup_marker} | 
    awk "{gsub(/^$backup_pool./,\"\",\$1); print\$1}" > $local_list
# If no local snapshots exist, we may need to send the entire filesystem, which we'll do later
# So, no error check here as a non-zero return code means the local filesystem or snaps are missing
# We'll catch this later on as a case where we need to send the entire filesystem

# Destroy the current backup marker snapshot on the remote system if it exists
grep -q ${current_backup_marker} $remote_list
if [ $? == 0 ]; then
  $ssh $remote $zfs destroy ${current_backup_marker} 
  if [ $? != 0 ]; then
    fatal_and_exit "-E- remote $zfs destroy on $hostname command failed" $mailto
  fi
fi

# Create the current backup marker snapshot on the remote system
$ssh $remote $zfs snapshot ${current_backup_marker}
if [ $? != 0 ]; then
  fatal_and_exit "-E- remote $zfs snapshot on $hostname command failed" $mailto
fi

# Check to see if the previous backup marker exists in the remote snapshot list.
# Check to see if the previous backup marker exists in the local snapshot list.
# If the previous backup markers exists, perform an incremental replicate. Else:
# 1) check to see if a common snapshot exists, and perform an incremental replicate.
# 2) if no common snapshot exists, destroy the local filesystem, and perform a full replicate.
grep -q ${previous_backup_marker} $remote_list
no_markers=$?
grep -q ${previous_backup_marker} $local_list
no_markers=$(($no_markers || $?))

if [ $no_markers == 0 ]; then
  # We found backup markers, incrementally send the new snaps

  # First, rollback the local backup pool to the previous backup marker in case the previous
  # backup was interrupted for some reason. If we don't do this, the zfs send -R command
  # below may complain about snaps already existing as it tries to resend from the 
  # previous backup marker again from a previously interrupted replicate.
  $zfs rollback -rf ${backup_pool}/${previous_backup_marker} 
  if [ $? != 0 ]; then
    sleep 120
    $zfs rollback -rf ${backup_pool}/${previous_backup_marker}
    if [ $? != 0 ]; then
      fatal_and_exit "-E- remote incremental $zfs rollback command failed on $hostname" $mailto
    fi
  fi
  # Now it should be safe to send the snaps
  replicate "$zfs send -Rc -I${previous_backup_marker} ${current_backup_marker}"
  if [ $? != 0 ]; then
    fatal_and_exit "-E- remote incremental $zfs send command failed on $hostname" $mailto
  fi
else
  # We didn't find any backup markers, next check to see if we have a common snapshot.

  # See what the most recent snapshot on the remote end is.
  latest=$(tail -n 1 $remote_list)

  # I did this to make sure that diff would always display the most recent common
  # Since we're keying off the context of the diff, we need to ensure we will get context
  # by injecting a known difference in case no others exist in the lists.
  echo bogus.remote >> $remote_list
  echo bogus.local  >> $local_list
  common=$(diff -u $remote_list $local_list | grep '^ ' | tail -n 1)

  if [[ -n "$common" ]]; then
    # We found a common snapshot, incrementally send the new snaps
    replicate "$zfs send -Rc -I${common/*@/@} ${current_backup_marker}"
    if [ $? != 0 ]; then
      fatal_and_exit "-E- remote incremental $zfs send command failed on $hostname" $mailto
    fi
  else
    # We did not find any markers or a common snapshot
    # At this point, we'll have to send the entire filesystem
    # Destroy the local filesystem if it exists before receving the full replicate
    zfs list ${backup_pool}/${remote_fs} > /dev/null 2>&1
    if [ $? == 0 ]; then
      if [[ $destroy_local_filesystem_on_full_replicate == 1 ]]; then
        $zfs destroy -r ${backup_pool}/${remote_fs}
        if [ $? != 0 ]; then
          fatal_and_exit "-E- remote full $zfs destroy command failed on $hostname" $mailto
        fi
      else
        echo "-W- We need to destroy a local filesystem before receiving a full stream."
	echo "    However, since the option is set to prevent this, skipping replicate operation."
        fatal_and_exit "unable to destroy local filesystem:\n$zfs destroy -r ${backup_pool}/${remote_fs} not able to run on $hostname" $mailto
      fi
    fi
    # Send the full filesystem
    replicate "$zfs send -Rc ${current_backup_marker}"
    if [ $? != 0 ]; then
      fatal_and_exit "-E- remote full $zfs send command failed on $hostname" $mailto
    fi
  fi
fi
 
# destroy the previous backup markers now that we've replicated past them
# don't check the return codes here because these may not exist, and that is ok
$zfs destroy ${backup_pool}/${previous_backup_marker} > /dev/null 2>&1
$ssh $remote $zfs destroy ${previous_backup_marker} > /dev/null 2>&1
sleep 1

# Rename the current backup marker to be the previous backup marker
$zfs rename ${backup_pool}/${current_backup_marker} ${backup_pool}/${previous_backup_marker}
if [ $? != 0 ]; then
  fatal_and_exit "-E- local $zfs rename command failed on $hostname" $mailto
fi
$ssh $remote $zfs rename ${current_backup_marker} ${previous_backup_marker}
if [ $? != 0 ]; then
  fatal_and_exit "-E- remote $zfs rename command failed on $hostname" $mailto
fi