#!/bin/sh
#-------------------------------------------------------------------------+
# Copyright (C) 2016 Matt Churchyard (churchers@gmail.com)
# All rights reserved
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted providing that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# vm recv ...
# recieve a guest that is being sent from another host
#
# @param string _name name of the new guest to recieve into
#
migration::recv(){
    local _name
    local _ds="default"
    local _conf _port _opt _start _stage _triple

    while getopts d:s12t _opt; do
        case $_opt in
            d) _ds="${OPTARG}" ;;
            s) _start="1" ;;
            1) _stage="1" ;;
            2) _stage="2" ;;
            t) _triple="1" ;;
            *) util::usage ;;
        esac
    done

    shift $((OPTIND - 1))
    _name="$1"

    [ -z "${_name}" ] && util::usage
    util::check_name "${_name}" || util::err "invalid virtual machine name - '${_name}'"
    datastore::get "${_ds}" || util::err "unable to load datastore - '${_ds}'"
    [ -z "${VM_DS_ZFS}" ] && util::err "${_ds} datastore must be on ZFS to support migration"
    [ -n "${_stage}" -a -n "${_triple}" ] && util::err "single stage and triple stage are mutually exclusive"

    # find a port to use
    vm::find_available_net_port "_port" "12000"
    [ -z "${_port}" ] && util::err "unable to allocate a port to recieve on"

    echo "Recieving guest into ${VM_DS_PATH}/${_name}"

    # STAGE 1
    # full send by default although user can specify an incremental snapshot
    # on sending side.
    [ -z "${_stage}" -o "${_stage}" = "1" ] && migration::__recv_snapshot "1"

    # STAGE 1b
    # only performed with the -t (triple stage) option
    # this is useful for large guests as it performs an incremental send while
    # the guest is still running. this should hopefully be much smaller than a full send
    # and so complete quicker, allowing the guest to be off for less time
    [ -n "${_triple}" ] && migration::__recv_snapshot "1b"

    # STAGE 2
    # this is the stage when the guest is shutdown
    # receive a snapshot, then start the guest if requested
    if [ -z "${_stage}" -o "${_stage}" = "2" ]; then
        migration::__recv_snapshot "2"

        # update config file
        echo "  * updating configuration file"
        zfs mount "${VM_DS_ZFS_DATASET}/${_name}"
        _conf=$(find "${VM_DS_PATH}/${_name}/" -name "*.conf" | awk -F"/" '{print $NF}')
        [ -z "${_conf}" ] && util::err_inline "unable to locate guest configuration file"
        mv "${VM_DS_PATH}/${_name}/${_conf}" "${VM_DS_PATH}/${_name}/${_name}.conf"

        if [ -n "${_start}" ]; then
            echo "  * attempting to start ${_name}"
            core::__start "${_name}"
            exit
        fi
    fi

    echo "  * done"
}

# receive a snapshot over the network
#
# @param string _state the migration stage to display
#
migration::__recv_snapshot(){
    local _stage="$1"

    echo "  * stage ${_stage}: waiting for snapshot on port ${_port}"
    socat -u "TCP-LISTEN:${_port}" EXEC:"zfs recv ${VM_DS_ZFS_DATASET}/${_name}" >/dev/null 2>&1
    [ $? -eq 0 ] || util::err_inline "error detected while recieving snapshot"
    echo "  * stage ${_stage}: complete"
}

# vm send ...
# send a guest to another system
#
# @param string _name name of the guest to send
# @param string _host host to send to
#
migration::send(){
    local _name _host _port _stage _triple _inc
    local _snap1 _snap2 _snap3 _state _running _pid _count=0
    local IFS=$'\n'

    while getopts i:12t _opt; do
        case $_opt in
            i) _inc="${OPTARG}" ;;
            1) _stage="1" ;;
            2) _stage="2" ;;
            t) _triple="1" ;;
            *) util::usage ;;
        esac
    done

    shift $((OPTIND - 1))
    _name="$1"
    _host="$2"

    [ -z "${_name}" -o -z "${_host}" ] && util::usage
    datastore::get_guest "${_name}" || util::err "unable to locate guest - '${_name}'"
    [ -z "${VM_DS_ZFS}" ] && util::err "${VM_DS_NAME} datastore must be on ZFS to support migration"
    [ -n "${_stage}" -a -n "${_triple}" ] && util::err "single stage and triple stage are mutually exclusive"
    [ "${_stage}" = "2" -a -z "${_inc}" ] && util::err "source snapshot must be given when running stage 2"

    # split host & port
    echo "${_host}" | egrep -iqs '^.+:[0-9]+$'
    [ $? -eq 0 ] || util::err "destination must be specified in host:port format"
    _port="${_host##*:}"
    _host="${_host%%:*}"

    # check compatability
    config::load "${VM_DS_PATH}/${_name}/${_name}.conf"
    migration::__check_compat

    # check if vm is running
    vm::confirm_stopped "${_name}" "1" >/dev/null 2>&1
    _state="$?"
    [ ${_state} -eq 2 ] && util::err "guest is powered up on another host"
    [ ${_state} -eq 1 ] && _running="1"

    # try to get pid if it is running
    if [ -n "${_running}" ]; then
        _pid=$(pgrep -fx "bhyve: ${_name}")
        [ -z "${_pid}" ] && util::err "guest seems to be running but can't find its pid"
    fi

    echo "Sending ${_name} to ${_host}"

    # STAGE 1
    # send the first snapshot
    if [ -z "${_stage}" -o "${_stage}" = "1" ]; then
        _snap1="$(date +'%Y%m%d%H%M%S')"
        echo "  * stage 1: taking snapshot - ${_snap1}"
        zfs snapshot -r "${VM_DS_ZFS_DATASET}/${_name}@${_snap1}" >/dev/null 2>&1
        [ $? -eq 0 ] || util::err_inline "error taking snapshot"

        # send the first snapshot
        migration::__send_snapshot "1" "${_snap1}" "${_inc}"

        # only wait if we have further stages
        if [ "${_stage}" != "1" ]; then
            echo "  * stage 1: giving time for remote socket to close"
            sleep 5
        fi
    fi

    # STAGE 1b
    # do it again if in triple stage
    if [ -n "${_triple}" ]; then
        _snap2="$(date +'%Y%m%d%H%M%S')"
        echo "  * stage 1b: taking snapshot - ${_snap2}"
        zfs snapshot -r "${VM_DS_ZFS_DATASET}/${_name}@${_snap2}" >/dev/null 2>&1
        [ $? -eq 0 ] || util::err_inline "error taking snapshot"

        # send the middle snapshot
        migration::__send_snapshot "1b" "${_snap2}" "${_snap1}"

        echo "  * stage 1b: giving time for remote socket to close"
        sleep 5
    fi

    # do we need to run stage2?
    if [ "${_stage}" = "1" ]; then
        echo "  * done"
        exit
    fi

    # if it's running, try and stop it
    if [ -n "${_running}" ]; then
        echo -n "  * stage 2: attempting to stop guest locally"

        kill ${_pid} >/dev/null 2>&1
        sleep 1
        kill ${_pid} >/dev/null 2>&1

        while [ ${_count} -lt 60 ]; do
            sleep 2
            kill -0 ${_pid} >/dev/null 2>&1 || break
            echo -n "."
            _count=$((_count + 1))
        done
        echo ""
    fi

    # has it stopped?
    kill -0 ${_pid} >/dev/null 2>&1 && util::err_inline "failed to stop guest"
    echo "  * stage 2: guest powered off"

    _snap3="$(date +'%Y%m%d%H%M%S')"
    echo "  * stage 2: taking snapshot - ${_snap3}"
    zfs snapshot -r "${VM_DS_ZFS_DATASET}/${_name}@${_snap3}" >/dev/null 2>&1
    [ $? -eq 0 ] || util::err_inline "error taking snapshot"

    # triple stage use snap 2
    # if single stage and snapshot given, use that
    # otherwise snap1
    if [ "${_triple}" = "1" ]; then
        migration::__send_snapshot "2" "${_snap3}" "${_snap2}"
    elif [ "${_stage}" = "2" ]; then
        migration::__send_snapshot "2" "${_snap3}" "${_inc}"
    else
        migration::__send_snapshot "2" "${_snap3}" "${_snap1}"
    fi

    echo "  * removing snapshots"
    [ -n "${_snap1}" ] && zfs destroy "${VM_DS_ZFS_DATASET}/${_name}@${_snap1}" >/dev/null 2>&1
    [ -n "${_snap2}" ] && zfs destroy "${VM_DS_ZFS_DATASET}/${_name}@${_snap2}" >/dev/null 2>&1
    [ -n "${_snap3}" ] && zfs destroy "${VM_DS_ZFS_DATASET}/${_name}@${_snap3}" >/dev/null 2>&1
    echo "  * done"
}

# send a snapshot to the remote host
#
# @param string _stage stage number to display
# @param string _snapshot the snapshot to send
# @param optional string _inc incremental source to use
#
migration::__send_snapshot(){
    local _stage="$1"
    local _snapshot="$2"
    local _inc="$3"

    if [ -n "${_inc}" ]; then
        echo "  * stage ${_stage}: sending ${VM_DS_ZFS_DATASET}/${_name}@${_snapshot} (incremental source ${_inc})"
        socat -u EXEC:"zfs send -Ri ${_inc} ${VM_DS_ZFS_DATASET}/${_name}@${_snapshot}" "TCP4:${_host}:${_port}" >/dev/null 2>&1
    else
        echo "  * stage ${_stage}: sending ${VM_DS_ZFS_DATASET}/${_name}@${_snapshot}"
        socat -u EXEC:"zfs send -R ${VM_DS_ZFS_DATASET}/${_name}@${_snapshot}" "TCP4:${_host}:${_port}" >/dev/null 2>&1
    fi

    [ $? -eq 0 ] || util::err_inline "error detected while sending snapshot"
    echo "  * stage ${_stage}: complete"
}

# see if a guest can be migrated.
# there are a few guest settings that are likely to
# cause the guest to break if it's moved to another host
#
migration::__check_compat(){
    local _setting _err _num=0

    # check pass through
    config::get "_setting" "passthru0"
    [ -n "${_setting}" ] && _err="pci pass-through enabled"

    # check for custom disks
    # file/zvol are under guest dataset and should go across ok
    # custom disks could be anywhere
    while true; do
        config::get "_setting" "disk${_num}_type"
        [ -z "${_setting}" ] && break
        [ "${_setting}" = "custom" ] && _err="custom disk(s) configured" && break
        _num=$((_num + 1))
    done

    [ -n "${_err}" ] && util::err "migration is not supported for this guest (${_err})"
}
