Improvements to batching logic:

- Avoid 60s timeout putting into self.__reqs
- Handle graceful stopping of all bg threads
- Allow cancelling of hanging streams
- Reraise bg_exceptions when they happen
- Align shutdown timeout with client-defined insert timeout
- Add mock tests of cancelling bidi streamsa

Author: tsmith023

mock_tests/conftest.py

@@ -1,7 +1,7 @@
 import json
 import time
 from concurrent import futures
-from typing import Generator, Mapping
+from typing import AsyncGenerator, Generator, Mapping
 
 import grpc
 import pytest
@@ -141,6 +141,16 @@ def weaviate_client(
     client.close()
 
 
+@pytest.fixture(scope="function")
+async def weaviate_client_async(
+    weaviate_mock: HTTPServer, start_grpc_server: grpc.Server
+) -> AsyncGenerator[weaviate.WeaviateAsyncClient, None]:
+    client = weaviate.use_async_with_local(port=MOCK_PORT, host=MOCK_IP, grpc_port=MOCK_PORT_GRPC)
+    await client.connect()
+    yield client
+    await client.close()
+
+
 @pytest.fixture(scope="function")
 def weaviate_timeouts_client(
     weaviate_timeouts_mock: HTTPServer, start_grpc_server: grpc.Server
@@ -368,3 +378,39 @@ def forbidden(
     service = MockForbiddenWeaviateService()
     weaviate_pb2_grpc.add_WeaviateServicer_to_server(service, start_grpc_server)
     return weaviate_client.collections.use("ForbiddenCollection")
+
+
+class MockWeaviateService(weaviate_pb2_grpc.WeaviateServicer):
+    def BatchStream(
+        self,
+        request_iterator: Generator[batch_pb2.BatchStreamRequest, None, None],
+        context: grpc.ServicerContext,
+    ) -> Generator[batch_pb2.BatchStreamReply, None, None]:
+        while True:
+            if context.is_active():
+                time.sleep(0.1)
+            else:
+                raise grpc.RpcError(grpc.StatusCode.DEADLINE_EXCEEDED, "Deadline exceeded")
+
+
+@pytest.fixture(scope="function")
+def stream_cancel(
+    weaviate_client: weaviate.WeaviateClient,
+    weaviate_mock: HTTPServer,
+    start_grpc_server: grpc.Server,
+) -> Generator[weaviate.collections.Collection, None, None]:
+    name = "StreamCancelCollection"
+    weaviate_mock.expect_request(f"/v1/schema/{name}").respond_with_response(
+        Response(status=404)
+    )  # skips __create_batch_reset vectorizer logic
+    weaviate_pb2_grpc.add_WeaviateServicer_to_server(MockWeaviateService(), start_grpc_server)
+    client = weaviate.connect_to_local(
+        port=MOCK_PORT,
+        host=MOCK_IP,
+        grpc_port=MOCK_PORT_GRPC,
+        additional_config=weaviate.classes.init.AdditionalConfig(
+            timeout=weaviate.classes.init.Timeout(insert=1)
+        ),
+    )
+    yield client.collections.use(name)
+    client.close()

mock_tests/test_connect.py

@@ -0,0 +1,25 @@
+import time
+import pytest
+import weaviate
+from weaviate.proto.v1 import batch_pb2
+
+
+def test_bidi_stream_cancel_sync(stream_cancel: weaviate.collections.Collection):
+    def gen():
+        time.sleep(10)
+        yield batch_pb2.BatchStreamRequest()
+
+    out, call = stream_cancel._connection.grpc_batch_stream(gen())
+    assert call.is_active()
+    call.cancel()
+    assert not call.is_active()
+    with pytest.raises(weaviate.exceptions.WeaviateBatchStreamError) as e:
+        next(out)
+    assert "StatusCode.CANCELLED(Locally cancelled by application!)" in e.value.message
+
+
+def test_batch_stream_hanging_server(stream_cancel: weaviate.collections.Collection):
+    with pytest.raises(weaviate.exceptions.WeaviateBatchStreamError) as e:
+        with stream_cancel.batch.stream() as batch:
+            batch.add_object()
+    assert "The server did not hangup its side of the stream gracefully in time" in e.value.message

weaviate/collections/batch/async_.py

@@ -14,7 +14,6 @@
 
 from weaviate.collections.batch.base import (
     GCP_STREAM_TIMEOUT,
-    SHUTDOWN_TIMEOUT,
     ObjectsBatchRequest,
     ReferencesBatchRequest,
     _BatchDataWrapper,
@@ -194,14 +193,16 @@ async def recv_wrapper() -> None:
 
     async def _wait(self) -> None:
         assert self.__bg_tasks is not None
-        deadline = time.time() + SHUTDOWN_TIMEOUT
+        # this is how long an insert will take to timeout for, so we wait at most this time +5s for the batch to finish after shutdown is initiated, in case the server never hangs up
+        shutdown_timeout = self.__connection.timeout_config.insert + 5
+        deadline = time.time() + shutdown_timeout
         while time.time() < deadline:
             if not self.__bg_tasks.any_alive():
                 break
             await asyncio.sleep(0.1)
         if self.__bg_tasks.any_alive():
             logger.warning(
-                f"Background batch tasks did not exit within {SHUTDOWN_TIMEOUT}s. "
+                f"Background batch tasks did not exit within {shutdown_timeout}s. "
                 f"Forcing shutdown. inflight_objs={len(self.__inflight_objs)}, "
                 f"inflight_refs={len(self.__inflight_refs)}, "
                 f"loop_alive={self.__bg_tasks.loop_alive()}, "
@@ -211,7 +212,7 @@ async def _wait(self) -> None:
             self.__bg_tasks.recv.cancel()
             self.__bg_tasks.loop.cancel()
         try:
-            await asyncio.wait_for(self.__bg_tasks.gather(), timeout=5)
+            await asyncio.wait_for(self.__bg_tasks.gather(), timeout=None)
         except asyncio.TimeoutError as e:
             raise WeaviateBatchStreamError(
                 "Background batch tasks did not terminate after forced shutdown."

weaviate/collections/batch/base.py

@@ -64,7 +64,6 @@
 GCP_STREAM_TIMEOUT = (
     160  # GCP connections have a max lifetime of 180s, leave 20s of buffer as safety
 )
-SHUTDOWN_TIMEOUT = 300  # time to wait for background threads to exit after shutdown is initiated, in seconds, in the event the server never hangs up
 
 
 class BatchRequest(ABC, Generic[TBatchInput, TBatchReturn]):
@@ -911,10 +910,10 @@ def recv_alive(self) -> bool:
             return self.recv.is_alive()
         return True  # not started yet so considered alive
 
-    def join(self, timeout: Optional[float] = None) -> None:
+    def join(self) -> None:
         """Join the background threads."""
-        self.loop.join(timeout=timeout)
-        self.recv.join(timeout=timeout)
+        self.loop.join()
+        self.recv.join()
 
 
 class _ClusterBatch:

weaviate/collections/batch/sync.py

@@ -10,7 +10,6 @@
 
 from weaviate.collections.batch.base import (
     GCP_STREAM_TIMEOUT,
-    SHUTDOWN_TIMEOUT,
     ObjectsBatchRequest,
     ReferencesBatchRequest,
     _BatchDataWrapper,
@@ -158,14 +157,16 @@ def _start(self) -> None:
                 )
 
     def _wait(self) -> None:
-        deadline = time.time() + SHUTDOWN_TIMEOUT
+        # this is how long an insert will take to timeout for, so we wait at most this time +5s for the batch to finish after shutdown is initiated, in case the server never hangs up
+        shutdown_timeout = self.__connection.timeout_config.insert + 5
+        deadline = time.time() + shutdown_timeout
         while time.time() < deadline:
             if not self.__any_threads_alive():
                 break
             time.sleep(0.1)
         if self.__any_threads_alive():
             logger.warning(
-                f"Background batch threads did not exit within {SHUTDOWN_TIMEOUT}s. "
+                f"Background batch threads did not exit within {shutdown_timeout}s. "
                 f"Forcing shutdown. inflight_objs={len(self.__inflight_objs)}, "
                 f"inflight_refs={len(self.__inflight_refs)}, "
                 f"loop_alive={self.__bg_threads.loop_alive()}, "
@@ -175,7 +176,7 @@ def _wait(self) -> None:
             self.__is_stopped.set()
             self.__cancel_active_stream()  # force __recv to exit by cancelling the stream
 
-        self.__bg_threads.join(timeout=5)
+        self.__bg_threads.join()
         if self.__any_threads_alive():
             raise WeaviateBatchStreamError(
                 "Background batch threads did not terminate after forced shutdown. "
@@ -193,10 +194,29 @@ def _wait(self) -> None:
             self.__results_for_wrapper.imported_shards
         )
 
+        if self.__bg_exception is not None:
+            if "StatusCode.CANCELLED(Locally cancelled by application!)" in str(
+                self.__bg_exception
+            ):
+                raise WeaviateBatchStreamError(
+                    "The server did not hangup its side of the stream gracefully in time"
+                )
+            raise self.__bg_exception
+
     def _shutdown(self) -> None:
         # Shutdown the current batch and wait for all requests to be finished
         self.__is_stopped.set()
 
+    def __put(self, req: _BatchStreamRequest | None):
+        while True:
+            try:
+                self.__reqs.put(req, timeout=1)
+                return True
+            except Full:
+                if self.__bg_exception is not None or self.__shutdown_loop.is_set():
+                    return False
+                return self.__put(req)
+
     def __loop(self) -> None:
         refresh_time: float = 0.01
         while self.__bg_exception is None and not self.__shutdown_loop.is_set():
@@ -238,21 +258,18 @@ def __loop(self) -> None:
                     if paused:
                         logger.info("Server is back up, resuming batching loop...")
                         paused = False
-                    try:
-                        self.__reqs.put(req, timeout=60)
-                    except Full as e:
-                        logger.warning(
-                            "Batch queue is blocked for more than 60 seconds. Exiting the loop"
-                        )
-                        self.__bg_exception = e
+                    if not self.__put(req):
+                        logger.info("Batch loop is shutting down, stopping putting requests...")
                         return
             elif (
                 self.__is_stopped.is_set()
                 and not self.__is_hungup.is_set()
                 and not self.__is_shutting_down.is_set()
                 and not self.__is_oom.is_set()
             ):
-                self.__reqs.put(None)
+                if not self.__put(None):
+                    logger.info("Batch loop is shutting down, stopping putting shutdown signal...")
+                    return
             time.sleep(refresh_time)
 
     def __generate_stream_requests(

weaviate/connect/v4.py

@@ -1024,7 +1024,7 @@ def grpc_batch_stream(
             request_iterator=requests,
             timeout=self.timeout_config.stream,
             metadata=self.grpc_headers(),
-        )()
+        )
 
         def generator():
             try: