esrally/driver/driver.py

# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# 	http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import asyncio
import collections
import concurrent.futures
import datetime
import itertools
import logging
import math
import multiprocessing
import queue
import sys
import threading
import time
from dataclasses import dataclass
from enum import Enum
from io import BytesIO
from typing import Callable, Optional

import thespian.actors

from esrally import (
    PROGRAM_NAME,
    actor,
    client,
    config,
    exceptions,
    metrics,
    paths,
    telemetry,
    track,
    types,
)
from esrally.client import delete_api_keys
from esrally.driver import runner, scheduler
from esrally.track import TrackProcessorRegistry, load_track, load_track_plugins
from esrally.utils import console, convert, net


##################################
#
# Messages sent between drivers
#
##################################
class PrepareBenchmark:
    """
    Initiates preparation steps for a benchmark. The benchmark should only be started after StartBenchmark is sent.
    """

    def __init__(self, config: types.Config, track):
        """
        :param config: Rally internal configuration object.
        :param track: The track to use.
        """
        self.config = config
        self.track = track


class StartBenchmark:
    pass


class Bootstrap:
    """
    Prompts loading of track code on new actors
    """

    def __init__(self, cfg: types.Config, worker_id=None):
        self.config = cfg
        self.worker_id = worker_id


class PrepareTrack:
    """
    Initiates preparation of a track.

    """

    def __init__(self, track):
        """
        :param track: The track to use.
        """
        self.track = track


class TrackPrepared:
    pass


class StartTaskLoop:
    def __init__(self, track_name, cfg: types.Config):
        self.track_name = track_name
        self.cfg = cfg


class DoTask:
    def __init__(self, task, cfg: types.Config):
        self.task = task
        self.cfg = cfg


@dataclass(frozen=True)
class WorkerTask:
    """
    Unit of work that should be completed by the low-level TaskExecutionActor
    """

    func: Callable
    params: dict


class ReadyForWork:
    pass


class WorkerIdle:
    pass


class PreparationComplete:
    def __init__(self, distribution_flavor, distribution_version, revision):
        self.distribution_flavor = distribution_flavor
        self.distribution_version = distribution_version
        self.revision = revision


class StartWorker:
    """
    Starts a worker.
    """

    def __init__(self, worker_id, config: types.Config, track, client_allocations, client_contexts):
        """
        :param worker_id: Unique (numeric) id of the worker.
        :param config: Rally internal configuration object.
        :param track: The track to use.
        :param client_allocations: A structure describing which clients need to run which tasks.
        :param client_contexts: A dict ``ClientContext`` objects keyed by client ID
        """
        self.worker_id = worker_id
        self.config = config
        self.track = track
        self.client_allocations = client_allocations
        self.client_contexts = client_contexts


class Drive:
    """
    Tells a load generator to drive (either after a join point or initially).
    """

    def __init__(self, client_start_timestamp):
        self.client_start_timestamp = client_start_timestamp


class CompleteCurrentTask:
    """
    Tells a load generator to prematurely complete its current task. This is used to model task dependencies for parallel tasks (i.e. if a
    specific task that is marked accordingly in the track finishes, it will also signal termination of all other tasks in the same parallel
    element).
    """


class UpdateSamples:
    """
    Used to send samples from a load generator node to the master.
    """

    def __init__(self, client_id, samples):
        self.client_id = client_id
        self.samples = samples


class JoinPointReached:
    """
    Tells the master that a load generator has reached a join point. Used for coordination across multiple load generators.
    """

    def __init__(self, worker_id, task):
        self.worker_id = worker_id
        # Using perf_counter here is fine even in the distributed case. Although we "leak" this value to other
        # machines, we will only ever interpret this value on the same machine (see `Drive` and the implementation
        # in `Driver#joinpoint_reached()`).
        self.worker_timestamp = time.perf_counter()
        self.task = task


class BenchmarkComplete:
    """
    Indicates that the benchmark is complete.
    """

    def __init__(self, metrics):
        self.metrics = metrics


class TaskFinished:
    def __init__(self, metrics, next_task_scheduled_in):
        self.metrics = metrics
        self.next_task_scheduled_in = next_task_scheduled_in


class DriverActor(actor.RallyActor):
    RESET_RELATIVE_TIME_MARKER = "reset_relative_time"

    WAKEUP_INTERVAL_SECONDS = 1

    # post-process request metrics every N seconds and send it to the metrics store
    POST_PROCESS_INTERVAL_SECONDS = 30

    """
    Coordinates all workers. This is actually only a thin actor wrapper layer around ``Driver`` which does the actual work.
    """

    def __init__(self):
        super().__init__()
        self.benchmark_actor = None
        self.driver = None
        self.status = "init"
        self.post_process_timer = 0
        self.cluster_details = {}

    def receiveMsg_PoisonMessage(self, poisonmsg, sender):
        self.logger.error("Main driver received a fatal indication from a load generator (%s). Shutting down.", poisonmsg.details)
        self.driver.close()
        self.send(self.benchmark_actor, actor.BenchmarkFailure("Fatal track or load generator indication", poisonmsg.details))

    def receiveMsg_BenchmarkFailure(self, msg, sender):
        self.logger.error("Main driver received a fatal exception from a load generator. Shutting down.")
        self.driver.close()
        self.send(self.benchmark_actor, msg)

    def receiveMsg_BenchmarkCancelled(self, msg, sender):
        self.logger.info("Main driver received a notification that the benchmark has been cancelled.")
        self.driver.close()
        self.send(self.benchmark_actor, msg)

    def receiveMsg_ActorExitRequest(self, msg, sender):
        self.logger.info("Main driver received ActorExitRequest and will terminate all load generators.")
        self.status = "exiting"

    def receiveMsg_ChildActorExited(self, msg, sender):
        # is it a worker?
        if msg.childAddress in self.driver.workers:
            worker_index = self.driver.workers.index(msg.childAddress)
            if self.status == "exiting":
                self.logger.debug("Worker [%d] has exited.", worker_index)
            else:
                self.logger.error("Worker [%d] has exited prematurely. Aborting benchmark.", worker_index)
                self.send(self.benchmark_actor, actor.BenchmarkFailure(f"Worker [{worker_index}] has exited prematurely."))
        else:
            self.logger.debug("A track preparator has exited.")

    def receiveUnrecognizedMessage(self, msg, sender):
        self.logger.debug("Main driver received unknown message [%s] (ignoring).", str(msg))

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_PrepareBenchmark(self, msg, sender):
        self.benchmark_actor = sender
        self.driver = Driver(self, msg.config)
        self.driver.prepare_benchmark(msg.track)

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_StartBenchmark(self, msg, sender):
        self.benchmark_actor = sender
        self.driver.start_benchmark()
        self.wakeupAfter(datetime.timedelta(seconds=DriverActor.WAKEUP_INTERVAL_SECONDS))

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_TrackPrepared(self, msg, track_preparation_actor):
        self.transition_when_all_children_responded(
            track_preparation_actor, msg, expected_status=None, new_status=None, transition=self._after_track_prepared
        )

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_JoinPointReached(self, msg, sender):
        self.driver.joinpoint_reached(msg.worker_id, msg.worker_timestamp, msg.task)

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_UpdateSamples(self, msg, sender):
        self.driver.update_samples(msg.samples)

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_WakeupMessage(self, msg, sender):
        if msg.payload == DriverActor.RESET_RELATIVE_TIME_MARKER:
            self.driver.reset_relative_time()
        elif not self.driver.finished():
            self.post_process_timer += DriverActor.WAKEUP_INTERVAL_SECONDS
            if self.post_process_timer >= DriverActor.POST_PROCESS_INTERVAL_SECONDS:
                self.post_process_timer = 0
                self.driver.post_process_samples()
            self.driver.update_progress_message()
            self.wakeupAfter(datetime.timedelta(seconds=DriverActor.WAKEUP_INTERVAL_SECONDS))

    def create_client(self, host, cfg: types.Config, worker_id):
        worker = self.createActor(Worker, targetActorRequirements=self._requirements(host))
        self.send(worker, Bootstrap(cfg, worker_id))
        return worker

    def start_worker(self, driver, worker_id, cfg: types.Config, track, allocations, client_contexts=None):
        self.send(driver, StartWorker(worker_id, cfg, track, allocations, client_contexts))

    def drive_at(self, driver, client_start_timestamp):
        self.send(driver, Drive(client_start_timestamp))

    def complete_current_task(self, driver):
        self.send(driver, CompleteCurrentTask())

    def on_task_finished(self, metrics, next_task_scheduled_in):
        if next_task_scheduled_in > 0:
            self.wakeupAfter(datetime.timedelta(seconds=next_task_scheduled_in), payload=DriverActor.RESET_RELATIVE_TIME_MARKER)
        else:
            self.driver.reset_relative_time()
        self.send(self.benchmark_actor, TaskFinished(metrics, next_task_scheduled_in))

    def _requirements(self, host):
        if host == "localhost":
            return {"coordinator": True}
        else:
            return {"ip": host}

    def prepare_track(self, hosts, cfg: types.Config, track):
        self.track = track
        self.logger.info("Starting prepare track process on hosts [%s]", hosts)
        self.children = [self._create_track_preparator(h) for h in hosts]
        msg = Bootstrap(cfg)
        for child in self.children:
            self.send(child, msg)

    @actor.no_retry("driver")  # pylint: disable=no-value-for-parameter
    def receiveMsg_ReadyForWork(self, msg, task_preparation_actor):
        msg = PrepareTrack(self.track)
        self.send(task_preparation_actor, msg)

    def _create_track_preparator(self, host):
        return self.createActor(TrackPreparationActor, targetActorRequirements=self._requirements(host))

    def _after_track_prepared(self):
        cluster_version = self.cluster_details["version"] if self.cluster_details else {}
        # manually compiled versions don't expose build_flavor but Rally expects a value in telemetry devices
        # we should default to trial/basic, but let's default to oss for now to avoid breaking the charts
        build_flavor = cluster_version.get("build_flavor", "oss")
        build_version = cluster_version.get("number", build_flavor)
        build_hash = cluster_version.get("build_hash", build_flavor)

        for child in self.children:
            self.send(child, thespian.actors.ActorExitRequest())
        self.children = []
        self.send(
            self.benchmark_actor,
            PreparationComplete(
                build_flavor,
                build_version,
                build_hash,
            ),
        )

    def on_benchmark_complete(self, metrics):
        self.send(self.benchmark_actor, BenchmarkComplete(metrics))


def load_local_config(coordinator_config) -> types.Config:
    cfg = config.auto_load_local_config(
        coordinator_config,
        additional_sections=[
            # only copy the relevant bits
            "track",
            "driver",
            "client",
            # due to distribution version...
            "mechanic",
            "telemetry",
        ],
    )
    # set root path (normally done by the main entry point)
    cfg.add(config.Scope.application, "node", "rally.root", paths.rally_root())
    return cfg


class TaskExecutionActor(actor.RallyActor):
    """
    This class should be used for long-running tasks, as it ensures they do not block the actor's messaging system
    """

    def __init__(self):
        super().__init__()
        self.pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
        self.executor_future = None
        self.wakeup_interval = 5
        self.task_preparation_actor = None
        self.logger = logging.getLogger(__name__)
        self.track_name = None
        self.cfg: Optional[types.Config] = None

    @actor.no_retry("task executor")  # pylint: disable=no-value-for-parameter
    def receiveMsg_StartTaskLoop(self, msg, sender):
        self.task_preparation_actor = sender
        self.track_name = msg.track_name
        self.cfg = load_local_config(msg.cfg)
        if self.cfg.opts("track", "test.mode.enabled"):
            self.wakeup_interval = 0.5
        track.load_track_plugins(self.cfg, self.track_name)
        self.send(self.task_preparation_actor, ReadyForWork())

    @actor.no_retry("task executor")  # pylint: disable=no-value-for-parameter
    def receiveMsg_DoTask(self, msg, sender):
        # actor can arbitrarily execute code based on these messages. if anyone besides our parent sends a task, ignore
        if sender != self.task_preparation_actor:
            msg = (
                f"TaskExecutionActor expected message from [{self.task_preparation_actor}]"
                " but the received the following from [{sender}]: {vars(msg)}"
            )
            raise exceptions.RallyError(msg)
        task = msg.task
        if self.executor_future is not None:
            msg = f"TaskExecutionActor received DoTask message [{vars(msg)}], but was already busy"
            raise exceptions.RallyError(msg)
        if task is None:
            self.send(self.task_preparation_actor, WorkerIdle())
        else:
            # this is a potentially long-running operation so we offload it a background thread so we don't block
            # the actor (e.g. logging works properly as log messages are forwarded timely).
            self.executor_future = self.pool.submit(task.func, **task.params)
            self.wakeupAfter(datetime.timedelta(seconds=self.wakeup_interval))

    @actor.no_retry("task executor")  # pylint: disable=no-value-for-parameter
    def receiveMsg_WakeupMessage(self, msg, sender):
        if self.executor_future is not None and self.executor_future.done():
            e = self.executor_future.exception(timeout=0)
            if e:
                self.logger.exception("Worker failed. Notifying parent...", exc_info=e)
                # the exception might be user-defined and not be on the load path of the original sender. Hence, it
                # cannot be deserialized on the receiver so we convert it here to a plain string.
                self.send(self.task_preparation_actor, actor.BenchmarkFailure("Error in task executor", str(e)))
            else:
                self.executor_future = None
                self.send(self.task_preparation_actor, ReadyForWork())
        else:
            self.wakeupAfter(datetime.timedelta(seconds=self.wakeup_interval))

    def receiveMsg_BenchmarkFailure(self, msg, sender):
        # sent by our no_retry infrastructure; forward to master
        self.send(self.task_preparation_actor, msg)


class TrackPreparationActor(actor.RallyActor):
    class Status(Enum):
        INITIALIZING = "initializing"
        PROCESSOR_RUNNING = "processor running"
        PROCESSOR_COMPLETE = "processor complete"

    def __init__(self):
        super().__init__()
        self.processors = queue.Queue()
        self.driver_actor = None
        self.logger.info("Track Preparator started")
        self.status = self.Status.INITIALIZING
        self.children = []
        self.tasks = []
        self.cfg: Optional[types.Config] = None
        self.data_root_dir = None
        self.track = None

    def receiveMsg_PoisonMessage(self, poisonmsg, sender):
        self.logger.error("Track Preparator received a fatal indication from a load generator (%s). Shutting down.", poisonmsg.details)
        self.send(self.driver_actor, actor.BenchmarkFailure("Fatal track preparation indication", poisonmsg.details))

    @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
    def receiveMsg_Bootstrap(self, msg, sender):
        self.driver_actor = sender
        # load node-specific config to have correct paths available
        self.cfg = load_local_config(msg.config)
        # this instance of load_track occurs once per host, so install dependencies if necessary
        load_track(self.cfg, install_dependencies=False)
        self.send(self.driver_actor, ReadyForWork())

    @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
    def receiveMsg_ActorExitRequest(self, msg, sender):
        self.logger.debug("ActorExitRequest received. Forwarding to children")
        for child in self.children:
            self.send(child, msg)

    @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
    def receiveMsg_BenchmarkFailure(self, msg, sender):
        # sent by our generic worker; forward to parent
        self.send(self.driver_actor, msg)

    @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
    def receiveMsg_PrepareTrack(self, msg, sender):
        assert self.cfg is not None
        self.data_root_dir = self.cfg.opts("benchmarks", "local.dataset.cache")
        tpr = TrackProcessorRegistry(self.cfg)
        self.track = msg.track
        self.logger.info("Preparing track [%s]", self.track.name)
        self.logger.info("Reloading track [%s] to ensure plugins are up-to-date.", self.track.name)
        # the track might have been loaded on a different machine (the coordinator machine) so we force a track
        # update to ensure we use the latest version of plugins.
        load_track(self.cfg)
        load_track_plugins(self.cfg, self.track.name, register_track_processor=tpr.register_track_processor, force_update=True)
        # we expect on_prepare_track can take a long time. seed a queue of tasks and delegate to child workers
        self.children = [self._create_task_executor() for _ in range(num_cores(self.cfg))]
        for processor in tpr.processors:
            self.processors.put(processor)
        self._seed_tasks(self.processors.get())
        self.send_to_children_and_transition(
            self, StartTaskLoop(self.track.name, self.cfg), self.Status.INITIALIZING, self.Status.PROCESSOR_RUNNING
        )

    def resume(self):
        assert self.cfg is not None
        if not self.processors.empty():
            self._seed_tasks(self.processors.get())
            self.send_to_children_and_transition(
                self, StartTaskLoop(self.track.name, self.cfg), self.Status.PROCESSOR_COMPLETE, self.Status.PROCESSOR_RUNNING
            )
        else:
            self.send(self.driver_actor, TrackPrepared())

    def _seed_tasks(self, processor):
        self.tasks = list(WorkerTask(func, params) for func, params in processor.on_prepare_track(self.track, self.data_root_dir))

    def _create_task_executor(self):
        return self.createActor(TaskExecutionActor)

    @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
    def receiveMsg_ReadyForWork(self, msg, task_execution_actor):
        assert self.cfg is not None
        if self.tasks:
            next_task = self.tasks.pop()
        else:
            next_task = None
        new_msg = DoTask(next_task, self.cfg)
        self.logger.debug("Track Preparator sending %s to %s", vars(new_msg), task_execution_actor)
        self.send(task_execution_actor, new_msg)

    @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
    def receiveMsg_WorkerIdle(self, msg, sender):
        self.transition_when_all_children_responded(sender, msg, self.Status.PROCESSOR_RUNNING, self.Status.PROCESSOR_COMPLETE, self.resume)


def num_cores(cfg: types.Config):
    return int(cfg.opts("system", "available.cores", mandatory=False, default_value=multiprocessing.cpu_count()))


ApiKey = collections.namedtuple("ApiKey", ["id", "secret"])


@dataclass
class ClientContext:
    client_id: int
    parent_worker_id: int
    api_key: Optional[ApiKey] = None


class Driver:
    def __init__(self, driver_actor, config: types.Config, es_client_factory_class=client.EsClientFactory):
        """
        Coordinates all workers. It is technology-agnostic, i.e. it does not know anything about actors. To allow us to hook in an actor,
        we provide a ``target`` parameter which will be called whenever some event has occurred. The ``target`` can use this to send
        appropriate messages.

        :param target: A target that will be notified of important events.
        :param config: The current config object.
        """
        self.logger = logging.getLogger(__name__)
        self.driver_actor = driver_actor
        self.config = config
        self.es_client_factory = es_client_factory_class
        self.default_sync_es_client = None
        self.track = None
        self.challenge = None
        self.metrics_store = None
        self.load_driver_hosts = []
        self.workers = []
        # which client ids are assigned to which workers?
        self.clients_per_worker = {}
        self.client_contexts = {}
        self.generated_api_key_ids = []

        self.progress_reporter = console.progress()
        self.progress_counter = 0
        self.quiet = False
        self.allocations = None
        self.raw_samples = []
        self.most_recent_sample_per_client = {}
        self.sample_post_processor = None

        self.number_of_steps = 0
        self.currently_completed = 0
        self.workers_completed_current_step = {}
        self.current_step = -1
        self.tasks_per_join_point = None
        self.complete_current_task_sent = False

        self.telemetry = None

    def create_es_clients(self):
        all_hosts = self.config.opts("client", "hosts").all_hosts
        distribution_version = self.config.opts("mechanic", "distribution.version", mandatory=False)
        distribution_flavor = self.config.opts("mechanic", "distribution.flavor", mandatory=False)
        es = {}
        for cluster_name, cluster_hosts in all_hosts.items():
            all_client_options = self.config.opts("client", "options").all_client_options
            cluster_client_options = dict(all_client_options[cluster_name])
            # Use retries to avoid aborts on long living connections for telemetry devices
            cluster_client_options["retry_on_timeout"] = True
            es[cluster_name] = self.es_client_factory(
                cluster_hosts, cluster_client_options, distribution_version=distribution_version, distribution_flavor=distribution_flavor
            ).create()
        return es

    def prepare_telemetry(self, es, enable, index_names, data_stream_names, build_hash, serverless_mode, serverless_operator):
        enabled_devices = self.config.opts("telemetry", "devices")
        telemetry_params = self.config.opts("telemetry", "params")
        log_root = paths.race_root(self.config)

        es_default = es["default"]

        if enable:
            devices = [
                telemetry.NodeStats(telemetry_params, es, self.metrics_store),
                telemetry.ExternalEnvironmentInfo(es_default, self.metrics_store),
                telemetry.ClusterEnvironmentInfo(es_default, self.metrics_store, build_hash),
                telemetry.JvmStatsSummary(es_default, self.metrics_store),
                telemetry.IndexStats(es_default, self.metrics_store),
                telemetry.MlBucketProcessingTime(es_default, self.metrics_store),
                telemetry.MasterNodeStats(telemetry_params, es_default, self.metrics_store),
                telemetry.SegmentStats(log_root, es_default),
                telemetry.CcrStats(telemetry_params, es, self.metrics_store),
                telemetry.RecoveryStats(telemetry_params, es, self.metrics_store),
                telemetry.ShardStats(telemetry_params, es, self.metrics_store),
                telemetry.TransformStats(telemetry_params, es, self.metrics_store),
                telemetry.SearchableSnapshotsStats(telemetry_params, es, self.metrics_store),
                telemetry.DataStreamStats(telemetry_params, es, self.metrics_store),
                telemetry.IngestPipelineStats(es, self.metrics_store),
                telemetry.DiskUsageStats(telemetry_params, es_default, self.metrics_store, index_names, data_stream_names),
                telemetry.BlobStoreStats(telemetry_params, es, self.metrics_store),
                telemetry.GeoIpStats(es_default, self.metrics_store),
            ]
        else:
            devices = []
        self.telemetry = telemetry.Telemetry(
            enabled_devices,
            devices=devices,
            serverless_mode=serverless_mode,
            serverless_operator=serverless_operator,
        )

    def wait_for_rest_api(self, es):
        es_default = es["default"]
        self.logger.info("Checking if REST API is available.")
        if client.wait_for_rest_layer(es_default, max_attempts=40):
            self.logger.info("REST API is available.")
        else:
            self.logger.error("REST API layer is not yet available. Stopping benchmark.")
            raise exceptions.SystemSetupError("Elasticsearch REST API layer is not available.")

    def retrieve_cluster_info(self, es):
        try:
            return es["default"].info()
        except BaseException:
            self.logger.exception("Could not retrieve cluster info on benchmark start")
            return None

    def retrieve_build_hash_from_nodes_info(self, es):
        try:
            nodes_info = es["default"].nodes.info(filter_path="**.build_hash")
            nodes = nodes_info["nodes"]
            # assumption: build hash is the same across all the nodes
            first_node_id = next(iter(nodes))
            return nodes[first_node_id]["build_hash"]
        except BaseException:
            self.logger.exception("Could not retrieve build hash from nodes info")
            return None

    def create_api_key(self, es, client_id):
        self.logger.debug("Creating ES API key for client [%s].", client_id)
        try:
            api_key = client.create_api_key(es, client_id)
            self.logger.debug("ES API key created for client [%s].", client_id)
            # Store the API key ID for deletion upon benchmark completion
            self.generated_api_key_ids.append(api_key["id"])
            return api_key
        except Exception as e:
            self.logger.error("Unable to create API keys. Stopping benchmark.")
            raise exceptions.SystemSetupError(e.message)

    def prepare_benchmark(self, t):
        self.track = t
        self.challenge = select_challenge(self.config, self.track)
        self.quiet = self.config.opts("system", "quiet.mode", mandatory=False, default_value=False)
        downsample_factor = int(self.config.opts("reporting", "metrics.request.downsample.factor", mandatory=False, default_value=1))
        self.metrics_store = metrics.metrics_store(cfg=self.config, track=self.track.name, challenge=self.challenge.name, read_only=False)

        self.sample_post_processor = SamplePostprocessor(
            self.metrics_store, downsample_factor, self.track.meta_data, self.challenge.meta_data
        )

        es_clients = self.create_es_clients()
        self.default_sync_es_client = es_clients["default"]

        skip_rest_api_check = self.config.opts("mechanic", "skip.rest.api.check")
        uses_static_responses = self.config.opts("client", "options").uses_static_responses
        serverless_mode = convert.to_bool(self.config.opts("driver", "serverless.mode", mandatory=False, default_value=False))
        serverless_operator = convert.to_bool(self.config.opts("driver", "serverless.operator", mandatory=False, default_value=False))
        build_hash = None
        if skip_rest_api_check:
            self.logger.info("Skipping REST API check as requested explicitly.")
        elif uses_static_responses:
            self.logger.info("Skipping REST API check as static responses are used.")
        else:
            if serverless_mode and not serverless_operator:
                self.logger.info("Skipping REST API check while targetting serverless cluster with a public user.")
            else:
                self.wait_for_rest_api(es_clients)
            self.driver_actor.cluster_details = self.retrieve_cluster_info(es_clients)
            if serverless_mode:
                # overwrite static serverless version number
                self.driver_actor.cluster_details["version"]["number"] = "serverless"
                if serverless_operator:
                    # overwrite build hash if running as operator
                    build_hash = self.retrieve_build_hash_from_nodes_info(es_clients)
                    self.logger.info("Retrieved actual build hash [%s] from serverless cluster.", build_hash)
                    self.driver_actor.cluster_details["version"]["build_hash"] = build_hash

        # Avoid issuing any requests to the target cluster when static responses are enabled. The results
        # are not useful and attempts to connect to a non-existing cluster just lead to exception traces in logs.
        self.prepare_telemetry(
            es_clients,
            enable=not uses_static_responses,
            index_names=self.track.index_names(),
            data_stream_names=self.track.data_stream_names(),
            build_hash=build_hash,
            serverless_mode=serverless_mode,
            serverless_operator=serverless_operator,
        )

        for host in self.config.opts("driver", "load_driver_hosts"):
            host_config = {
                # for simplicity we assume that all benchmark machines have the same specs
                "cores": num_cores(self.config)
            }
            if host != "localhost":
                host_config["host"] = net.resolve(host)
            else:
                host_config["host"] = host

            self.load_driver_hosts.append(host_config)

        self.driver_actor.prepare_track([h["host"] for h in self.load_driver_hosts], self.config, self.track)

    def start_benchmark(self):
        self.logger.info("Benchmark is about to start.")
        # ensure relative time starts when the benchmark starts.
        self.reset_relative_time()
        self.logger.info("Attaching cluster-level telemetry devices.")
        self.telemetry.on_benchmark_start()
        self.logger.info("Cluster-level telemetry devices are now attached.")

        allocator = Allocator(self.challenge.schedule)
        self.allocations = allocator.allocations
        self.number_of_steps = len(allocator.join_points) - 1
        self.tasks_per_join_point = allocator.tasks_per_joinpoint

        self.logger.info(
            "Benchmark consists of [%d] steps executed by [%d] clients.",
            self.number_of_steps,
            len(self.allocations),  # type: ignore[arg-type]  # TODO remove the below ignore when introducing type hints
        )
        # avoid flooding the log if there are too many clients
        if allocator.clients < 128:
            self.logger.debug("Allocation matrix:\n%s", "\n".join([str(a) for a in self.allocations]))

        create_api_keys = self.config.opts("client", "options").all_client_options["default"].get("create_api_key_per_client", None)
        worker_assignments = calculate_worker_assignments(self.load_driver_hosts, allocator.clients)
        worker_id = 0
        for assignment in worker_assignments:
            host = assignment["host"]
            for clients in assignment["workers"]:
                # don't assign workers without any clients
                if len(clients) > 0:
                    self.logger.debug("Allocating worker [%d] on [%s] with [%d] clients.", worker_id, host, len(clients))
                    worker = self.driver_actor.create_client(host, self.config, worker_id)

                    client_allocations = ClientAllocations()
                    worker_client_contexts = {}
                    for client_id in clients:
                        client_allocations.add(client_id, self.allocations[client_id])
                        self.clients_per_worker[client_id] = worker_id
                        client_context = ClientContext(client_id=client_id, parent_worker_id=worker_id)

                        if create_api_keys:
                            resp = self.create_api_key(self.default_sync_es_client, client_id)
                            client_context.api_key = ApiKey(id=resp["id"], secret=resp["api_key"])

                        worker_client_contexts[client_id] = client_context
                        self.client_contexts[worker_id] = worker_client_contexts
                    self.driver_actor.start_worker(
                        worker, worker_id, self.config, self.track, client_allocations, client_contexts=worker_client_contexts
                    )
                    self.workers.append(worker)
                    worker_id += 1

        self.update_progress_message()

    def joinpoint_reached(self, worker_id, worker_local_timestamp, task_allocations):
        self.currently_completed += 1
        self.workers_completed_current_step[worker_id] = (worker_local_timestamp, time.perf_counter())
        self.logger.debug(
            "[%d/%d] workers reached join point [%d/%d].",
            self.currently_completed,
            len(self.workers),
            self.current_step + 1,
            self.number_of_steps,
        )
        if self.currently_completed == len(self.workers):
            self.logger.info("All workers completed their tasks until join point [%d/%d].", self.current_step + 1, self.number_of_steps)
            # we can go on to the next step
            self.currently_completed = 0
            self.complete_current_task_sent = False
            # make a copy and reset early to avoid any race conditions from clients that reach a join point already while we are sending...
            workers_curr_step = self.workers_completed_current_step
            self.workers_completed_current_step = {}
            self.update_progress_message(task_finished=True)
            # clear per step
            self.most_recent_sample_per_client = {}
            self.current_step += 1

            self.logger.debug("Postprocessing samples...")
            self.post_process_samples()
            if self.finished():
                self.telemetry.on_benchmark_stop()
                self.logger.info("All steps completed.")
                # Some metrics store implementations return None because no external representation is required.
                # pylint: disable=assignment-from-none
                m = self.metrics_store.to_externalizable(clear=True)
                self.logger.debug("Closing metrics store...")
                self.metrics_store.close()
                # immediately clear as we don't need it anymore and it can consume a significant amount of memory
                self.metrics_store = None
                if self.generated_api_key_ids:
                    self.logger.debug("Deleting auto-generated client API keys...")
                    try:
                        delete_api_keys(self.default_sync_es_client, self.generated_api_key_ids)
                    except exceptions.RallyError:
                        console.warn(
                            "Unable to delete auto-generated API keys. You may need to manually delete them. "
                            "Please check the logs for details."
                        )
                self.logger.debug("Sending benchmark results...")
                self.driver_actor.on_benchmark_complete(m)
            else:
                self.move_to_next_task(workers_curr_step)
        else:
            self.may_complete_current_task(task_allocations)

    def move_to_next_task(self, workers_curr_step):
        if self.config.opts("track", "test.mode.enabled"):
            # don't wait if test mode is enabled and start the next task immediately.
            waiting_period = 0
        else:
            # start the next task in one second (relative to master's timestamp)
            #
            # Assumption: We don't have a lot of clock skew between reaching the join point and sending the next task
            #             (it doesn't matter too much if we're a few ms off).
            waiting_period = 1.0
        # Some metrics store implementations return None because no external representation is required.
        # pylint: disable=assignment-from-none
        m = self.metrics_store.to_externalizable(clear=True)
        self.driver_actor.on_task_finished(m, waiting_period)
        # Using a perf_counter here is fine also in the distributed case as we subtract it from `master_received_msg_at` making it
        # a relative instead of an absolute value.
        start_next_task = time.perf_counter() + waiting_period
        for worker_id, worker in enumerate(self.workers):
            worker_ended_task_at, master_received_msg_at = workers_curr_step[worker_id]
            worker_start_timestamp = worker_ended_task_at + (start_next_task - master_received_msg_at)
            self.logger.debug(
                "Scheduling next task for worker id [%d] at their timestamp [%f] (master timestamp [%f])",
                worker_id,
                worker_start_timestamp,
                start_next_task,
            )
            self.driver_actor.drive_at(worker, worker_start_timestamp)

    def may_complete_current_task(self, task_allocations):
        any_joinpoints_completing_parent = [a for a in task_allocations if a.task.any_task_completes_parent]
        joinpoints_completing_parent = [a for a in task_allocations if a.task.preceding_task_completes_parent]

        # If 'completed-by' is set to 'any', then we *do* want to check for completion by
        # any client and *not* wait until the remaining runner has completed. This way the 'parallel' task will exit
        # on the completion of _any_ client for any task, i.e. given a contrived track with two tasks to execute inside
        # a parallel block:
        #   * parallel:
        #     * bulk-1, with clients 8
        #     * bulk-2, with clients: 8
        #
        # 1. Both 'bulk-1' and 'bulk-2' execute in parallel
        # 2. 'bulk-1' client[0]'s runner is first to complete and reach the next joinpoint successfully
        # 3. 'bulk-1' will now cause the parent task to complete and _not_ wait for all 8 clients' runner to complete,
        # or for 'bulk-2' at all
        #
        # The reasoning for the distinction between 'any_joinpoints_completing_parent' & 'joinpoints_completing_parent'
        # is to simplify logic, otherwise we'd need to implement some form of state machine involving actor-to-actor
        # communication.

        if len(any_joinpoints_completing_parent) > 0 and not self.complete_current_task_sent:
            self.logger.info(
                "Any task before join point [%s] is able to complete the parent structure. Telling all clients to exit immediately.",
                any_joinpoints_completing_parent[0].task,
            )

            self.complete_current_task_sent = True
            for worker in self.workers:
                self.driver_actor.complete_current_task(worker)

        # If we have a specific 'completed-by' task specified, then we want to make sure that all clients for that task
        # are able to complete their runners as expected before completing the parent
        elif len(joinpoints_completing_parent) > 0 and not self.complete_current_task_sent:
            # while this list could contain multiple items, it should always be the same task (but multiple
            # different clients) so any item is sufficient.
            current_join_point = joinpoints_completing_parent[0].task
            self.logger.info(
                "Tasks before join point [%s] are able to complete the parent structure. Checking "
                "if all [%d] clients have finished yet.",
                current_join_point,
                len(current_join_point.clients_executing_completing_task),
            )
            pending_client_ids = []
            for client_id in current_join_point.clients_executing_completing_task:
                # We assume that all clients have finished if their corresponding worker has finished
                worker_id = self.clients_per_worker[client_id]
                if worker_id not in self.workers_completed_current_step:
                    pending_client_ids.append(client_id)

            # are all clients executing said task already done? if so we need to notify the remaining clients
            if len(pending_client_ids) == 0:
                # As we are waiting for other clients to finish, we would send this message over and over again.
                # Hence we need to memorize whether we have already sent it for the current step.
                self.complete_current_task_sent = True
                self.logger.info("All affected clients have finished. Notifying all clients to complete their current tasks.")
                for worker in self.workers:
                    self.driver_actor.complete_current_task(worker)
            else:
                if len(pending_client_ids) > 32:
                    self.logger.info("[%d] clients did not yet finish.", len(pending_client_ids))
                else:
                    self.logger.info("Client id(s) [%s] did not yet finish.", ",".join(map(str, pending_client_ids)))

    def reset_relative_time(self):
        self.logger.debug("Resetting relative time of request metrics store.")
        self.metrics_store.reset_relative_time()

    def finished(self):
        return self.current_step == self.number_of_steps

    def close(self):
        self.progress_reporter.finish()
        if self.metrics_store and self.metrics_store.opened:
            self.metrics_store.close()

    def update_samples(self, samples):
        if len(samples) > 0:
            self.raw_samples += samples
            # We need to check all samples, they will be from different clients
            for s in samples:
                self.most_recent_sample_per_client[s.client_id] = s

    def update_progress_message(self, task_finished=False):
        if not self.quiet and self.current_step >= 0:
            tasks = ",".join([t.name for t in self.tasks_per_join_point[self.current_step]])

            if task_finished:
                total_progress = 1.0
            else:
                # we only count clients which actually contribute to progress. If clients are executing tasks eternally in a parallel
                # structure, we should not count them. The reason is that progress depends entirely on the client(s) that execute the
                # task that is completing the parallel structure.
                progress_per_client = [
                    s.percent_completed for s in self.most_recent_sample_per_client.values() if s.percent_completed is not None
                ]

                num_clients = max(len(progress_per_client), 1)
                total_progress = sum(progress_per_client) / num_clients
            self.progress_reporter.print("Running %s" % tasks, "[%3d%% done]" % (round(total_progress * 100)))
            if task_finished:
                self.progress_reporter.finish()

    def post_process_samples(self):
        # we do *not* do this here to avoid concurrent updates (actors are single-threaded) but rather to make it clear that we use
        # only a snapshot and that new data will go to a new sample set.
        raw_samples = self.raw_samples
        self.raw_samples = []
        self.sample_post_processor(raw_samples)


class SamplePostprocessor:
    def __init__(self, metrics_store, downsample_factor, track_meta_data, challenge_meta_data):
        self.logger = logging.getLogger(__name__)
        self.metrics_store = metrics_store
        self.track_meta_data = track_meta_data
        self.challenge_meta_data = challenge_meta_data
        self.throughput_calculator = ThroughputCalculator()
        self.downsample_factor = downsample_factor

    def __call__(self, raw_samples):
        if len(raw_samples) == 0:
            return
        total_start = time.perf_counter()
        start = total_start
        final_sample_count = 0
        for idx, sample in enumerate(raw_samples):
            if idx % self.downsample_factor == 0:
                final_sample_count += 1
                client_id_meta_data = {"client_id": sample.client_id}
                meta_data = self.merge(
                    self.track_meta_data,
                    self.challenge_meta_data,
                    sample.operation_meta_data,
                    sample.task.meta_data,
                    sample.request_meta_data,
                    client_id_meta_data,
                )

                self.metrics_store.put_value_cluster_level(
                    name="latency",
                    value=convert.seconds_to_ms(sample.latency),
                    unit="ms",
                    task=sample.task.name,
                    operation=sample.operation_name,
                    operation_type=sample.operation_type,
                    sample_type=sample.sample_type,
                    absolute_time=sample.absolute_time,
                    relative_time=sample.relative_time,
                    meta_data=meta_data,
                )

                self.metrics_store.put_value_cluster_level(
                    name="service_time",
                    value=convert.seconds_to_ms(sample.service_time),
                    unit="ms",
                    task=sample.task.name,
                    operation=sample.operation_name,
                    operation_type=sample.operation_type,
                    sample_type=sample.sample_type,
                    absolute_time=sample.absolute_time,
                    relative_time=sample.relative_time,
                    meta_data=meta_data,
                )

                self.metrics_store.put_value_cluster_level(
                    name="processing_time",
                    value=convert.seconds_to_ms(sample.processing_time),
                    unit="ms",
                    task=sample.task.name,
                    operation=sample.operation_name,
                    operation_type=sample.operation_type,
                    sample_type=sample.sample_type,
                    absolute_time=sample.absolute_time,
                    relative_time=sample.relative_time,
                    meta_data=meta_data,
                )

                for timing in sample.dependent_timings:
                    self.metrics_store.put_value_cluster_level(
                        name="service_time",
                        value=convert.seconds_to_ms(timing.service_time),
                        unit="ms",
                        task=timing.task.name,
                        operation=timing.operation_name,
                        operation_type=timing.operation_type,
                        sample_type=timing.sample_type,
                        absolute_time=timing.absolute_time,
                        relative_time=timing.relative_time,
                        meta_data=self.merge(timing.request_meta_data, client_id_meta_data),
                    )

        end = time.perf_counter()
        self.logger.debug("Storing latency and service time took [%f] seconds.", (end - start))
        start = end
        aggregates = self.throughput_calculator.calculate(raw_samples)
        end = time.perf_counter()
        self.logger.debug("Calculating throughput took [%f] seconds.", (end - start))
        start = end
        for task, samples in aggregates.items():
            meta_data = self.merge(self.track_meta_data, self.challenge_meta_data, task.operation.meta_data, task.meta_data)
            for absolute_time, relative_time, sample_type, throughput, throughput_unit in samples:
                self.metrics_store.put_value_cluster_level(
                    name="throughput",
                    value=throughput,
                    unit=throughput_unit,
                    task=task.name,
                    operation=task.operation.name,
                    operation_type=task.operation.type,
                    sample_type=sample_type,
                    absolute_time=absolute_time,
                    relative_time=relative_time,
                    meta_data=meta_data,
                )
        end = time.perf_counter()
        self.logger.debug("Storing throughput took [%f] seconds.", (end - start))
        start = end
        # this will be a noop for the in-memory metrics store. If we use an ES metrics store however, this will ensure that we already send
        # the data and also clear the in-memory buffer. This allows users to see data already while running the benchmark. In cases where
        # it does not matter (i.e. in-memory) we will still defer this step until the end.
        #
        # Don't force refresh here in the interest of short processing times. We don't need to query immediately afterwards so there is
        # no need for frequent refreshes.
        self.metrics_store.flush(refresh=False)
        end = time.perf_counter()
        self.logger.debug("Flushing the metrics store took [%f] seconds.", (end - start))
        self.logger.debug(
            "Postprocessing [%d] raw samples (downsampled to [%d] samples) took [%f] seconds in total.",
            len(raw_samples),
            final_sample_count,
            (end - total_start),
        )

    def merge(self, *args):
        result = {}
        for arg in args:
            if arg is not None:
                result.update(arg)
        return result


def calculate_worker_assignments(host_configs, client_count):
    """
    Assigns clients to workers on the provided hosts.

    :param host_configs: A list of dicts where each dict contains the host name (key: ``host``) and the number of
                         available CPU cores (key: ``cores``).
    :param client_count: The number of clients that should be used at most.
    :return: A list of dicts containing the host (key: ``host``) and a list of workers (key ``workers``). Each entry
             in that list contains another list with the clients that should be assigned to these workers.
    """
    assignments = []
    client_idx = 0
    host_count = len(host_configs)
    clients_per_host = math.ceil(client_count / host_count)
    remaining_clients = client_count
    for host_config in host_configs:
        # the last host might not need to simulate as many clients as the rest of the hosts as we eagerly
        # assign clients to hosts.
        clients_on_this_host = min(clients_per_host, remaining_clients)
        assignment = {
            "host": host_config["host"],
            "workers": [],
        }
        assignments.append(assignment)

        workers_on_this_host = host_config["cores"]
        clients_per_worker = [0] * workers_on_this_host

        # determine how many clients each worker should simulate
        for c in range(clients_on_this_host):
            clients_per_worker[c % workers_on_this_host] += 1

        # assign client ids to workers
        for client_count_for_worker in clients_per_worker:
            worker_assignment = []
            assignment["workers"].append(worker_assignment)
            for c in range(client_idx, client_idx + client_count_for_worker):
                worker_assignment.append(c)
            client_idx += client_count_for_worker

        remaining_clients -= clients_on_this_host

    assert remaining_clients == 0

    return assignments


ClientAllocation = collections.namedtuple("ClientAllocation", ["client_id", "task"])


class ClientAllocations:
    def __init__(self):
        self.allocations = []

    def add(self, client_id, tasks):
        self.allocations.append({"client_id": client_id, "tasks": tasks})

    def is_joinpoint(self, task_index):
        return all(isinstance(t.task, JoinPoint) for t in self.tasks(task_index))

    def tasks(self, task_index, remove_empty=True):
        current_tasks = []
        for allocation in self.allocations:
            tasks_at_index = allocation["tasks"][task_index]
            if remove_empty and tasks_at_index is not None:
                current_tasks.append(ClientAllocation(allocation["client_id"], tasks_at_index))
        return current_tasks


class Worker(actor.RallyActor):
    """
    The actual worker that applies load against the cluster(s).

    It will also regularly send measurements to the master node so it can consolidate them.
    """

    WAKEUP_INTERVAL_SECONDS = 5

    def __init__(self):
        super().__init__()
        self.driver_actor = None
        self.worker_id = None
        self.config: Optional[types.Config] = None
        self.track = None
        self.client_allocations = None
        self.client_contexts = None
        self.current_task_index = 0
        self.next_task_index = 0
        self.on_error = None
        self.pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
        # cancellation via future does not work, hence we use our own mechanism with a shared variable and polling
        self.cancel = threading.Event()
        # used to indicate that we want to prematurely consider this completed. This is *not* due to cancellation
        # but a regular event in a benchmark and used to model task dependency of parallel tasks.
        self.complete = threading.Event()
        self.executor_future = None
        self.sampler = None
        self.start_driving = False
        self.wakeup_interval = Worker.WAKEUP_INTERVAL_SECONDS
        self.sample_queue_size = None

    @actor.no_retry("worker")  # pylint: disable=no-value-for-parameter
    def receiveMsg_Bootstrap(self, msg, sender):
        self.driver_actor = sender
        self.worker_id = msg.worker_id
        # load node-specific config to have correct paths available
        self.config = load_local_config(msg.config)
        load_track(self.config, install_dependencies=False)
        self.logger.debug("Worker[%d] has Python load path %s after bootstrap.", self.worker_id, sys.path)

    @actor.no_retry("worker")  # pylint: disable=no-value-for-parameter
    def receiveMsg_StartWorker(self, msg, sender):
        assert self.config is not None
        self.logger.info("Worker[%d] is about to start.", msg.worker_id)
        self.on_error = self.config.opts("driver", "on.error")
        self.sample_queue_size = int(self.config.opts("reporting", "sample.queue.size", mandatory=False, default_value=1 << 20))
        self.track = msg.track
        track.set_absolute_data_path(self.config, self.track)
        self.client_allocations = msg.client_allocations
        self.client_contexts = msg.client_contexts
        self.current_task_index = 0
        self.cancel.clear()
        # we need to wake up more often in test mode
        if self.config.opts("track", "test.mode.enabled"):
            self.wakeup_interval = 0.5
        runner.register_default_runners(self.config)
        if self.track.has_plugins:
            track.load_track_plugins(self.config, self.track.name, runner.register_runner, scheduler.register_scheduler)
        self.drive()

    @actor.no_retry("worker")  # pylint: disable=no-value-for-parameter
    def receiveMsg_Drive(self, msg, sender):
        sleep_time = datetime.timedelta(seconds=msg.client_start_timestamp - time.perf_counter())
        self.logger.debug(
            "Worker[%d] is continuing its work at task index [%d] on [%f], that is in [%s].",
            self.worker_id,
            self.current_task_index,
            msg.client_start_timestamp,
            sleep_time,
        )
        self.start_driving = True
        self.wakeupAfter(sleep_time)

    @actor.no_retry("worker")  # pylint: disable=no-value-for-parameter
    def receiveMsg_CompleteCurrentTask(self, msg, sender):
        # finish now ASAP. Remaining samples will be sent with the next WakeupMessage. We will also need to skip to the next
        # JoinPoint. But if we are already at a JoinPoint at the moment, there is nothing to do.
        if self.at_joinpoint():
            self.logger.info(
                "Worker[%s] has received CompleteCurrentTask but is currently at join point at index [%d]. Ignoring.",
                str(self.worker_id),
                self.current_task_index,
            )
        else:
            self.logger.info(
                "Worker[%s] has received CompleteCurrentTask. Completing tasks at index [%d].", str(self.worker_id), self.current_task_index
            )
            self.complete.set()

    @actor.no_retry("worker")  # pylint: disable=no-value-for-parameter
    def receiveMsg_WakeupMessage(self, msg, sender):
        # it would be better if we could send ourselves a message at a specific time, simulate this with a boolean...
        if self.start_driving:
            self.start_driving = False
            self.drive()
        else:
            current_samples = self.send_samples()
            if self.cancel.is_set():
                self.logger.info("Worker[%s] has detected that benchmark has been cancelled. Notifying master...", str(self.worker_id))
                self.send(self.driver_actor, actor.BenchmarkCancelled())
            elif self.executor_future is not None and self.executor_future.done():
                e = self.executor_future.exception(timeout=0)
                if e:
                    self.logger.exception(
                        "Worker[%s] has detected a benchmark failure. Notifying master...", str(self.worker_id), exc_info=e
                    )
                    # the exception might be user-defined and not be on the load path of the master driver. Hence, it cannot be
                    # deserialized on the receiver so we convert it here to a plain string.
                    self.send(self.driver_actor, actor.BenchmarkFailure(f"Error in load generator [{self.worker_id}]", str(e)))
                else:
                    self.logger.debug("Worker[%s] is ready for the next task.", str(self.worker_id))
                    self.executor_future = None
                    self.drive()
            else:
                if current_samples and len(current_samples) > 0:
                    most_recent_sample = current_samples[-1]
                    if most_recent_sample.percent_completed is not None:
                        self.logger.debug(
                            "Worker[%s] is executing [%s] (%.2f%% complete).",
                            str(self.worker_id),
                            most_recent_sample.task,
                            most_recent_sample.percent_completed * 100.0,
                        )
                    else:
                        # TODO: This could be misleading given that one worker could execute more than one task...
                        self.logger.debug(
                            "Worker[%s] is executing [%s] (dependent eternal task).", str(self.worker_id), most_recent_sample.task
                        )
                else:
                    self.logger.debug("Worker[%s] is executing (no samples).", str(self.worker_id))
                self.wakeupAfter(datetime.timedelta(seconds=self.wakeup_interval))

    def receiveMsg_ActorExitRequest(self, msg, sender):
        self.logger.debug("Worker[%s] has received ActorExitRequest.", str(self.worker_id))
        if self.executor_future is not None and self.executor_future.running():
            self.cancel.set()
        self.pool.shutdown()
        self.logger.debug("Worker[%s] is exiting due to ActorExitRequest.", str(self.worker_id))

    def receiveMsg_BenchmarkFailure(self, msg, sender):
        # sent by our no_retry infrastructure; forward to master
        self.send(self.driver_actor, msg)

    def receiveUnrecognizedMessage(self, msg, sender):
        self.logger.debug("Worker[%d] received unknown message [%s] (ignoring).", self.worker_id, str(msg))

    def drive(self):
        assert self.config is not None
        task_allocations = self.current_tasks_and_advance()
        # skip non-tasks in the task list
        while len(task_allocations) == 0:
            task_allocations = self.current_tasks_and_advance()

        if self.at_joinpoint():
            self.logger.debug("Worker[%d] reached join point at index [%d].", self.worker_id, self.current_task_index)
            # clients that don't execute tasks don't need to care about waiting
            if self.executor_future is not None:
                self.executor_future.result()
            self.send_samples()
            self.cancel.clear()
            self.complete.clear()
            self.executor_future = None
            self.sampler = None
            self.send(self.driver_actor, JoinPointReached(self.worker_id, task_allocations))
        else:
            # There may be a situation where there are more (parallel) tasks than workers. If we were asked to complete all tasks, we not
            # only need to complete actively running tasks but actually all scheduled tasks until we reach the next join point.
            if self.complete.is_set():
                self.logger.info(
                    "Worker[%d] skips tasks at index [%d] because it has been asked to complete all tasks until next join point.",
                    self.worker_id,
                    self.current_task_index,
                )
            else:
                self.logger.debug("Worker[%d] is executing tasks at index [%d].", self.worker_id, self.current_task_index)
                self.sampler = Sampler(start_timestamp=time.perf_counter(), buffer_size=self.sample_queue_size)
                executor = AsyncIoAdapter(
                    self.config,
                    self.track,
                    task_allocations,
                    self.sampler,
                    self.cancel,
                    self.complete,
                    self.on_error,
                    self.client_contexts,
                    self.worker_id,
                )

                self.executor_future = self.pool.submit(executor)
                self.wakeupAfter(datetime.timedelta(seconds=self.wakeup_interval))

    def at_joinpoint(self):
        return self.client_allocations.is_joinpoint(self.current_task_index)

    def current_tasks_and_advance(self):
        self.current_task_index = self.next_task_index
        current = self.client_allocations.tasks(self.current_task_index)
        self.next_task_index += 1
        self.logger.debug("Worker[%d] is at task index [%d].", self.worker_id, self.current_task_index)
        return current

    def send_samples(self):
        if self.sampler:
            samples = self.sampler.samples
            if len(samples) > 0:
                self.send(self.driver_actor, UpdateSamples(self.worker_id, samples))
            return samples
        return None


class Sampler:
    """
    Encapsulates management of gathered samples.
    """

    def __init__(self, start_timestamp, buffer_size=16384):
        self.start_timestamp = start_timestamp
        self.q = queue.Queue(maxsize=buffer_size)
        self.logger = logging.getLogger(__name__)

    def add(
        self,
        task,
        client_id,
        sample_type,
        meta_data,
        absolute_time,
        request_start,
        latency,
        service_time,
        processing_time,
        throughput,
        ops,
        ops_unit,
        time_period,
        percent_completed,
        dependent_timing=None,
    ):
        try:
            self.q.put_nowait(
                Sample(
                    client_id,
                    absolute_time,
                    request_start,
                    self.start_timestamp,
                    task,
                    sample_type,
                    meta_data,
                    latency,
                    service_time,
                    processing_time,
                    throughput,
                    ops,
                    ops_unit,
                    time_period,
                    percent_completed,
                    dependent_timing,
                )
            )
        except queue.Full:
            self.logger.warning("Dropping sample for [%s] due to a full sampling queue.", task.operation.name)

    @property
    def samples(self):
        samples = []
        try:
            while True:
                samples.append(self.q.get_nowait())
        except queue.Empty:
            pass
        return samples


class Sample:
    def __init__(
        self,
        client_id,
        absolute_time,
        request_start,
        task_start,
        task,
        sample_type,
        request_meta_data,
        latency,
        service_time,
        processing_time,
        throughput,
        total_ops,
        total_ops_unit,
        time_period,
        percent_completed,
        dependent_timing=None,
        operation_name=None,
        operation_type=None,
    ):
        self.client_id = client_id
        self.absolute_time = absolute_time
        self.request_start = request_start
        self.task_start = task_start
        self.task = task
        self.sample_type = sample_type
        self.request_meta_data = request_meta_data
        self.latency = latency
        self.service_time = service_time
        self.processing_time = processing_time
        self.throughput = throughput
        self.total_ops = total_ops
        self.total_ops_unit = total_ops_unit
        self.time_period = time_period
        self._dependent_timing = dependent_timing
        self._operation_name = operation_name
        self._operation_type = operation_type
        # may be None for eternal tasks!
        self.percent_completed = percent_completed

    @property
    def operation_name(self):
        return self._operation_name if self._operation_name else self.task.operation.name

    @property
    def operation_type(self):
        return self._operation_type if self._operation_type else self.task.operation.type

    @property
    def operation_meta_data(self):
        return self.task.operation.meta_data

    @property
    def relative_time(self):
        return self.request_start - self.task_start

    @property
    def dependent_timings(self):
        if self._dependent_timing:
            for t in self._dependent_timing:
                timing = t.pop("dependent_timing")
                meta_data = self._merge(self.request_meta_data, t)
                yield Sample(
                    self.client_id,
                    timing["absolute_time"],
                    timing["request_start"],
                    self.task_start,
                    self.task,
                    self.sample_type,
                    meta_data,
                    0,
                    timing["service_time"],
                    0,
                    0,
                    self.total_ops,
                    self.total_ops_unit,
                    self.time_period,
                    self.percent_completed,
                    None,
                    timing["operation"],
                    timing["operation-type"],
                )

    def __repr__(self, *args, **kwargs):
        return (
            f"[{self.absolute_time}; {self.relative_time}] [client [{self.client_id}]] [{self.task}] "
            f"[{self.sample_type}]: [{self.latency}s] request latency, [{self.service_time}s] service time, "
            f"[{self.total_ops} {self.total_ops_unit}]"
        )

    def _merge(self, *args):
        result = {}
        for arg in args:
            if arg is not None:
                result.update(arg)
        return result


def select_challenge(config: types.Config, t):
    challenge_name = config.opts("track", "challenge.name")
    selected_challenge = t.find_challenge_or_default(challenge_name)

    if not selected_challenge:
        raise exceptions.SystemSetupError(
            "Unknown challenge [%s] for track [%s]. You can list the available tracks and their "
            "challenges with %s list tracks." % (challenge_name, t.name, PROGRAM_NAME)
        )
    return selected_challenge


class ThroughputCalculator:
    class TaskStats:
        """
        Stores per task numbers needed for throughput calculation in between multiple calculations.
        """

        def __init__(self, bucket_interval, sample_type, start_time):
            self.unprocessed = []
            self.total_count = 0
            self.interval = 0
            self.bucket_interval = bucket_interval
            # the first bucket is complete after one bucket interval is over
            self.bucket = bucket_interval
            self.sample_type = sample_type
            self.has_samples_in_sample_type = False
            # start relative to the beginning of our (calculation) time slice.
            self.start_time = start_time

        @property
        def throughput(self):
            return self.total_count / self.interval

        def maybe_update_sample_type(self, current_sample_type):
            if self.sample_type < current_sample_type:
                self.sample_type = current_sample_type
                self.has_samples_in_sample_type = False

        def update_interval(self, absolute_sample_time):
            self.interval = max(absolute_sample_time - self.start_time, self.interval)

        def can_calculate_throughput(self):
            return self.interval > 0 and self.interval >= self.bucket

        def can_add_final_throughput_sample(self):
            return self.interval > 0 and not self.has_samples_in_sample_type

        def finish_bucket(self, new_total):
            self.unprocessed = []
            self.total_count = new_total
            self.has_samples_in_sample_type = True
            self.bucket = int(self.interval) + self.bucket_interval

    def __init__(self):
        self.task_stats = {}

    def calculate(self, samples, bucket_interval_secs=1):
        """
        Calculates global throughput based on samples gathered from multiple load generators.

        :param samples: A list containing all samples from all load generators.
        :param bucket_interval_secs: The bucket interval for aggregations.
        :return: A global view of throughput samples.
        """

        samples_per_task = {}
        # first we group all samples by task (operation).
        for sample in samples:
            k = sample.task
            if k not in samples_per_task:
                samples_per_task[k] = []
            samples_per_task[k].append(sample)

        global_throughput = {}
        # with open("raw_samples_new.csv", "a") as sample_log:
        # print("client_id,absolute_time,relative_time,operation,sample_type,total_ops,time_period", file=sample_log)
        for k, v in samples_per_task.items():
            task = k
            if task not in global_throughput:
                global_throughput[task] = []
            # sort all samples by time
            if task in self.task_stats:
                samples = itertools.chain(v, self.task_stats[task].unprocessed)
            else:
                samples = v
            current_samples = sorted(samples, key=lambda s: s.absolute_time)

            # Calculate throughput based on service time if the runner does not provide one, otherwise use it as is and
            # only transform the values into the expected structure.
            first_sample = current_samples[0]
            if first_sample.throughput is None:
                task_throughput = self.calculate_task_throughput(task, current_samples, bucket_interval_secs)
            else:
                task_throughput = self.map_task_throughput(current_samples)
            global_throughput[task].extend(task_throughput)

        return global_throughput

    def calculate_task_throughput(self, task, current_samples, bucket_interval_secs):
        task_throughput = []

        if task not in self.task_stats:
            first_sample = current_samples[0]
            self.task_stats[task] = ThroughputCalculator.TaskStats(
                bucket_interval=bucket_interval_secs,
                sample_type=first_sample.sample_type,
                start_time=first_sample.absolute_time - first_sample.time_period,
            )
        current = self.task_stats[task]
        count = current.total_count
        last_sample = None
        for sample in current_samples:
            last_sample = sample
            # print("%d,%f,%f,%s,%s,%d,%f" %
            #       (sample.client_id, sample.absolute_time, sample.relative_time, sample.operation, sample.sample_type,
            #        sample.total_ops, sample.time_period), file=sample_log)

            # once we have seen a new sample type, we stick to it.
            current.maybe_update_sample_type(sample.sample_type)

            # we need to store the total count separately and cannot update `current.total_count` immediately here
            # because we would count all raw samples in `unprocessed` twice. Hence, we'll only update
            # `current.total_count` when we have calculated a new throughput sample.
            count += sample.total_ops
            current.update_interval(sample.absolute_time)

            if current.can_calculate_throughput():
                current.finish_bucket(count)
                task_throughput.append(
                    (
                        sample.absolute_time,
                        sample.relative_time,
                        current.sample_type,
                        current.throughput,
                        # we calculate throughput per second
                        f"{sample.total_ops_unit}/s",
                    )
                )
            else:
                current.unprocessed.append(sample)

        # also include the last sample if we don't have one for the current sample type, even if it is below the bucket
        # interval (mainly needed to ensure we show throughput data in test mode)
        if last_sample is not None and current.can_add_final_throughput_sample():
            current.finish_bucket(count)
            task_throughput.append(
                (
                    last_sample.absolute_time,
                    last_sample.relative_time,
                    current.sample_type,
                    current.throughput,
                    f"{last_sample.total_ops_unit}/s",
                )
            )

        return task_throughput

    def map_task_throughput(self, current_samples):
        throughput = []
        for sample in current_samples:
            throughput.append(
                (
                    sample.absolute_time,
                    sample.relative_time,
                    sample.sample_type,
                    sample.throughput,
                    f"{sample.total_ops_unit}/s",
                )
            )
        return throughput


class AsyncIoAdapter:
    def __init__(self, cfg: types.Config, track, task_allocations, sampler, cancel, complete, abort_on_error, client_contexts, worker_id):
        self.cfg = cfg
        self.track = track
        self.task_allocations = task_allocations
        self.sampler = sampler
        self.cancel = cancel
        self.complete = complete
        self.abort_on_error = abort_on_error
        self.client_contexts = client_contexts
        self.parent_worker_id = worker_id
        self.profiling_enabled = self.cfg.opts("driver", "profiling")
        self.assertions_enabled = self.cfg.opts("driver", "assertions")
        self.debug_event_loop = self.cfg.opts("system", "async.debug", mandatory=False, default_value=False)
        self.logger = logging.getLogger(__name__)

    def __call__(self, *args, **kwargs):
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            loop.set_debug(self.debug_event_loop)
            loop.set_exception_handler(self._logging_exception_handler)
            asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(self.run())
        finally:
            loop.close()

    def _logging_exception_handler(self, loop, context):
        self.logger.error("Uncaught exception in event loop: %s", context)

    async def run(self):
        def es_clients(client_id, all_hosts, all_client_options, distribution_version, distribution_flavor):
            es = {}
            context = self.client_contexts.get(client_id)
            api_key = context.api_key
            for cluster_name, cluster_hosts in all_hosts.items():
                es[cluster_name] = client.EsClientFactory(
                    cluster_hosts,
                    all_client_options[cluster_name],
                    distribution_version=distribution_version,
                    distribution_flavor=distribution_flavor,
                ).create_async(api_key=api_key, client_id=client_id)
            return es

        if self.assertions_enabled:
            self.logger.info("Task assertions enabled")
        runner.enable_assertions(self.assertions_enabled)

        clients = []
        awaitables = []
        # A parameter source should only be created once per task - it is partitioned later on per client.
        params_per_task = {}
        for client_id, task_allocation in self.task_allocations:
            task = task_allocation.task
            if task not in params_per_task:
                param_source = track.operation_parameters(self.track, task)
                params_per_task[task] = param_source
            schedule = schedule_for(task_allocation, params_per_task[task])
            es = es_clients(
                client_id,
                self.cfg.opts("client", "hosts").all_hosts,
                self.cfg.opts("client", "options"),
                self.cfg.opts("mechanic", "distribution.version", mandatory=False),
                self.cfg.opts("mechanic", "distribution.flavor", mandatory=False),
            )
            clients.append(es)
            async_executor = AsyncExecutor(
                client_id, task, schedule, es, self.sampler, self.cancel, self.complete, task.error_behavior(self.abort_on_error)
            )
            final_executor = AsyncProfiler(async_executor) if self.profiling_enabled else async_executor
            awaitables.append(final_executor())
        task_names = [t.task.task.name for t in self.task_allocations]
        self.logger.info("Worker[%s] executing tasks: %s", self.parent_worker_id, task_names)
        run_start = time.perf_counter()
        try:
            _ = await asyncio.gather(*awaitables)
        finally:
            run_end = time.perf_counter()
            self.logger.info(
                "Worker[%s] finished executing tasks %s in %f seconds", self.parent_worker_id, task_names, (run_end - run_start)
            )
            await asyncio.get_event_loop().shutdown_asyncgens()
            shutdown_asyncgens_end = time.perf_counter()
            self.logger.debug("Total time to shutdown asyncgens: %f seconds.", (shutdown_asyncgens_end - run_end))
            for c in clients:
                for es in c.values():
                    await es.close()
            transport_close_end = time.perf_counter()
            self.logger.debug("Total time to close transports: %f seconds.", (transport_close_end - shutdown_asyncgens_end))


class AsyncProfiler:
    def __init__(self, target):
        """
        :param target: The actual executor which should be profiled.
        """
        self.target = target
        self.profile_logger = logging.getLogger("rally.profile")

    async def __call__(self, *args, **kwargs):
        # initialize lazily, we don't need it in the majority of cases
        # pylint: disable=import-outside-toplevel
        import io as python_io

        import yappi

        yappi.start()
        try:
            return await self.target(*args, **kwargs)
        finally:
            yappi.stop()
            s = python_io.StringIO()
            yappi.get_func_stats().print_all(
                out=s, columns={0: ("name", 140), 1: ("ncall", 8), 2: ("tsub", 8), 3: ("ttot", 8), 4: ("tavg", 8)}
            )

            profile = "\n=== Profile START ===\n"
            profile += s.getvalue()
            profile += "=== Profile END ==="
            self.profile_logger.info(profile)


class AsyncExecutor:
    def __init__(self, client_id, task, schedule, es, sampler, cancel, complete, on_error):
        """
        Executes tasks according to the schedule for a given operation.

        :param task: The task that is executed.
        :param schedule: The schedule for this task.
        :param es: Elasticsearch client that will be used to execute the operation.
        :param sampler: A container to store raw samples.
        :param cancel: A shared boolean that indicates we need to cancel execution.
        :param complete: A shared boolean that indicates we need to prematurely complete execution.
        :param on_error: A string specifying how the load generator should behave on errors.
        """
        self.client_id = client_id
        self.task = task
        self.op = task.operation
        self.schedule_handle = schedule
        self.es = es
        self.sampler = sampler
        self.cancel = cancel
        self.complete = complete
        self.on_error = on_error
        self.logger = logging.getLogger(__name__)

    async def __call__(self, *args, **kwargs):
        any_task_completes_parent = self.task.any_completes_parent
        task_completes_parent = self.task.completes_parent
        total_start = time.perf_counter()
        # lazily initialize the schedule
        self.logger.debug("Initializing schedule for client id [%s].", self.client_id)
        schedule = self.schedule_handle()
        # Start the schedule's timer early so the warmup period is independent of any deferred start due to ramp-up
        self.schedule_handle.start()
        rampup_wait_time = self.schedule_handle.ramp_up_wait_time
        if rampup_wait_time:
            self.logger.debug("client id [%s] waiting [%.2f]s for ramp-up.", self.client_id, rampup_wait_time)
            await asyncio.sleep(rampup_wait_time)

        self.logger.debug("Entering main loop for client id [%s].", self.client_id)
        # noinspection PyBroadException
        try:
            async for expected_scheduled_time, sample_type, percent_completed, runner, params in schedule:
                if self.cancel.is_set():
                    self.logger.info("User cancelled execution.")
                    break
                absolute_expected_schedule_time = total_start + expected_scheduled_time
                throughput_throttled = expected_scheduled_time > 0
                if throughput_throttled:
                    rest = absolute_expected_schedule_time - time.perf_counter()
                    if rest > 0:
                        await asyncio.sleep(rest)

                absolute_processing_start = time.time()
                processing_start = time.perf_counter()
                self.schedule_handle.before_request(processing_start)
                with self.es["default"].new_request_context() as request_context:
                    total_ops, total_ops_unit, request_meta_data = await execute_single(runner, self.es, params, self.on_error)
                    request_start = request_context.request_start
                    request_end = request_context.request_end

                processing_end = time.perf_counter()
                service_time = request_end - request_start
                processing_time = processing_end - processing_start
                time_period = request_end - total_start
                self.schedule_handle.after_request(processing_end, total_ops, total_ops_unit, request_meta_data)
                # Allow runners to override the throughput calculation in very specific circumstances. Usually, Rally
                # assumes that throughput is the "amount of work" (determined by the "weight") per unit of time
                # (determined by the elapsed time period). However, in certain cases (e.g. shard recovery or other
                # long running operations where there is a dedicated stats API to determine progress), it is
                # advantageous if the runner calculates throughput directly. The following restrictions apply:
                #
                # * Only one client must call that runner (when throughput is calculated, it is aggregated across
                #   all clients but if the runner provides it, we take the value as is).
                # * The runner should be rate-limited as each runner call will result in one throughput sample.
                #
                throughput = request_meta_data.pop("throughput", None)
                # Do not calculate latency separately when we run unthrottled. This metric is just confusing then.
                latency = request_end - absolute_expected_schedule_time if throughput_throttled else service_time
                # If this task completes the parent task we should *not* check for completion by another client but
                # instead continue until our own runner has completed. We need to do this because the current
                # worker (process) could run multiple clients that execute the same task. We do not want all clients to
                # finish this task as soon as the first of these clients has finished but rather continue until the last
                # client has finished that task.
                if task_completes_parent:
                    completed = runner.completed
                else:
                    completed = self.complete.is_set() or runner.completed
                # last sample should bump progress to 100% if externally completed.
                if completed:
                    progress = 1.0
                elif runner.percent_completed:
                    progress = runner.percent_completed
                else:
                    progress = percent_completed

                self.sampler.add(
                    self.task,
                    self.client_id,
                    sample_type,
                    request_meta_data,
                    absolute_processing_start,
                    request_start,
                    latency,
                    service_time,
                    processing_time,
                    throughput,
                    total_ops,
                    total_ops_unit,
                    time_period,
                    progress,
                    request_meta_data.pop("dependent_timing", None),
                )

                if completed:
                    self.logger.info("Task [%s] is considered completed due to external event.", self.task)
                    break
        except BaseException as e:
            self.logger.exception("Could not execute schedule")
            raise exceptions.RallyError(f"Cannot run task [{self.task}]: {e}") from None
        finally:
            # Actively set it if this task completes its parent
            if task_completes_parent:
                self.logger.info(
                    "Task [%s] completes parent. Client id [%s] is finished executing it and signals completion.",
                    self.task,
                    self.client_id,
                )
                self.complete.set()
            elif any_task_completes_parent:
                self.logger.info(
                    "Task [%s] completes parent. Client id [%s] is finished executing it and signals completion of all "
                    "remaining clients, immediately.",
                    self.task,
                    self.client_id,
                )
                self.complete.set()


async def execute_single(runner, es, params, on_error):
    """
    Invokes the given runner once and provides the runner's return value in a uniform structure.

    :return: a triple of: total number of operations, unit of operations, a dict of request meta data (may be None).
    """
    # pylint: disable=import-outside-toplevel
    import elasticsearch

    fatal_error = False
    try:
        async with runner:
            return_value = await runner(es, params)
        if isinstance(return_value, tuple) and len(return_value) == 2:
            total_ops, total_ops_unit = return_value
            request_meta_data = {"success": True}
        elif isinstance(return_value, dict):
            total_ops = return_value.pop("weight", 1)
            total_ops_unit = return_value.pop("unit", "ops")
            request_meta_data = return_value
            if "success" not in request_meta_data:
                request_meta_data["success"] = True
        else:
            total_ops = 1
            total_ops_unit = "ops"
            request_meta_data = {"success": True}
    except elasticsearch.TransportError as e:
        # we *specifically* want to distinguish connection refused (a node died?) from connection timeouts
        # pylint: disable=unidiomatic-typecheck
        if type(e) is elasticsearch.ConnectionError:
            fatal_error = True

        total_ops = 0
        total_ops_unit = "ops"
        request_meta_data = {"success": False, "error-type": "transport"}
        # For the 'errors' attribute, errors are ordered from
        # most recently raised (index=0) to least recently raised (index=N)
        #
        # If an HTTP status code is available with the error it
        # will be stored under 'status'. If HTTP headers are available
        # they are stored under 'headers'.
        if e.errors:
            if hasattr(e.errors[0], "status"):
                request_meta_data["http-status"] = e.errors[0].status
        # connection timeout errors don't provide a helpful description
        if isinstance(e, elasticsearch.ConnectionTimeout):
            request_meta_data["error-description"] = "network connection timed out"
        else:
            error_description = e.message
            request_meta_data["error-description"] = error_description
    except elasticsearch.ApiError as e:
        total_ops = 0
        total_ops_unit = "ops"
        request_meta_data = {"success": False, "error-type": "api"}
        error_message = ""

        # Some runners return a raw response, causing the 'error' property to be a string literal of the bytes/BytesIO object,
        # we should avoid bubbling that up
        # e.g. ApiError(413, '<_io.BytesIO object at 0xffffaf146a70>')
        if isinstance(e.body, bytes):
            # could be an empty body
            if error_body := e.body.decode("utf-8"):
                error_message = error_body
            else:
                # to be consistent with an empty 'e.error'
                error_message = str(None)
        elif isinstance(e.body, BytesIO):
            # could be an empty body
            if error_body := e.body.read().decode("utf-8"):
                error_message = error_body
            else:
                # to be consistent with an empty 'e.error'
                error_message = str(None)
        # fallback to 'error' property if the body isn't bytes/BytesIO
        else:
            if isinstance(e.error, bytes):
                error_message = e.error.decode("utf-8")
            elif isinstance(e.error, BytesIO):
                error_message = e.error.read().decode("utf-8")
            else:
                # if the 'error' is empty, we get back str(None)
                error_message = e.error

        if isinstance(e.info, bytes):
            error_info = e.info.decode("utf-8")
        elif isinstance(e.info, BytesIO):
            error_info = e.info.read().decode("utf-8")
        else:
            error_info = e.info

        if error_info:
            error_message += f" ({error_info})"

        error_description = error_message

        request_meta_data["error-description"] = error_description
        if e.status_code:
            request_meta_data["http-status"] = e.status_code

    except KeyError as e:
        logging.getLogger(__name__).exception("Cannot execute runner [%s]; most likely due to missing parameters.", str(runner))
        msg = "Cannot execute [%s]. Provided parameters are: %s. Error: [%s]." % (str(runner), list(params.keys()), str(e))
        raise exceptions.SystemSetupError(msg)

    if not request_meta_data["success"]:
        if on_error == "abort" or fatal_error:
            msg = "Request returned an error. Error type: %s" % request_meta_data.get("error-type", "Unknown")

            if description := request_meta_data.get("error-description"):
                msg += f", Description: {description}"

            if http_status := request_meta_data.get("http-status"):
                msg += f", HTTP Status: {http_status}"

            raise exceptions.RallyAssertionError(msg)

    return total_ops, total_ops_unit, request_meta_data


class JoinPoint:
    def __init__(self, id, clients_executing_completing_task=None, any_task_completes_parent=None):
        """

        :param id: The join point's id.
        :param clients_executing_completing_task: An array of client indices which execute a task that can prematurely complete its parent
        element. Provide 'None' or an empty array if no task satisfies this predicate.
        """
        if clients_executing_completing_task is None:
            clients_executing_completing_task = []
        if any_task_completes_parent is None:
            any_task_completes_parent = []
        self.id = id
        self.any_task_completes_parent = any_task_completes_parent
        self.clients_executing_completing_task = clients_executing_completing_task
        self.num_clients_executing_completing_task = len(clients_executing_completing_task)
        self.preceding_task_completes_parent = self.num_clients_executing_completing_task > 0

    def __hash__(self):
        return hash(self.id)

    def __eq__(self, other):
        return isinstance(other, type(self)) and self.id == other.id

    def __repr__(self, *args, **kwargs):
        return "JoinPoint(%s)" % self.id


class TaskAllocation:
    def __init__(self, task, client_index_in_task, global_client_index, total_clients):
        """

        :param task: The current task which is always a leaf task.
        :param client_index_in_task: The task-specific index for the allocated client.
        :param global_client_index:  The globally unique index for the allocated client across
                                     all concurrently executed tasks.
        :param total_clients: The total number of clients executing tasks concurrently.
        """
        self.task = task
        self.client_index_in_task = client_index_in_task
        self.global_client_index = global_client_index
        self.total_clients = total_clients

    def __hash__(self):
        return hash(self.task) ^ hash(self.global_client_index)

    def __eq__(self, other):
        return isinstance(other, type(self)) and self.task == other.task and self.global_client_index == other.global_client_index

    def __repr__(self, *args, **kwargs):
        return (
            f"TaskAllocation [{self.client_index_in_task}/{self.task.clients}] for {self.task} "
            f"and [{self.global_client_index}/{self.total_clients}] in total"
        )


class Allocator:
    """
    Decides which operations runs on which client and how to partition them.
    """

    def __init__(self, schedule):
        self.schedule = schedule

    @property
    def allocations(self):
        """
        Calculates an allocation matrix consisting of two dimensions. The first dimension is the client. The second dimension are the task
         this client needs to run. The matrix shape is rectangular (i.e. it is not ragged). There are three types of entries in the matrix:

          1. Normal tasks: They need to be executed by a client.
          2. Join points: They are used as global coordination points which all clients need to reach until the benchmark can go on. They
                          indicate that a client has to wait until the master signals it can go on.
          3. `None`: These are inserted by the allocator to keep the allocation matrix rectangular. Clients have to skip `None` entries
                     until one of the other entry types are encountered.

        :return: An allocation matrix with the structure described above.
        """
        max_clients = self.clients
        allocations = [None] * max_clients
        for client_index in range(max_clients):
            allocations[client_index] = []
        join_point_id = 0
        # start with an artificial join point to allow master to coordinate that all clients start at the same time
        next_join_point = JoinPoint(join_point_id)
        for client_index in range(max_clients):
            allocations[client_index].append(next_join_point)
        join_point_id += 1

        for task in self.schedule:
            start_client_index = 0
            clients_executing_completing_task = []
            any_task_completes_parent = []
            for sub_task in task:
                for client_index in range(start_client_index, start_client_index + sub_task.clients):
                    # this is the actual client that will execute the task. It may differ from the logical one in case we over-commit (i.e.
                    # more tasks than actually available clients)
                    physical_client_index = client_index % max_clients
                    if sub_task.completes_parent:
                        clients_executing_completing_task.append(physical_client_index)
                    elif sub_task.any_completes_parent:
                        any_task_completes_parent.append(physical_client_index)

                    ta = TaskAllocation(
                        task=sub_task,
                        client_index_in_task=client_index - start_client_index,
                        global_client_index=client_index,
                        # if task represents a parallel structure this is the total number of clients
                        # executing sub-tasks concurrently.
                        total_clients=task.clients,
                    )
                    allocations[physical_client_index].append(ta)
                start_client_index += sub_task.clients
            # uneven distribution between tasks and clients, e.g. there are 5 (parallel) tasks but only 2 clients. Then, one of them
            # executes three tasks, the other one only two. So we need to fill in a `None` for the second one.
            if start_client_index % max_clients > 0:
                # pin the index range to [0, max_clients). This simplifies the code below.
                start_client_index = start_client_index % max_clients
                for client_index in range(start_client_index, max_clients):
                    allocations[client_index].append(None)

            # let all clients join after each task, then we go on
            next_join_point = JoinPoint(join_point_id, clients_executing_completing_task, any_task_completes_parent)
            for client_index in range(max_clients):
                allocations[client_index].append(next_join_point)
            join_point_id += 1
        return allocations

    @property
    def join_points(self):
        """
        :return: A list of all join points for this allocations.
        """
        return [allocation for allocation in self.allocations[0] if isinstance(allocation, JoinPoint)]

    @property
    def tasks_per_joinpoint(self):
        """

        Calculates a flat list of all tasks that are run in between join points.

        Consider the following schedule (2 clients):

        1. task1 and task2 run by both clients in parallel
        2. join point
        3. task3 run by client 1
        4. join point

        The results in: [{task1, task2}, {task3}]

        :return: A list of sets containing all tasks.
        """
        tasks = []
        current_tasks = set()

        allocs = self.allocations
        # assumption: the shape of allocs is rectangular (i.e. each client contains the same number of elements)
        for idx in range(0, len(allocs[0])):
            for client in range(0, self.clients):
                allocation = allocs[client][idx]
                if isinstance(allocation, TaskAllocation):
                    current_tasks.add(allocation.task)
                elif isinstance(allocation, JoinPoint) and len(current_tasks) > 0:
                    tasks.append(current_tasks)
                    current_tasks = set()

        return tasks

    @property
    def clients(self):
        """
        :return: The maximum number of clients involved in executing the given schedule.
        """
        max_clients = 1
        for task in self.schedule:
            max_clients = max(max_clients, task.clients)
        return max_clients


#######################################
#
# Scheduler related stuff
#
#######################################


# Runs a concrete schedule on one worker client
# Needs to determine the runners and concrete iterations per client.
def schedule_for(task_allocation, parameter_source):
    """
    Calculates a client's schedule for a given task.

    :param task_allocation: The task allocation that should be executed by this schedule.
    :param parameter_source: The parameter source that should be used for this task.
    :return: A generator for the operations the given client needs to perform for this task.
    """
    logger = logging.getLogger(__name__)
    task = task_allocation.task
    op = task.operation
    sched = scheduler.scheduler_for(task)

    # We cannot use the global client index here because we need to support parallel execution of tasks
    # with multiple clients. Consider the following scenario:
    #
    # * Clients 0-3 bulk index into indexA
    # * Clients 4-7 bulk index into indexB
    #
    # Now we need to ensure that we start partitioning parameters correctly in both cases. And that means we
    # need to start from (client) index 0 in both cases instead of 0 for indexA and 4 for indexB.
    client_index = task_allocation.client_index_in_task

    # guard all logging statements with the client index and only emit them for the first client. This information is
    # repetitive and may cause issues in thespian with many clients (an excessive number of actor messages is sent).
    if client_index == 0:
        logger.debug("Choosing [%s] for [%s].", sched, task)
    runner_for_op = runner.runner_for(op.type)
    params_for_op = parameter_source.partition(client_index, task.clients)
    if hasattr(sched, "parameter_source"):
        if client_index == 0:
            logger.debug("Setting parameter source [%s] for scheduler [%s]", params_for_op, sched)
        sched.parameter_source = params_for_op

    if requires_time_period_schedule(task, runner_for_op, params_for_op):
        warmup_time_period = task.warmup_time_period if task.warmup_time_period else 0
        if client_index == 0:
            logger.info(
                "Creating time-period based schedule with [%s] distribution for [%s] with a warmup period of [%s] "
                "seconds and a time period of [%s] seconds.",
                task.schedule,
                task.name,
                str(warmup_time_period),
                str(task.time_period),
            )
        loop_control = TimePeriodBased(warmup_time_period, task.time_period)
    else:
        warmup_iterations = task.warmup_iterations if task.warmup_iterations else 0
        if task.iterations:
            iterations = task.iterations
        elif params_for_op.infinite:
            # this is usually the case if the parameter source provides a constant
            iterations = 1
        else:
            iterations = None
        if client_index == 0:
            logger.debug(
                "Creating iteration-count based schedule with [%s] distribution for [%s] with [%s] warmup "
                "iterations and [%s] iterations.",
                task.schedule,
                task.name,
                str(warmup_iterations),
                str(iterations),
            )
        loop_control = IterationBased(warmup_iterations, iterations)

    if client_index == 0:
        if loop_control.infinite:
            logger.debug("Parameter source will determine when the schedule for [%s] terminates.", task.name)
        else:
            logger.debug("%s schedule will determine when the schedule for [%s] terminates.", str(loop_control), task.name)

    return ScheduleHandle(task_allocation, sched, loop_control, runner_for_op, params_for_op)


def requires_time_period_schedule(task, task_runner, params):
    if task.warmup_time_period is not None or task.time_period is not None:
        return True
    # user has explicitly requested iterations
    if task.warmup_iterations is not None or task.iterations is not None:
        return False
    # the runner determines completion
    if task_runner.completed is not None:
        return True
    # If the parameter source ends after a finite amount of iterations, we will run with a time-based schedule
    return not params.infinite


class ScheduleHandle:
    def __init__(self, task_allocation, sched, task_progress_control, runner, params):
        """
        Creates a generator that will yield individual task invocations for the provided schedule.

        :param task_allocation: The task allocation for which the schedule is generated.
        :param sched: The scheduler for this task.
        :param task_progress_control: Controls how and how often this generator will loop.
        :param runner: The runner for a given operation.
        :param params: The parameter source for a given operation.
        :return: A generator for the corresponding parameters.
        """
        self.task_allocation = task_allocation
        self.operation_type = task_allocation.task.operation.type
        self.sched = sched
        self.task_progress_control = task_progress_control
        self.runner = runner
        self.params = params
        # TODO: Can we offload the parameter source execution to a different thread / process? Is this too heavy-weight?
        # from concurrent.futures import ThreadPoolExecutor
        # import asyncio
        # self.io_pool_exc = ThreadPoolExecutor(max_workers=1)
        # self.loop = asyncio.get_event_loop()

    @property
    def ramp_up_wait_time(self):
        """
        :return: the number of seconds to wait until this client should start so load can gradually ramp-up.
        """
        ramp_up_time_period = self.task_allocation.task.ramp_up_time_period
        if ramp_up_time_period:
            return ramp_up_time_period * (self.task_allocation.global_client_index / self.task_allocation.total_clients)
        else:
            return 0

    def start(self):
        self.task_progress_control.start()

    def before_request(self, now):
        self.sched.before_request(now)

    def after_request(self, now, weight, unit, request_meta_data):
        self.sched.after_request(now, weight, unit, request_meta_data)

    def params_with_operation_type(self):
        p = self.params.params()
        p.update({"operation-type": self.operation_type})
        return p

    async def __call__(self):
        next_scheduled = 0
        if self.task_progress_control.infinite:
            param_source_knows_progress = hasattr(self.params, "percent_completed")
            while True:
                try:
                    next_scheduled = self.sched.next(next_scheduled)
                    # does not contribute at all to completion. Hence, we cannot define completion.
                    percent_completed = self.params.percent_completed if param_source_knows_progress else None
                    # current_params = await self.loop.run_in_executor(self.io_pool_exc, self.params.params)
                    yield (
                        next_scheduled,
                        self.task_progress_control.sample_type,
                        percent_completed,
                        self.runner,
                        self.params_with_operation_type(),
                    )
                    self.task_progress_control.next()
                except StopIteration:
                    return
        else:
            while not self.task_progress_control.completed:
                try:
                    next_scheduled = self.sched.next(next_scheduled)
                    # current_params = await self.loop.run_in_executor(self.io_pool_exc, self.params.params)
                    yield (
                        next_scheduled,
                        self.task_progress_control.sample_type,
                        self.task_progress_control.percent_completed,
                        self.runner,
                        self.params_with_operation_type(),
                    )
                    self.task_progress_control.next()
                except StopIteration:
                    return


class TimePeriodBased:
    def __init__(self, warmup_time_period, time_period):
        self._warmup_time_period = warmup_time_period
        self._time_period = time_period
        if warmup_time_period is not None and time_period is not None:
            self._duration = self._warmup_time_period + self._time_period
        else:
            self._duration = None
        self._start = None
        self._now = None

    def start(self):
        self._now = time.perf_counter()
        self._start = self._now

    @property
    def _elapsed(self):
        return self._now - self._start

    @property
    def sample_type(self):
        return metrics.SampleType.Warmup if self._elapsed < self._warmup_time_period else metrics.SampleType.Normal

    @property
    def infinite(self):
        return self._time_period is None

    @property
    def percent_completed(self):
        return self._elapsed / self._duration

    @property
    def completed(self):
        return self._now >= (self._start + self._duration)

    def next(self):
        self._now = time.perf_counter()

    def __str__(self):
        return "time-period-based"


class IterationBased:
    def __init__(self, warmup_iterations, iterations):
        self._warmup_iterations = warmup_iterations
        self._iterations = iterations
        if warmup_iterations is not None and iterations is not None:
            self._total_iterations = self._warmup_iterations + self._iterations
            if self._total_iterations == 0:
                raise exceptions.RallyAssertionError("Operation must run at least for one iteration.")
        else:
            self._total_iterations = None
        self._it = None

    def start(self):
        self._it = 0

    @property
    def sample_type(self):
        return metrics.SampleType.Warmup if self._it < self._warmup_iterations else metrics.SampleType.Normal

    @property
    def infinite(self):
        return self._iterations is None

    @property
    def percent_completed(self):
        return (self._it + 1) / self._total_iterations

    @property
    def completed(self):
        return self._it >= self._total_iterations

    def next(self):
        self._it += 1

    def __str__(self):
        return "iteration-count-based"