Implement consistent retry policy when calling Elasticsearch

connectors/config.py

@@ -59,8 +59,9 @@ def _default_config():
                 "max_concurrency": 5,
                 "chunk_max_mem_size": 5,
                 "concurrent_downloads": 10,
-                "max_retries": 3,
             },
+            "max_retries": 5,
+            "retry_timeout": 10,
             "retry_on_timeout": True,
             "request_timeout": 120,
             "max_wait_duration": 120,

connectors/es/client.py

@@ -9,6 +9,7 @@
 import time
 from enum import Enum
 
+from elastic_transport import ConnectionTimeout
 from elastic_transport.client_utils import url_to_node_config
 from elasticsearch import ApiError, AsyncElasticsearch, ConflictError
 from elasticsearch import (
@@ -17,7 +18,11 @@
 
 from connectors import __version__
 from connectors.logger import logger, set_extra_logger
-from connectors.utils import CancellableSleeps
+from connectors.utils import (
+    CancellableSleeps,
+    RetryStrategy,
+    time_to_sleep_between_retries,
+)
 
 
 class License(Enum):
@@ -41,6 +46,10 @@ def __init__(self, config):
             use_default_ports_for_scheme=True,
         )
         self._sleeps = CancellableSleeps()
+        self._retrier = TransientElasticsearchRetrier(
+            logger, config.get("max_retries", 5), config.get("retry_timeout", 10)
+        )
+
         options = {
             "hosts": [self.host],
             "request_timeout": config.get("request_timeout", 120),
@@ -100,7 +109,9 @@ async def has_active_license_enabled(self, license_):
             Tuple: (boolean if `license_` is enabled and not expired, actual license Elasticsearch is using)
         """
 
-        license_response = await self.client.license.get()
+        license_response = await self._retrier.execute_with_retry(
+            self.client.license.get
+        )
         license_info = license_response.get("license", {})
         is_expired = license_info.get("status", "").lower() == "expired"
 
@@ -125,23 +136,9 @@ async def has_active_license_enabled(self, license_):
         )
 
     async def close(self):
+        await self._retrier.close()
         await self.client.close()
 
-    async def ping(self):
-        try:
-            await self.client.info()
-        except ApiError as e:
-            logger.error(f"The server returned a {e.status_code} code")
-            if e.info is not None and "error" in e.info and "reason" in e.info["error"]:
-                logger.error(e.info["error"]["reason"])
-            return False
-        except ElasticConnectionError as e:
-            logger.error("Could not connect to the server")
-            if e.message is not None:
-                logger.error(e.message)
-            return False
-        return True
-
     async def wait(self):
         backoff = self.initial_backoff_duration
         start = time.time()
@@ -162,6 +159,80 @@ async def wait(self):
         await self.close()
         return False
 
+    async def ping(self):
+        try:
+            await self.client.info()
+        except ApiError as e:
+            logger.error(f"The server returned a {e.status_code} code")
+            if e.info is not None and "error" in e.info and "reason" in e.info["error"]:
+                logger.error(e.info["error"]["reason"])
+            return False
+        except ElasticConnectionError as e:
+            logger.error("Could not connect to the server")
+            if e.message is not None:
+                logger.error(e.message)
+            return False
+        return True
+
+
+class RetryInterruptedError(Exception):
+    pass
+
+
+class TransientElasticsearchRetrier:
+    def __init__(
+        self,
+        logger_,
+        max_retries,
+        retry_timeout,
+        retry_strategy=RetryStrategy.LINEAR_BACKOFF,
+    ):
+        self._logger = logger_
+        self._sleeps = CancellableSleeps()
+        self._keep_retrying = True
+        self._max_retries = max_retries
+        self._retry_timeout = retry_timeout
+        self._retry_strategy = retry_strategy
+
+    async def close(self):
+        self._sleeps.cancel()
+        self._keep_retrying = False
+
+    async def _sleep(self, retry):
+        time_to_sleep = time_to_sleep_between_retries(
+            self._retry_strategy, self._retry_timeout, retry
+        )
+        self._logger.debug(f"Attempt {retry}: sleeping for {time_to_sleep}")
+        await self._sleeps.sleep(time_to_sleep)
+
+    async def execute_with_retry(self, func):
+        retry = 0
+        while self._keep_retrying and retry < self._max_retries:
+            retry += 1
+            try:
+                result = await func()
+
+                return result
+            except ConnectionTimeout:
+                self._logger.debug(f"Attempt {retry}: connection timeout")
+
+                if retry >= self._max_retries:
+                    raise
+            except ApiError as e:
+                self._logger.debug(
+                    f"Attempt {retry}: api error with status {e.status_code}"
+                )
+
+                if e.status_code != 429:
+                    raise
+                if retry >= self._max_retries:
+                    raise
+
+            await self._sleep(retry)
+
+        msg = "Retry operation was interrupted"
+        raise RetryInterruptedError(msg)
+
 
 def with_concurrency_control(retries=3):
     def wrapper(func):

connectors/es/management_client.py

@@ -4,6 +4,8 @@
 # you may not use this file except in compliance with the Elastic License 2.0.
 #
 
+from functools import partial
+
 from elasticsearch import (
     NotFoundError as ElasticNotFoundError,
 )
@@ -39,34 +41,48 @@ async def ensure_exists(self, indices=None):
 
         for index in indices:
             logger.debug(f"Checking index {index}")
-            if not await self.client.indices.exists(index=index):
-                await self.client.indices.create(index=index)
+            if not await self._retrier.execute_with_retry(
+                partial(self.client.indices.exists, index=index)
+            ):
+                await self._retrier.execute_with_retry(
+                    partial(self.client.indices.create, index=index)
+                )
                 logger.debug(f"Created index {index}")
 
     async def create_content_index(self, search_index_name, language_code):
         settings = Settings(language_code=language_code, analysis_icu=False).to_hash()
         mappings = Mappings.default_text_fields_mappings(is_connectors_index=True)
 
-        return await self.client.indices.create(
-            index=search_index_name, mappings=mappings, settings=settings
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.indices.create,
+                index=search_index_name,
+                mappings=mappings,
+                settings=settings,
+            )
         )
 
     async def ensure_content_index_mappings(self, index, mappings):
         # open = Match open, non-hidden indices. Also matches any non-hidden data stream.
         # Content indices are always non-hidden.
-        response = await self.client.indices.get_mapping(index=index)
+        response = await self._retrier.execute_with_retry(
+            partial(self.client.indices.get_mapping, index=index)
+        )
 
         existing_mappings = response[index].get("mappings", {})
         if len(existing_mappings) == 0:
             if mappings:
                 logger.debug(
                     "Index %s has no mappings or it's empty. Adding mappings...", index
                 )
-                await self.client.indices.put_mapping(
-                    index=index,
-                    dynamic=mappings.get("dynamic", False),
-                    dynamic_templates=mappings.get("dynamic_templates", []),
-                    properties=mappings.get("properties", {}),
+                await self._retrier.execute_with_retry(
+                    partial(
+                        self.client.indices.put_mapping,
+                        index=index,
+                        dynamic=mappings.get("dynamic", False),
+                        dynamic_templates=mappings.get("dynamic_templates", []),
+                        properties=mappings.get("properties", {}),
+                    )
                 )
                 logger.debug("Successfully added mappings for index %s", index)
             else:
@@ -82,34 +98,62 @@ async def ensure_ingest_pipeline_exists(
         self, pipeline_id, version, description, processors
     ):
         try:
-            await self.client.ingest.get_pipeline(id=pipeline_id)
+            await self._retrier.execute_with_retry(
+                partial(self.client.ingest.get_pipeline, id=pipeline_id)
+            )
         except ElasticNotFoundError:
-            await self.client.ingest.put_pipeline(
-                id=pipeline_id,
-                version=version,
-                description=description,
-                processors=processors,
+            await self._retrier.execute_with_retry(
+                partial(
+                    self.client.ingest.put_pipeline,
+                    id=pipeline_id,
+                    version=version,
+                    description=description,
+                    processors=processors,
+                )
             )
 
     async def delete_indices(self, indices):
-        await self.client.indices.delete(index=indices, ignore_unavailable=True)
+        await self._retrier.execute_with_retry(
+            partial(self.client.indices.delete, index=indices, ignore_unavailable=True)
+        )
 
     async def clean_index(self, index_name):
-        return await self.client.delete_by_query(
-            index=index_name, body={"query": {"match_all": {}}}, ignore_unavailable=True
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.delete_by_query,
+                index=index_name,
+                body={"query": {"match_all": {}}},
+                ignore_unavailable=True,
+            )
         )
 
     async def list_indices(self):
-        return await self.client.indices.stats(index="search-*")
+        return await self._retrier.execute_with_retry(
+            partial(self.client.indices.stats, index="search-*")
+        )
 
     async def index_exists(self, index_name):
-        return await self.client.indices.exists(index=index_name)
+        return await self._retrier.execute_with_retry(
+            partial(self.client.indices.exists, index=index_name)
+        )
 
     async def upsert(self, _id, index_name, doc):
-        await self.client.index(
-            id=_id,
-            index=index_name,
-            document=doc,
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.index,
+                id=_id,
+                index=index_name,
+                document=doc,
+            )
+        )
+
+    async def bulk_insert(self, operations, pipeline):
+        return await self._retrier.execute_with_retry(
+            partial(
+                self.client.bulk,
+                operations=operations,
+                pipeline=pipeline,
+            )
         )
 
     async def yield_existing_documents_metadata(self, index):

connectors/es/sink.py

@@ -43,7 +43,6 @@
     aenumerate,
     get_size,
     iso_utc,
-    retryable,
 )
 
 __all__ = ["SyncOrchestrator"]
@@ -130,20 +129,16 @@ def _bulk_op(self, doc, operation=OP_INDEX):
 
     @tracer.start_as_current_span("_bulk API call", slow_log=1.0)
     async def _batch_bulk(self, operations, stats):
-        @retryable(retries=self.max_retires)
-        async def _bulk_api_call():
-            return await self.client.client.bulk(
-                operations=operations, pipeline=self.pipeline["name"]
-            )
-
         # TODO: treat result to retry errors like in async_streaming_bulk
         task_num = len(self.bulk_tasks)
 
         if self._logger.isEnabledFor(logging.DEBUG):
             self._logger.debug(
                 f"Task {task_num} - Sending a batch of {len(operations)} ops -- {get_mb_size(operations)}MiB"
             )
-        res = await _bulk_api_call()
+
+        # TODO: retry 429s for individual items here
+        res = await self.client.bulk_insert(operations, self.pipeline["name"])
         if res.get("errors"):
             for item in res["items"]:
                 for op, data in item.items():