Force verification to make at most one request

Author: sethmlarson

elasticsearch/_async/transport.py

@@ -27,6 +27,7 @@
     ConnectionError,
     ConnectionTimeout,
     ElasticsearchWarning,
+    NotElasticsearchError,
     SerializationError,
     TransportError,
 )
@@ -117,6 +118,10 @@ async def _async_init(self):
         self.loop = get_running_loop()
         self.kwargs["loop"] = self.loop
 
+        # Set our 'verified_once' implementation to one that
+        # works with 'asyncio' instead of 'threading'
+        self._verified_once = Once()
+
         # Now that we have a loop we can create all our HTTP connections...
         self.set_connections(self.hosts)
         self.seed_connections = list(self.connection_pool.connections[:])
@@ -332,8 +337,17 @@ async def perform_request(self, method, url, headers=None, params=None, body=Non
         )
 
         # Before we make the actual API call we verify the Elasticsearch instance.
-        if not self._verified_elasticsearch:
-            await self._do_verify_elasticsearch(headers=headers, timeout=timeout)
+        if self._verified_elasticsearch is None:
+            await self._verified_once.call(
+                self._do_verify_elasticsearch, headers=headers, timeout=timeout
+            )
+
+        # If '_verified_elasticsearch' is False we know we're not connected to Elasticsearch.
+        if self._verified_elasticsearch is False:
+            raise NotElasticsearchError(
+                "The client noticed that the server is not Elasticsearch "
+                "and we do not support this unknown product"
+            )
 
         for attempt in range(self.max_retries + 1):
             connection = self.get_connection()
@@ -471,7 +485,20 @@ async def _do_verify_elasticsearch(self, headers, timeout):
             raise error
 
         # Check the information we got back from the index request.
-        _verify_elasticsearch(info_headers, info_response)
+        self._verified_elasticsearch = _verify_elasticsearch(
+            info_headers, info_response
+        )
+
+
+class Once:
+    """Simple class which forces an async function to only execute once."""
+
+    def __init__(self):
+        self._lock = asyncio.Lock()
+        self._called = False
 
-        # If we made it through the above call this config is verified.
-        self._verified_elasticsearch = True
+    async def call(self, func, *args, **kwargs):
+        async with self._lock:
+            if not self._called:
+                self._called = True
+                await func(*args, **kwargs)

elasticsearch/compat.py

@@ -70,6 +70,17 @@ def to_bytes(x, encoding="ascii"):
 except (ImportError, AttributeError):
     pass
 
+try:
+    from threading import Lock
+except ImportError:  # Python <3.7 isn't guaranteed to have threading support.
+
+    class Lock:
+        def __enter__(self):
+            pass
+
+        def __exit__(self, *_):
+            pass
+
 
 __all__ = [
     "string_types",

elasticsearch/compat.pyi

@@ -16,7 +16,7 @@
 #  under the License.
 
 import sys
-from typing import Callable, Tuple, Type, Union
+from typing import Callable, ContextManager, Tuple, Type, Union
 
 PY2: bool
 string_types: Tuple[type, ...]
@@ -25,6 +25,8 @@ to_str: Callable[[Union[str, bytes]], str]
 to_bytes: Callable[[Union[str, bytes]], bytes]
 reraise_exceptions: Tuple[Type[Exception], ...]
 
+Lock: ContextManager[None]
+
 if sys.version_info[0] == 2:
     from itertools import imap as map
     from urllib import quote as quote

elasticsearch/transport.py

@@ -22,6 +22,7 @@
 from platform import python_version
 
 from ._version import __versionstr__
+from .compat import Lock
 from .connection import Urllib3HttpConnection
 from .connection_pool import ConnectionPool, DummyConnectionPool, EmptyConnectionPool
 from .exceptions import (
@@ -204,9 +205,22 @@ def __init__(
         if http_client_meta:
             self._client_meta += (http_client_meta,)
 
-        # Flag which is set after verifying that we're
-        # connected to Elasticsearch.
-        self._verified_elasticsearch = False
+        # Tri-state flag that describes what state the verification
+        # of whether we're connected to an Elasticsearch cluster or not.
+        # The three states are:
+        # - 'None': Means we've either not started the verification process
+        #   or that the verification is in progress. '_verified_once' ensures
+        #   that multiple requests don't kick off multiple verification processes.
+        # - 'True': Means we've verified that we're talking to Elasticsearch or
+        #   that we can't rule out Elasticsearch due to auth issues. A warning
+        #   will be raised if we receive 401/403.
+        # - 'False': Means we've discovered we're not talking to Elasticsearch,
+        #   should raise an error in this case for every request.
+        self._verified_elasticsearch = None
+
+        # Ensures that the ES verification request only fires once and that
+        # all requests block until this request returns back.
+        self._verified_once = Once()
 
     def add_connection(self, host):
         """
@@ -391,7 +405,17 @@ def perform_request(self, method, url, headers=None, params=None, body=None):
         )
 
         # Before we make the actual API call we verify the Elasticsearch instance.
-        self._do_verify_elasticsearch(headers=headers, timeout=timeout)
+        if self._verified_elasticsearch is None:
+            self._verified_once.call(
+                self._do_verify_elasticsearch, headers=headers, timeout=timeout
+            )
+
+        # If '_verified_elasticsearch' is False we know we're not connected to Elasticsearch.
+        if self._verified_elasticsearch is False:
+            raise NotElasticsearchError(
+                "The client noticed that the server is not Elasticsearch "
+                "and we do not support this unknown product"
+            )
 
         for attempt in range(self.max_retries + 1):
             connection = self.get_connection()
@@ -513,7 +537,7 @@ def _do_verify_elasticsearch(self, headers, timeout):
         error we instead emit an 'ElasticsearchWarning'.
         """
         # Product check has already been done, no need to do again.
-        if self._verified_elasticsearch:
+        if self._verified_elasticsearch is not None:
             return
 
         headers = {header.lower(): value for header, value in (headers or {}).items()}
@@ -566,10 +590,9 @@ def _do_verify_elasticsearch(self, headers, timeout):
             raise error
 
         # Check the information we got back from the index request.
-        _verify_elasticsearch(info_headers, info_response)
-
-        # If we made it through the above call this config is verified.
-        self._verified_elasticsearch = True
+        self._verified_elasticsearch = _verify_elasticsearch(
+            info_headers, info_response
+        )
 
 
 def _verify_elasticsearch(headers, response):
@@ -616,7 +639,20 @@ def _verify_elasticsearch(headers, response):
         # 7.14+ and there's a bad 'X-Elastic-Product' HTTP header
         or ((7, 14, 0) <= version_number and bad_product_header)
     ):
-        raise NotElasticsearchError(
-            "The client noticed that the server is not Elasticsearch "
-            "and we do not support this unknown product"
-        )
+        return False
+
+    return True
+
+
+class Once:
+    """Simple class which forces a function to only execute once."""
+
+    def __init__(self):
+        self._lock = Lock()
+        self._called = False
+
+    def call(self, func, *args, **kwargs):
+        with self._lock:
+            if not self._called:
+                self._called = True
+                func(*args, **kwargs)

test_elasticsearch/test_async/test_transport.py

@@ -33,6 +33,7 @@
     AuthorizationException,
     ConnectionError,
     ElasticsearchWarning,
+    NotElasticsearchError,
     TransportError,
 )
 
@@ -714,3 +715,111 @@ async def test_verify_elasticsearch_skips_on_auth_errors(self, exception_cls):
                 },
             ),
         ]
+
+    async def test_multiple_requests_verify_elasticsearch_success(self, event_loop):
+        t = AsyncTransport(
+            [
+                {
+                    "data": '{"version":{"number":"7.13.0","build_flavor":"default"},"tagline":"You Know, for Search"}',
+                    "delay": 1,
+                }
+            ],
+            connection_class=DummyConnection,
+        )
+
+        results = []
+        completed_at = []
+
+        async def request_task():
+            try:
+                results.append(await t.perform_request("GET", "/_search"))
+            except Exception as e:
+                results.append(e)
+            completed_at.append(event_loop.time())
+
+        # Execute a bunch of requests concurrently.
+        tasks = []
+        start_time = event_loop.time()
+        for _ in range(10):
+            tasks.append(event_loop.create_task(request_task()))
+        await asyncio.gather(*tasks)
+        end_time = event_loop.time()
+
+        # Exactly 10 results completed
+        assert len(results) == 10
+
+        # No errors in the results
+        assert all(isinstance(result, dict) for result in results)
+
+        # Assert that this took longer than 2 seconds but less than 2.1 seconds
+        duration = end_time - start_time
+        assert 2 <= duration <= 2.1
+
+        # Assert that every result came after ~2 seconds, no fast completions.
+        assert all(
+            2 <= completed_time - start_time <= 2.1 for completed_time in completed_at
+        )
+
+        # Assert that the cluster is "verified"
+        assert t._verified_elasticsearch
+
+        # See that the first request is always 'GET /' for ES check
+        calls = t.connection_pool.connections[0].calls
+        assert calls[0][0] == ("GET", "/")
+
+        # The rest of the requests are 'GET /_search' afterwards
+        assert all(call[0][:2] == ("GET", "/_search") for call in calls[1:])
+
+    async def test_multiple_requests_verify_elasticsearch_errors(self, event_loop):
+        t = AsyncTransport(
+            [
+                {
+                    "data": '{"version":{"number":"7.13.0","build_flavor":"default"},"tagline":"BAD TAGLINE"}',
+                    "delay": 1,
+                }
+            ],
+            connection_class=DummyConnection,
+        )
+
+        results = []
+        completed_at = []
+
+        async def request_task():
+            try:
+                results.append(await t.perform_request("GET", "/_search"))
+            except Exception as e:
+                results.append(e)
+            completed_at.append(event_loop.time())
+
+        # Execute a bunch of requests concurrently.
+        tasks = []
+        start_time = event_loop.time()
+        for _ in range(10):
+            tasks.append(event_loop.create_task(request_task()))
+        await asyncio.gather(*tasks)
+        end_time = event_loop.time()
+
+        # Exactly 10 results completed
+        assert len(results) == 10
+
+        # All results were errors
+        assert all(isinstance(result, NotElasticsearchError) for result in results)
+
+        # Assert that one request was made but not 2 requests.
+        duration = end_time - start_time
+        assert 1 <= duration <= 1.1
+
+        # Assert that every result came after ~1 seconds, no fast completions.
+        assert all(
+            1 <= completed_time - start_time <= 1.1 for completed_time in completed_at
+        )
+
+        # Assert that the cluster is definitely not Elasticsearch
+        assert t._verified_elasticsearch is False
+
+        # See that the first request is always 'GET /' for ES check
+        calls = t.connection_pool.connections[0].calls
+        assert calls[0][0] == ("GET", "/")
+
+        # The rest of the requests are 'GET /_search' afterwards
+        assert all(call[0][:2] == ("GET", "/_search") for call in calls[1:])