Check for distinct replicas on riak cluster nodes for a given n_val

riak-vnode-check.py

#! /usr/bin/env python
"""
Check for replicas on distinct nodes within a riak cluster for a given n_val

https://gist.github.com/millerdev/b6f50640d5fc2519ae8c9c2de3c64af2
"""
import re
from argparse import ArgumentParser
from subprocess import check_output

def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        '-n', '--n_val', metavar='N', type=int, default=3,
        help='n_val (default: 3)',
    )
    n_val = parser.parse_args().n_val
    ring_size = get_ring_size()
    print("n_val: {}".format(n_val))
    print("ring size: {}".format(ring_size))

    ref_partitions = list(range(ring_size)) * 2
    all_partitions = set()
    for node in get_nodes():
        partitions = get_partitions(node)
        indistinct = get_indistinct_replicas(partitions, n_val, ref_partitions)
        if indistinct:
            print("{} replicas are not distinct: {}".format(node, indistinct))
        else:
            print("{} looks good".format(node))
        all_partitions.update(partitions)

    partition_set = set(ref_partitions)
    if all_partitions != partition_set:
        unexpected = all_partitions - partition_set
        if unexpected:
            print("unexpected partitions: {}".format(list(unexpected)))
        notfound = partition_set - all_partitions
        if notfound:
            print("partitions not found: {}".format(list(notfound)))

def get_ring_size():
    output = check_output(["riak-admin", "cluster", "partition-count"])
    match = re.match(r"Cluster-wide partition-count: (\d+)", output)
    if not match:
        raise RuntimeError("bad partition count: {}".format(output))
    return int(match.group(1))

def get_nodes():
    node_name = r"([\w_.\-]+@[\w_.\-]+) +\|"
    output = check_output(["riak-admin", "cluster", "status"])
    return re.findall(node_name, output)

def get_partitions(node):
    part_expr = r"^\| primary \|[\w ]+\| *(\d+) *\|$"
    command = ["riak-admin", "cluster", "partitions", "--node", node]
    output = check_output(command)
    return [int(m) for m in re.findall(part_expr, output, re.MULTILINE)]

def get_indistinct_replicas(partitions, n_val, ref_partitions):
    indistinct = set()
    local_set = set(partitions)
    for id_ in partitions:
        for i in range(1, n_val):
            if ref_partitions[id_ + i] in local_set:
                indistinct.add(id_)
                indistinct.add(ref_partitions[id_ + i])
    return sorted(indistinct)

if __name__ == "__main__":
    main()