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)
