Update storage tools with last changes from storage branch.

docs/configuration.rst

@@ -140,6 +140,8 @@ remote servers. Available options are:
 * ``storage.max_connections`` represents the maximum number of client connections to be made against the same server or
   bucket. The default value is 8.
 
+* ``storage.max_workers`` indicates the maximum number of worker threads used for making storage files transfers.
+
 * ``storage.mirror_files`` is used to provide a json file that specify the mapping for mirrors URLs resolution.
   Example::
 

esrally/storage/_adapter.py

@@ -16,15 +16,16 @@
 # under the License.
 from __future__ import annotations
 
+import datetime
 import importlib
 import logging
 import threading
 from abc import ABC, abstractmethod
-from collections.abc import Iterable
-from typing import NamedTuple, Protocol, runtime_checkable
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
 
-from esrally.config import Config
 from esrally.storage._range import NO_RANGE, RangeSet
+from esrally.types import Config
 
 LOG = logging.getLogger(__name__)
 
@@ -40,33 +41,27 @@ def write(self, data: bytes) -> None:
         pass
 
 
-@runtime_checkable
-class Readable(Protocol):
-
-    def read(self, size: int = -1) -> bytes:
-        pass
-
-
-class Head(NamedTuple):
-    url: str
+@dataclass
+class Head:
+    url: str | None = None
     content_length: int | None = None
-    accept_ranges: bool = False
+    accept_ranges: bool | None = None
     ranges: RangeSet = NO_RANGE
     document_length: int | None = None
     crc32c: str | None = None
+    date: datetime.datetime | None = None
 
     @classmethod
     def create(
         cls,
-        url: str,
+        url: str | None = None,
         content_length: int | None = None,
         accept_ranges: bool | None = None,
         ranges: RangeSet = NO_RANGE,
         document_length: int | None = None,
         crc32c: str | None = None,
+        date: datetime.datetime | None = None,
     ) -> Head:
-        if accept_ranges is None:
-            accept_ranges = bool(ranges)
         if content_length is None and ranges:
             content_length = ranges.size
         if document_length is None and not ranges:
@@ -78,21 +73,24 @@ def create(
             ranges=ranges,
             document_length=document_length,
             crc32c=crc32c,
+            date=date,
         )
 
+    def check(self, other: Head) -> None:
+        for field in ("url", "content_length", "accept_ranges", "ranges", "document_length", "crc32c", "date"):
+            want = getattr(self, field)
+            got = getattr(other, field)
+            if not ({got, want} & {None, NO_RANGE}) and got != want:
+                raise ValueError(f"unexpected '{field}': got {got}, want {want}")
+
 
 class Adapter(ABC):
     """Base class for storage class client implementation"""
 
-    # A commas separated list of URL prefixes used to associate an adapter implementation to a remote file URL.
-    # This value will be overridden by `Adapter` subclasses to be consumed by `AdapterRegistry` class.
-    # Example:
-    #   ```
-    #   class HTTPAdapter(Adapter):
-    #       # The value will serve to associate any URL with "https" scheme to `HTTPAdapter` subclass.
-    #       __adapter_URL_prefixes__ = "http://, https://"
-    #   ```
-    __adapter_URL_prefixes__: str = ""
+    @classmethod
+    def match_url(cls, url: str) -> str:
+        """It returns a canonical URL in case this adapter accepts the URL, None otherwise."""
+        raise NotImplementedError
 
     @classmethod
     def from_config(cls, cfg: Config) -> Adapter:
@@ -114,33 +112,34 @@ def head(self, url: str) -> Head:
         """
 
     @abstractmethod
-    def get(self, url: str, stream: Writable, ranges: RangeSet = NO_RANGE) -> Head:
+    def get(self, url: str, stream: Writable, head: Head | None = None) -> Head:
         """It downloads a remote bucket object to a local file path.
 
         :param url: it represents the URL of the remote file object.
         :param stream: it represents the local file stream where to write data to.
-        :param ranges: it represents the portion of the file to transfer (it must be empty or a continuous range).
+        :param head: it allows to specify optional parameters:
+            - range: portion of the file to transfer (it must be empty or a continuous range).
+            - document_length: the number of bytes to transfer.
+            - crc32c the CRC32C checksum of the file.
+            - date: the date the file has been modified.
         :raises ServiceUnavailableError: in case on temporary service failure.
         """
 
 
 ADAPTER_CLASS_NAMES = ",".join(
     [
+        "esrally.storage._tracks:TracksRepositoryAdapter",
+        "esrally.storage._s3:S3Adapter",
         "esrally.storage._http:HTTPAdapter",
     ]
 )
 
 
-class AdapterClassEntry(NamedTuple):
-    url_prefix: str
-    cls: type[Adapter]
-
-
 class AdapterRegistry:
     """AdapterClassRegistry allows to register classes of adapters to be selected according to the target URL."""
 
     def __init__(self, cfg: Config) -> None:
-        self._classes: list[AdapterClassEntry] = []
+        self._classes: list[type[Adapter]] = []
         self._adapters: dict[type[Adapter], Adapter] = {}
         self._lock = threading.Lock()
         self._cfg = cfg
@@ -155,29 +154,38 @@ def from_config(cls, cfg: Config) -> AdapterRegistry:
         )
         for spec in adapters_specs:
             module_name, class_name = spec.split(":")
-            module = importlib.import_module(module_name)
-            obj = getattr(module, class_name)
+            try:
+                module = importlib.import_module(module_name)
+            except ModuleNotFoundError:
+                LOG.exception("unable to import module '%s'.", module_name)
+                continue
+            try:
+                obj = getattr(module, class_name)
+            except AttributeError:
+                raise ValueError("Invalid Adapter class name: '{class_name}'.")
             if not isinstance(obj, type) or not issubclass(obj, Adapter):
                 raise TypeError(f"'{obj}' is not a valid subclass of Adapter")
-            registry.register_class(obj, obj.__adapter_URL_prefixes__.split(","))
+            registry.register_class(obj)
         return registry
 
-    def register_class(self, cls: type[Adapter], prefixes: Iterable[str]) -> type[Adapter]:
+    def register_class(self, cls: type[Adapter]) -> type[Adapter]:
         with self._lock:
-            for p in prefixes:
-                self._classes.append(AdapterClassEntry(p.strip(), cls))
-            # The list of adapter classes is kept sorted from the longest prefix to the shorter to ensure that matching
-            # a shorter URL prefix will never hide matching a longer one.
-            self._classes.sort(key=lambda e: len(e.url_prefix), reverse=True)
+            self._classes.append(cls)
         return cls
 
-    def get(self, url: str) -> Adapter:
+    def get(self, url: str) -> tuple[Adapter, str]:
         with self._lock:
-            for e in self._classes:
-                if url.startswith(e.url_prefix):
-                    adapter = self._adapters.get(e.cls)
-                    if adapter is None:
-                        adapter = e.cls.from_config(self._cfg)
-                    self._adapters[e.cls] = adapter
-                    return adapter
-        raise ValueError(f"No adapter found for url '{url}'")
+            for cls in self._classes:
+                try:
+                    actual_url = cls.match_url(url)
+                    break
+                except NotImplementedError:
+                    continue
+            else:
+                raise ValueError(f"No adapter found for url '{url}'")
+
+        adapter = self._adapters.get(cls)
+        if adapter is None:
+            adapter = cls.from_config(self._cfg)
+        self._adapters[cls] = adapter
+        return adapter, actual_url

esrally/storage/_client.py

@@ -25,48 +25,51 @@
 from random import Random
 from typing import NamedTuple
 
-from esrally import config
+from esrally import types
 from esrally.storage._adapter import (
-    Adapter,
     AdapterRegistry,
     Head,
     ServiceUnavailableError,
     Writable,
 )
 from esrally.storage._mirror import MirrorList
-from esrally.storage._range import NO_RANGE, RangeSet
 from esrally.utils import pretty
 from esrally.utils.threads import WaitGroup, WaitGroupLimitError
 
 LOG = logging.getLogger(__name__)
 
 MIRRORS_FILES = "~/.rally/storage-mirrors.json"
 MAX_CONNECTIONS = 8
+RANDOM = Random(time.monotonic_ns())
 
 
-class Client(Adapter):
+class Client:
     """It handles client instances allocation allowing reusing pre-allocated instances from the same thread."""
 
     @classmethod
-    def from_config(cls, cfg: config.Config) -> Client:
+    def from_config(
+        cls, cfg: types.Config, adapters: AdapterRegistry | None = None, mirrors: MirrorList | None = None, random: Random | None = None
+    ) -> Client:
+        if adapters is None:
+            adapters = AdapterRegistry.from_config(cfg)
+        if mirrors is None:
+            mirrors = MirrorList.from_config(cfg)
+        if random is None:
+            random_seed = cfg.opts(section="storage", key="storage.random_seed", default_value=None, mandatory=False)
+            if random_seed is None:
+                random = RANDOM
+            else:
+                random = Random(random_seed)
         max_connections = int(cfg.opts(section="storage", key="storage.max_connections", default_value=MAX_CONNECTIONS, mandatory=False))
-        random_seed = cfg.opts(section="storage", key="storage.random_seed", default_value=None, mandatory=False)
-        random = None
-        if random_seed is not None:
-            random = Random(random_seed)
-        return cls(
-            max_connections=max_connections, mirrors=MirrorList.from_config(cfg), adapters=AdapterRegistry.from_config(cfg), random=random
-        )
+        return cls(adapters=adapters, mirrors=mirrors, random=random, max_connections=max_connections)
 
     def __init__(
         self,
         adapters: AdapterRegistry,
         mirrors: MirrorList,
-        random: Random | None = None,
+        random: Random,
         max_connections: int = MAX_CONNECTIONS,
     ):
-        if random is None:
-            random = Random(time.time())
         self._adapters: AdapterRegistry = adapters
         self._cached_heads: dict[str, tuple[Head | Exception, float]] = {}
         self._connections: dict[str, WaitGroup] = defaultdict(lambda: WaitGroup(max_count=max_connections))
@@ -90,7 +93,7 @@ def head(self, url: str, ttl: float | None = None) -> Head:
                     # cached value or error is enough recent to be used.
                     return _head_or_raise(value)
 
-        adapter = self._adapters.get(url)
+        adapter, url = self._adapters.get(url)
         try:
             value = adapter.head(url)
         except Exception as ex:
@@ -107,14 +110,17 @@ def head(self, url: str, ttl: float | None = None) -> Head:
 
         return _head_or_raise(value)
 
-    def resolve(
-        self, url: str, document_length: int | None = None, crc32c: str | None = None, accept_ranges: bool = False, ttl: float = 60.0
-    ) -> Iterator[Head]:
+    @property
+    def adapters(self) -> AdapterRegistry:
+        return self._adapters
+
+    def resolve(self, url: str, check: Head | None, ttl: float = 60.0) -> Iterator[Head]:
         """It looks up mirror list for given URL and yield mirror heads.
         :param url: the remote file URL at its mirrored source location.
-        :param document_length: if not none it will filter out mirrors which file has an unexpected document lengths.
-        :param crc32c: if not none it will filter out mirrors which file has an unexpected crc32c checksum.
-        :param accept_ranges: if True it will filter out mirrors that are not supporting ranges.
+        :param check: extra parameters to mach remote heads.
+            - document_length: if not none it will filter out mirrors which file has an unexpected document lengths.
+            - crc32c: if not none it will filter out mirrors which file has an unexpected crc32c checksum.
+            - accept_ranges: if True it will filter out mirrors that are not supporting ranges.
         :param ttl: the time to live value (in seconds) to use for cached heads retrieval.
         :return: iterator over mirror URLs
         """
@@ -143,43 +149,43 @@ def resolve(
 
         for u in urls:
             try:
-                head = self.head(u, ttl=ttl)
-            except Exception:
+                got = self.head(u, ttl=ttl)
+            except Exception as ex:
                 # The exception is already logged by head method before caching it.
+                LOG.error("Failed to fetch remote head for file: %s, %s", u, ex)
                 continue
-            if document_length is not None and head.document_length is not None and document_length != head.document_length:
-                LOG.debug("unexpected document length: %r, got %d, want %d", url, head.document_length, document_length)
-                continue
-            if crc32c is not None and head.crc32c is not None and head.crc32c != crc32c:
-                LOG.debug("unexpected crc32c checksum: %r, got %d, want %d", url, head.crc32c, crc32c)
-                continue
-            if accept_ranges and not head.accept_ranges:
-                LOG.debug("it doesn't accept ranges: %r", url)
-                continue
-            yield head
-
-    def get(
-        self, url: str, stream: Writable, ranges: RangeSet = NO_RANGE, document_length: int | None = None, crc32c: str | None = None
-    ) -> Head:
+            if check is not None:
+                try:
+                    check.check(got)
+                except ValueError as ex:
+                    LOG.debug("unexpected mirrored file (url='%s'): %s", url, ex)
+                    continue
+            yield got
+
+    def get(self, url: str, stream: Writable, head: Head | None = None) -> Head:
         """It downloads a remote bucket object to a local file path.
 
         :param url: the URL of the remote file.
         :param stream: the destination file stream where to write data to.
-        :param document_length: the document length of the file to transfer.
-        :param crc32c: the crc32c checksum of the file to transfer.
-        :param ranges: the portion of the file to transfer.
+        :param head: extra params for getting the file:
+            - document_length: the document length of the file to transfer.
+            - crc32c: the crc32c checksum of the file to transfer.
+            - ranges: the portion of the file to transfer.
         :raises ServiceUnavailableError: in case on temporary service failure.
         """
-        for head in self.resolve(url, accept_ranges=bool(ranges), document_length=document_length, crc32c=crc32c):
-            connections = self._server_connections(head.url)
+        for got in self.resolve(url, check=head):
+            if got.url is None:
+                LOG.error("resolved mirror URL is None: %s", url)
+                continue
+            connections = self._server_connections(got.url)
             try:
                 connections.add(1)
             except WaitGroupLimitError:
                 LOG.debug("connection limit exceeded: url='%s'", url)
                 continue
-            adapter = self._adapters.get(head.url)
+            adapter, url = self._adapters.get(got.url)
             try:
-                return adapter.get(head.url, stream, ranges)
+                return adapter.get(url, stream, head=head)
             except ServiceUnavailableError as ex:
                 LOG.debug("service unavailable error received: url='%s' %s", url, ex)
                 with self._lock:
@@ -217,7 +223,8 @@ def monitor(self):
         for url, connection in sorted(connections.items()):
             latency = self._average_latency(url)
             infos.append(f"- '{url}' count={connection.count} latency={pretty.duration(latency)}")
-        LOG.info("Active client connection(s):\n  %s", "\n  ".join(infos))
+        if len(infos) > 0:
+            LOG.info("Active client connection(s):\n  %s", "\n  ".join(infos))
 
 
 class ServerStats(NamedTuple):