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)))
