use context manager to run the timeout background tasks

Author: miguelgrinberg

elasticsearch/_async/helpers.py

@@ -33,6 +33,7 @@
     Union,
 )
 
+from ..compat import safe_task
 from ..exceptions import ApiError, NotFoundError, TransportError
 from ..helpers.actions import (
     _TYPE_BULK_ACTION,
@@ -94,28 +95,24 @@ async def get_items() -> None:
             try:
                 async for item in actions:
                     await item_queue.put(item)
-            except Exception:
+            finally:
                 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
+        async with safe_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
 
     ret = chunker.flush()
     if ret:

elasticsearch/compat.py

@@ -15,10 +15,13 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import asyncio
+from contextlib import contextmanager, asynccontextmanager
 import inspect
 import os
 import sys
 from pathlib import Path
+from threading import Thread
 from typing import Tuple, Type, Union
 
 string_types: Tuple[Type[str], Type[bytes]] = (str, bytes)
@@ -76,9 +79,46 @@ def warn_stacklevel() -> int:
     return 0
 
 
+@contextmanager
+def safe_thread(target, *args, **kwargs):
+    """Run a thread within a context manager block.
+
+    The thread is automatically joined when the block ends. If the thread raised
+    an exception, it is raised in the caller's context.
+    """
+    captured_exception = None
+
+    def run():
+        try:
+            target(*args, **kwargs)
+        except BaseException as exc:
+            nonlocal captured_exception
+            captured_exception = exc
+
+    thread = Thread(target=run)
+    thread.start()
+    yield
+    thread.join()
+    if captured_exception:
+        raise captured_exception
+
+
+@asynccontextmanager
+async def safe_task(coro):
+    """Run a background task within a context manager block.
+
+    The task is awaited when the block ends.
+    """
+    task = asyncio.create_task(coro)
+    yield task
+    await task
+
+
 __all__ = [
     "string_types",
     "to_str",
     "to_bytes",
     "warn_stacklevel",
+    "safe_thread",
+    "safe_task",
 ]

elasticsearch/helpers/actions.py

@@ -39,7 +39,7 @@
 from elastic_transport import OpenTelemetrySpan
 
 from .. import Elasticsearch
-from ..compat import to_bytes
+from ..compat import to_bytes, safe_thread
 from ..exceptions import ApiError, NotFoundError, TransportError
 from ..serializer import Serializer
 from .errors import BulkIndexError, ScanError
@@ -266,35 +266,28 @@ def _chunk_actions(
         )
 
         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))
+            finally:
+                # make sure we signal the end even if there is an exception
+                item_queue.put((BulkMeta.done, None))
 
-        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()
+        with safe_thread(get_items):
+            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:
+                    break
+                ret = chunker.feed(action, data)
+                if ret:
+                    yield ret
 
     ret = chunker.flush()
     if ret: