blob: ae7da6ced19cc5dc4af809b2db6c112f4f31ebc5 [file] [log] [blame]
#!/usr/bin/env python
"""
Generate the partitions json file from the $OC* environment variables
Usage: onos-gen-partitions output_file [num_nodes] [-e]
If output file is not provided, the json is written to stdout.
"""
from os import environ
from collections import deque, OrderedDict
import re
import json
import sys
import hashlib
convert = lambda text: int(text) if text.isdigit() else text.lower()
alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
def get_OC_vars():
vars = []
for var in environ:
if re.match(r"OC[0-9]+", var):
vars.append(var)
return sorted(vars, key=alphanum_key)
def get_nodes(vars, port=9876):
node = lambda k: { 'id': k, 'ip': k, 'port': port }
return [ node(environ[v]) for v in vars ]
def generate_base_partition(nodes):
return {
'id': 0,
'members': nodes
}
def generate_extended_partitions_scaling(nodes, k, partitions=3, equal=False):
l = deque(nodes)
perms = []
for i in range(1, partitions + 1):
if equal:
members = list(l)
else:
members = list(l)[:k]
part = {
'id': i,
'members': members
}
perms.append(part)
l.rotate(-2)
return perms
if __name__ == '__main__':
vars = get_OC_vars()
# NOTE: likely prone to errors
nodes = get_nodes(vars)
num = None
equal = False
if len(sys.argv) >= 3:
num = int(sys.argv[2])
try:
equal = "-e" in sys.argv[3]
except:
equal = False
if num:
nodes = nodes[:num]
base_partition = generate_base_partition([v.get('id') for v in nodes])
extended_partitions = generate_extended_partitions_scaling([v.get('id') for v in nodes],
3, equal=equal)
partitions = []
partitions.append(base_partition)
partitions.extend(extended_partitions)
name = hash("HAScaling")
data = {
'name': name,
'nodes': nodes,
'partitions': partitions
}
output = json.dumps(data, indent=4)
if len(sys.argv) >= 2:
filename = sys.argv[1]
with open(filename, 'w') as f:
f.write(output)
else:
print output