blob: f5184f292503462b8c00c443435d75a57eef34fd [file] [log] [blame]
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -08001#!/usr/bin/env python
2"""
Thiago Santos720a1622016-09-20 14:03:34 -03003usage: onos-gen-partitions [-h] [-s PARTITION_SIZE] [-n NUM_PARTITIONS]
4 [filename] [node_ip [node_ip ...]]
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -08005
Thiago Santos720a1622016-09-20 14:03:34 -03006Generate the partitions json file given a list of IPs or from the $OC*
7environment variables.
8
9positional arguments:
10 filename File to write output to. If none is provided, output
11 is written to stdout.
12 node_ip IP Address(es) of the node(s) in the cluster. If no
13 IPs are given, will use the $OC* environment
14 variables. NOTE: these arguemnts are only processed
15 after the filename argument.
16
17optional arguments:
18 -h, --help show this help message and exit
19 -s PARTITION_SIZE, --partition-size PARTITION_SIZE
20 Number of nodes per partition. Note that partition
21 sizes smaller than 3 are not fault tolerant. Defaults
22 to 3.
23 -n NUM_PARTITIONS, --num-partitions NUM_PARTITIONS
24 Number of partitions. Defaults to the number of nodes
25 in the cluster.
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080026"""
Brian O'Connor6e192432015-02-26 15:17:23 -080027
28from os import environ
Thiago Santos720a1622016-09-20 14:03:34 -030029from collections import deque
30import argparse
Brian O'Connor6e192432015-02-26 15:17:23 -080031import re
32import json
Thomas Vachuskab56b9172015-11-24 11:24:23 -080033import hashlib
Brian O'Connor6e192432015-02-26 15:17:23 -080034
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080035convert = lambda text: int(text) if text.isdigit() else text.lower()
36alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
Brian O'Connor6e192432015-02-26 15:17:23 -080037
38def get_OC_vars():
39 vars = []
40 for var in environ:
41 if re.match(r"OC[0-9]+", var):
42 vars.append(var)
43 return sorted(vars, key=alphanum_key)
44
Brian O'Connorfb8a2752016-05-16 16:48:23 -070045def get_nodes(ips=None, port=9876):
Madan Jampaniec1df022015-10-13 21:23:03 -070046 node = lambda k: { 'id': k, 'ip': k, 'port': port }
Brian O'Connorfb8a2752016-05-16 16:48:23 -070047 if not ips:
48 ips = [ environ[v] for v in get_OC_vars() ]
49 return [ node(v) for v in ips ]
Brian O'Connor6e192432015-02-26 15:17:23 -080050
Thiago Santos720a1622016-09-20 14:03:34 -030051def generate_partitions(nodes, k, n):
Brian O'Connor6e192432015-02-26 15:17:23 -080052 l = deque(nodes)
Madan Jampaniec1df022015-10-13 21:23:03 -070053 perms = []
Thiago Santos720a1622016-09-20 14:03:34 -030054 for i in range(1, n+1):
Madan Jampaniec1df022015-10-13 21:23:03 -070055 part = {
Madan Jampaniab7e7cd2016-01-14 14:02:32 -080056 'id': i,
Madan Jampaniec1df022015-10-13 21:23:03 -070057 'members': list(l)[:k]
58 }
59 perms.append(part)
Brian O'Connor6e192432015-02-26 15:17:23 -080060 l.rotate(-1)
Madan Jampaniec1df022015-10-13 21:23:03 -070061 return perms
Brian O'Connor6e192432015-02-26 15:17:23 -080062
63if __name__ == '__main__':
Thiago Santos720a1622016-09-20 14:03:34 -030064 parser = argparse.ArgumentParser(
65 description="Generate the partitions json file given a list of IPs or from the $OC* environment variables.")
66 parser.add_argument(
67 '-s', '--partition-size', type=int, default=3,
68 help="Number of nodes per partition. Note that partition sizes smaller than 3 are not fault tolerant. Defaults to 3." )
69 parser.add_argument(
70 '-n', '--num-partitions', type=int,
71 help="Number of partitions. Defaults to the number of nodes in the cluster." )
72 # TODO: make filename and nodes independent. This will break backwards compatibility with existing usage.
73 parser.add_argument(
74 'filename', metavar='filename', type=str, nargs='?',
75 help='File to write output to. If none is provided, output is written to stdout.')
76 parser.add_argument(
77 'nodes', metavar='node_ip', type=str, nargs='*',
78 help='IP Address(es) of the node(s) in the cluster. If no IPs are given, ' +
79 'will use the $OC* environment variables. NOTE: these arguemnts' +
80 ' are only processed after the filename argument.')
81
82 args = parser.parse_args()
83 filename = args.filename
84 partition_size = args.partition_size
85 nodes = get_nodes(args.nodes)
86 num_partitions = args.num_partitions
87 if not num_partitions:
88 num_partitions = len(nodes)
89
90 partitions = generate_partitions([v.get('id') for v in nodes], partition_size, num_partitions)
Brian O'Connorfb8a2752016-05-16 16:48:23 -070091 m = hashlib.sha256()
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080092 for node in nodes:
Brian O'Connorfb8a2752016-05-16 16:48:23 -070093 m.update(node['ip'])
94 name = int(m.hexdigest()[:8], base=16) # 32-bit int based on SHA256 digest
Madan Jampaniec1df022015-10-13 21:23:03 -070095 data = {
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080096 'name': name,
Brian O'Connor6e192432015-02-26 15:17:23 -080097 'nodes': nodes,
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080098 'partitions': partitions
Brian O'Connor6e192432015-02-26 15:17:23 -080099 }
100 output = json.dumps(data, indent=4)
101
Thiago Santos720a1622016-09-20 14:03:34 -0300102 if filename:
Brian O'Connor6e192432015-02-26 15:17:23 -0800103 with open(filename, 'w') as f:
104 f.write(output)
105 else:
106 print output