ZFS pool comparison

zfs_pooldiff.py

# Copyright (C) 2016  Hugo Geoffroy
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from __future__ import print_function
from collections import defaultdict
from itertools import chain
from sys import argv, exit, stderr

# List of supported VDEV types
VDEVS = "mirror raidz raidz1 raidz2 raidz3 log cache spare".split()

def parse(text):
    """Parser for ZFS virtual device specifications, as given to `zpool create`.
    
    >>> parse("mirror a b c g mirror d e f h log mirror g h i")
    [(('mirror',), 0, ('a', 'b', 'c', 'g')),
     (('mirror',), 1, ('d', 'e', 'f', 'h')),
     (('log', 'mirror'), 0, ('g', 'h', 'i'))]
    """
    # This would be easier to write w
    words = text.split()
    typestack, devstack = [], []
    counts = defaultdict(int)
    for index, word in enumerate(words):
        if word not in VDEVS:                         # STAGE 1
            # not a vdev type, append to the devices stack for use in stage 2
            devstack.append(word)
        if word in VDEVS or index == len(words) - 1:  # STAGE 2
            # need to consume to the devices stack, and yield current state of type stack
            if devstack:
                _typestack = tuple(typestack)
                groupindex = counts[_typestack]
                counts[_typestack] += 1
                yield (_typestack, groupindex), set(devstack)
                devstack = []
            else:
                assert not typestack or typestack[-1] == "log", "unexpected device identifier"
        if word in VDEVS:                             # STAGE 3
            # append to typestack, following zpool vdev grammar
            if word != "mirror":
                if typestack and typestack[-1] != "log":
                    typestack = []
                typestack.append(word)
            elif not typestack or typestack[-1] != "mirror":
                typestack.append(word)

def compare(name, source, target,
            attach_opts='',
            detach_opts='',
            add_opts='',
            remove_opts='',
            replace_opts='',):
    """Compare two ZFS virtual device specifications, generating the commands required to apply the changes.
    
    >>> current = "X mirror a b log mirror Y Z"
    >>> target = "mirror a b c d mirror  e f g h log mirror i j k spare m n o cache p q r s"
    >>> '\n'.join(compare("pool", current, target))
    zpool remove  pool X
    zpool add  pool mirror e g h f
    zpool add  pool spare o m n
    zpool add  pool cache s r q p
    zpool replace  pool Y
    zpool replace  pool Z
    zpool attach  pool j
    zpool attach  pool b
    zpool attach  pool b
    """
    # STAGE 1 : Parse source & target, swap same-devices group indexes
    source_spec = defaultdict(set, parse(source))
    target_spec = dict(parse(target))
    for key, vdevs in list(target_spec.items()):
        for _key, _vdevs in list(source_spec.items()):
            if _vdevs == vdevs and key != _key:
                source_spec[key], source_spec[_key] = source_spec[_key], source_spec[key]
    # STAGE 2 : Determine removed/inserted devices
    removals = []
    for key, vdevs in source_spec.items():
        for vdev in vdevs:
            if vdev not in target_spec.get(key, ()):
                removals.append((key, vdev))
    insertions = []
    for key, vdevs in target_spec.items():
        for vdev in vdevs:
            if vdev not in source_spec[key]:
                insertions.append((key, vdev))              
    # STAGE 3 : Transform matching insertions/removals in replacements
    replacements = []
    for key, removed in list(removals):
        for _key, inserted in list(insertions):
            if key == _key:
                replacements.append((removed, inserted))
                removals.remove((key, removed))
                insertions.remove((key, inserted))
                break
    # STAGE 4 : Transform mirror insertions/removals into attachments/detachments
    attachments = []
    for key, inserted in list(insertions):
        types, index = key
        if "mirror" in types:
            others = list(vdev for vdev in source_spec[key] if vdev != inserted)
            if others:
                attachments.append((inserted, others[0]))
                insertions.remove((key, inserted))
                source_spec[key].add(inserted)
    detachments = []
    for key, removed in list(removals):
        types, index = key
        if "mirror" in types:
            others = list(vdev for vdev in source_spec[key] if vdev != removed)
            if others:
                detachments.append((removed))
                removals.remove((key, removed))
                source_spec[key].remove(removed)
    # STAGE 5 : Produce zpool commands
    adds = defaultdict(list)
    for (types, index), inserted in insertions:
        adds[types, index].append(inserted)
    for (types, index), added in adds.items():
        yield "zpool add {} {} {}".format(add_opts, name, ' '.join(chain(types, added)))
    for removed, inserted in replacements:
        yield "zpool replace {} {} {} {}".format(replace_opts, name, removed, inserted)
    for attached, other in attachments:
        yield "zpool attach {} {} {} {}".format(attach_opts, name, other, attached)
    for detached in detachments:
        yield "zpool detach {} {} {}".format(detach_opts, name, detached)
    for (types, index), removed in removals:
        yield "zpool remove {} {} {}".format(remove_opts, name, removed)
        
def main():
    if len(argv) < 4:
        print("\n".join([
            "# zfs-pooldiff --- ZFS pool comparison",
            "# ",
            "# Usage: zfs-pooldiff pool_name \"old_vdevs\" \"new_vdevs\"",
        ]), file=stderr)
        exit(1)
    name, current, target = argv[-3:]
    for command in compare(name, current, target):
        print(command)

if __name__ == '__main__':
    main()