blob: f7ae14df351d66b243709bf7fad903f8b14e6a2a [file] [log] [blame]
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -08001#!/usr/bin/env python
2"""
Brian O'Connor6e192432015-02-26 15:17:23 -08003 Generate the partitions json file from the $OC* environment variables
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -08004
Madan Jampani80244e52016-05-17 20:03:43 -07005 Usage: onos-gen-partitions [output file] [node_ip ...]
Brian O'Connor6e192432015-02-26 15:17:23 -08006 If output file is not provided, the json is written to stdout.
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -08007"""
Brian O'Connor6e192432015-02-26 15:17:23 -08008
9from os import environ
10from collections import deque, OrderedDict
11import re
12import json
13import sys
Thomas Vachuskab56b9172015-11-24 11:24:23 -080014import hashlib
Brian O'Connor6e192432015-02-26 15:17:23 -080015
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080016convert = lambda text: int(text) if text.isdigit() else text.lower()
17alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
Brian O'Connor6e192432015-02-26 15:17:23 -080018
19def get_OC_vars():
20 vars = []
21 for var in environ:
22 if re.match(r"OC[0-9]+", var):
23 vars.append(var)
24 return sorted(vars, key=alphanum_key)
25
Brian O'Connorfb8a2752016-05-16 16:48:23 -070026def get_nodes(ips=None, port=9876):
Madan Jampaniec1df022015-10-13 21:23:03 -070027 node = lambda k: { 'id': k, 'ip': k, 'port': port }
Brian O'Connorfb8a2752016-05-16 16:48:23 -070028 if not ips:
29 ips = [ environ[v] for v in get_OC_vars() ]
30 return [ node(v) for v in ips ]
Brian O'Connor6e192432015-02-26 15:17:23 -080031
Madan Jampaniab7e7cd2016-01-14 14:02:32 -080032def generate_base_partition(nodes):
33 return {
34 'id': 0,
35 'members': nodes
36 }
37
38def generate_extended_partitions(nodes, k):
Brian O'Connor6e192432015-02-26 15:17:23 -080039 l = deque(nodes)
Madan Jampaniec1df022015-10-13 21:23:03 -070040 perms = []
Brian O'Connor6e192432015-02-26 15:17:23 -080041 for i in range(1, len(nodes)+1):
Madan Jampaniec1df022015-10-13 21:23:03 -070042 part = {
Madan Jampaniab7e7cd2016-01-14 14:02:32 -080043 'id': i,
Madan Jampaniec1df022015-10-13 21:23:03 -070044 'members': list(l)[:k]
45 }
46 perms.append(part)
Brian O'Connor6e192432015-02-26 15:17:23 -080047 l.rotate(-1)
Madan Jampaniec1df022015-10-13 21:23:03 -070048 return perms
Brian O'Connor6e192432015-02-26 15:17:23 -080049
50if __name__ == '__main__':
Madan Jampani80244e52016-05-17 20:03:43 -070051 nodes = get_nodes(sys.argv[2:])
Madan Jampaniab7e7cd2016-01-14 14:02:32 -080052 base_partition = generate_base_partition([v.get('id') for v in nodes])
53 extended_partitions = generate_extended_partitions([v.get('id') for v in nodes], 3)
54 partitions = []
55 partitions.append(base_partition)
56 partitions.extend(extended_partitions)
Brian O'Connorfb8a2752016-05-16 16:48:23 -070057 m = hashlib.sha256()
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080058 for node in nodes:
Brian O'Connorfb8a2752016-05-16 16:48:23 -070059 m.update(node['ip'])
60 name = int(m.hexdigest()[:8], base=16) # 32-bit int based on SHA256 digest
Madan Jampaniec1df022015-10-13 21:23:03 -070061 data = {
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080062 'name': name,
Brian O'Connor6e192432015-02-26 15:17:23 -080063 'nodes': nodes,
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080064 'partitions': partitions
Brian O'Connor6e192432015-02-26 15:17:23 -080065 }
66 output = json.dumps(data, indent=4)
67
68 if len(sys.argv) == 2:
69 filename = sys.argv[1]
70 with open(filename, 'w') as f:
71 f.write(output)
72 else:
73 print output