#!/usr/bin/python


import os.path
import os
import sys
import optparse

sched_path = '/proc/sys/kernel/sched_domain/'
cpu_dir_path = 'cpu'
sd_dir_path = 'domain'
flags_file_path = 'flags'
cpu_sys_path = '/sys/devices/system/cpu/'

verbose = False

license = "tunesd changes the scheduling domain flags to modify the schedulers balancing behaviour\n"\
"Copyright (C) 2010  Markus Pargmann (mpargman <AT> allfex.org)\n\n"\
\
"This program is free software: you can redistribute it and/or modify\n"\
"it under the terms of the GNU General Public License as published by\n"\
"the Free Software Foundation, either version 3 of the License, or\n"\
"(at your option) any later version.\n"\
\
"This program is distributed in the hope that it will be useful,\n"\
"but WITHOUT ANY WARRANTY; without even the implied warranty of\n"\
"MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n"\
"GNU General Public License for more details.\n\n"\
\
"You should have received a copy of the GNU General Public License\n"\
"along with this program.  If not, see <http://www.gnu.org/licenses/>."
usage_descr = "usage: %prog [options] [[<CPU>]:[<DOMAIN>]:[<FLAG>]=-1/0/1 [...]]"
version = "kernels 2.6.38"
description = "tunesd changes the scheduling domain flags to modify the schedulers balancing behaviour. Without an argument the program will print the domain hierarchy and the flags. If CPU or DOMAIN are empty, the program assumes all CPUs or DOMAINs are meant. FLAG is one of the flags listed when executing the program with the -f option. 0 means clear the flag, 1 means set the flag. -1 means clear all flags, be careful with this option, it can have a huge impact on the performance of your system."
flaglist = \
"SD_LOAD_BALANCE		0x0001	/* Do load balancing on this domain. */\n"\
"SD_BALANCE_NEWIDLE	0x0002	/* Balance when about to become idle */\n"\
"SD_BALANCE_EXEC		0x0004	/* Balance on exec */\n"\
"SD_BALANCE_FORK		0x0008	/* Balance on fork, clone */\n"\
"SD_BALANCE_WAKE		0x0010  /* Balance on wakeup */\n"\
"SD_WAKE_AFFINE		0x0020	/* Wake task to waking CPU */\n"\
"SD_PREFER_LOCAL		0x0040  /* Prefer to keep tasks local to this domain */\n"\
"SD_SHARE_CPUPOWER	0x0080	/* Domain members share cpu power */\n"\
"SD_POWERSAVINGS_BALANCE	0x0100	/* Balance for power savings */\n"\
"SD_SHARE_PKG_RESOURCES	0x0200	/* Domain members share cpu pkg resources */\n"\
"SD_SERIALIZE		0x0400	/* Only a single load balancing instance */\n"\
"SD_ASYM_PACKING		0x0800  /* Place busy groups earlier in the domain */\n"\
"SD_PREFER_SIBLING	0x1000	/* Prefer to place tasks in a sibling domain */"

flag_trans = {}
flag_trans['SD_LOAD_BALANCE'] = 0x0001
flag_trans['SD_BALANCE_NEWIDLE'] = 0x0002
flag_trans['SD_BALANCE_EXEC'] = 0x0004
flag_trans['SD_BALANCE_FORK'] = 0x0008
flag_trans['SD_BALANCE_WAKE'] = 0x0010
flag_trans['SD_WAKE_AFFINE'] = 0x0020
flag_trans['SD_PREFER_LOCAL'] = 0x0040
flag_trans['SD_SHARE_CPUPOWER'] = 0x0080
flag_trans['SD_POWERSAVINGS_BALANCE'] = 0x0100
flag_trans['SD_SHARE_PKG_RESOURCES'] = 0x0200
flag_trans['SD_SERIALIZE'] = 0x0400
flag_trans['SD_ASYM_PACKING'] = 0x0800
flag_trans['SD_PREFER_SIBLING'] = 0x1000

def vb_print(a):
	if verbose:
		print(a)

def cpu_dir(cpu):
	return sched_path + cpu_dir_path + str(cpu) + '/'

def sd_dir(cpu, sd):
	return cpu_dir(cpu) + sd_dir_path + str(sd) + '/'

def flags_file(cpu, sd):
	return sd_dir(cpu, sd) + flags_file_path

def check_support():
	if os.path.isdir(sd_dir(0,0)):
		return True
	return False

def read_flags(cpu, sd):
	f = open(flags_file(cpu, sd), 'r')
	result = int(f.read().strip())
	f.close()
	return result

def write_flags(cpu, sd, flags):
	f = open(flags_file(cpu, sd), 'w')
	f.write(str(flags))
	f.close()

def set_flags(cpu, sd, flags):
	val = read_flags(cpu, sd)
	val = val | flags
	write_flags(cpu, sd, val)

def unset_flags(cpu, sd, flags):
	val = read_flags(cpu, sd)
	val = val & (~flags)
	write_flags(cpu, sd, val)

def clear_flags(cpu, sd):
	write_flags(cpu, sd, 0)

def cpu_request(cpu, sd, flags, operation):
	if sd == -1:
		for i in os.listdir(cpu_dir(cpu)):
			sdid = int(i[len(sd_dir_path):])
			if operation == 1:
				set_flags(cpu, sdid, flags)
			elif operation == 0:
				unset_flags(cpu, sdid, flags)
			else:
				clear_flags(cpu, sdid)
	else:
		if operation == 1:
			set_flags(cpu, sd, flags)
		elif operation == 0:
			unset_flags(cpu, sd, flags)
		else:
			clear_flags(cpu, sd)


def process_request(cpu, sd, flags, operation):
	if cpu == -1:
		for i in os.listdir(sched_path):
			cpuid = int(i[len(cpu_dir_path):])
			cpu_request(cpuid, sd, flags, operation)
	else:
		cpu_request(cpu, sd, flags, operation)

def translate_flag(flagstr):
	return flag_trans[flagstr]

def get_flagstr(flags):
	s = []
	for k,v in flag_trans.items():
		if flags & v != 0:
			s.append(k)
	return s

def parse_args(args):
	for i in args:
		ids = i.split(':')
		if len(ids) != 3:
			print("Error: Argument '" + i + "' is not in the format [<CPU>]:[<SD>]:<FLAG>=<OPERATION> . Continuing work")
			continue
		cpuid = -1
		sdid = -1
		if ids[0] != '':
			cpuid = int(ids[0])
		if ids[1] != '':
			sdid = int(ids[1])
		op = ids[2].split('=')
		flag = translate_flag(op[0])
		if len(op) != 2 or (int(op[1]) != 0 and int(op[1]) != 1 and int(op[1]) != -1):
			print("Error: Argument '" + i + "' is not in the format [<CPU>]:[<SD>]:<FLAG>=<OPERATION> . Continuing work")
			continue
		process_request(cpuid, sdid, flag, int(op[1]))

class domain:
	def __init__(self, cpus = None):
		self.childs = []
		if cpus == None:
			self.cpus = set()
		else:
			self.cpus = cpus

def append_domain(tree, domain):
	for i in tree.childs:
		if domain.cpus.issubset(i.cpus):
			append_domain(i, domain)
			return
	tree.childs.append(domain)

def parse_cpulist(filepath):
	vb_print('parsing file ' + filepath)
	f = open(filepath, 'r')
	s = f.read().strip()
	f.close()
	cpus = []
	for i in s.split(','):
		if '-' in i:
			r = i.split('-')
			cpus.extend(range(int(r[0]), int(r[1]) + 1))
		else:
			cpus.append(int(i))
	vb_print('read cpu list: ' + s)
	vb_print('parsed to be: ' + ','.join(str(i) for i in cpus))
	return cpus

def create_domains():
	domains = []
	sibl_done = set()
	core_done = set()
	for cpu in os.listdir(cpu_sys_path):
		if cpu == 'cpufreq' or cpu == 'cpuidle' or os.path.isfile(cpu_sys_path + cpu) or cpu[:3] != 'cpu':
			continue
		cpuid = int(cpu[3:])
		if cpuid not in sibl_done:
			siblings = parse_cpulist(cpu_sys_path + cpu + '/topology/thread_siblings_list')
			if len(siblings) != 1:
				domains.append(domain(set(siblings)))
			for i in siblings:
				sibl_done.add(i)
		if cpuid not in core_done:
			cores = parse_cpulist(cpu_sys_path + cpu + '/topology/core_siblings_list')
			if len(cores) != 1:
				domains.append(domain(set(cores)))
			for i in cores:
				core_done.add(i)
		vb_print('sibl_done: ' + ','.join(str(i) for i in sibl_done))
		vb_print('core_done: ' + ','.join(str(i) for i in core_done))

	vb_print('found ' + str(len(domains)) + ' different domains')
	tree = domain(sibl_done)

	for i in range(0, len(domains)):
		max_k = -1
		max_v = -1
		for k in range(0, len(domains)):
			if len(domains[k].cpus) > max_v:
				max_k = k
				max_v = len(domains[k].cpus)
		append_domain(tree, domains[max_k])
		del domains[max_k]

	if len(tree.childs) == 1:
		tree = tree.childs[0]
	return tree

def get_lvl(lvl):
	if lvl == 0:
		return ''
	s = ''
	for i in range(0, 2*lvl - 1):
		s += "    "
	s += "`---"
	return s

def dump_tree(tree, lvl):
	if len(tree.childs) == 0:
		print(get_lvl(lvl) + 'Domain CPUs:' + str(tree.cpus))
		return
	print(get_lvl(lvl) + 'Domain')
	for i in tree.childs:
		dump_tree(i, lvl + 1)

def dump_flags():
	for cpu in os.listdir(sched_path):
		cpuid = int(cpu[len(cpu_dir_path):])
		print(cpu)
		for sd in os.listdir(cpu_dir(cpuid)):
			sdid = int(sd[len(sd_dir_path):])
			flags = read_flags(cpuid, sdid)
			flagstr = get_flagstr(flags)
			print("\t" + sd)
			for flag in flagstr:
				print("\t\t" + flag)

def dump_hierarchy():
	tree = create_domains()
	print("")
	print("Domain Hierarchy:")
	dump_tree(tree, 0)
	print("")
	print("Domain flags:")
	dump_flags()
	return

if not check_support():
	print("ERROR: There are no scheduling domains. Did you enable kernel support? CONFIG_SCHED_DEBUG and CONFIG_SYSCTL must be set. Also make sure this is no single CPU system")
	sys.exit(1)


opts = optparse.OptionParser(description=description, usage = usage_descr, version = version)

opts.add_option('-v', '--verbose', action='store_true', dest='verbose', default=False, help='Verbose output-mode')
opts.add_option('-l', '--license', action='store_true', dest='license', default=False, help='License information')
opts.add_option('-f', '--flags', action='store_true', dest='listflags', default=False, help='List available flags')
(options, arguments) = opts.parse_args()

if options.license:
	print(license)
	sys.exit(0)
if options.listflags:
	print(flaglist)
	sys.exit(0)

verbose = options.verbose

if len(arguments) != 0:
	parse_args(arguments)

dump_hierarchy()