bulk timeouts

Author: miguelgrinberg

elasticsearch/_async/helpers.py

@@ -39,6 +39,8 @@
     _TYPE_BULK_ACTION_BODY,
     _TYPE_BULK_ACTION_HEADER,
     _TYPE_BULK_ACTION_HEADER_AND_BODY,
+    _TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY,
+    BulkMeta,
     _ActionChunker,
     _process_bulk_chunk_error,
     _process_bulk_chunk_success,
@@ -54,9 +56,10 @@
 
 
 async def _chunk_actions(
-    actions: AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY],
+    actions: AsyncIterable[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY],
     chunk_size: int,
     max_chunk_bytes: int,
+    flush_after_seconds: Optional[float],
     serializer: Serializer,
 ) -> AsyncIterable[
     Tuple[
@@ -76,10 +79,44 @@ async def _chunk_actions(
     chunker = _ActionChunker(
         chunk_size=chunk_size, max_chunk_bytes=max_chunk_bytes, serializer=serializer
     )
-    async for action, data in actions:
-        ret = chunker.feed(action, data)
-        if ret:
-            yield ret
+
+    if not flush_after_seconds:
+        async for action, data in actions:
+            ret = chunker.feed(action, data)
+            if ret:
+                yield ret
+    else:
+        item_queue: asyncio.Queue[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY] = (
+            asyncio.Queue()
+        )
+
+        async def get_items() -> None:
+            try:
+                async for item in actions:
+                    await item_queue.put(item)
+            except Exception:
+                await item_queue.put((BulkMeta.done, None))
+                raise
+            await item_queue.put((BulkMeta.done, None))
+
+        item_getter_job = asyncio.create_task(get_items())
+
+        timeout: Optional[float] = flush_after_seconds
+        while True:
+            try:
+                action, data = await asyncio.wait_for(item_queue.get(), timeout=timeout)
+                timeout = flush_after_seconds
+            except asyncio.TimeoutError:
+                action, data = BulkMeta.flush, None
+                timeout = None
+
+            if action is BulkMeta.done:
+                break
+            ret = chunker.feed(action, data)
+            if ret:
+                yield ret
+        await item_getter_job
+
     ret = chunker.flush()
     if ret:
         yield ret
@@ -162,6 +199,7 @@ async def async_streaming_bulk(
     actions: Union[Iterable[_TYPE_BULK_ACTION], AsyncIterable[_TYPE_BULK_ACTION]],
     chunk_size: int = 500,
     max_chunk_bytes: int = 100 * 1024 * 1024,
+    flush_after_seconds: Optional[float] = None,
     raise_on_error: bool = True,
     expand_action_callback: Callable[
         [_TYPE_BULK_ACTION], _TYPE_BULK_ACTION_HEADER_AND_BODY
@@ -234,7 +272,7 @@ async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
     ]
     bulk_actions: List[bytes]
     async for bulk_data, bulk_actions in _chunk_actions(
-        map_actions(), chunk_size, max_chunk_bytes, serializer
+        map_actions(), chunk_size, max_chunk_bytes, flush_after_seconds, serializer
     ):
         for attempt in range(max_retries + 1):
             to_retry: List[bytes] = []

elasticsearch/helpers/actions.py

@@ -16,10 +16,11 @@
 #  under the License.
 
 import logging
+import queue
 import time
 from enum import Enum
 from operator import methodcaller
-from queue import Queue
+from threading import Thread
 from typing import (
     Any,
     Callable,
@@ -48,16 +49,23 @@
 
 class BulkMeta(Enum):
     flush = 1
+    done = 2
 
 
 BULK_FLUSH = BulkMeta.flush
 
-_TYPE_BULK_ACTION = Union[bytes, str, Dict[str, Any], BulkMeta]
+_TYPE_BULK_ACTION = Union[bytes, str, Dict[str, Any]]
 _TYPE_BULK_ACTION_HEADER = Dict[str, Any]
-_TYPE_BULK_ACTION_HEADER_AND_META = Union[Dict[str, Any], BulkMeta]
 _TYPE_BULK_ACTION_BODY = Union[None, bytes, Dict[str, Any]]
 _TYPE_BULK_ACTION_HEADER_AND_BODY = Tuple[
-    _TYPE_BULK_ACTION_HEADER_AND_META, _TYPE_BULK_ACTION_BODY
+    _TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY
+]
+
+_TYPE_BULK_ACTION_WITH_META = Union[bytes, str, Dict[str, Any], BulkMeta]
+_TYPE_BULK_ACTION_HEADER_WITH_META = Union[Dict[str, Any], BulkMeta]
+_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY = Union[
+    Tuple[_TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY],
+    Tuple[BulkMeta, Any],
 ]
 
 
@@ -71,9 +79,6 @@ def expand_action(data: _TYPE_BULK_ACTION) -> _TYPE_BULK_ACTION_HEADER_AND_BODY:
     if isinstance(data, (bytes, str)):
         return {"index": {}}, to_bytes(data, "utf-8")
 
-    if isinstance(data, BulkMeta):
-        return data, {}
-
     # make sure we don't alter the action
     data = data.copy()
     op_type: str = data.pop("_op_type", "index")
@@ -151,7 +156,9 @@ def __init__(
         ] = []
 
     def feed(
-        self, action: _TYPE_BULK_ACTION_HEADER_AND_META, data: _TYPE_BULK_ACTION_BODY
+        self,
+        action: _TYPE_BULK_ACTION_HEADER_WITH_META,
+        data: _TYPE_BULK_ACTION_BODY,
     ) -> Optional[
         Tuple[
             List[
@@ -224,9 +231,10 @@ def flush(
 
 
 def _chunk_actions(
-    actions: Iterable[_TYPE_BULK_ACTION_HEADER_AND_BODY],
+    actions: Iterable[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY],
     chunk_size: int,
     max_chunk_bytes: int,
+    flush_after_seconds: Optional[float],
     serializer: Serializer,
 ) -> Iterable[
     Tuple[
@@ -246,10 +254,48 @@ def _chunk_actions(
     chunker = _ActionChunker(
         chunk_size=chunk_size, max_chunk_bytes=max_chunk_bytes, serializer=serializer
     )
-    for action, data in actions:
-        ret = chunker.feed(action, data)
-        if ret:
-            yield ret
+
+    if not flush_after_seconds:
+        for action, data in actions:
+            ret = chunker.feed(action, data)
+            if ret:
+                yield ret
+    else:
+        item_queue: queue.Queue[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY] = (
+            queue.Queue()
+        )
+
+        def get_items() -> None:
+            ret = None
+            try:
+                for item in actions:
+                    item_queue.put(item)
+            except BaseException as exc:
+                ret = exc
+            item_queue.put((BulkMeta.done, ret))
+
+        item_getter_job = Thread(target=get_items)
+        item_getter_job.start()
+
+        timeout: Optional[float] = flush_after_seconds
+        while True:
+            try:
+                action, data = item_queue.get(timeout=timeout)
+                timeout = flush_after_seconds
+            except queue.Empty:
+                action, data = BulkMeta.flush, None
+                timeout = None
+
+            if action is BulkMeta.done:
+                if isinstance(data, BaseException):
+                    raise data
+                break
+            ret = chunker.feed(action, data)
+            if ret:
+                yield ret
+
+        item_getter_job.join()
+
     ret = chunker.flush()
     if ret:
         yield ret
@@ -376,9 +422,10 @@ def _process_bulk_chunk(
 
 def streaming_bulk(
     client: Elasticsearch,
-    actions: Iterable[_TYPE_BULK_ACTION],
+    actions: Iterable[_TYPE_BULK_ACTION_WITH_META],
     chunk_size: int = 500,
     max_chunk_bytes: int = 100 * 1024 * 1024,
+    flush_after_seconds: Optional[float] = None,
     raise_on_error: bool = True,
     expand_action_callback: Callable[
         [_TYPE_BULK_ACTION], _TYPE_BULK_ACTION_HEADER_AND_BODY
@@ -440,6 +487,13 @@ def streaming_bulk(
 
         serializer = client.transport.serializers.get_serializer("application/json")
 
+        def expand_action_with_meta(
+            data: _TYPE_BULK_ACTION_WITH_META,
+        ) -> _TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY:
+            if isinstance(data, BulkMeta):
+                return data, None
+            return expand_action_callback(data)
+
         bulk_data: List[
             Union[
                 Tuple[_TYPE_BULK_ACTION_HEADER],
@@ -448,9 +502,10 @@ def streaming_bulk(
         ]
         bulk_actions: List[bytes]
         for bulk_data, bulk_actions in _chunk_actions(
-            map(expand_action_callback, actions),
+            map(expand_action_with_meta, actions),
             chunk_size,
             max_chunk_bytes,
+            flush_after_seconds,
             serializer,
         ):
             for attempt in range(max_retries + 1):
@@ -572,6 +627,7 @@ def parallel_bulk(
     thread_count: int = 4,
     chunk_size: int = 500,
     max_chunk_bytes: int = 100 * 1024 * 1024,
+    flush_after_seconds: Optional[float] = None,
     queue_size: int = 4,
     expand_action_callback: Callable[
         [_TYPE_BULK_ACTION], _TYPE_BULK_ACTION_HEADER_AND_BODY
@@ -611,7 +667,7 @@ def _setup_queues(self) -> None:
             super()._setup_queues()  # type: ignore[misc]
             # The queue must be at least the size of the number of threads to
             # prevent hanging when inserting sentinel values during teardown.
-            self._inqueue: Queue[
+            self._inqueue: queue.Queue[
                 Tuple[
                     List[
                         Union[
@@ -620,7 +676,7 @@ def _setup_queues(self) -> None:
                     ],
                     List[bytes],
                 ]
-            ] = Queue(max(queue_size, thread_count))
+            ] = queue.Queue(max(queue_size, thread_count))
             self._quick_put = self._inqueue.put
 
     with client._otel.helpers_span("helpers.parallel_bulk") as otel_span:
@@ -640,7 +696,11 @@ def _setup_queues(self) -> None:
                     )
                 ),
                 _chunk_actions(
-                    expanded_actions, chunk_size, max_chunk_bytes, serializer
+                    expanded_actions,
+                    chunk_size,
+                    max_chunk_bytes,
+                    flush_after_seconds,
+                    serializer,
                 ),
             ):
                 yield from result

test_elasticsearch/test_helpers.py

@@ -18,6 +18,7 @@
 import pickle
 import threading
 import time
+from typing import Optional
 from unittest import mock
 
 import pytest
@@ -85,7 +86,6 @@ def setup_method(self, _):
     def test_expand_action(self):
         assert helpers.expand_action({}) == ({"index": {}}, {})
         assert helpers.expand_action({"key": "val"}) == ({"index": {}}, {"key": "val"})
-        assert helpers.expand_action(helpers.BULK_FLUSH) == (helpers.BULK_FLUSH, {})
 
     def test_expand_action_actions(self):
         assert helpers.expand_action(
@@ -157,30 +157,44 @@ def test__source_metadata_or_source(self):
             {"_source": {"key2": "val2"}, "key": "val", "_op_type": "update"}
         ) == ({"update": {}}, {"key2": "val2"})
 
-    def test_chunks_are_chopped_by_byte_size(self):
+    @pytest.mark.parametrize("flush_seconds", [None, 10])
+    def test_chunks_are_chopped_by_byte_size(self, flush_seconds: Optional[float]):
         assert 100 == len(
-            list(helpers._chunk_actions(self.actions, 100000, 1, JSONSerializer()))
+            list(
+                helpers._chunk_actions(
+                    self.actions, 100000, 1, flush_seconds, JSONSerializer()
+                )
+            )
         )
 
-    def test_chunks_are_chopped_by_chunk_size(self):
+    @pytest.mark.parametrize("flush_seconds", [None, 10])
+    def test_chunks_are_chopped_by_chunk_size(self, flush_seconds: Optional[float]):
         assert 10 == len(
-            list(helpers._chunk_actions(self.actions, 10, 99999999, JSONSerializer()))
+            list(
+                helpers._chunk_actions(
+                    self.actions, 10, 99999999, flush_seconds, JSONSerializer()
+                )
+            )
         )
 
-    def test_chunks_are_chopped_by_byte_size_properly(self):
+    @pytest.mark.parametrize("flush_seconds", [None, 10])
+    def test_chunks_are_chopped_by_byte_size_properly(
+        self, flush_seconds: Optional[float]
+    ):
         max_byte_size = 170
         chunks = list(
             helpers._chunk_actions(
-                self.actions, 100000, max_byte_size, JSONSerializer()
+                self.actions, 100000, max_byte_size, flush_seconds, JSONSerializer()
             )
         )
         assert 25 == len(chunks)
         for chunk_data, chunk_actions in chunks:
             chunk = b"".join(chunk_actions)
             assert len(chunk) <= max_byte_size
 
-    def test_chunks_are_chopped_by_flush(self):
-        flush = helpers.expand_action(helpers.BULK_FLUSH)
+    @pytest.mark.parametrize("flush_seconds", [None, 10])
+    def test_chunks_are_chopped_by_flush(self, flush_seconds: Optional[float]):
+        flush = (helpers.BULK_FLUSH, None)
         actions = (
             self.actions[:3]
             + [flush] * 2  # two consecutive flushes after 3 items
@@ -189,7 +203,11 @@ def test_chunks_are_chopped_by_flush(self):
             + self.actions[4:]
             + [flush]  # flush at the end
         )
-        chunks = list(helpers._chunk_actions(actions, 100, 99999999, JSONSerializer()))
+        chunks = list(
+            helpers._chunk_actions(
+                actions, 100, 99999999, flush_seconds, JSONSerializer()
+            )
+        )
         assert 3 == len(chunks)
         assert len(chunks[0][0]) == 3
         assert len(chunks[0][1]) == 6