Add support for synchronous Transport sniffing

Author: sethmlarson

elastic_transport/__init__.py

@@ -48,6 +48,7 @@
     NodeConfig,
     QueryParams,
     RequestOptions,
+    SniffOptions,
 )
 from ._node import AiohttpHttpNode, BaseNode, RequestsHttpNode, Urllib3HttpNode
 from ._node_pool import NodePool, NodeSelector, RandomSelector, RoundRobinSelector
@@ -95,6 +96,7 @@
     "SerializationError",
     "Serializer",
     "ServiceUnavailableError",
+    "SniffOptions",
     "TextSerializer",
     "TlsError",
     "TooManyRequestsError",

elastic_transport/_async_transport.py

@@ -168,22 +168,23 @@ async def perform_request(
 
             except TransportError as e:
                 retry = False
+                node_failure = e.status not in NOT_DEAD_NODE_HTTP_STATUSES
                 if isinstance(e, ConnectionTimeout):
                     retry = self.retry_on_timeout
+                    node_failure = True
                 elif isinstance(e, ConnectionError):
                     retry = True
+                    node_failure = True
                 elif e.status in self.retry_on_status:
                     retry = 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 retry:
-                    try:
-                        # only mark as dead if we are retrying
-                        if e.status not in NOT_DEAD_NODE_HTTP_STATUSES:
-                            await self.mark_dead(node)
-                    except TransportError:
-                        # If sniffing on failure, it could fail too. Catch the
-                        # exception not to interrupt the retries.
-                        pass
                     # raise exception on last retry
                     if attempt == self.max_retries:
                         e.errors = tuple(errors)
@@ -199,10 +200,6 @@ async def perform_request(
                 self.node_pool.mark_live(node)
                 return response, data
 
-    async def mark_dead(self, node: BaseNode) -> None:
-        """Marks a node as dead and optionally starts sniffing for additional nodes if enabled"""
-        self.node_pool.mark_dead(node)
-
     async def close(self) -> None:
         """
         Explicitly closes all nodes in the transport's pool

elastic_transport/_compat.py

@@ -16,7 +16,9 @@
 #  under the License.
 
 import asyncio
+import inspect
 import sys
+from pathlib import Path
 from urllib.parse import quote as _quote
 from urllib.parse import urlencode, urlparse
 
@@ -67,6 +69,50 @@ def __enter__(self) -> None:
         def __exit__(self, *_) -> None:
             pass
 
+        def acquire(self, blocking: bool = True) -> bool:
+            return True
+
+        def release(self) -> None:
+            pass
+
+
+def warn_stacklevel() -> int:
+    """Dynamically determine warning stacklevel for warnings based on the call stack"""
+    try:
+        # Grab the root module from the current module '__name__'
+        module_name = __name__.partition(".")[0]
+        module_path = Path(sys.modules[module_name].__file__)
+
+        # If the module is a folder we're looking at
+        # subdirectories, otherwise we're looking for
+        # an exact match.
+        module_is_folder = module_path.name == "__init__.py"
+        if module_is_folder:
+            module_path = module_path.parent
+
+        # Look through frames until we find a file that
+        # isn't a part of our module, then return that stacklevel.
+        for level, frame in enumerate(inspect.stack()):
+            # Garbage collecting frames
+            frame_filename = Path(frame.filename)
+            del frame
+
+            if (
+                # If the module is a folder we look at subdirectory
+                module_is_folder
+                and module_path not in frame_filename.parents
+            ) or (
+                # Otherwise we're looking for an exact match.
+                not module_is_folder
+                and module_path != frame_filename
+            ):
+                return level
+    except KeyError:
+        pass
+    except Exception:
+        return 2
+    return 0
+
 
 __all__ = [
     "get_running_loop",

elastic_transport/_models.py

@@ -210,24 +210,24 @@ def _normalize_key(self, key: str) -> str:
 class ApiResponseMeta:
     """Metadata that is returned from Transport.perform_request()"""
 
-    #: Node which handled the request
-    node: "NodeConfig"
-
-    #: Number of seconds from start of request to start of response
-    duration: float
+    #: HTTP status code
+    status: int
 
     #: HTTP version being used
     http_version: str
 
-    #: HTTP status code
-    status: int
-
     #: HTTP headers
     headers: HttpHeaders
 
+    #: Number of seconds from start of request to start of response
+    duration: float
+
+    #: Node which handled the request
+    node: "NodeConfig"
+
     #: Extras that can be set to anything, typically used by third-parties.
     #: Third-party keys should start with an underscore and prefix.
-    _extras: Dict[str, Any] = field(default_factory=dict)
+    _extras: Dict[str, Any] = field(default_factory=dict, repr=False)
 
     @property
     def mimetype(self) -> Optional[str]:
@@ -388,6 +388,14 @@ def __hash__(self) -> int:
         )
 
 
+@dataclass()
+class SniffOptions:
+    """Options which are passed to Transport.sniff_callback"""
+
+    is_initial_sniff: bool
+    sniff_timeout: Optional[float]
+
+
 @dataclass(frozen=True, repr=True)
 class RequestOptions:
     """Options which can be passed per request to the Transport"""

elastic_transport/_node/_http_aiohttp.py

@@ -24,7 +24,7 @@
 import warnings
 from typing import Tuple
 
-from .._compat import get_running_loop
+from .._compat import get_running_loop, warn_stacklevel
 from .._exceptions import ConnectionError, ConnectionTimeout, SecurityWarning, TlsError
 from .._models import ApiResponseMeta, HttpHeaders, NodeConfig
 from ..client_utils import DEFAULT, client_meta_version, normalize_headers
@@ -77,7 +77,7 @@ def __init__(self, config: NodeConfig):
                 if config.ssl_show_warn:
                     warnings.warn(
                         f"Connecting to {self.base_url!r} using TLS with verify_certs=False is insecure",
-                        stacklevel=2,
+                        stacklevel=warn_stacklevel(),
                         category=SecurityWarning,
                     )
 

elastic_transport/_node/_http_requests.py

@@ -23,6 +23,7 @@
 
 import urllib3
 
+from .._compat import warn_stacklevel
 from .._exceptions import ConnectionError, ConnectionTimeout, SecurityWarning, TlsError
 from .._models import ApiResponseMeta, HttpHeaders, NodeConfig
 from ..client_utils import DEFAULT, client_meta_version
@@ -97,7 +98,7 @@ def __init__(self, config: NodeConfig):
         ):
             warnings.warn(
                 f"Connecting to {self.base_url!r} using TLS with verify_certs=False is insecure",
-                stacklevel=2,
+                stacklevel=warn_stacklevel(),
                 category=SecurityWarning,
             )
 

elastic_transport/_node/_http_urllib3.py

@@ -25,6 +25,7 @@
 from urllib3.exceptions import ConnectTimeoutError, ReadTimeoutError
 from urllib3.util.retry import Retry
 
+from .._compat import warn_stacklevel
 from .._exceptions import ConnectionError, ConnectionTimeout, SecurityWarning, TlsError
 from .._models import ApiResponseMeta, HttpHeaders, NodeConfig
 from ..client_utils import DEFAULT, client_meta_version
@@ -87,7 +88,7 @@ def __init__(self, config: NodeConfig):
                 if config.ssl_show_warn:
                     warnings.warn(
                         f"Connecting to {self.base_url!r} using TLS with verify_certs=False is insecure",
-                        stacklevel=2,
+                        stacklevel=warn_stacklevel(),
                         category=SecurityWarning,
                     )
                 else:

elastic_transport/_node_pool.py

@@ -289,7 +289,7 @@ def add(self, node_config: NodeConfig) -> None:
                 # Start the node as dead because 'dead_nodes' is thread-safe.
                 # The node will be resurrected on the next call to .get()
                 self.dead_consecutive_failures[node.config] = 0
-                self.dead_nodes.put((0.0, node))
+                self.dead_nodes.put((time.time(), node))
 
     def remove(self, node_config: NodeConfig) -> None:
         # Can't mark a seed node as removed.