def testSpeedControl(self): settings = Settings() counter_server = CounterServer(self.dir_, settings) self.serv = SpeedControlServer(self.dir_, settings, rpc_server=self.rpc_server, counter_server=counter_server) self.cli = SpeedControlClient(self.serv, 'localhost', 0) self.assertEqual(self.cli.require(10), (10, 0)) self.serv.set_speed(100) almost = lambda left, right: left[0] == right[0] \ and abs(left[1] - right[1]) < 0.1 self.assertTrue(almost(self.cli.require(100), (100, 0.5))) self.assertTrue(almost(self.cli.require(1), (0, 0.5))) self.serv.set_instance_speed(50) self.serv.calc_spans() self.serv.reset() self.assertTrue(almost(self.cli.require(50), (50, 1.1))) self.assertTrue(almost(self.cli.require(1), (0, 1.1))) self.serv.set_adaptive(True) addr = 'localhost#0' self.serv.counter_server.acc(addr, 'banned_start', 1100) self.serv.counter_server.acc(addr, 'banned_end', 1500) self.serv.counter_server.acc(addr, 'normal_start', 1000) self.serv.counter_server.acc(addr, 'normal_end', 1100) self.serv.counter_server.acc(addr, 'normal_pages', 200) self.serv.counter_server.inc(addr, 'pages', 500) self.serv.counter_server.inc(addr, 'secs', 1000) self.serv.calc_spans() self.serv.reset() self.assertTrue(almost(self.cli.require(24), (24, 0.5))) self.assertTrue(almost(self.cli.require(1), (0, 0.5)))
def init(self): with self.lock: if self.inited: return self.log_file = os.path.join(self.working_dir, 'job.log') self.logger = self.logger or get_logger(name='cola_task', filename=self.log_file, server=self.master_ip) for i in range(self.n_tasks): self.counter_clients[i] = CounterClient(self.counter_server, app_name=self.job_name) self.budget_clients[i] = BudgetApplyClient( self.budget_server, app_name=self.job_name) self.speed_clients[i] = SpeedControlClient( self.speed_server, self.ip, self.task_start_id + i, app_name=self.job_name) self.init_tasks() self._init_counter_sync() self._init_idle_status_checker() self.inited = True
class Test(unittest.TestCase): def setUp(self): self.port = random.randint(10000, 30000) self.addr = 'localhost:%s' % self.port self.rpc_server = ColaRPCServer(('localhost', self.port)) thd = threading.Thread(target=self.rpc_server.serve_forever) thd.setDaemon(True) thd.start() self.dir_ = tempfile.mkdtemp() def tearDown(self): try: self.rpc_server.shutdown() finally: shutil.rmtree(self.dir_) def testBudgetApply(self): self.serv = BudgetApplyServer(self.dir_, Settings(), rpc_server=self.rpc_server) self.cli1 = BudgetApplyClient(self.serv) self.cli2 = BudgetApplyClient('localhost:%s'%self.port) try: self.serv.set_budgets(90) self.assertEqual(self.cli1.apply(50), 50) self.assertEqual(self.cli2.apply(50), 40) self.cli1.finish(50) self.assertEqual(50, self.serv.finished) self.cli2.finish(50) self.assertEqual(90, self.serv.finished) self.cli1.error(10) self.assertEqual(90, self.serv.applied) self.serv.finished = 0 self.cli2.error(10) self.assertEqual(80, self.serv.applied) finally: self.serv.shutdown() def testCounter(self): self.serv = CounterServer(self.dir_, Settings(), rpc_server=self.rpc_server) self.cli1 = CounterClient(self.serv) self.cli2 = CounterClient('localhost:%s'%self.port) try: self.cli1.global_inc('pages', 10) self.cli2.global_inc('pages', 2) self.assertEqual(self.cli1.get_global_inc('pages'), 10) self.assertEqual(self.cli2.get_global_inc('pages'), 2) self.assertEqual(self.serv.inc_counter.get('global', 'pages', 0), 0) self.cli1.sync() self.cli2.sync() self.assertEqual(self.cli1.get_global_inc('pages'), None) self.assertEqual(self.cli2.get_global_inc('pages'), None) self.assertEqual(self.serv.inc_counter.get('global', 'pages'), 12) self.cli1.local_inc(self.addr, 0, 'pages', 100) self.assertEqual(self.cli1.get_local_inc(self.addr, 0, 'pages'), 100) self.cli1.sync() self.assertEqual(self.serv.inc_counter.get('%s#%s'%(self.addr, 0), 'pages'), 100) self.cli1.global_acc('normal', 100) self.cli2.global_acc('normal', 'test') self.assertEqual(self.cli1.get_global_acc('normal'), [100,]) self.assertEqual(self.cli2.get_global_acc('normal'), ['test',]) self.cli1.sync() self.cli2.sync() self.assertEqual(self.cli1.get_global_acc('normal'), None) self.assertEqual(self.cli2.get_global_acc('normal'), None) self.assertEqual(self.serv.acc_counter.get('global', 'normal'), [100, 'test']) self.cli2.local_acc(self.addr, 1, 'normal', 100) self.assertEqual(self.cli2.get_local_acc(self.addr, 1, 'normal'), [100, ]) self.cli2.sync() self.assertEqual(self.serv.acc_counter.get('%s#%s'%(self.addr, 1), 'normal'), [100, ]) finally: self.serv.shutdown() def testSpeedControl(self): settings = Settings() counter_server = CounterServer(self.dir_, settings) self.serv = SpeedControlServer(self.dir_, settings, rpc_server=self.rpc_server, counter_server=counter_server) self.cli = SpeedControlClient(self.serv, 'localhost', 0) self.assertEqual(self.cli.require(10), (10, 0)) self.serv.set_speed(100) almost = lambda left, right: left[0] == right[0] \ and abs(left[1] - right[1]) < 0.1 self.assertTrue(almost(self.cli.require(100), (100, 0.5))) self.assertTrue(almost(self.cli.require(1), (0, 0.5))) self.serv.set_instance_speed(50) self.serv.calc_spans() self.serv.reset() self.assertTrue(almost(self.cli.require(50), (50, 1.1))) self.assertTrue(almost(self.cli.require(1), (0, 1.1))) self.serv.set_adaptive(True) addr = 'localhost#0' self.serv.counter_server.acc(addr, 'banned_start', 1100) self.serv.counter_server.acc(addr, 'banned_end', 1500) self.serv.counter_server.acc(addr, 'normal_start', 1000) self.serv.counter_server.acc(addr, 'normal_end', 1100) self.serv.counter_server.acc(addr, 'normal_pages', 200) self.serv.counter_server.inc(addr, 'pages', 500) self.serv.counter_server.inc(addr, 'secs', 1000) self.serv.calc_spans() self.serv.reset() self.assertTrue(almost(self.cli.require(24), (24, 0.5))) self.assertTrue(almost(self.cli.require(1), (0, 0.5)))