Remove redundant stream cancellation logic, align async with sync impl

Author: tsmith023

.github/workflows/main.yaml

@@ -22,14 +22,13 @@ env:
   WEAVIATE_128: 1.28.16
   WEAVIATE_129: 1.29.11
   WEAVIATE_130: 1.30.22
-  WEAVIATE_131: 1.31.20
-  WEAVIATE_132: 1.32.23
-  WEAVIATE_133: 1.33.10
-  WEAVIATE_134: 1.34.5
-  WEAVIATE_135: 1.35.16-efdedfa
-  WEAVIATE_136: 1.36.9-d905e6c
-  WEAVIATE_137: 1.37.0-rc.0-b313954.amd64
-
+  WEAVIATE_131: 1.31.22
+  WEAVIATE_132: 1.32.27
+  WEAVIATE_133: 1.33.18
+  WEAVIATE_134: 1.34.19
+  WEAVIATE_135: 1.35.15
+  WEAVIATE_136: 1.36.6-8edcf08.amd64
+  WEAVIATE_137: 1.37.0-dev-29d5c87.amd64
 
 jobs:
   lint-and-format:

mock_tests/test_connect.py

@@ -115,7 +115,6 @@ def __init__(
         self.__oom_wait_time = 300
 
         self.__shutdown_loop = asyncio.Event()
-        self.__sent_sentinel = asyncio.Event()
 
         self.__objs_cache_lock = asyncio.Lock()
         self.__objs_cache: dict[str, BatchObject] = {}
@@ -195,34 +194,12 @@ async def _wait(self) -> None:
         assert self.__bg_tasks is not None
         # 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"Forcing shutdown. inflight_objs={len(self.__inflight_objs)}, "
-                f"inflight_refs={len(self.__inflight_refs)}, "
-                f"loop_alive={self.__bg_tasks.loop_alive()}, "
-                f"recv_alive={self.__bg_tasks.recv_alive()}"
-            )
-            self.__shutdown_loop.set()  # force __loop to exit
-            self.__bg_tasks.recv.cancel()
-            self.__bg_tasks.loop.cancel()
         try:
-            await asyncio.wait_for(self.__bg_tasks.gather(), timeout=None)
+            await asyncio.wait_for(self.__bg_tasks.gather(), timeout=shutdown_timeout)
         except asyncio.TimeoutError as e:
             raise WeaviateBatchStreamError(
                 "Background batch tasks did not terminate after forced shutdown."
             ) from e
-        if self.__bg_tasks.any_alive():
-            raise WeaviateBatchStreamError(
-                "Background batch tasks did not terminate after forced shutdown. "
-                f"loop_alive={self.__bg_tasks.loop_alive()}, "
-                f"recv_alive={self.__bg_tasks.recv_alive()}"
-            )
 
         # copy the results to the public results
         self.__results_for_wrapper_backup.results = self.__results_for_wrapper.results
@@ -237,6 +214,15 @@ async def _wait(self) -> None:
     async def _shutdown(self) -> None:
         self.__is_stopped.set()
 
+    async def __put(self, req: _BatchStreamRequest | None):
+        try:
+            await asyncio.wait_for(self.__reqs.put(req), timeout=1)
+            return True
+        except asyncio.TimeoutError:
+            if self.__bg_exception is not None or self.__shutdown_loop.is_set():
+                return False
+            return await self.__put(req)
+
     async def __loop(self) -> None:
         refresh_time: float = 0.01
         while self.__bg_exception is None and not self.__shutdown_loop.is_set():
@@ -278,23 +264,18 @@ async def __loop(self) -> None:
                     if paused:
                         logger.info("Server is back up, resuming batching loop...")
                         paused = False
-                    try:
-                        await asyncio.wait_for(self.__reqs.put(req), timeout=60)
-                    except asyncio.TimeoutError 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 new requests...")
                         return
             elif (
                 self.__is_stopped.is_set()
-                and not self.__sent_sentinel.is_set()
                 and not self.__is_hungup.is_set()
                 and not self.__is_shutting_down.is_set()
                 and not self.__is_oom.is_set()
             ):
-                await self.__reqs.put(None)
-                self.__sent_sentinel.set()
+                await self.__put(None)
+                logger.info("Sent sentinel, stopping batch loop...")
+                return
             await asyncio.sleep(refresh_time)
 
     def __generate_stream_requests(
@@ -347,10 +328,7 @@ def request_maker():
         if len(request.data.objects.values) > 0 or len(request.data.references.values) > 0:
             yield _BatchStreamRequest(request, uuids, beacons)
 
-    async def __send(
-        self,
-    ) -> AsyncGenerator[batch_pb2.BatchStreamRequest, None]:
-        self.__sent_sentinel.clear()
+    async def __send(self) -> AsyncGenerator[batch_pb2.BatchStreamRequest, None]:
         yield batch_pb2.BatchStreamRequest(
             start=batch_pb2.BatchStreamRequest.Start(
                 consistency_level=self.__batch_grpc._consistency_level,
@@ -393,14 +371,13 @@ async def __send(
         logger.info("Batch send thread exiting due to exception...")
 
     async def __recv(self) -> None:
-        stream = self.__batch_grpc.astream(
-            connection=self.__connection,
-            requests=self.__send(),
-        )
         self.__is_renewing_stream.clear()
         self.__is_shutting_down.clear()
         self.__is_hungup.clear()
-        async for message in stream:
+        async for message in self.__batch_grpc.astream(
+            connection=self.__connection,
+            requests=self.__send(),
+        ):
             if message.HasField("started"):
                 logger.info("Batch stream started successfully")
 

weaviate/collections/batch/async_.py

@@ -910,10 +910,10 @@ def recv_alive(self) -> bool:
             return self.recv.is_alive()
         return True  # not started yet so considered alive
 
-    def join(self) -> None:
+    def join(self, timeout: float | None = None) -> None:
         """Join the background threads."""
-        self.loop.join()
-        self.recv.join()
+        self.loop.join(timeout=timeout)
+        self.recv.join(timeout=timeout)
 
 
 class _ClusterBatch:

weaviate/collections/batch/base.py

@@ -66,7 +66,6 @@ def __init__(
 
         self.__connection = connection
         self.__is_gcp_on_wcd = connection._connection_params.is_gcp_on_wcd()
-        self.__stream_start: Optional[float] = None
         self.__is_renewing_stream = threading.Event()
         self.__consistency_level: ConsistencyLevel = consistency_level or ConsistencyLevel.QUORUM
         self.__batch_size = 100
@@ -124,26 +123,6 @@ def number_errors(self) -> int:
     def __all_threads_alive(self) -> bool:
         return self.__bg_threads.is_alive()
 
-    def __any_threads_alive(self) -> bool:
-        return self.__bg_threads.any_alive()
-
-    def __set_active_stream(self, call: Call) -> None:
-        with self.__stream_lock:
-            self.__active_stream = call
-
-    def __clear_active_stream(self) -> None:
-        with self.__stream_lock:
-            self.__active_stream = None
-
-    def __cancel_active_stream(self) -> bool:
-        with self.__stream_lock:
-            stream = self.__active_stream
-
-        if stream is None:
-            return False
-
-        return stream.cancel()
-
     def _start(self) -> None:
         self.__start_bg_threads()
         logger.info("Provisioned stream to the server for batch processing")
@@ -159,30 +138,12 @@ def _start(self) -> None:
     def _wait(self) -> None:
         # 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"Forcing shutdown. inflight_objs={len(self.__inflight_objs)}, "
-                f"inflight_refs={len(self.__inflight_refs)}, "
-                f"loop_alive={self.__bg_threads.loop_alive()}, "
-                f"recv_alive={self.__bg_threads.recv_alive()}"
-            )
-            self.__shutdown_loop.set()  # force __loop to exit
-            self.__is_stopped.set()
-            self.__cancel_active_stream()  # force __recv to exit by cancelling the stream
-
-        self.__bg_threads.join()
-        if self.__any_threads_alive():
+        try:
+            self.__bg_threads.join(shutdown_timeout)
+        except TimeoutError as e:
             raise WeaviateBatchStreamError(
-                "Background batch threads did not terminate after forced shutdown. "
-                f"loop_alive={self.__bg_threads.loop_alive()}, "
-                f"recv_alive={self.__bg_threads.recv_alive()}"
-            )
+                "Background batch threads did not terminate after forced shutdown."
+            ) from e
 
         # copy the results to the public results
         self.__results_for_wrapper_backup.results = self.__results_for_wrapper.results
@@ -194,15 +155,6 @@ 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()
@@ -267,9 +219,9 @@ def __loop(self) -> None:
                 and not self.__is_shutting_down.is_set()
                 and not self.__is_oom.is_set()
             ):
-                if not self.__put(None):
-                    logger.info("Batch loop is shutting down, stopping putting shutdown signal...")
-                    return
+                self.__put(None)
+                logger.info("Sent sentinel, stopping batch loop...")
+                return
             time.sleep(refresh_time)
 
     def __generate_stream_requests(
@@ -370,16 +322,13 @@ def __send(
         logger.info("Batch send thread exiting due to exception...")
 
     def __recv(self) -> None:
-        gen, call = self.__batch_grpc.stream(
-            connection=self.__connection,
-            requests=self.__send(),
-        )
-        self.__set_active_stream(call)
-
         self.__is_renewing_stream.clear()
         self.__is_shutting_down.clear()
         self.__is_hungup.clear()
-        for message in gen:
+        for message in self.__batch_grpc.stream(
+            connection=self.__connection,
+            requests=self.__send(),
+        ):
             if message.HasField("started"):
                 logger.info("Batch stream started successfully")
 

weaviate/collections/batch/sync.py

@@ -1018,27 +1018,22 @@ def grpc_batch_objects(
     def grpc_batch_stream(
         self,
         requests: Generator[batch_pb2.BatchStreamRequest, None, None],
-    ) -> tuple[Generator[batch_pb2.BatchStreamReply, None, None], Call]:
+    ) -> Generator[batch_pb2.BatchStreamReply, None, None]:
         assert self.grpc_stub is not None
-        call = self.grpc_stub.BatchStream(
-            request_iterator=requests,
-            timeout=self.timeout_config.stream,
-            metadata=self.grpc_headers(),
-        )
-
-        def generator():
-            try:
-                for msg in call:
-                    yield msg
-            except RpcError as e:
-                error = cast(Call, e)
-                if error.code() == StatusCode.PERMISSION_DENIED:
-                    raise InsufficientPermissionsError(error)
-                if error.code() == StatusCode.ABORTED:
-                    raise _BatchStreamShutdownError()
-                raise WeaviateBatchStreamError(f"{error.code()}({error.details()})")
-
-        return generator(), call
+        try:
+            for msg in self.grpc_stub.BatchStream(
+                request_iterator=requests,
+                timeout=self.timeout_config.stream,
+                metadata=self.grpc_headers(),
+            ):
+                yield msg
+        except RpcError as e:
+            error = cast(Call, e)
+            if error.code() == StatusCode.PERMISSION_DENIED:
+                raise InsufficientPermissionsError(error)
+            if error.code() == StatusCode.ABORTED:
+                raise _BatchStreamShutdownError()
+            raise WeaviateBatchStreamError(f"{error.code()}({error.details()})")
 
     def grpc_batch_delete(
         self, request: batch_delete_pb2.BatchDeleteRequest