Ensure that sleep tasks are timed (#1121)

With this commit we ensure that also the `sleep` task is included in
time measurements. Regular runners in Rally issue requests against
Elasticsearch and time measurement is done via callbacks in the request
lifecycle but the `sleep` runner is different because it only waits for
a certain time without any calls.

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

esrally/driver/runner.py

@@ -1650,7 +1650,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"

tests/driver/runner_test.py

@@ -3329,6 +3329,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")
@@ -3340,6 +3342,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)