blob: e69bcc178955d04d6c2e05f084bd667323ab7061 [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 Jampani597282a2016-06-04 09:23:17 -070032def generate_partitions(nodes, k):
Brian O'Connor6e192432015-02-26 15:17:23 -080033 l = deque(nodes)
Madan Jampaniec1df022015-10-13 21:23:03 -070034 perms = []
Brian O'Connor6e192432015-02-26 15:17:23 -080035 for i in range(1, len(nodes)+1):
Madan Jampaniec1df022015-10-13 21:23:03 -070036 part = {
Madan Jampaniab7e7cd2016-01-14 14:02:32 -080037 'id': i,
Madan Jampaniec1df022015-10-13 21:23:03 -070038 'members': list(l)[:k]
39 }
40 perms.append(part)
Brian O'Connor6e192432015-02-26 15:17:23 -080041 l.rotate(-1)
Madan Jampaniec1df022015-10-13 21:23:03 -070042 return perms
Brian O'Connor6e192432015-02-26 15:17:23 -080043
44if __name__ == '__main__':
Madan Jampani80244e52016-05-17 20:03:43 -070045 nodes = get_nodes(sys.argv[2:])
Madan Jampani597282a2016-06-04 09:23:17 -070046 partitions = generate_partitions([v.get('id') for v in nodes], 3)
Brian O'Connorfb8a2752016-05-16 16:48:23 -070047 m = hashlib.sha256()
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080048 for node in nodes:
Brian O'Connorfb8a2752016-05-16 16:48:23 -070049 m.update(node['ip'])
50 name = int(m.hexdigest()[:8], base=16) # 32-bit int based on SHA256 digest
Madan Jampaniec1df022015-10-13 21:23:03 -070051 data = {
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080052 'name': name,
Brian O'Connor6e192432015-02-26 15:17:23 -080053 'nodes': nodes,
Aaron Kruglikoveb0ae4e2015-11-10 19:16:16 -080054 'partitions': partitions
Brian O'Connor6e192432015-02-26 15:17:23 -080055 }
56 output = json.dumps(data, indent=4)
57
58 if len(sys.argv) == 2:
59 filename = sys.argv[1]
60 with open(filename, 'w') as f:
61 f.write(output)
62 else:
63 print output