zreplicate

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import optparse
import os
import time
import itertools
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from zfslib import children_first, parents_first, chronosorted
from zfslib import Dataset, Pool, Snapshot, PoolSet, ZFSConnection
from zfslib import stderr

# ========== common code =================================


def simplify(x):
    '''Take a list of tuples where each tuple is in form [v1,v2,...vn]
    and then coalesce all tuples tx and ty where tx[v1] equals ty[v2],
    preserving v3...vn of tx and discarding v3...vn of ty.

    m = [
    (1,2,"one"),
    (2,3,"two"),
    (3,4,"three"),
    (8,9,"three"),
    (4,5,"four"),
    (6,8,"blah"),
    ]
    simplify(x) -> [[1, 5, 'one'], [6, 9, 'blah']]
    '''
    y = list(x)
    if len(x) < 2:
        return y
    for idx, o in enumerate(list(y)):
        for idx2, p in enumerate(list(y)):
            if idx == idx2:
                continue
            if o and p and o[0] == p[1]:
                y[idx] = None
                y[idx2] = list(p)
                y[idx2][0] = p[0]
                y[idx2][1] = o[1]
    return [n for n in y if n is not None]


def uniq(seq, idfun=None):
    '''Makes a sequence 'unique' in the style of UNIX command uniq'''
    # order preserving
    if idfun is None:
        def idfun(x):
            return x
    seen = {}
    result = []
    for item in seq:
        marker = idfun(item)
        # in old Python versions:
        # if seen.has_key(marker)
        # but in new ones:
        if marker in seen:
            continue
        seen[marker] = 1
        result.append(item)
    return result

#===================== configuration =====================

parser = optparse.OptionParser(
    "usage: %prog [-onv] [-b BUFSIZE] [-l RATELIMIT] <srcdatasetname> <dstdatasetname>")
parser.add_option(
    '-o', '--progress', action='store_true', dest='progress', default=False,
    help='show progress (depends on the executabilty of the \'bar\' a.k.a. \'clpbar\' program, or the \'pv\' program) (default: %default)')
parser.add_option(
    '-l', '--rate-limit', action='store', dest='ratelimit', default=-1, type="int",
    help='rate limit in bytes per second (requires --progress) (default: %default which means no limit)')
parser.add_option('-n', '--dry-run', action='store_true', dest='dryrun',
                  default=False, help='don\'t actually manipulate any file systems')
parser.add_option(
    '-b', '--bufsize', action='store', dest='bufsize', default=-1,
    help='buffer size in bytes for network operations (default: %default i.e. OS default)')
parser.add_option(
    '-C', '--force-compression', action='store_true', dest='force_compression', default=False,
    help='forcibly enable gzip compression for transfers via SSH (default: %default i.e. obey SSH configuration)')
parser.add_option('-v', '--verbose', action='store_true',
                  dest='verbose', default=False, help='be verbose (default: %default)')
parser.add_option(
    '-t', '--trust', action='store_true', dest='trust', default=False,
    help='automatically trust unknown or mismatched SSH keys of remote hosts (default: %default)')
opts, args = parser.parse_args(sys.argv[1:])

try:
    bufsize = int(opts.bufsize)
    assert bufsize >= 16384 or bufsize == -1
except (ValueError, AssertionError), e:
    parser.error("bufsize must be an integer greater than 16384")

if len(args) == 2:
    try:
        source_host, source_dataset_name = args[0].split(":", 1)
    except ValueError:
        source_host, source_dataset_name = "localhost", args[0]
    try:
        destination_host, destination_dataset_name = args[1].split(":", 1)
    except ValueError:
        destination_host, destination_dataset_name = "localhost", args[1]
else:
    parser.error("arguments are wrong")

if opts.ratelimit != -1:
    if opts.ratelimit < 1024:
        parser.error("rate limit (%s) must be higher than 1024 bytes per second" %
                     opts.ratelimit)
    if not opts.progress:
        parser.error("to apply a rate limit, you must specify --progress too")


def verbose_stderr(*args, **kwargs):
    if opts.verbose:
        stderr(*args, **kwargs)

#===================== end configuration =================

# ================ start program algorithm ===================

src_conn = ZFSConnection(
    source_host, trust=opts.trust, dataset=source_dataset_name)
dst_conn = ZFSConnection(
    destination_host, trust=opts.trust, dataset=destination_dataset_name)

verbose_stderr("Replicating dataset %s:%s into %s:%s..." % (
    source_host, source_dataset_name,
    destination_host, destination_dataset_name))

verbose_stderr("Assessing that the source dataset exists... %s" %
               source_dataset_name)
try:
    source_dataset = src_conn.pools.lookup(source_dataset_name)
except KeyError:
    stderr(
        "Error: the source dataset does not exist.  Backup cannot continue.")
    sys.exit(2)


verbose_stderr("Assessing that the destination dataset exists... %s" %
               destination_dataset_name)
try:
    dst_pools = dst_conn.pools
    destination_dataset = dst_pools.lookup(destination_dataset_name)
except KeyError:
    stderr(
        "Error: the destination dataset does not exist.  Backup cannot continue.")
    sys.exit(2)

# it is time to determine which datasets need to be synced
# we walk the entire dataset structure, and sync snapshots recursively


def recursive_replicate(s, d):
    sched = []

    # we first collect all snapshot names, to later see if they are on both
    # sides, one side, or what
    all_snapshots = []
    if s:
        all_snapshots.extend(s.get_snapshots())
    if d:
        all_snapshots.extend(d.get_snapshots())
    all_snapshots = [y[1] for y in chronosorted([(
        x.get_creation(), x.name) for x in all_snapshots])]
    snapshot_pairs = []
    for snap in all_snapshots:
        try:
            ssnap = s.get_snapshot(snap)
        except (KeyError, AttributeError):
            ssnap = None
        try:
            dsnap = d.get_snapshot(snap)
        except (KeyError, AttributeError):
            dsnap = None
        # if the source snapshot exists and is not already in the table of snapshots
        # then pair it up with its destination snapshot (if it exists) or None
        # and add it to the table of snapshots
        if ssnap and not snap in [x[0].name for x in snapshot_pairs]:
            snapshot_pairs.append((ssnap, dsnap))

    # now we have a list of all snapshots, paired up by name, and in chronological order
    # (it's quadratic complexity, but who cares)
    # now we need to find the snapshot pair that happens to be the the most
    # recent common pair
    found_common_pair = False
    for idx, (m, n) in enumerate(snapshot_pairs):
        if m and n and m.name == n.name:
            found_common_pair = idx

    # we have combed through the snapshot pairs
    # time to check what the latest common pair is
    if not s.get_snapshots():
        # well, no snapshots in source, do nothing
        pass
    elif found_common_pair is False:
        # no snapshot is in common, problem!
        # theoretically destroying destination dataset and resyncing it recursively would work
        # but this requires work in the optimizer that comes later
        sched.append(("full", s, d, None, snapshot_pairs[-1][0]))
    elif found_common_pair == len(snapshot_pairs) - 1:
        # the latest snapshot of both datasets that is common to both, is the latest snapshot in the source
        # we have nothing to do here because the datasets are "in sync"
        pass
    else:
        # the source dataset has more recent snapshots, not present in the destination dataset
        # we need to transfer those
        snapshots_to_transfer = [x[
            0] for x in snapshot_pairs[found_common_pair:]]
        for n, x in enumerate(snapshots_to_transfer):
            if n == 0:
                continue
            sched.append(("incremental", s, d, snapshots_to_transfer[n-1], x))

    # now let's apply the same argument to the children
    children_sched = []
    for c in [x for x in s.children if not isinstance(x, Snapshot)]:
        try:
            cd = d.get_child(c.name)
        except (KeyError, AttributeError):
            cd = None
        children_sched.extend(recursive_replicate(c, cd))

    # and return our schedule of operations to the parent
    return sched + children_sched

operation_schedule = recursive_replicate(source_dataset, destination_dataset)


def optimize(operation_schedule):
    # now let's optimize the operation schedule
    # the optimization is quite basic
    # step 1: if a snapshot is scheduled to be synced, and its parent has the same snapshot to be synced, we skip it
    # step 2: coalesce operations

    # this optimizer removes superfluous operations and consolidates the
    # remaining ones

    # pass 0: (unimplemented)
    # if any incremental transfers are scheduled, but full transfers are scheduled for their parents
    # then drop them
    # this should not happen so far because the recursive_replicate
    # function should never generate that combination of operations

    # pass 1:
    # group operations based on the source and destination snapshot names
    # look up each operation in its corresponding group
    # if the parent dataset is listed in the corresponding group
    # drop it like a mad beat
    def pass1(sched):

        # first step is to build a dictionary of operations
        # grouped by source snapshot name and destination snapshot name
        operations_by_srcsnap_to_dstsnap = {}
        operations_by_dstsnap = {}
        for op, src, dst, srcs, dsts in filter(lambda x: x[0] == 'incremental', sched):
            srcname = src.get_path()
            opsummary = srcs.name + "->" + dsts.name
            if opsummary not in operations_by_srcsnap_to_dstsnap:
                operations_by_srcsnap_to_dstsnap[opsummary] = []
            operations_by_srcsnap_to_dstsnap[
                opsummary].append((op, src, dst, srcs, dsts))
        for op, src, dst, srcs, dsts in sched:
            srcname = src.get_path()
            opsummary = dsts.name
            if opsummary not in operations_by_dstsnap:
                operations_by_dstsnap[opsummary] = []
            operations_by_dstsnap[opsummary].append((op, src, dst, srcs, dsts))

        # second step is to iterate through the operations and look them up in
        # their group
        optimized = []
        for op, src, dst, srcs, dsts in sched:
            if op == "incremental":
                opsummary = srcs.name + "->" + dsts.name
                commonprefix = os.path.commonprefix(
                    [
                        x[1].get_path() for x in operations_by_srcsnap_to_dstsnap[opsummary]
                        if x[1].get_path().startswith(src.get_path())
                        or src.get_path().startswith(x[1].get_path())
                    ]
                )
                if len(commonprefix) < len(src.get_path()):
                    # the parent is scheduled to experience the same operation
                    # as the one evaluated right now
                    continue  # then ignore it and go straight to evaluating the next structure
            elif op == "full":
                opsummary = dsts.name
                commonprefix = os.path.commonprefix(
                    [
                        x[1].get_path() for x in operations_by_dstsnap[opsummary]
                        if x[1].get_path().startswith(src.get_path())
                        or src.get_path().startswith(x[1].get_path())
                    ]
                )
                if len(commonprefix) < len(src.get_path()):
                    # the parent is scheduled to send the snapshot that is
                    # being evaluated right now
                    continue  # then ignore it and go straight to evaluating the next structure
            else:
                assert 0, "Not reached"
            # print "Operation",op,src,opsummary,"accepted"
            optimized.append((op, src, dst, srcs, dsts))

        return optimized

    # pass 2:
    # coalesce operations based on contiguousness
    def pass2(sched):

        # first step is to build a dictionary of operations
        # grouped by source dataset
        operations_by_src = {}
        for op, src, dst, srcs, dsts in sched:
            if src not in operations_by_src:
                operations_by_src[src] = []
            operations_by_src[src].append((op, src, dst, srcs, dsts))

        # second step is to iterate through it
        # this requires that the operations in sched be sorted
        for src, operations in operations_by_src.items():
            # transpose operations so srcs and dsts are the first two
            operations = [(
                srcs, dsts, op, src, dst) for op, src, dst, srcs, dsts in operations]
            operations = simplify(operations)
            operations = [(
                op, src, dst, srcs, dsts) for srcs, dsts, op, src, dst in operations]
            operations_by_src[src] = operations

        # flatten dictionary back into a list, preserving the order of the
        # original schedule
        datasets = uniq([src for op, src, dst, srcs, dsts in sched])
        sched = [v for x in datasets for v in operations_by_src[x]]
        return sched

    # pass 3:
    # sort operations so incremental child operations happen before
    # incremental parent operations
    def pass3(sched):
        return sorted(sched, key=lambda k: -(k[1].get_path().count('/')))

    operation_schedule = pass1(operation_schedule)
    operation_schedule = pass2(operation_schedule)
    operation_schedule = pass3(operation_schedule)
    return operation_schedule


def recursive_cleanup(s, d):
    all_snapshots = []

    if s:
        all_snapshots.extend(s.get_snapshots())
    if d:
        all_snapshots.extend(d.get_snapshots())

    all_snapshots = [y[1] for y in chronosorted([(
        x.get_creation(), x.name) for x in all_snapshots])]

    snapshot_pairs = []
    for snap in all_snapshots:
        try:
            ssnap = s.get_snapshot(snap)
        except (KeyError, AttributeError):
            ssnap = None
        try:
            dsnap = d.get_snapshot(snap)
        except (KeyError, AttributeError):
            dsnap = None
        # if the source snapshot exists and is not already in the table of snapshots
        # then pair it up with its destination snapshot (if it exists) or None
        # and add it to the table of snapshots
        if ssnap and not snap in [x[0].name for x in snapshot_pairs]:
            snapshot_pairs.append((ssnap, dsnap))

    found_common_pair = False
    for idx, (m, n) in enumerate(snapshot_pairs):
        if m and n and m.name == n.name:
            found_common_pair = idx

    if found_common_pair:
        print "Found common pair:"
        print snapshot_pairs[found_common_pair]
    else:
        print "No common pair.  Need to do full sync."

    src_snaps = []
    dest_snaps = []
    if s:
        src_snaps = s.get_snapshots()
    if d:
        dest_snaps = d.get_snapshots()

    if src_snaps and dest_snaps:
        junk = set([n.name for n in dest_snaps]) - set(
            [n.name for n in src_snaps])

        for j in junk:
            try:
                jdsnap = d.get_snapshot(j)
            except:
                pass
            if jdsnap:
                verbose_stderr("Deleting snapshot on destination: %s" %
                               jdsnap.get_path())
                verbose_stderr("Result %s" % dst_conn.destroy(jdsnap.get_path()))

    # now let's apply the same process to the children
    for c in [x for x in s.children if not isinstance(x, Snapshot)]:
        try:
            cd = d.get_child(c.name)
        except (KeyError, AttributeError):
            cd = None
        recursive_cleanup(c, cd)


optimized_operation_schedule = optimize(operation_schedule)

send_opts = ["-R"]
receive_opts = []
if opts.dryrun:
    receive_opts.append("-n")
if opts.verbose:
    send_opts.append("-v")
    receive_opts.append("-v")

# a little debugging code right here
# import subprocess
# _oldpopen = subprocess.Popen
# def mypopen(*args,**kwargs):
    # print args[0]
    # return _oldpopen(*args,**kwargs)
# subprocess.Popen = mypopen

verbose_stderr("=================================")

# Clean that up snapshots on destination not found on source.  This is a
# bug in zfs/send receive where it cannot destroy or rename on the
# destination."
verbose_stderr("Cleaning up snapshots on destination")
recursive_cleanup(source_dataset, destination_dataset)
verbose_stderr("Cleaning up snapshots on destination")

verbose_stderr("=================================")

for op, src, dst, srcs, dsts in optimized_operation_schedule:
    # assert 0, (op,src,dst,srcs,dsts)

    source_dataset_path = src.get_path()
    if srcs:
        this_send_opts = ["-I", "@" + srcs.name]
    else:
        this_send_opts = []

    if dst:
        destination_dataset_path = dst.get_path()
    else:
        commonbase = os.path.commonprefix(
            [src.get_path(), source_dataset.get_path()])
        remainder = src.get_path()[len(commonbase):]
        destination_dataset_path = destination_dataset.get_path() + remainder
    destination_snapshot_path = dsts.get_path()

    verbose_stderr("Recursively replicating %s to %s" %
                   (source_dataset_path, destination_dataset_path))
    verbose_stderr("Base snapshot available in destination: %s" % srcs)
    verbose_stderr("Target snapshot available in source:    %s" % dsts)

    # finagle the send() part of transfer by
    # manually selecting incremental sending of intermediate snapshots
    # rather than sending of differential snapshots
    # which would happen if we used the fromsnapshot= parameter to transfer()

    src_conn.transfer(
        dst_conn,
        destination_snapshot_path,
        destination_dataset_path,
        showprogress=opts.progress,
        compression=opts.force_compression,
        ratelimit=opts.ratelimit,
        bufsize=bufsize,
        send_opts=send_opts + this_send_opts,
        receive_opts=receive_opts
    )
    verbose_stderr("=================================")

verbose_stderr("Replication complete.")