Repository: elastic/rally

Issue:
Allow Tracks to Declare Additional Dependencies
If track code would require additional libraries to be installed, it would be good for them to be able to declare them.

I propose we extend the Track Registry concept to allow for additional dependencies, for example:

**track.py**
```python
def register(registry):
  registry.register_dependency("numpy", "==1.21.4")
```

Which would be calling...
```python
def register_dependency(library, version_string):
```

...and would invoke the runtime-local `pip` in-code, in the fashion mentioned [in the pip docs](https://pip.pypa.io/en/latest/user_guide/#using-pip-from-your-program) (as a `subprocess`).

Potential goblins:
1. I'm not sure whether we should try to perform an uninstall of any additional libraries after the race ends (or on any failure, or whatever)
2. Registry invocations are re-run in every worker process, due to `load_track` being used as a way to ensure all Actors have the track code classloaded.  We would likely only want to run `pip` on the first call per `DriverActor`, so some check for the library's existence (and compliance with the `version_string`) is appropriate

Discussion:
pquentin: What is the status of this? As far as I can tell the tests pass, the failure is a transient one.

@elasticmachine please run tests
DJRickyB: @pquentin The failure seems legit. I will fix the test, document the feature, then re-request your review when it's ready.
DJRickyB: Reviewers, I've removed the "experimental" designation from the description and ":internal" from the labels. Our incubation period/process for experimental features is not well-defined, and I think bringing one more thing out of experimental state without a defined process just adds to our team's considerable tech debt.
michaelbaamonde: Late to the party, but I'll take a look at this today.

Patch:
diff --git a/docs/command_line_reference.rst b/docs/command_line_reference.rst
index 82bf03ea..34d79f19 100644
--- a/docs/command_line_reference.rst
+++ b/docs/command_line_reference.rst
@@ -181,7 +181,7 @@ Selects the track repository that Rally should use to resolve tracks. By default
 ``track-revision``
 ~~~~~~~~~~~~~~~~~~
 
-Selects a specific revision in the track repository. By default, Rally will choose the most appropriate branch on its own but in some cases it is necessary to specify a certain commit. This is mostly needed when testing whether a change in performance has occurred due to a change in the workload. 
+Selects a specific revision in the track repository. By default, Rally will choose the most appropriate branch on its own but in some cases it is necessary to specify a certain commit. This is mostly needed when testing whether a change in performance has occurred due to a change in the workload.
 
 ``track``
 ~~~~~~~~~
@@ -656,7 +656,7 @@ Default value: ``timeout:60`` (applies any time ``timeout`` is not specified)
 Rally recognizes the following client options in addition:
 
 * ``max_connections``: By default, Rally will choose the maximum allowed number of connections automatically (equal to the number of simulated clients but at least 256 connections). With this property it is possible to override that logic but a minimum of 256 is enforced internally.
-* ``enable_cleanup_closed`` (default: ``false``): In some cases, SSL connections might not be properly closed and the number of open connections increases as a result. When this client option is set to ``true``, the Elasticsearch client will check and forcefully close these connections.
+* ``enable_cleanup_closed`` (default: ``true``): In some cases, `Elasticsearch does not properly close SSL connections <https://github.com/elastic/elasticsearch/issues/76642>`_ and the number of open connections increases as a result. When this client option is set to ``true``, the Elasticsearch client will check and forcefully close these connections.
 * ``static_responses``: The path to a JSON file containing path patterns and the corresponding responses. When this value is set to ``true``, Rally will not send requests to Elasticsearch but return static responses as specified by the file. This is useful to diagnose performance issues in Rally itself. See below for a specific example.
 
 **Examples**
diff --git a/docs/track.rst b/docs/track.rst
index bb3a8d4e..7a98a334 100644
--- a/docs/track.rst
+++ b/docs/track.rst
@@ -126,17 +126,18 @@ We recommend the following path:
 Anatomy of a track
 ==================
 
-A track JSON file consists of the following sections:
-
-* indices
-* templates
-* data-streams
-* composable-templates
-* component-templates
-* corpora
-* operations
-* schedule
-* challenges
+A track JSON file can include the following sections:
+
+* :ref:`indices <track_indices>`
+* :ref:`templates <track_templates>`
+* :ref:`data-streams <track_data_streams>`
+* :ref:`composable-templates <track_composable_templates>`
+* :ref:`component-templates <track_component_templates>`
+* :ref:`corpora <track_corpora>`
+* :ref:`challenge(s) <track_challenge>`
+* :ref:`schedule <track_schedule>`
+* :ref:`operations <track_operations>`
+* :ref:`dependencies <track_dependencies>`
 
 In the ``indices`` and ``templates`` sections you define the relevant indices and index templates. These sections are optional but recommended if you want to create indices and index templates with the help of Rally. The index templates here represent the `legacy Elasticsearch index templates <https://www.elastic.co/guide/en/elasticsearch/reference/7.9/indices-templates-v1.html>`_ which have been deprecated in Elasticsearch 7.9. Users should refer to the ``composable-templates`` and ``component-templates`` for new tracks.
 
@@ -200,6 +201,9 @@ If the ``meta`` structure contains the same key on different elements, more spec
 
 E.g. a key defined on a task, will override the same key defined on a challenge. All properties defined within the merged ``meta`` structure, will get copied into each metrics record.
 
+
+.. _track_indices:
+
 indices
 .......
 
@@ -221,6 +225,9 @@ Example::
         }
     ]
 
+
+.. _track_templates:
+
 templates
 .........
 
@@ -242,6 +249,9 @@ Example::
         }
     ]
 
+
+.. _track_data_streams:
+
 data-streams
 ............
 
@@ -259,6 +269,9 @@ Example::
         }
     ]
 
+
+.. _track_composable_templates:
+
 composable-templates
 ....................
 
@@ -282,6 +295,9 @@ Example::
         }
     ]
 
+
+.. _track_component_templates:
+
 component-templates
 ....................
 
@@ -301,6 +317,9 @@ Example::
         }
     ]
 
+
+.. _track_corpora:
+
 corpora
 .......
 
@@ -2945,6 +2964,30 @@ The following meta data is always returned:
 * ``weight``: The number of fetched pages. Should always equal to the ``pages`` parameter.
 * ``unit``: The unit in which to interpret ``weight``. Always "ops".
 
+
+.. _track_dependencies:
+
+dependencies
+............
+
+If your track code requires additional Python dependencies at runtime (for instance, for code included in :ref:`custom runners <adding_tracks_custom_runners>` or :ref:`parameter sources <adding_tracks_custom_param_sources>`) you can declare pip-installable packages in your track's ``dependencies`` section.
+``dependencies`` is an array of ``pip install`` targets, which Rally will manage when the benchmark begins::
+
+  "dependencies": [
+    "pytoml",
+    "git+https://github.com/elastic/apm-agent-python@main",
+    "eland<8.0.0"
+  ]
+
+
+Rally will manage the installation of the libraries in the target directory ``$RALLY_HOME/libs``, erasing its contents between races.
+
+.. warning::
+
+    Because modules will be loaded in the same environment as ``esrally`` itself, conflicts with Rally and its own dependencies are possible with the usage of Track Dependencies.
+    **It is highly recommended to use pinned versions and to carefully test dependencies for their impact on your benchmarks.**
+
+
 Examples
 ========
 
diff --git a/esrally/async_connection.py b/esrally/async_connection.py
index 5ce270ae..09b77d2d 100644
--- a/esrally/async_connection.py
+++ b/esrally/async_connection.py
@@ -206,7 +206,7 @@ class AIOHttpConnection(elasticsearch.AIOHttpConnection):
         )
 
         self._trace_configs = [trace_config] if trace_config else None
-        self._enable_cleanup_closed = kwargs.get("enable_cleanup_closed", False)
+        self._enable_cleanup_closed = kwargs.get("enable_cleanup_closed", True)
 
         static_responses = kwargs.get("static_responses")
         self.use_static_responses = static_responses is not None
diff --git a/esrally/driver/driver.py b/esrally/driver/driver.py
index 2fa16ed9..255a8014 100644
--- a/esrally/driver/driver.py
+++ b/esrally/driver/driver.py
@@ -469,7 +469,8 @@ class TrackPreparationActor(actor.RallyActor):
     def receiveMsg_Bootstrap(self, msg, sender):
         # load node-specific config to have correct paths available
         self.cfg = load_local_config(msg.config)
-        load_track(self.cfg)
+        # this instance of load_track occurs once per host, so install dependencies if necessary
+        load_track(self.cfg, install_dependencies=True)
         self.send(sender, ReadyForWork())
 
     @actor.no_retry("track preparator")  # pylint: disable=no-value-for-parameter
diff --git a/esrally/paths.py b/esrally/paths.py
index 763119f0..f3c61594 100644
--- a/esrally/paths.py
+++ b/esrally/paths.py
@@ -41,6 +41,10 @@ def install_root(cfg=None):
     return os.path.join(races_root(cfg), install_id)
 
 
+def libs():
+    return os.path.join(rally_confdir(), "libs")
+
+
 def logs():
     """
     :return: The absolute path to the directory that contains Rally's log file.
diff --git a/esrally/racecontrol.py b/esrally/racecontrol.py
index 85529a42..950beca6 100644
--- a/esrally/racecontrol.py
+++ b/esrally/racecontrol.py
@@ -191,7 +191,7 @@ class BenchmarkCoordinator:
             if specified_version < min_es_version:
                 raise exceptions.SystemSetupError(f"Cluster version must be at least [{min_es_version}] but was [{distribution_version}]")
 
-        self.current_track = track.load_track(self.cfg)
+        self.current_track = track.load_track(self.cfg, install_dependencies=False)
         self.track_revision = self.cfg.opts("track", "repository.revision", mandatory=False)
         challenge_name = self.cfg.opts("track", "challenge.name")
         self.current_challenge = self.current_track.find_challenge_or_default(challenge_name)
diff --git a/esrally/resources/track-schema.json b/esrally/resources/track-schema.json
index a4b4c9a7..10112327 100644
--- a/esrally/resources/track-schema.json
+++ b/esrally/resources/track-schema.json
@@ -585,6 +585,14 @@
         "$ref": "#/definitions/challenge"
       }
     },
+    "dependencies": {
+      "type": "array",
+      "minItems": 0,
+      "uniqueItems": true,
+      "items": {
+        "type": "string"
+      }
+    },
     "challenge": {
       "$ref": "#/definitions/challenge"
     },
diff --git a/esrally/track/loader.py b/esrally/track/loader.py
index 18b9548f..fdfc7e30 100644
--- a/esrally/track/loader.py
+++ b/esrally/track/loader.py
@@ -19,6 +19,8 @@ import json
 import logging
 import os
 import re
+import shutil
+import subprocess
 import sys
 import tempfile
 import urllib.error
@@ -29,7 +31,7 @@ import jsonschema
 import tabulate
 from jinja2 import meta
 
-from esrally import PROGRAM_NAME, config, exceptions, time, version
+from esrally import PROGRAM_NAME, config, exceptions, paths, time, version
 from esrally.track import params, track
 from esrally.utils import collections, console, convert, io, modules, net, opts, repo
 
@@ -180,7 +182,7 @@ def track_info(cfg):
             console.println("")
 
 
-def load_track(cfg):
+def load_track(cfg, install_dependencies=False):
     """
 
     Loads a track
@@ -189,15 +191,43 @@ def load_track(cfg):
     :return: The loaded track.
     """
     repo = track_repo(cfg)
-    return _load_single_track(cfg, repo, repo.track_name)
+    return _load_single_track(cfg, repo, repo.track_name, install_dependencies)
 
 
-def _load_single_track(cfg, track_repository, track_name):
+def _install_dependencies(dependencies):
+    def _cleanup():
+        # fully destructive is fine, we only allow one Rally to run at a time and we will rely on the pip cache for download caching
+        console.info("Cleaning track dependency directory...")
+        logging.info("Cleaning track dependency directory [%s]...", paths.libs())
+        shutil.rmtree(paths.libs(), onerror=_trap)
+
+    def _trap(function, path, exc_info):
+        logging.exception("Failed to clean up [%s] with [%s]", path, function, exc_info=True)
+        raise exceptions.SystemSetupError(f"Unable to clean [{paths.libs()}]. See Rally log for more information.")
+
+    if dependencies:
+        _cleanup()
+        log_path = os.path.join(paths.logs(), "dependency.log")
+        console.info(f"Installing track dependencies [{', '.join(dependencies)}]")
+        try:
+            with open(log_path, "ab") as install_log:
+                subprocess.check_call(
+                    [sys.executable, "-m", "pip", "install", *dependencies, "--upgrade", "--target", paths.libs()],
+                    stdout=install_log,
+                    stderr=install_log,
+                )
+        except subprocess.CalledProcessError:
+            raise exceptions.SystemSetupError(f"Installation of track dependencies failed. See [{install_log.name}] for more information.")
+
+
+def _load_single_track(cfg, track_repository, track_name, install_dependencies=False):
     try:
         track_dir = track_repository.track_dir(track_name)
         reader = TrackFileReader(cfg)
         current_track = reader.read(track_name, track_repository.track_file(track_name), track_dir)
         tpr = TrackProcessorRegistry(cfg)
+        if install_dependencies:
+            _install_dependencies(current_track.dependencies)
         has_plugins = load_track_plugins(cfg, track_name, register_track_processor=tpr.register_track_processor)
         current_track.has_plugins = has_plugins
         for processor in tpr.processors:
@@ -213,7 +243,15 @@ def _load_single_track(cfg, track_repository, track_name):
         raise
 
 
-def load_track_plugins(cfg, track_name, register_runner=None, register_scheduler=None, register_track_processor=None, force_update=False):
+def load_track_plugins(
+    cfg,
+    track_name,
+    register_runner=None,
+    register_scheduler=None,
+    register_track_processor=None,
+    install_dependencies=False,
+    force_update=False,
+):
     """
     Loads plugins that are defined for the current track (as specified by the configuration).
 
@@ -222,6 +260,7 @@ def load_track_plugins(cfg, track_name, register_runner=None, register_scheduler
     :param register_runner: An optional function where runners can be registered.
     :param register_scheduler: An optional function where custom schedulers can be registered.
     :param register_track_processor: An optional function where track processors can be registered.
+    :param install_dependencies: If set to ``True``, install declared dependencies from the track.py. Defaults to ``False``.
     :param force_update: If set to ``True`` this ensures that the track is first updated from the remote repository.
                          Defaults to ``False``.
     :return: True iff this track defines plugins and they have been loaded.
@@ -1083,6 +1122,9 @@ class TrackPluginReader:
         return self.loader.can_load()
 
     def load(self):
+        # get dependent libraries installed in a prior step. ensure dir exists to make sure loading works correctly.
+        os.makedirs(paths.libs(), exist_ok=True)
+        sys.path.insert(0, paths.libs())
         root_module = self.loader.load()
         try:
             # every module needs to have a register() method
@@ -1158,6 +1200,7 @@ class TrackSpecificationReader:
         ]
         corpora = self._create_corpora(self._r(track_specification, "corpora", mandatory=False, default_value=[]), indices, data_streams)
         challenges = self._create_challenges(track_specification)
+        dependencies = self._r(track_specification, "dependencies", mandatory=False)
         # at this point, *all* track params must have been referenced in the templates
         return track.Track(
             name=self.name,
@@ -1170,6 +1213,7 @@ class TrackSpecificationReader:
             composable_templates=composable_templates,
             component_templates=component_templates,
             corpora=corpora,
+            dependencies=dependencies,
             root=self.base_path,
         )
 
diff --git a/esrally/track/track.py b/esrally/track/track.py
index c9f6b218..1305a197 100644
--- a/esrally/track/track.py
+++ b/esrally/track/track.py
@@ -433,6 +433,7 @@ class Track:
         composable_templates=None,
         component_templates=None,
         corpora=None,
+        dependencies=None,
         has_plugins=False,
         root=None,
     ):
@@ -462,6 +463,7 @@ class Track:
         self.templates = templates if templates else []
         self.composable_templates = composable_templates if composable_templates else []
         self.component_templates = component_templates if component_templates else []
+        self.dependencies = dependencies if dependencies else []
         self.has_plugins = has_plugins
         self.root = root
 
diff --git a/esrally/utils/modules.py b/esrally/utils/modules.py
index a7d1b9a7..05bbe75a 100644
--- a/esrally/utils/modules.py
+++ b/esrally/utils/modules.py
@@ -60,7 +60,6 @@ class ComponentLoader:
     def _load_component(self, component_name, module_dirs):
         # precondition: A module with this name has to exist provided that the caller has called #can_load() before.
         root_module_name = "%s.%s" % (component_name, self.component_entry_point)
-
         for p in self._modules(module_dirs, component_name):
             self.logger.debug("Loading module [%s]", p)
             m = importlib.import_module(p)
diff --git a/it/resources/track_with_dependency/track.json b/it/resources/track_with_dependency/track.json
new file mode 100644
index 00000000..5388160e
--- /dev/null
+++ b/it/resources/track_with_dependency/track.json
@@ -0,0 +1,12 @@
+{
+    "schedule": [
+        {
+          "operation": {
+            "operation-type": "no-op"
+          }
+        }
+    ],
+    "dependencies": [
+      "pytoml"
+    ]
+}
diff --git a/it/resources/track_with_dependency/track.py b/it/resources/track_with_dependency/track.py
new file mode 100644
index 00000000..661ce026
--- /dev/null
+++ b/it/resources/track_with_dependency/track.py
@@ -0,0 +1,15 @@
+async def noop(es, params):
+    # pylint: disable=import-outside-toplevel,import-error
+    import pytoml as toml
+
+    ret_val = {}
+    # no-ops is our correctness marker
+    toml_values = ["weight = 1", 'unit = "no-ops"']
+    await es.cluster.health()
+    for toml_value in toml_values:
+        ret_val.update(toml.loads(toml_value))
+    return ret_val
+
+
+def register(registry):
+    registry.register_runner("no-op", noop, async_runner=True)


Test patch:
diff --git a/it/dependencies_test.py b/it/dependencies_test.py
new file mode 100644
index 00000000..8da82618
--- /dev/null
+++ b/it/dependencies_test.py
@@ -0,0 +1,31 @@
+# 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 os
+
+import it
+
+
+@it.random_rally_config
+def test_track_dependencies(cfg):
+    port = 19200
+    it.wait_until_port_is_free(port_number=port)
+    dist_version = it.DISTRIBUTIONS[-1]
+    cwd = os.path.dirname(__file__)
+    track_path = os.path.join(cwd, "resources", "track_with_dependency")
+    # workaround for MacOS and Python deficiency. http://sealiesoftware.com/blog/archive/2017/6/5/Objective-C_and_fork_in_macOS_1013.html
+    os.environ["OBJC_DISABLE_INITIALIZE_FORK_SAFETY"] = "YES"
+    assert it.race(cfg, f"--distribution-version={dist_version} --track-path={track_path} --car=4gheap,basic-license") == 0
diff --git a/tests/client_test.py b/tests/client_test.py
index a5e5eb3d..583bb346 100644
--- a/tests/client_test.py
+++ b/tests/client_test.py
@@ -28,6 +28,7 @@ import pytest
 import urllib3.exceptions
 
 from esrally import client, doc_link, exceptions
+from esrally.async_connection import AIOHttpConnection
 from esrally.utils import console
 from tests import run_async
 
@@ -349,3 +350,15 @@ class TestRestLayer:
         )
         with pytest.raises(exceptions.SystemSetupError, match="Could not connect to cluster via https. Is this an https endpoint?"):
             client.wait_for_rest_layer(es, max_attempts=3)
+
+
+class TestAsyncConnection:
+    @run_async
+    async def test_enable_cleanup_close(self):
+        connection = AIOHttpConnection()
+        # pylint: disable=protected-access
+        assert connection._enable_cleanup_closed is True
+
+        connection = AIOHttpConnection(enable_cleanup_closed=False)
+        # pylint: disable=protected-access
+        assert connection._enable_cleanup_closed is False
diff --git a/tests/driver/runner_test.py b/tests/driver/runner_test.py
index 12c85ebb..e091fe42 100644
--- a/tests/driver/runner_test.py
+++ b/tests/driver/runner_test.py
@@ -5795,13 +5795,13 @@ class TestComposite:
         assert len(timings) == 3
 
         assert timings[0]["operation"] == "initial-call"
-        assert math.isclose(timings[0]["service_time"], 0.1, rel_tol=0.05)
+        assert timings[0]["service_time"] == pytest.approx(0.1, abs=0.1)
 
         assert timings[1]["operation"] == "stream-a"
-        assert math.isclose(timings[1]["service_time"], 0.2, rel_tol=0.05)
+        assert timings[1]["service_time"] == pytest.approx(0.2, abs=0.1)
 
         assert timings[2]["operation"] == "stream-b"
-        assert math.isclose(timings[2]["service_time"], 0.1, rel_tol=0.05)
+        assert timings[2]["service_time"] == pytest.approx(0.1, abs=0.1)
 
         # common properties
         for timing in timings: