diff --git a/esrally/client.py b/esrally/client.py index ff92b7553..39a08f86b 100644 --- a/esrally/client.py +++ b/esrally/client.py @@ -148,14 +148,10 @@ def loads(self, s): return super().loads(s) async def on_request_start(session, trace_config_ctx, params): - meta = RallyAsyncElasticsearch.request_context.get() - # this can happen if multiple requests are sent on the wire for one logical request (e.g. scrolls) - if "request_start" not in meta: - meta["request_start"] = time.perf_counter() + RallyAsyncElasticsearch.on_request_start() async def on_request_end(session, trace_config_ctx, params): - meta = RallyAsyncElasticsearch.request_context.get() - meta["request_end"] = time.perf_counter() + RallyAsyncElasticsearch.on_request_end() trace_config = aiohttp.TraceConfig() trace_config.on_request_start.append(on_request_start) @@ -175,6 +171,18 @@ def init_request_context(self): RallyAsyncElasticsearch.request_context.set(ctx) return ctx + @classmethod + def on_request_start(cls): + meta = RallyAsyncElasticsearch.request_context.get() + # this can happen if multiple requests are sent on the wire for one logical request (e.g. scrolls) + if "request_start" not in meta: + meta["request_start"] = time.perf_counter() + + @classmethod + def on_request_end(cls): + meta = RallyAsyncElasticsearch.request_context.get() + meta["request_end"] = time.perf_counter() + def return_raw_response(self): ctx = RallyAsyncElasticsearch.request_context.get() ctx["raw_response"] = True diff --git a/esrally/driver/runner.py b/esrally/driver/runner.py index ca5b390ac..f1aa928ed 100644 --- a/esrally/driver/runner.py +++ b/esrally/driver/runner.py @@ -1630,7 +1630,11 @@ class Sleep(Runner): """ async def __call__(self, es, params): - await asyncio.sleep(mandatory(params, "duration", "sleep")) + es.on_request_start() + try: + await asyncio.sleep(mandatory(params, "duration", "sleep")) + finally: + es.on_request_end() def __repr__(self, *args, **kwargs): return "sleep" diff --git a/tests/driver/runner_test.py b/tests/driver/runner_test.py index b5a69be85..599776819 100644 --- a/tests/driver/runner_test.py +++ b/tests/driver/runner_test.py @@ -3328,6 +3328,8 @@ async def test_missing_parameter(self, sleep, es): await r(es, params={}) self.assertEqual(0, es.call_count) + self.assertEqual(1, es.on_request_start.call_count) + self.assertEqual(1, es.on_request_end.call_count) self.assertEqual(0, sleep.call_count) @mock.patch("elasticsearch.Elasticsearch") @@ -3339,6 +3341,8 @@ async def test_sleep(self, sleep, es): await r(es, params={"duration": 4.3}) self.assertEqual(0, es.call_count) + self.assertEqual(1, es.on_request_start.call_count) + self.assertEqual(1, es.on_request_end.call_count) sleep.assert_called_once_with(4.3)