From 0b370fc155741d309164aad447bd17734ca9c370 Mon Sep 17 00:00:00 2001 From: Mark Veidemanis Date: Fri, 12 Aug 2022 22:27:49 +0100 Subject: [PATCH] Improve channel allocation and write basic tests for it --- api/views.py | 2 +- modules/chankeep.py | 41 +++++++++++----- runtest.sh | 3 ++ tests/test_chankeep.py | 109 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 143 insertions(+), 12 deletions(-) create mode 100755 runtest.sh create mode 100644 tests/test_chankeep.py diff --git a/api/views.py b/api/views.py index 2af0241..d5795a2 100644 --- a/api/views.py +++ b/api/views.py @@ -522,4 +522,4 @@ class API(object): name = f"{net}{num}" if name not in main.IRCPool.keys(): return dumps({"success": False, "reason": f"relay {num} not on {net}"}) - return dumps({"nickname": main.IRCPool[name].nickname}) \ No newline at end of file + return dumps({"nickname": main.IRCPool[name].nickname}) diff --git a/modules/chankeep.py b/modules/chankeep.py index 89b8e60..6865f28 100644 --- a/modules/chankeep.py +++ b/modules/chankeep.py @@ -61,7 +61,7 @@ def getChanFree(net, new): return (chanfree, chanlimits.pop()) -def emptyChanAllocate(net, flist, relay, new): +def emptyChanAllocate(net, flist, new): chanfree = getChanFree(net, new) if not chanfree: return @@ -102,20 +102,21 @@ def emptyChanAllocate(net, flist, relay, new): flist = flist[:sum_free] debug(f"emptyChanAllocate() flist truncated to {sum_free}, length nis now {len(flist)}") trace(f"emptyChanAllocate() best effort allocation: {flist}") + newlist = list(flist) for i in chanfree[0].keys(): for x in range(chanfree[0][i]): - if not len(flist): + if not len(newlist): break if i in allocated.keys(): - allocated[i].append(flist.pop()) + allocated[i].append(newlist.pop()) else: - allocated[i] = [flist.pop()] + allocated[i] = [newlist.pop()] return allocated -def populateChans(net, clist, relay, new): +def populateChans(net, clist, new): # divided = array_split(clist, relay) - allocated = emptyChanAllocate(net, clist, relay, new) + allocated = emptyChanAllocate(net, clist, new) if not allocated: return for i in allocated.keys(): @@ -148,7 +149,7 @@ def minifyChans(net, listinfo): return listinfo -def keepChannels(net, listinfo, mean, sigrelay, relay): +def keepChannels(net, listinfo, mean, sigrelay, relay, chanlimit): listinfo = minifyChans(net, listinfo) if not listinfo: return @@ -159,23 +160,27 @@ def keepChannels(net, listinfo, mean, sigrelay, relay): if not sigrelay <= main.config["ChanKeep"]["MaxRelay"]: error("Network %s is too big to cover: %i relays required" % (net, sigrelay)) return + num_instances = len(getActiveRelays(net)) + max_chans = chanlimit * num_instances if coverAll: needed = relay - len(getActiveRelays(net)) debug(f"keepChannels() coverAll asking to provision {needed} relays for {net} relay:{relay}") newNums = modules.provision.provisionMultipleRelays(net, needed) flist = [i[0] for i in listinfo] - populateChans(net, flist, relay, newNums) + chosen = sorted(flist, reverse=True, key=lambda x: x[1])[:max_chans] + populateChans(net, chosen, newNums) else: needed = sigrelay - len(getActiveRelays(net)) debug(f"keepChannels() NOT coverAll asking to provision {needed} relays for {net} sigrelay:{sigrelay}") newNums = modules.provision.provisionMultipleRelays(net, needed) siglist = [i[0] for i in listinfo if int(i[1]) > mean] - populateChans(net, siglist, sigrelay, newNums) + chosen = sorted(siglist, reverse=True, key=lambda x: x[1])[:max_chans] + populateChans(net, chosen, newNums) notifyJoin(net) def joinSingle(net, channel): - eca = emptyChanAllocate(net, [channel], None, []) + eca = emptyChanAllocate(net, [channel], []) if not eca: return False if not len(eca.keys()) == 1: @@ -272,8 +277,22 @@ def _initialList(net, num, listinfo, chanlimit): p.execute() debug("List parsing completed on %s" % net) - keepChannels(net, listinfo, mean, sigrelay, relay) + keepChannels(net, listinfo, mean, sigrelay, relay, chanlimit) + + # return (listinfo, mean, sigrelay, relay) def initialList(net, num, listinfo, chanlimit): deferToThread(_initialList, net, num, deepcopy(listinfo), chanlimit) + + +def chankeep_handler(net, num, listinfo, chanlimit): + """ + Handle a channel keep request. + :param net: + :param num: + :param listinfo: + :param chanlimit: + :return: + """ + listinfo, mean, sigrelay, relay = _initialList(net, num, listinfo, chanlimit) \ No newline at end of file diff --git a/runtest.sh b/runtest.sh new file mode 100755 index 0000000..1ee7505 --- /dev/null +++ b/runtest.sh @@ -0,0 +1,3 @@ +#!/bin/sh +#pre-commit run -a +python -m unittest discover -s tests -p 'test_*.py' diff --git a/tests/test_chankeep.py b/tests/test_chankeep.py new file mode 100644 index 0000000..f35c653 --- /dev/null +++ b/tests/test_chankeep.py @@ -0,0 +1,109 @@ +import logging +from unittest import TestCase +from unittest.mock import MagicMock, patch +from random import randint +from modules import chankeep +from math import ceil +import heapq + +class TestChanKeep(TestCase): + def setUp(self): + self.net = "testnet" + self.num = 1 + self.chanlimit = 100 + chankeep.main.initConf() + chankeep.main.r = MagicMock() + chankeep.main.g = MagicMock() + chankeep.main.g.pipeline = MagicMock() + chankeep.main.config["ChanKeep"]["Provision"] = False + + self.listinfo = self.generate_listinfo() + self.chan_name_list = [x[0] for x in self.listinfo] + self.chan_member_list = [x[1] for x in self.listinfo] + + def generate_listinfo(self, ranges=None): + """ + Create a fake listinfo. + Where #channel has 192 users, and #channel2 has 188 users. + listinfo = [["#channel", 192], ["#channel2", 188]] + """ + if not ranges: + ranges = [[100, 5, 10], [400, 100, 200], [2, 500, 1000]] + listinfo = [] + for channum, min, max in ranges: + for i in range(channum): + chan_name = f"#num-{channum}-{i}" + chan_users = randint(min, max) + listinfo.append([chan_name, chan_users]) + return listinfo + + def percent_diff(self, a, b): + return (abs(b - a) / a) * 100.0 + + def test_alt_listinfo(self): + # We're looking for a perc of 1000-1100 + # And a sigrelay of 2 + # We only want those 10 big channels + instances = 1 + chanlimit = 5 + max_chans = instances * chanlimit + listinfo = self.generate_listinfo(ranges=[[1000, 1, 2], [200, 400, 800], [10, 1000, 2000]]) + listinfo_num = [x[1] for x in listinfo] + + listlength = len(listinfo) + cumul = 0 + try: + cumul += sum(int(i[1]) for i in listinfo) + except TypeError: + return + mean = round(cumul / listlength, 2) + siglength = 0 + insiglength = 0 + sigcumul = 0 + insigcumul = 0 + for i in listinfo: + if int(i[1]) > mean: + siglength += 1 + sigcumul += int(i[1]) + elif int(i[1]) < mean: + insiglength += 1 + insigcumul += int(i[1]) + + sigrelay = ceil(siglength / chanlimit) + relay = ceil(listlength / chanlimit) + print(f"len:{listlength} cumul:{cumul} mean:{mean} siglength:{siglength} insiglength:{insiglength} sigrelay:{sigrelay} relay:{relay} sigcumul:{sigcumul} insigcumul:{insigcumul}") + # We want a return between 1000 and 1100 + + list_insig = [x for x in listinfo_num if x < mean] + list_sig = [x for x in listinfo if x[1] > mean] + chosen = sorted(list_sig, reverse=True, key=lambda x: x[1])[:max_chans] + print("CHOSEN", chosen) + self.assertEqual(len(chosen), 5) + + @patch("modules.chankeep.keepChannels") + def test__initialList(self, keepchannels): + chankeep._initialList(self.net, self.num, self.listinfo, self.chanlimit) + net, passed_list, mean, sigrelay, relay, chanlimit = keepchannels.call_args_list[0][0] + self.assertEqual(net, self.net) + self.assertEqual(passed_list, self.listinfo) + self.assertEqual(chanlimit, self.chanlimit) + # print(net, mean, sigrelay, relay) + + @patch("modules.chankeep.getChanFree") + def test_empty_chan_allocate(self, getchanfree): + getchanfree.return_value = ({1: 600}, 600) # pretend we have 600 channels free + eca = chankeep.emptyChanAllocate(self.net, self.chan_name_list, []) + self.assertEqual(len(eca), 1) + num = list(eca.keys())[0] + chans = eca[list(eca.keys())[0]] + self.assertEqual(num, self.num) + self.assertCountEqual(chans, self.chan_name_list) + + getchanfree.return_value = ({1: 100}, 10) + eca = chankeep.emptyChanAllocate(self.net, self.chan_name_list, []) + self.assertEqual(len(eca), 1) + num = list(eca.keys())[0] + chans = eca[list(eca.keys())[0]] + self.assertEqual(num, self.num) + self.assertEqual(len(chans), 100) + #self.assertCountEqual(chans, self.chan_name_list) \ No newline at end of file