zsync

#!/usr/bin/env python

# Licensed under an MIT-style License:
#
# Copyright (c) 2012 Fog Creek Software, Inc.
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


# Changelog
# ---------
# 2011-02-13  Tim Stewart  <tim@stoo.org>
#     * Initial Release


from __future__ import print_function

import os
import sys
import syslog
import shlex
import subprocess
import signal
import re
import textwrap
import select
import datetime
from datetime import datetime
from contextlib import contextmanager


class NonZeroExitStatus(Exception): pass
class InvalidDestinationError(Exception): pass
class NoSuchDatasetError(Exception): pass
class LockfileExists(Exception): pass
class SendNoSuchDataset(Exception): pass


class Zfs(object):
    """This class represents the overall ZFS system.

    It can be used to get a list of all ZFS pools and datasets on the
    system."""

    @classmethod
    def _datasets(cls, dset_type, filter, cmd_prefix="", cmd_postfix=""):
        """Internal implementation of dataset selection code."""
        cmd = '%s zfs list -t %s -r -Ho name,creation -s creation %s' % (
            cmd_prefix, dset_type, cmd_postfix)
        stdout, stderr = cls._run_cmd(cmd)
        lines = [line for line in stdout.splitlines() if line.strip()]
        results = []
        filter_re = re.compile(filter)
        for line in lines:
            cols = line.split('\t')
            name = cols[0]
            cdate = datetime.strptime(cols[1], '%a %b %d %H:%M %Y')

            if filter_re.match(name):
                results.append(Dataset(name, create=cdate))
        return results

    @classmethod
    def _run_cmd(cls, args):
        args_list = shlex.split(args)
        p = subprocess.Popen(args_list, stdin=None,
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        if p.returncode:
            raise NonZeroExitStatus(stderr.strip())
        return (stdout, stderr)

    @classmethod
    def pools(cls):
        """Return a list of ZFS pools."""
        stdout, stderr = cls._run_cmd('zpool list -Ho name')
        return [Zpool(l) for l in stdout.strip().splitlines()]

    @classmethod
    def datasets(cls, filter=".*", cmd_prefix=""):
        """Return a list of datasets in all pools."""
        return cls._datasets('all', filter, cmd_prefix)

    @classmethod
    def filesystems(cls, filter=".*", cmd_prefix=""):
        """Return a list of all datasets that are filesystems in all pools."""
        return cls._datasets('filesystem', filter, cmd_prefix)

    @classmethod
    def snapshots(cls, filter=".*", cmd_prefix=""):
        """Return a list of all datasets that are snapshots in all pools."""
        return cls._datasets('snapshot', filter, cmd_prefix)

    @classmethod
    def volumes(cls, filter=".*", cmd_prefix=""):
        """Return a list of all datasets that are volumes in all pools."""
        return cls._datasets('volume', filter, cmd_prefix)


class Zpool(object):
    """Represents an individual ZFS pool.

    Can be used to perform pool-wide operations and to get a list of
    datasets within a pool."""

    def __init__(self, name):
        self._name = name

    @property
    def name(self):
        """The ZFS pool's name"""
        return self._name

    def datasets(self, filter=".*", cmd_prefix=""):
        """Return a list of datasets in all pools."""
        return Zfs._datasets('all', filter, cmd_prefix, self.name)

    def filesystems(self, filter=".*", cmd_prefix=""):
        """Return a list of all datasets that are filesystems in all pools."""
        return Zfs._datasets('filesystem', filter, cmd_prefix, self.name)

    def snapshots(self, filter=".*", cmd_prefix=""):
        """Return a list of all datasets that are snapshots in all pools."""
        return Zfs._datasets('snapshot', filter, cmd_prefix, self.name)

    def volumes(self, filter=".*", cmd_prefix=""):
        """Return a list of all datasets that are volumes in all pools."""
        return Zfs._datasets('all', filter, cmd_prefix, self.name)


class Dataset(object):
    """Represents a dataset within a ZFS pool.

    It could be a filesystem, a volume, or a snapshot.  Allows access
    to properties of the dataset."""

    def __init__(self, name, cmd_prefix="", **attrs):
        """Initialize a new Dataset.

        If we are provided only a name then fill in the attributes
        ourselves.  Otherwise, fill with the values provided.  Only
        the ZFS classes should use **attrs directly in this
        constructor."""
        self._name = name
        self._cmd_prefix = cmd_prefix
        self._attrs = attrs
        if not attrs:
            dsets = Zfs._datasets('all', '^%s$' % (name), cmd_prefix)
            if not dsets:
                raise NoSuchDatasetError
            if len(dsets) > 1:
                raise Exception("Unexpected logic bug")

    @property
    def name(self):
        """The dataset's name."""
        return self._name

    @property
    def attrs(self):
        """Dictionary of dataset attributes."""
        return self._attrs

    @property
    def snapshots(self):
        """A list of snapshots based on this dataset."""
        return Zfs._datasets('snapshot', '%s@.*'
                             % (self.name), self._cmd_prefix)

    def __str__(self):
        return 'dset:%s  create:%s' % (self.name, self.attrs['create'])


@contextmanager
def zfs_recv_lockfile(remhost, remfs):
    """ Lock file is created based on the destination host and
    dataset, to allow for concurrent replications when safe"""
    lockfile = '/tmp/zfs-replicate-%s-%s.lock' \
               % (remhost.split('.')[0], remfs.replace('/', '_'))
    lockfile_fd = 0
    try:
        lockfile_fd = os.open(lockfile, os.O_EXCL|os.O_CREAT)
        os.close(lockfile_fd)
        yield
    except OSError:
        raise LockfileExists
    finally:
        if lockfile_fd:
            os.unlink(lockfile)


def log(msg, log=True, stdout=False, stderr=False):
    for line in msg.split('\n'):
        if stdout:
            print(line)
            sys.stdout.flush()
        if stderr:
            print('ERROR:  %s' % line, file=sys.stderr)
            sys.stderr.flush()
        if log:
            if stderr:
                syslog.syslog('ERROR:  %s' % line)
            else:
                syslog.syslog('INFO:   %s' % line)


def split_dest(dest):
    m = re.match('^([^:]+):([^:]+)$', dest)
    if not m:
        raise InvalidDestinationError
    return m.group(1), m.group(2)


def run_repl(begin_snap, end_snap, dst_host, dst_fs, verbose=False,
             dryrun=False):
    """Perform replication to dst_host into dst_fs.

    If begin_snap is None, perform a full replication, otherwise,
    perform an incremental send/receive between begin_snap and
    end_snap."""

    send_args = ''
    if begin_snap:
        send_args = send_args + ' -I @%s' % (begin_snap.name.split('@')[1])
    send_cmd = 'zfs send -v%s %s' % (send_args, end_snap.name)
    recv_cmd = 'ssh %s zfs receive -v -F %s' % (dst_host, dst_fs)
    send_list = shlex.split(send_cmd)
    recv_list = shlex.split(recv_cmd)

    if dryrun:
        log('[DRY RUN] Not running command: %s | %s' %
            (send_cmd, recv_cmd), stdout=True)
    else:
        log('Running command: %s | %s' % (send_cmd, recv_cmd), stdout=verbose)

    if not dryrun:
        send_p = subprocess.Popen(send_list,
                                  stdout=subprocess.PIPE,
                                  stderr=subprocess.PIPE)
        recv_p = subprocess.Popen(recv_list,
                                  stdin=send_p.stdout,
                                  stdout=subprocess.PIPE,
                                  stderr=subprocess.PIPE)
        send_p.stdout.close() # Ensures EPIPE is properly sent to send_p

        no_such_dset = ''
        no_such_dset_RE = "^warning: cannot send '(.*)': no such pool or dataset$"

        send_err_EOF = False
        recv_out_EOF = False
        recv_err_EOF = False
        while True:
            rready, wready, xready = select.select(
                [send_p.stderr, recv_p.stdout, recv_p.stderr], [], [])
            buf = ''
            for fd in rready:
                buf = fd.readline().rstrip()
                if fd is send_p.stderr:
                    if buf:
                        log('  SEND (stderr): %s' % (buf), stdout=verbose)
                        dset = re.search(no_such_dset_RE, buf)
                        if dset:
                            no_such_dset = dset.group(1)
                    else:
                        send_err_EOF = True

                elif fd is recv_p.stdout:
                    if buf:
                        log('  RECV (stdout): %s' % (buf), stdout=verbose)
                    else:
                        recv_out_EOF = True

                elif fd is recv_p.stderr:
                    if buf:
                        log('  RECV (stderr): %s' % (buf), stdout=verbose)
                    else:
                        recv_err_EOF = True

            if send_err_EOF and recv_out_EOF and recv_err_EOF:
                break

        send_retcode = send_p.wait()
        recv_retcode = recv_p.wait()

        if no_such_dset:
            raise SendNoSuchDataset(no_such_dset)

        if send_retcode:
            raise NonZeroExitStatus(send_cmd)
        if recv_retcode:
            raise NonZeroExitStatus(recv_cmd)

    return


def rm_remote_dsets(dsetlist, dst_host, verbose=False, dryrun=False):
    """Remove remote datasets"""

    for dset in dsetlist:
        cmd = 'ssh %s zfs destroy %s' % (dst_host, dset.name)
        cmd_list = shlex.split(cmd)

        if dryrun:
            log('[DRY RUN] Not running command: %s' % (cmd), stdout=True)
        else:
            log('Running command: %s' % (cmd), stdout=verbose)

        if not dryrun:
            p = subprocess.Popen(cmd_list,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE)
            out_EOF = False
            err_EOF = False
            while True:
                rready, wready, xready = select.select(
                    [p.stdout, p.stderr], [], [])
                buf = ''
                for fd in rready:
                    buf = fd.readline().rstrip()
                    if fd is p.stdout:
                        if buf:
                            log('  REMOTE (stdout): %s' % (buf), stdout=verbose)
                        else:
                            out_EOF = True

                    elif fd is p.stderr:
                        if buf:
                            log('  REMOTE (stderr): %s' % (buf), stdout=verbose)
                        else:
                            err_EOF = True

                if out_EOF and err_EOF:
                    break

            retcode = p.wait()

            if retcode:
                raise NonZeroExitStatus(cmd)

    return


def replicate(fs, remhost, remfs, options):
    # Validate source dataset
    try:
        fs_dset = Dataset(fs)
    except NoSuchDatasetError:
        log('No such local dataset: %s' % (fs), stderr=True)
        sys.exit(1)
    if len(fs_dset.snapshots) == 0:
        log('No snapshots for local dataset: %s' % (fs), stderr=True)
        sys.exit(1)

    local_snaps = fs_dset.snapshots
    local_snaps.reverse()
    if options.end:
        try:
            end_snap = Dataset('%s@%s' % (fs, options.end))
        except NoSuchDatasetError:
            log('No such remote dataset: %s@%s' % (fs, options.end),
                stderr=True)
            sys.exit(1)
        local_text = '(specified): '
    else:
        end_snap = local_snaps[0]
        staging_snap = local_snaps[-1]
        local_text = '(newest):    '


    # Validate destination filesystem
    remfs_dset = None
    try:
        remfs_dset = Dataset(remfs, 'ssh %s' % (remhost))
    except NoSuchDatasetError:
        try:
            rslash = remfs.rfind('/')
            if rslash < 0:
                raise NoSuchDatasetError
            Dataset(remfs[:rslash], 'ssh %s' % (remhost))
        except NoSuchDatasetError:
            log('Remote dataset or parent does not exist: %s' % (remfs),
                stderr=True)
            sys.exit(1)
    except NonZeroExitStatus, e:
        log('Error running ZFS command: %s' % (e), stderr=True)
        sys.exit(1)

    # Find the newest remote snapshot that also exists in the local
    # list
    begin_snap = None
    rem_text = remfs
    repl_text = 'Full replication, no snapshots on remote host'
    rem_snaps_todel = []
    if remfs_dset and len(remfs_dset.snapshots) > 0:
        rem_snaps = remfs_dset.snapshots
        rem_snaps.reverse()
        for rsnap in rem_snaps:
            matches = [lsnap for lsnap in local_snaps
                       if lsnap.name.split('@')[1] == rsnap.name.split('@')[1]]
            if matches:
                begin_snap = matches[0]
                break
            rem_snaps_todel.append(rsnap)

        if not begin_snap:
            log(textwrap.fill(textwrap.dedent('''\
            No common snapshots between local and remote datasets.
            You will need to perform a full replication.  Please
            destroy the remote dataset to begin.'''), 70), stderr=True)
            sys.exit(1)
        rem_text = begin_snap.name
        repl_text = 'Incremental update since last replication'

    # Report lag and exit if we're in lag mode
    if options.truelag or options.snaplag:
        if begin_snap:
            if options.truelag:
                prefix = "truelag: "
                td = datetime.now() - begin_snap.attrs['create']
            else:
                prefix = "snaplag: "
                td = end_snap.attrs['create'] - begin_snap.attrs['create']
            minutes = int(td.seconds/60 + td.days*24*60)
            log("%s%s minutes (dset:%s, rsnap:%s, lsnap:%s)"
                % (prefix, minutes, fs,
                   begin_snap.name.split('@')[1],
                   end_snap.name.split('@')[1]),
                stdout=True)
            sys.exit(0)
        else:
            log("Dataset not yet replicated, cannot calculate lag",
                stderr=True)
            sys.exit(1)

    # Are snapshots the same on both sides?
    if begin_snap and \
           end_snap.name.split('@')[1] == begin_snap.name.split('@')[1]:
        log('Remote snapshot of %s is up-to-date, no work to do' %
            end_snap.name, stdout=(not options.cron))
        sys.exit(0)

    try:
        with zfs_recv_lockfile(remhost, remfs):
            # Print out validated parameters
            log(textwrap.dedent('''\
            === Replication parameters validated ===
            Local snapshot %s %s
            Remote dataset, host:        %s on %s
            Replication type:            %s
            Remote starting dataset:     %s\
            ''' % (local_text, end_snap.name, remfs, remhost,
                   repl_text, rem_text)), stdout=options.verbose)

            # Snapshot analysis
            if begin_snap:
                log(textwrap.dedent("""\
                === Incremental Snapshot Analysis ===
                  `-->' indicates starting snapshot on src and dest
                  `DEL' indicates remote snapshots to be destroyed
                Local snapshot candidates (newest first):"""),
                    stdout=options.verbose)
                for snap in local_snaps:
                    if snap.name.split('@')[1] == begin_snap.name.split('@')[1]:
                        log('--> %s' % (snap), stdout=options.verbose)
                        break
                    else:
                        log('    %s' % (snap), stdout=options.verbose)
                log('Remote snapshot candidates (newest first)',
                    stdout=options.verbose)
                for snap in rem_snaps:
                    if snap.name.split('@')[1] == begin_snap.name.split('@')[1]:
                        log('--> %s' % (snap), stdout=options.verbose)
                        break
                    elif snap in rem_snaps_todel:
                        log('DEL %s' % (snap), stdout=options.verbose)
                    else:
                        log('    %s' % (snap), stdout=options.verbose)

            if options.dryrun:
                log(textwrap.dedent('''\
                ==============================================
                DRY RUN - Replication commands will not be run
                ==============================================\
                '''), stdout=options.verbose)

            if rem_snaps_todel:
                log('Removing remote snapshots...', stdout=options.verbose)
                rm_remote_dsets(rem_snaps_todel, remhost, options.verbose,
                                options.dryrun)

            start_time = datetime.now()
            log('=== Starting replication for %s at %s'
                % (fs, start_time.strftime('%c')), stdout=options.verbose)

            if not begin_snap:
                # Perform a staging replication if necessary.  This is
                # only necessary when doing a full replication, since
                # we don't get -I's behavior of sending all
                # intermediate snapshots.  Here we "stage" to the
                # first snapshot and then do a regular incremental
                # outside of the if statement.
                run_repl(None, staging_snap, remhost, remfs,
                          options.verbose, options.dryrun)
                begin_snap = staging_snap

            # If begin_snap and end_snap have the same name, then the
            # above staging replication got us fully up-to-date, so
            # skip the follow-up incremental
            if begin_snap.name.split('@')[1] != end_snap.name.split('@')[1]:
                run_repl(begin_snap, end_snap, remhost, remfs,
                         options.verbose, options.dryrun)

            end_time = datetime.now()
            log('=== Replication complete for %s at %s'
                % (fs, end_time.strftime('%c')), stdout=options.verbose)

    except LockfileExists:
        log("Lock file exists for %s:%s, receive already in progress" %
            (remhost.split('.')[0], remfs), stderr=(not options.cron))
        if options.cron:
            sys.exit(0)
        sys.exit(1)


def main():
    # Set up signal handling
    try:
        os.setpgid(0, 0)
    except OSError:
        # We shouldn't stop exeuction if we can't create our own
        # process group
        pass
    for sig in [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM]:
        signal.signal(sig, sig_handler)

    from optparse import OptionParser
    usage = textwrap.dedent("""\
    Usage: %prog [-hclLnv] [-e snapshot] filesystem dest
    """)
    usage += '\n'
    usage += textwrap.fill(textwrap.dedent("""\
    Replicate filesystem to dest using zfs send and receive commands
    (unless -l or -L is specified).  `filesystem' should be a ZFS
    dataset and `dest' should be of the form:
    """), 76)
    usage += "\n\n   remote_host:filesystem\n\n"
    usage += textwrap.fill(textwrap.dedent("""\
    SSH will be used to connect to `remote_host' and receive the
    stream into the given dataset.  Activity will be logged to syslog,
    and to standard output if -v is used.
    """), 76)
    parser = OptionParser(usage=usage)
    parser.add_option('-c', '--cron', dest='cron',
                      action="store_true", default=False,
                      help="use when running via cron: suppress -v "
                           "output when no work is to be done and do "
                           "not complain about lockfile presence")
    parser.add_option('-e', '--end', dest='end',
                      help="use specified snapshot as the replication target, "
                      "rather than choosing the newest snapshot available")
    parser.add_option('-l', '--truelag', dest='truelag', action="store_true",
                      help="do not replicate filesystems; instead, report "
                      "the lag time in minutes between the newest remote "
                      "snapshot and the current time to standard output.  "
                      "Useful for monitoring systems")
    parser.add_option('-L', '--snaplag', dest='snaplag', action="store_true",
                      help="do not replicate filesystems; instead, report "
                      "the lag time in minutes between the newest local and "
                      "remote snapshots to standard output.  Useful "
                      "for monitoring systems")
    parser.add_option('-n', '--dryrun', dest='dryrun',
                      action="store_true", default=False,
                      help="show steps without running actual commands, "
                      "implies --verbose")
    parser.add_option('-v', '--verbose', dest='verbose', action="store_true",
                      default=False, help="report progress on standard "
                      "output as well as syslog")
    options, args = parser.parse_args()
    if not args:
        parser.print_help()
        sys.exit(0)
    if len(args) != 2:
        parser.error("Please supply filesystem and dest arguments")

    if options.dryrun:
        options.verbose = True
    
    fs = args[0]
    try:
        remhost, remfs = split_dest(args[1])
    except InvalidDestinationError:
        parser.error("Invalid dest specification")

    # Options look good, set up logging
    syslog.openlog("zsync", syslog.LOG_PID, syslog.LOG_LOCAL0)

    # Begin replication
    done = False
    while not done:
        try:
            replicate(fs, remhost, remfs, options)
            done = True
        except SendNoSuchDataset, e:
            log(textwrap.fill(textwrap.dedent("""\
            WARNING: Source snapshot (%s) disappeared during
            replication run.  `zfs send -R ...' creates a list of
            snapshots to send when the command is started and will
            fail mid-transfer if a snapshot disappears.  This happens
            often with hourly snapshots dropping off then end of the
            scheduled snapshot rotation.  Replication will be
            restarted--you can probably ignore this warning.
            """ % (e)), 70), stdout=options.verbose)


def sig_handler(sig, frame):
    # Ignore the signal we're about to send out to our process group,
    # since it will come back to us
    signal.signal(sig, signal.SIG_IGN)
    os.killpg(os.getpgid(0), sig)
    log(textwrap.fill(textwrap.dedent('''\
    Replication aborted.  Snapshots that have already transferred will
    remain on the destination host; only the in-progress snapshot will
    have to be transferred again.  You should be able to restart later
    where you left off.
    '''), 70), stderr=True)
    sys.exit(1)


if __name__ == "__main__":
    main()