elastic_transport/_async_transport.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 logging
from typing import (
    Any,
    Awaitable,
    Callable,
    Collection,
    List,
    Mapping,
    Optional,
    Tuple,
    Type,
    Union,
)

from ._compat import await_if_coro
from ._exceptions import (
    ConnectionError,
    ConnectionTimeout,
    SniffingError,
    TransportError,
)
from ._models import DEFAULT, DefaultType, HttpHeaders, NodeConfig, SniffOptions
from ._node import AiohttpHttpNode, BaseAsyncNode
from ._node_pool import NodePool, NodeSelector
from ._otel import OpenTelemetrySpan
from ._serializer import Serializer
from ._transport import (
    DEFAULT_CLIENT_META_SERVICE,
    NOT_DEAD_NODE_HTTP_STATUSES,
    Transport,
    TransportApiResponse,
    validate_sniffing_options,
)
from .client_utils import resolve_default

_logger = logging.getLogger("elastic_transport.transport")


class AsyncTransport(Transport):
    """
    Encapsulation of transport-related to logic. Handles instantiation of the
    individual nodes as well as creating a node pool to hold them.

    Main interface is the :meth:`elastic_transport.Transport.perform_request` method.
    """

    def __init__(
        self,
        node_configs: List[NodeConfig],
        node_class: Union[str, Type[BaseAsyncNode]] = AiohttpHttpNode,
        node_pool_class: Type[NodePool] = NodePool,
        randomize_nodes_in_pool: bool = True,
        node_selector_class: Optional[Union[str, Type[NodeSelector]]] = None,
        dead_node_backoff_factor: Optional[float] = None,
        max_dead_node_backoff: Optional[float] = None,
        serializers: Optional[Mapping[str, Serializer]] = None,
        default_mimetype: str = "application/json",
        max_retries: int = 3,
        retry_on_status: Collection[int] = (429, 502, 503, 504),
        retry_on_timeout: bool = False,
        sniff_on_start: bool = False,
        sniff_before_requests: bool = False,
        sniff_on_node_failure: bool = False,
        sniff_timeout: Optional[float] = 0.5,
        min_delay_between_sniffing: float = 10.0,
        sniff_callback: Optional[
            Callable[
                ["AsyncTransport", "SniffOptions"],
                Union[List[NodeConfig], Awaitable[List[NodeConfig]]],
            ]
        ] = None,
        meta_header: bool = True,
        client_meta_service: Tuple[str, str] = DEFAULT_CLIENT_META_SERVICE,
    ):
        """
        :arg node_configs: List of 'NodeConfig' instances to create initial set of nodes.
        :arg node_class: subclass of :class:`~elastic_transport.BaseNode` to use
            or the name of the Connection (ie 'urllib3', 'requests')
        :arg node_pool_class: subclass of :class:`~elastic_transport.NodePool` to use
        :arg randomize_nodes_in_pool: Set to false to not randomize nodes within the pool.
            Defaults to true.
        :arg node_selector_class: Class to be used to select nodes within
            the :class:`~elastic_transport.NodePool`.
        :arg dead_node_backoff_factor: Exponential backoff factor to calculate the amount
            of time to timeout a node after an unsuccessful API call.
        :arg max_dead_node_backoff: Maximum amount of time to timeout a node after an
            unsuccessful API call.
        :arg serializers: optional dict of serializer instances that will be
            used for deserializing data coming from the server. (key is the mimetype)
        :arg max_retries: Maximum number of retries for an API call.
            Set to 0 to disable retries. Defaults to ``0``.
        :arg retry_on_status: set of HTTP status codes on which we should retry
            on a different node. defaults to ``(429, 502, 503, 504)``
        :arg retry_on_timeout: should timeout trigger a retry on different
            node? (default ``False``)
        :arg sniff_on_start: If ``True`` will sniff for additional nodes as soon
            as possible, guaranteed before the first request.
        :arg sniff_on_node_failure: If ``True`` will sniff for additional nodees
            after a node is marked as dead in the pool.
        :arg sniff_before_requests: If ``True`` will occasionally sniff for additional
            nodes as requests are sent.
        :arg sniff_timeout: Timeout value in seconds to use for sniffing requests.
            Defaults to 1 second.
        :arg min_delay_between_sniffing: Number of seconds to wait between calls to
            :meth:`elastic_transport.Transport.sniff` to avoid sniffing too frequently.
            Defaults to 10 seconds.
        :arg sniff_callback: Function that is passed a :class:`elastic_transport.Transport` and
            :class:`elastic_transport.SniffOptions` and should do node discovery and
            return a list of :class:`elastic_transport.NodeConfig` instances or a coroutine
            that returns the list.
        """

        # Since we don't pass all the sniffing options to super().__init__()
        # we want to validate the sniffing options here too.
        validate_sniffing_options(
            node_configs=node_configs,
            sniff_on_start=sniff_on_start,
            sniff_before_requests=sniff_before_requests,
            sniff_on_node_failure=sniff_on_node_failure,
            sniff_callback=sniff_callback,
        )

        super().__init__(
            node_configs=node_configs,
            node_class=node_class,
            node_pool_class=node_pool_class,
            randomize_nodes_in_pool=randomize_nodes_in_pool,
            node_selector_class=node_selector_class,
            dead_node_backoff_factor=dead_node_backoff_factor,
            max_dead_node_backoff=max_dead_node_backoff,
            serializers=serializers,
            default_mimetype=default_mimetype,
            max_retries=max_retries,
            retry_on_status=retry_on_status,
            retry_on_timeout=retry_on_timeout,
            sniff_timeout=sniff_timeout,
            min_delay_between_sniffing=min_delay_between_sniffing,
            meta_header=meta_header,
            client_meta_service=client_meta_service,
        )

        self._sniff_on_start = sniff_on_start
        self._sniff_before_requests = sniff_before_requests
        self._sniff_on_node_failure = sniff_on_node_failure
        self._sniff_timeout = sniff_timeout
        self._sniff_callback = sniff_callback  # type: ignore
        self._sniffing_task: Optional["asyncio.Task[Any]"] = None
        self._last_sniffed_at = 0.0

        # We set this to 'None' here but it'll never be None by the
        # time it's needed. Gets set within '_async_call()' which should
        # precede all logic within async calls.
        self._loop: asyncio.AbstractEventLoop = None  # type: ignore[assignment]

        # AsyncTransport doesn't require a thread lock for
        # sniffing. Uses '_sniffing_task' instead.
        self._sniffing_lock = None  # type: ignore[assignment]

    async def perform_request(  # type: ignore[override, return]
        self,
        method: str,
        target: str,
        *,
        body: Optional[Any] = None,
        headers: Union[Mapping[str, Any], DefaultType] = DEFAULT,
        max_retries: Union[int, DefaultType] = DEFAULT,
        retry_on_status: Union[Collection[int], DefaultType] = DEFAULT,
        retry_on_timeout: Union[bool, DefaultType] = DEFAULT,
        request_timeout: Union[Optional[float], DefaultType] = DEFAULT,
        client_meta: Union[Tuple[Tuple[str, str], ...], DefaultType] = DEFAULT,
        otel_span: Union[OpenTelemetrySpan, DefaultType] = DEFAULT,
    ) -> TransportApiResponse:
        """
        Perform the actual request. Retrieve a node from the node
        pool, pass all the information to it's perform_request method and
        return the data.

        If an exception was raised, mark the node as failed and retry (up
        to ``max_retries`` times).

        If the operation was successful and the node used was previously
        marked as dead, mark it as live, resetting it's failure count.

        :arg method: HTTP method to use
        :arg target: HTTP request target
        :arg body: body of the request, will be serialized using serializer and
            passed to the node
        :arg headers: Additional headers to send with the request.
        :arg max_retries: Maximum number of retries before giving up on a request.
            Set to ``0`` to disable retries.
        :arg retry_on_status: Collection of HTTP status codes to retry.
        :arg retry_on_timeout: Set to true to retry after timeout errors.
        :arg request_timeout: Amount of time to wait for a response to fail with a timeout error.
        :arg client_meta: Extra client metadata key-value pairs to send in the client meta header.
        :arg otel_span: OpenTelemetry span used to add metadata to the span.
        :returns: Tuple of the :class:`elastic_transport.ApiResponseMeta` with the deserialized response.
        """
        await self._async_call()

        if headers is DEFAULT:
            request_headers = HttpHeaders()
        else:
            request_headers = HttpHeaders(headers)
        max_retries = resolve_default(max_retries, self.max_retries)
        retry_on_timeout = resolve_default(retry_on_timeout, self.retry_on_timeout)
        retry_on_status = resolve_default(retry_on_status, self.retry_on_status)
        otel_span = resolve_default(otel_span, OpenTelemetrySpan(None))

        if self.meta_header:
            request_headers["x-elastic-client-meta"] = ",".join(
                f"{k}={v}"
                for k, v in self._transport_client_meta
                + resolve_default(client_meta, ())
            )

        # Serialize the request body to bytes based on the given mimetype.
        request_body: Optional[bytes]
        if body is not None:
            if "content-type" not in request_headers:
                raise ValueError(
                    "Must provide a 'Content-Type' header to requests with bodies"
                )
            request_body = self.serializers.dumps(
                body, mimetype=request_headers["content-type"]
            )
            otel_span.set_db_statement(request_body)
        else:
            request_body = None

        # Errors are stored from (oldest->newest)
        errors: List[Exception] = []

        for attempt in range(max_retries + 1):
            # If we sniff before requests are made we want to do so before
            # 'node_pool.get()' is called so our sniffed nodes show up in the pool.
            if self._sniff_before_requests:
                await self.sniff(False)

            retry = False
            node_failure = False
            last_response: Optional[TransportApiResponse] = None
            node: BaseAsyncNode = self.node_pool.get()  # type: ignore[assignment]
            start_time = self._loop.time()
            try:
                otel_span.set_node_metadata(node.host, node.port, node.base_url, target)
                resp = await node.perform_request(
                    method,
                    target,
                    body=request_body,
                    headers=request_headers,
                    request_timeout=request_timeout,
                )
                _logger.info(
                    "%s %s%s [status:%s duration:%.3fs]"
                    % (
                        method,
                        node.base_url,
                        target,
                        resp.meta.status,
                        self._loop.time() - start_time,
                    )
                )

                if method != "HEAD":
                    body = self.serializers.loads(resp.body, resp.meta.mimetype)
                else:
                    body = None

                if resp.meta.status in retry_on_status:
                    retry = True
                    # Keep track of the last response we see so we can return
                    # it in case the retried request returns with a transport error.
                    last_response = TransportApiResponse(resp.meta, body)

            except TransportError as e:
                _logger.info(
                    "%s %s%s [status:%s duration:%.3fs]"
                    % (
                        method,
                        node.base_url,
                        target,
                        "N/A",
                        self._loop.time() - start_time,
                    )
                )

                if isinstance(e, ConnectionTimeout):
                    retry = retry_on_timeout
                    node_failure = True
                elif isinstance(e, ConnectionError):
                    retry = True
                    node_failure = True

                # If the error was determined to be a node failure
                # we mark it dead in the node pool to allow for
                # other nodes to be retried.
                if node_failure:
                    self.node_pool.mark_dead(node)

                    if self._sniff_on_node_failure:
                        try:
                            await self.sniff(False)
                        except TransportError:
                            # If sniffing on failure, it could fail too. Catch the
                            # exception not to interrupt the retries.
                            pass

                if not retry or attempt >= max_retries:
                    # Since we're exhausted but we have previously
                    # received some sort of response from the API
                    # we should forward that along instead of the
                    # transport error. Likely to be more actionable.
                    if last_response is not None:
                        return last_response

                    e.errors = tuple(errors)
                    raise
                else:
                    _logger.warning(
                        "Retrying request after failure (attempt %d of %d)",
                        attempt,
                        max_retries,
                        exc_info=e,
                    )
                    errors.append(e)

            else:
                # If we got back a response we need to check if that status
                # is indicative of a healthy node even if it's a non-2XX status
                if (
                    200 <= resp.meta.status < 299
                    or resp.meta.status in NOT_DEAD_NODE_HTTP_STATUSES
                ):
                    self.node_pool.mark_live(node)
                else:
                    self.node_pool.mark_dead(node)

                    if self._sniff_on_node_failure:
                        try:
                            await self.sniff(False)
                        except TransportError:
                            # If sniffing on failure, it could fail too. Catch the
                            # exception not to interrupt the retries.
                            pass

                # We either got a response we're happy with or
                # we've exhausted all of our retries so we return it.
                if not retry or attempt >= max_retries:
                    return TransportApiResponse(resp.meta, body)
                else:
                    _logger.warning(
                        "Retrying request after non-successful status %d (attempt %d of %d)",
                        resp.meta.status,
                        attempt,
                        max_retries,
                    )

    async def sniff(self, is_initial_sniff: bool = False) -> None:  # type: ignore[override]
        await self._async_call()
        task = self._create_sniffing_task(is_initial_sniff)

        # Only block on the task if this is the initial sniff.
        # Otherwise we do the sniffing in the background.
        if is_initial_sniff and task:
            await task

    async def close(self) -> None:  # type: ignore[override]
        """
        Explicitly closes all nodes in the transport's pool
        """
        node: BaseAsyncNode
        for node in self.node_pool.all():  # type: ignore[assignment]
            await node.close()

    def _should_sniff(self, is_initial_sniff: bool) -> bool:
        """Decide if we should sniff or not. _async_init() must be called
        before using this function.The async implementation doesn't have a lock.
        """
        if is_initial_sniff:
            return True

        # Only start a new sniff if the previous run is completed.
        if self._sniffing_task:
            if not self._sniffing_task.done():
                return False
            # If there was a previous run we collect the sniffing task's
            # result as it could have failed with an exception.
            self._sniffing_task.result()

        return (
            self._loop.time() - self._last_sniffed_at
            >= self._min_delay_between_sniffing
        )

    def _create_sniffing_task(
        self, is_initial_sniff: bool
    ) -> Optional["asyncio.Task[Any]"]:
        """Creates a sniffing task if one should be created and returns the task if created."""
        task = None
        if self._should_sniff(is_initial_sniff):
            _logger.info("Started sniffing for additional nodes")
            # 'self._sniffing_task' is unset within the task implementation.
            task = self._loop.create_task(self._sniffing_task_impl(is_initial_sniff))
            self._sniffing_task = task
        return task

    async def _sniffing_task_impl(self, is_initial_sniff: bool) -> None:
        """Implementation of the sniffing task"""
        previously_sniffed_at = self._last_sniffed_at
        try:
            self._last_sniffed_at = self._loop.time()
            options = SniffOptions(
                is_initial_sniff=is_initial_sniff, sniff_timeout=self._sniff_timeout
            )
            assert self._sniff_callback is not None
            node_configs = await await_if_coro(self._sniff_callback(self, options))
            if not node_configs and is_initial_sniff:
                raise SniffingError(
                    "No viable nodes were discovered on the initial sniff attempt"
                )

            prev_node_pool_size = len(self.node_pool)
            for node_config in node_configs:
                self.node_pool.add(node_config)

            # Do some math to log which nodes are new/existing
            sniffed_nodes = len(node_configs)
            new_nodes = sniffed_nodes - (len(self.node_pool) - prev_node_pool_size)
            existing_nodes = sniffed_nodes - new_nodes
            _logger.debug(
                "Discovered %d nodes during sniffing (%d new nodes, %d already in pool)",
                sniffed_nodes,
                new_nodes,
                existing_nodes,
            )

        # If sniffing failed for any reason we
        # want to allow retrying immediately.
        except BaseException:
            self._last_sniffed_at = previously_sniffed_at
            raise

    async def _async_call(self) -> None:
        """Async constructor which is called on the first call to perform_request()
        because we're not guaranteed to be within an active asyncio event loop
        when __init__() is called.
        """
        if self._loop is not None:
            return  # Call at most once!
        self._loop = asyncio.get_running_loop()
        if self._sniff_on_start:
            await self.sniff(True)