def testCPULimitForFlows(self): """This tests that the client actions are limited properly.""" result = {} client_mock = action_mocks.CPULimitClientMock(result) client_mock = flow_test_lib.MockClient(self.client_id, client_mock, token=self.token) client_mock.EnableResourceUsage(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) worker_obj = worker.GRRWorker(token=self.token) flow.GRRFlow.StartFlow(client_id=self.client_id, flow_name=flow_test_lib.CPULimitFlow.__name__, cpu_limit=1000, network_bytes_limit=10000, token=self.token) self._Process([client_mock], worker_obj) self.assertEqual(result["cpulimit"], [1000, 980, 960]) self.assertEqual(result["networklimit"], [10000, 9000, 8000]) return result
def testLimitPropagation(self): """This tests that client actions are limited properly.""" client_mock = action_mocks.CPULimitClientMock( user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000], runtime_us=[rdfvalue.Duration.From(1, rdfvalue.SECONDS)]) flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=10000, runtime_limit=rdfvalue.Duration.From( 5, rdfvalue.SECONDS)) self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000]) self.assertEqual(client_mock.storage["runtimelimit"], [ rdfvalue.Duration.From(5, rdfvalue.SECONDS), rdfvalue.Duration.From(4, rdfvalue.SECONDS), rdfvalue.Duration.From(3, rdfvalue.SECONDS), ])
def testLimitPropagation(self): """This tests that client actions are limited properly.""" client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=10000) self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000])
def testCPULimitForFlows(self): """This tests that the client actions are limited properly.""" client_mock = action_mocks.CPULimitClientMock( user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) flow_test_lib.TestFlowHelper( flow_test_lib.CPULimitFlow.__name__, client_mock, token=self.token, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=10000) self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000])
def testNetworkLimitExceeded(self): """This tests that the network limit for flows is working.""" client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) with test_lib.SuppressLogs(): flow_id = flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=1500, check_flow_errors=False) rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(rdf_flow.flow_state, "ERROR") self.assertIn("bytes limit exceeded", rdf_flow.error_message)
def testRuntimeLimitExceeded(self): client_mock = action_mocks.CPULimitClientMock( user_cpu_usage=[1], system_cpu_usage=[1], network_usage=[1], runtime_us=[rdfvalue.Duration.From(4, rdfvalue.SECONDS)]) with test_lib.SuppressLogs(): flow_id = flow_test_lib.StartAndRunFlow( flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, runtime_limit=rdfvalue.Duration.From(9, rdfvalue.SECONDS), check_flow_errors=False) rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(rdf_flow.flow_state, "ERROR") self.assertIn("Runtime limit exceeded", rdf_flow.error_message)
def RunFlow(self, flow_name, **kwargs): result = {} client_mock = action_mocks.CPULimitClientMock(result) client_mock = flow_test_lib.MockClient(self.client_id, client_mock, token=self.token) worker_mock = ResourcedWorker(check_flow_errors=True, token=self.token) flow.GRRFlow.StartFlow(client_id=self.client_id, flow_name=flow_name, token=self.token, **kwargs) while True: client_processed = client_mock.Next() flows_run = [] for flow_run in worker_mock.Next(): flows_run.append(flow_run) if client_processed == 0 and not flows_run: break return result
def testCPULimitForHunts(self): worker_obj = worker.GRRWorker(token=self.token) client_ids = ["C.%016X" % i for i in xrange(10, 20)] result = {} client_mocks = [] for client_id in client_ids: client_mock = action_mocks.CPULimitClientMock(result) client_mock = flow_test_lib.MockClient( rdf_client.ClientURN(client_id), client_mock, token=self.token) client_mock.EnableResourceUsage(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) client_mocks.append(client_mock) flow_runner_args = rdf_flows.FlowRunnerArgs( flow_name=flow_test_lib.CPULimitFlow.__name__) with implementation.GRRHunt.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=flow_runner_args, cpu_limit=5000, per_client_cpu_limit=10000, network_bytes_limit=1000000, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() implementation.GRRHunt.StartClients(hunt.session_id, client_ids[:1]) self._Process(client_mocks, worker_obj) implementation.GRRHunt.StartClients(hunt.session_id, client_ids[1:2]) self._Process(client_mocks, worker_obj) implementation.GRRHunt.StartClients(hunt.session_id, client_ids[2:3]) self._Process(client_mocks, worker_obj) # The limiting factor here is the overall hunt limit of 5000 cpu # seconds. Clients that finish should decrease the remaining quota # and the following clients should get the reduced quota. self.assertEqual(result["cpulimit"], [ 5000.0, 4980.0, 4960.0, 4940.0, 4920.0, 4900.0, 4880.0, 4860.0, 4840.0 ]) self.assertEqual(result["networklimit"], [ 1000000L, 999000L, 998000L, 997000L, 996000L, 995000L, 994000L, 993000L, 992000L ]) result.clear() with implementation.GRRHunt.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=flow_runner_args, per_client_cpu_limit=3000, per_client_network_limit_bytes=3000000, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() implementation.GRRHunt.StartClients(hunt.session_id, client_ids[:1]) self._Process(client_mocks, worker_obj) implementation.GRRHunt.StartClients(hunt.session_id, client_ids[1:2]) self._Process(client_mocks, worker_obj) implementation.GRRHunt.StartClients(hunt.session_id, client_ids[2:3]) self._Process(client_mocks, worker_obj) # This time, the per client limit is 3000s / 3000000 bytes. Every # client should get the same limit. self.assertEqual(result["cpulimit"], [ 3000.0, 2980.0, 2960.0, 3000.0, 2980.0, 2960.0, 3000.0, 2980.0, 2960.0 ]) self.assertEqual(result["networklimit"], [ 3000000, 2999000, 2998000, 3000000, 2999000, 2998000, 3000000, 2999000, 2998000 ]) result.clear() for client_mock in client_mocks: client_mock.EnableResourceUsage(user_cpu_usage=[500], system_cpu_usage=[500], network_usage=[1000000]) with implementation.GRRHunt.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=flow_runner_args, per_client_cpu_limit=3000, cpu_limit=5000, per_client_network_limit_bytes=3000000, network_bytes_limit=5000000, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() implementation.GRRHunt.StartClients(hunt.session_id, client_ids[:1]) self._Process(client_mocks, worker_obj) implementation.GRRHunt.StartClients(hunt.session_id, client_ids[1:2]) self._Process(client_mocks, worker_obj) implementation.GRRHunt.StartClients(hunt.session_id, client_ids[2:3]) self._Process(client_mocks, worker_obj) # The first client gets the full per client limit of 3000s, and # uses all of it. The hunt has a limit of just 5000 total so the # second client gets started with a limit of 2000. It can only run # two of three states, the last client will not be started at all # due to out of quota. self.assertEqual(result["cpulimit"], [3000.0, 2000.0, 1000.0, 2000.0, 1000.0]) self.assertEqual(result["networklimit"], [3000000, 2000000, 1000000, 2000000, 1000000]) errors = list(hunt.GetClientsErrors()) self.assertEqual(len(errors), 2) # Client side out of cpu. self.assertIn("CPU limit exceeded", errors[0].log_message) # Server side out of cpu. self.assertIn("Out of CPU quota", errors[1].backtrace)