Add safety hatch in-case server doesn't clos its side of the stream in time (300s)

Author: tsmith023

weaviate/collections/batch/async_.py

@@ -14,6 +14,7 @@
 
 from weaviate.collections.batch.base import (
     GCP_STREAM_TIMEOUT,
+    SHUTDOWN_TIMEOUT,
     ObjectsBatchRequest,
     ReferencesBatchRequest,
     _BatchDataWrapper,
@@ -182,8 +183,20 @@ async def recv_wrapper() -> None:
             loop=loop,
         )
 
-    async def _wait(self):
+    async def _wait(self) -> None:
         assert self.__bg_tasks is not None
+        deadline = time.time() + SHUTDOWN_TIMEOUT
+        while time.time() < deadline:
+            if not self.__bg_tasks.all_alive():
+                break
+            await asyncio.sleep(0.1)
+        if self.__bg_tasks.all_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)}"
+            )
+            self.__shutdown_loop.set()  # force __loop to exit
         await self.__bg_tasks.gather()
 
         # copy the results to the public results

weaviate/collections/batch/base.py

@@ -64,6 +64,7 @@
 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]):

weaviate/collections/batch/sync.py

@@ -9,6 +9,7 @@
 
 from weaviate.collections.batch.base import (
     GCP_STREAM_TIMEOUT,
+    SHUTDOWN_TIMEOUT,
     ObjectsBatchRequest,
     ReferencesBatchRequest,
     _BatchDataWrapper,
@@ -135,6 +136,19 @@ def _start(self) -> None:
                 )
 
     def _wait(self) -> None:
+        deadline = time.time() + SHUTDOWN_TIMEOUT
+        while time.time() < deadline:
+            if not self.__bg_threads.is_alive():
+                break
+            time.sleep(0.1)
+        if self.__bg_threads.is_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)}"
+            )
+            self.__shutdown_loop.set()  # force __loop to exit
+
         self.__bg_threads.join()
 
         # copy the results to the public results