From 185e4d33d79f934672a3b7196de0848eea7c882e Mon Sep 17 00:00:00 2001
From: po-et <42566883+po-et@users.noreply.github.com>
Date: Tue, 20 Oct 2026 09:36:56 +0800
Subject: [PATCH] fix(batch): isolate failed sync stream recovery

---
 mock_tests/test_batch_recovery.py  | 432 +++++++++++++++++++++++++++++
 weaviate/collections/batch/sync.py | 199 ++++++++++---
 weaviate/connect/v4.py             |   2 +-
 3 files changed, 590 insertions(+), 43 deletions(-)
 create mode 100644 mock_tests/test_batch_recovery.py

diff --git a/mock_tests/test_batch_recovery.py b/mock_tests/test_batch_recovery.py
new file mode 100644
index 000000000..42850f1ee
--- /dev/null
+++ b/mock_tests/test_batch_recovery.py
@@ -0,0 +1,432 @@
+import json
+import threading
+import time
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from copy import deepcopy
+
+import grpc
+import pytest
+from grpc_health.v1 import health_pb2, health_pb2_grpc
+from werkzeug.wrappers import Response
+
+import weaviate
+from weaviate.classes.init import AdditionalConfig, Timeout
+from weaviate.classes.config import ConsistencyLevel
+from weaviate.proto.v1 import batch_pb2, search_get_pb2, weaviate_pb2_grpc
+
+from .conftest import MOCK_IP, MOCK_PORT_GRPC
+from .mock_data import mock_class
+
+
+def wait_until(predicate, timeout=6):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if predicate():
+            return
+        time.sleep(0.01)
+    assert predicate(), "condition did not become true before deadline"
+
+
+class RecoveryService(weaviate_pb2_grpc.WeaviateServicer):
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.stream_count = 0
+        self.initial_streams = 2
+        self.terminated = {}
+        self.warm = [threading.Event(), threading.Event()]
+        self.fault_go = threading.Event()
+        self.recovery_entered = threading.Event()
+        self.release = threading.Event()
+        self.targets = {1}
+        self.phase = "before_ack"
+        self.fault_barrier = threading.Barrier(1)
+        self.received = Counter()
+        self.metadata = []
+        self.deadlines = []
+        self.retry_failures = 0
+        self.retry_failed = threading.Event()
+        self.unhealthy = False
+        self.node_wait = threading.Event()
+
+    def Search(self, request, context):
+        return search_get_pb2.SearchReply()
+
+    def BatchStream(self, request_iterator, context):
+        with self.lock:
+            self.stream_count += 1
+            stream_id = self.stream_count
+            ended = self.terminated[stream_id] = threading.Event()
+            self.metadata.append(dict(context.invocation_metadata()))
+            self.deadlines.append(context.time_remaining())
+        context.add_callback(ended.set)
+        if stream_id > self.initial_streams:
+            self.recovery_entered.set()
+            if not self.release.wait(8):
+                context.abort(grpc.StatusCode.DEADLINE_EXCEEDED, "test recovery barrier timed out")
+            if self.retry_failures:
+                self.retry_failures -= 1
+                self.retry_failed.set()
+                context.abort(grpc.StatusCode.UNAVAILABLE, "temporarily unavailable")
+        yield batch_pb2.BatchStreamReply(started=batch_pb2.BatchStreamReply.Started())
+        for request in request_iterator:
+            if request.HasField("data"):
+                uuids = [obj.uuid for obj in request.data.objects.values]
+                beacons = [
+                    f"weaviate://localhost/{ref.from_collection}/{ref.from_uuid}/{ref.name}"
+                    for ref in request.data.references.values
+                ]
+                with self.lock:
+                    self.received.update(uuids)
+                    self.received.update(beacons)
+                fault = stream_id in self.targets and self.fault_go.is_set()
+                if fault:
+                    self.fault_barrier.wait(5)
+                if fault and self.phase == "before_ack":
+                    context.abort(grpc.StatusCode.UNAVAILABLE, "Socket closed")
+                yield batch_pb2.BatchStreamReply(
+                    acks=batch_pb2.BatchStreamReply.Acks(uuids=uuids, beacons=beacons)
+                )
+                if fault and self.phase == "after_ack":
+                    context.abort(grpc.StatusCode.UNAVAILABLE, "Socket closed")
+                yield batch_pb2.BatchStreamReply(
+                    results=batch_pb2.BatchStreamReply.Results(
+                        successes=[
+                            batch_pb2.BatchStreamReply.Results.Success(uuid=uuid) for uuid in uuids
+                        ]
+                        + [
+                            batch_pb2.BatchStreamReply.Results.Success(beacon=beacon)
+                            for beacon in beacons
+                        ]
+                    )
+                )
+                if fault:
+                    context.abort(grpc.StatusCode.UNAVAILABLE, "Socket closed")
+                if stream_id <= self.initial_streams:
+                    self.warm[stream_id - 1].set()
+            if request.HasField("stop"):
+                return
+
+
+class RecoveryHealth(health_pb2_grpc.HealthServicer):
+    def __init__(self):
+        self.status = None
+        self.calls = 0
+
+    def Check(self, request, context):
+        self.calls += 1
+        if self.status is not None:
+            context.abort(self.status, "optional health service unavailable")
+        return health_pb2.HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVING)
+
+
+@pytest.fixture
+def health_service():
+    return RecoveryHealth()
+
+
+@pytest.fixture
+def start_grpc_server(health_service):
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        server = grpc.server(executor)
+        health_pb2_grpc.add_HealthServicer_to_server(health_service, server)
+        assert server.add_insecure_port(f"[::]:{MOCK_PORT_GRPC}") == MOCK_PORT_GRPC
+        server.start()
+        try:
+            yield server
+        finally:
+            server.stop(0).wait(5)
+
+
+@pytest.fixture
+def recovery_client(httpserver, start_grpc_server):
+    service = RecoveryService()
+    metadata_calls = []
+
+    def metadata_handler(request):
+        metadata_calls.append(time.monotonic())
+        # Expose the old shared-connection reconnect without sleeping arbitrarily.
+        if service.fault_go.is_set():
+            service.recovery_entered.set()
+            if not service.release.wait(8):
+                return Response("test metadata barrier timed out", status=503)
+        return Response(json.dumps({"version": "1.39.0"}), content_type="application/json")
+
+    httpserver.expect_request("/v1/meta").respond_with_handler(metadata_handler)
+
+    def nodes_handler(request):
+        if service.unhealthy:
+            service.node_wait.set()
+        status = "UNHEALTHY" if service.unhealthy else "HEALTHY"
+        return Response(
+            json.dumps({"nodes": [{"status": status}, {"status": "HEALTHY"}]}),
+            content_type="application/json",
+        )
+
+    httpserver.expect_request("/v1/nodes").respond_with_handler(nodes_handler)
+    for name in ["Something", "Second"]:
+        config = deepcopy(mock_class)
+        config["class"] = name
+        httpserver.expect_request(f"/v1/schema/{name}").respond_with_json(config)
+    weaviate_pb2_grpc.add_WeaviateServicer_to_server(service, start_grpc_server)
+    client = weaviate.connect_to_local(
+        host=MOCK_IP,
+        port=httpserver.port,
+        grpc_port=MOCK_PORT_GRPC,
+        skip_init_checks=True,
+        auth_credentials=weaviate.auth.AuthApiKey("local-test-key"),
+        headers={"X-Recovery-Test": "preserved"},
+        additional_config=AdditionalConfig(timeout=Timeout(init=3, query=2, insert=2, stream=15)),
+    )
+    try:
+        yield client, service, metadata_calls
+    finally:
+        service.release.set()
+        client.close()
+        wait_until(lambda: not [t for t in threading.enumerate() if t.name.startswith("BgBatch")])
+
+
+@pytest.mark.parametrize("phase", ["before_ack", "after_ack", "after_result"])
+@pytest.mark.parametrize("targets", [{1}, {1, 2}])
+@pytest.mark.timeout(35)
+def test_recovery_preserves_shared_connection_and_replays_unfinished_objects(
+    recovery_client, phase, targets
+):
+    client, service, metadata_calls = recovery_client
+    service.phase = phase
+    service.targets = targets
+    service.fault_barrier = threading.Barrier(len(targets))
+    collections = [client.collections.use("Something"), client.collections.use("Second")]
+    expected = [set(), set()]
+    with ExitStack() as contexts:
+        batches = []
+        for i, collection in enumerate(collections):
+            batch = contexts.enter_context(collection.batch.stream())
+            batches.append(batch)
+            expected[i].add(batch.add_object({"names": ["warmup"]}))
+            assert service.warm[i].wait(5)
+        assert collections[0].query.fetch_objects(limit=1).objects == []
+        try:
+            service.fault_go.set()
+            fault_uuids = {}
+            for stream_id in targets:
+                uuid = batches[stream_id - 1].add_object({"names": ["fault"]})
+                expected[stream_id - 1].add(uuid)
+                fault_uuids[stream_id] = uuid
+            assert service.recovery_entered.wait(6)
+            assert client.is_connected()
+            assert collections[0].query.fetch_objects(limit=1).objects == []
+            assert len(metadata_calls) == 1
+            if 2 not in targets:
+                assert not service.terminated[2].is_set()
+                expected[1].add(batches[1].add_object({"names": ["unaffected"]}))
+        finally:
+            service.release.set()
+        wait_until(lambda: service.stream_count == 2 + len(targets))
+        for i, batch in enumerate(batches):
+            expected[i].add(batch.add_object({"names": ["after recovery"]}))
+
+    for i, collection in enumerate(collections):
+        assert set(collection.batch.results.objs.uuids.values()) == expected[i], service.received
+        assert not collection.batch.failed_objects
+    for uuid in fault_uuids.values():
+        assert service.received[str(uuid)] == (1 if phase == "after_result" else 2)
+    assert all(headers["x-recovery-test"] == "preserved" for headers in service.metadata)
+    assert all(headers["authorization"] == "Bearer local-test-key" for headers in service.metadata)
+    # gRPC's wire deadline rounds slightly, so allow one second of resolution.
+    assert all(0 < deadline < 16 for deadline in service.deadlines), service.deadlines
+    assert client.is_connected()
+    assert collections[0].query.fetch_objects(limit=1).objects == []
+
+
+@pytest.mark.timeout(30)
+def test_close_during_stream_recovery_does_not_reopen_client(recovery_client):
+    client, service, metadata_calls = recovery_client
+    service.retry_failures = 99
+    service.initial_streams = 1
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.fault_go.set()
+        batch.add_object({"names": ["fault"]})
+        try:
+            assert service.recovery_entered.wait(6)
+        finally:
+            service.release.set()
+        assert service.retry_failed.wait(5)
+        client.close()
+        wait_until(lambda: not [t for t in threading.enumerate() if t.name.startswith("BgBatch")])
+        assert not client.is_connected()
+        assert len(metadata_calls) == 1
+    assert not client.is_connected()
+
+
+@pytest.mark.timeout(30)
+def test_consecutive_failures_replay_each_unfinished_object_once(recovery_client):
+    client, service, metadata_calls = recovery_client
+    service.initial_streams = 1
+    service.retry_failures = 2
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        warmup = batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.fault_go.set()
+        fault = batch.add_object({"names": ["fault"]})
+        try:
+            assert service.recovery_entered.wait(6)
+        finally:
+            service.release.set()
+        wait_until(lambda: service.stream_count == 4, timeout=10)
+        after = batch.add_object({"names": ["after recovery"]})
+    assert set(collection.batch.results.objs.uuids.values()) == {warmup, fault, after}
+    assert service.received[str(fault)] == 2, service.received
+    assert service.received[str(after)] == 1, service.received
+    assert not collection.batch.failed_objects
+    assert len(metadata_calls) == 1
+    assert client.is_connected()
+
+
+@pytest.mark.parametrize("phase", ["before_ack", "after_ack", "after_result"])
+@pytest.mark.timeout(25)
+def test_recovery_replays_only_unfinished_references(recovery_client, phase):
+    client, service, metadata_calls = recovery_client
+    service.initial_streams = 1
+    service.phase = phase
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        source = batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.fault_go.set()
+        batch.add_reference(source, "related", source)
+        try:
+            assert service.recovery_entered.wait(6)
+        finally:
+            service.release.set()
+        wait_until(lambda: service.stream_count == 2)
+        batch.add_reference(source, "other", source)
+    beacon = f"weaviate://localhost/Something/{source}/related"
+    assert service.received[beacon] == (1 if phase == "after_result" else 2), service.received
+    assert service.received[f"weaviate://localhost/Something/{source}/other"] == 1
+    assert not collection.batch.failed_references
+    assert not collection.batch.results.refs.errors
+    assert len(metadata_calls) == 1
+
+
+@pytest.mark.timeout(25)
+def test_recovery_drains_when_context_exits_after_last_object(recovery_client):
+    client, service, _ = recovery_client
+    service.initial_streams = 1
+    service.release.set()
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        warmup = batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.fault_go.set()
+        fault = batch.add_object({"names": ["fault"]})
+    assert set(collection.batch.results.objs.uuids.values()) == {warmup, fault}
+    assert service.received[str(fault)] == 2
+
+
+@pytest.mark.timeout(40)
+def test_repeated_failures_stop_only_failed_batch_threads(recovery_client):
+    client, service, metadata_calls = recovery_client
+    service.initial_streams = 1
+    service.retry_failures = 99
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.fault_go.set()
+        batch.add_object({"names": ["fault"]})
+        try:
+            assert service.recovery_entered.wait(6)
+        finally:
+            service.release.set()
+        wait_until(
+            lambda: not [t for t in threading.enumerate() if t.name.startswith("BgBatch")],
+            timeout=25,
+        )
+        with pytest.raises(weaviate.exceptions.WeaviateBatchStreamError, match="unavailable"):
+            batch.add_object({"names": ["cannot be accepted"]})
+    assert service.stream_count == 6
+    assert client.is_connected()
+    assert collection.query.fetch_objects(limit=1).objects == []
+    assert len(metadata_calls) == 1
+
+
+@pytest.mark.timeout(30)
+def test_queued_objects_and_references_have_one_replay_owner(recovery_client):
+    client, service, _ = recovery_client
+    service.initial_streams = 1
+    service.retry_failures = 2
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        source = batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.fault_go.set()
+        fault = batch.add_object({"names": ["fault"]})
+        queued = []
+        try:
+            assert service.recovery_entered.wait(6)
+            for i in range(12):
+                queued.append(batch.add_object({"names": [f"queued {i}"]}))
+                batch.add_reference(source, f"related{i}", source)
+        finally:
+            service.release.set()
+        wait_until(lambda: service.stream_count == 4, timeout=10)
+    assert set(collection.batch.results.objs.uuids.values()) == {source, fault, *queued}
+    assert service.received[str(fault)] == 2, service.received
+    assert all(service.received[str(uuid)] == 1 for uuid in queued), service.received
+    for i in range(12):
+        assert service.received[f"weaviate://localhost/Something/{source}/related{i}"] == 1
+    assert not collection.batch.failed_objects
+    assert not collection.batch.failed_references
+
+
+@pytest.mark.timeout(20)
+def test_close_interrupts_wait_for_unhealthy_nodes(recovery_client):
+    client, service, metadata_calls = recovery_client
+    service.initial_streams = 1
+    collection = client.collections.use("Something").with_consistency_level(ConsistencyLevel.ALL)
+    with collection.batch.stream() as batch:
+        batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        service.unhealthy = True
+        service.fault_go.set()
+        batch.add_object({"names": ["fault"]})
+        assert service.node_wait.wait(5)
+        assert client.is_connected()
+        assert collection.query.fetch_objects(limit=1).objects == []
+        client.close()
+        wait_until(lambda: not [t for t in threading.enumerate() if t.name.startswith("BgBatch")])
+    assert not client.is_connected()
+    assert service.stream_count == 1
+    assert len(metadata_calls) == 1
+
+
+@pytest.mark.parametrize(
+    "status", [grpc.StatusCode.UNIMPLEMENTED, grpc.StatusCode.PERMISSION_DENIED]
+)
+@pytest.mark.timeout(20)
+def test_recovery_respects_skip_init_checks(recovery_client, health_service, status):
+    client, service, metadata_calls = recovery_client
+    health_service.status = status
+    service.initial_streams = 1
+    service.release.set()
+    collection = client.collections.use("Something")
+    with collection.batch.stream() as batch:
+        warmup = batch.add_object({"names": ["warmup"]})
+        assert service.warm[0].wait(5)
+        assert collection.query.fetch_objects(limit=1).objects == []
+        service.fault_go.set()
+        fault = batch.add_object({"names": ["fault"]})
+        wait_until(lambda: service.recovery_entered.is_set() or health_service.calls > 0)
+        if health_service.calls:
+            client.close()
+    assert health_service.calls == 0
+    assert set(collection.batch.results.objs.uuids.values()) == {warmup, fault}
+    assert service.received[str(fault)] == 2
+    assert len(metadata_calls) == 1
+    assert client.is_connected()
diff --git a/weaviate/collections/batch/sync.py b/weaviate/collections/batch/sync.py
index 6cf8c1edc..d6a1c0019 100644
--- a/weaviate/collections/batch/sync.py
+++ b/weaviate/collections/batch/sync.py
@@ -1,10 +1,12 @@
 import threading
 import time
 import uuid as uuid_package
+from collections import deque
 from concurrent.futures import ThreadPoolExecutor
 from queue import Empty, Full, Queue
-from typing import Generator, List, Optional, Set, Union
+from typing import Generator, List, Optional, Set, Union, cast
 
+from grpc import Call, RpcError, StatusCode
 from pydantic import ValidationError
 
 from weaviate.collections.batch.base import (
@@ -40,6 +42,7 @@
     WeaviateBatchFailedToReestablishStreamError,
     WeaviateBatchStreamError,
     WeaviateBatchValidationError,
+    WeaviateClosedClientError,
     WeaviateGRPCUnavailableError,
     WeaviateStartUpError,
 )
@@ -89,6 +92,7 @@ def __init__(
         self.__results_lock = threading.Lock()
 
         self.__bg_exception: Optional[Exception] = None
+        self.__consecutive_retries = 0
         self.__is_oom = threading.Event()
         self.__is_shutting_down = threading.Event()
         self.__is_hungup = threading.Event()
@@ -105,10 +109,15 @@ def __init__(
         self.__acks_lock = threading.Lock()
         self.__inflight_objs: set[str] = set()
         self.__inflight_refs: set[str] = set()
+        self.__unresolved_objs: set[str] = set()
+        self.__unresolved_refs: set[str] = set()
 
         # maxsize=1 so that __loop does not run faster than generator for __recv
         # thereby using too much buffer in case of server-side shutdown
         self.__reqs: Queue[Optional[_BatchStreamRequest]] = Queue(maxsize=1)
+        self.__send_lock = threading.Lock()
+        self.__replay_requests: deque[_BatchStreamRequest] = deque()
+        self.__send_finished = False
 
     @property
     def number_errors(self) -> int:
@@ -200,6 +209,8 @@ def __loop(self) -> None:
                             logger.info("Server is shutting down, pausing batching loop...")
                             paused = True
                         time.sleep(refresh_time)
+                        if self.__bg_exception is not None:
+                            return
                         if time.time() - start > self.__oom_wait_time:
                             raise WeaviateBatchFailedToReestablishStreamError(
                                 f"Batch stream was not re-established within {self.__oom_wait_time} seconds after an OOM message. Terminating batch."
@@ -255,7 +266,7 @@ def request_maker():
             uuids.add(obj.uuid)
 
         for reference in references:
-            ref = self.__batch_grpc.grpc_reference(reference._to_internal())
+            ref = self.__batch_grpc.grpc_reference(reference.model_copy()._to_internal())
             ref_size = ref.ByteSize() + per_object_overhead
 
             if total_size + ref_size >= self.__batch_grpc.grpc_max_msg_size:
@@ -273,6 +284,7 @@ def request_maker():
 
     def __send(
         self,
+        stopped: threading.Event,
     ) -> Generator[batch_pb2.BatchStreamRequest, None, None]:
         yield batch_pb2.BatchStreamRequest(
             start=batch_pb2.BatchStreamRequest.Start(
@@ -280,7 +292,7 @@ def __send(
             ),
         )
         stream_start = time.time()
-        while self.__bg_exception is None:
+        while self.__bg_exception is None and not stopped.is_set():
             if self.__is_gcp_on_wcd:
                 assert stream_start is not None, "stream_start should be set for GCP streams"
                 if time.time() - stream_start > GCP_STREAM_TIMEOUT:
@@ -291,16 +303,28 @@ def __send(
                     yield batch_pb2.BatchStreamRequest(stop=batch_pb2.BatchStreamRequest.Stop())
                     return
             try:
-                req = self.__reqs.get(timeout=1)
+                with self.__send_lock:
+                    if stopped.is_set():
+                        return
+                    if self.__replay_requests:
+                        req = self.__replay_requests.popleft()
+                    elif self.__send_finished:
+                        req = None
+                    else:
+                        req = self.__reqs.get(timeout=0.1)
+                    if req is not None:
+                        with self.__acks_lock:
+                            self.__inflight_objs.update(req.uuids)
+                            self.__inflight_refs.update(req.beacons)
+                            self.__unresolved_objs.update(req.uuids)
+                            self.__unresolved_refs.update(req.beacons)
                 if req is None:
+                    self.__send_finished = True
                     logger.info(
                         "Batching finished, stopping and closing the client-side of the stream"
                     )
                     yield batch_pb2.BatchStreamRequest(stop=batch_pb2.BatchStreamRequest.Stop())
                     return
-                with self.__acks_lock:
-                    self.__inflight_objs.update(req.uuids)
-                    self.__inflight_refs.update(req.beacons)
                 yield req.proto
                 continue
             except Empty:
@@ -318,14 +342,25 @@ def __send(
                 )
         logger.info("Batch send thread exiting due to exception...")
 
-    def __recv(self) -> None:
+    def __stream(self) -> Generator[batch_pb2.BatchStreamReply, None, None]:
+        stopped = threading.Event()
+        try:
+            yield from self.__batch_grpc.stream(
+                connection=self.__connection,
+                requests=self.__send(stopped),
+            )
+        finally:
+            # gRPC may still be waiting in next(requests) after the RPC ends.
+            # Stop that iterator before it can take requests for its replacement.
+            stopped.set()
+            with self.__send_lock:
+                pass  # Wait for any queue read already in progress before replaying.
+
+    def __recv(self) -> bool:
         self.__is_renewing_stream.clear()
         self.__is_shutting_down.clear()
         self.__is_hungup.clear()
-        for message in self.__batch_grpc.stream(
-            connection=self.__connection,
-            requests=self.__send(),
-        ):
+        for message in self.__stream():
             if message.HasField("started"):
                 logger.info("Batch stream started successfully")
 
@@ -347,6 +382,12 @@ def __recv(self) -> None:
                     self.__inflight_refs.difference_update(message.acks.beacons)
 
             if message.HasField("results"):
+                with self.__acks_lock:
+                    for entry in [*message.results.errors, *message.results.successes]:
+                        if entry.HasField("uuid"):
+                            self.__unresolved_objs.discard(entry.uuid)
+                        if entry.HasField("beacon"):
+                            self.__unresolved_refs.discard(entry.beacon)
                 result_objs = BatchObjectReturn()
                 result_refs = BatchReferenceReturn()
                 failed_objs: List[ErrorObject] = []
@@ -420,6 +461,8 @@ def __recv(self) -> None:
                     self.__results_for_wrapper.results.refs += result_refs
                     self.__results_for_wrapper.failed_objects.extend(failed_objs)
                     self.__results_for_wrapper.failed_references.extend(failed_refs)
+                if message.results.errors or message.results.successes:
+                    self.__consecutive_retries = 0
 
             if message.HasField("out_of_memory"):
                 logger.info(
@@ -427,6 +470,9 @@ def __recv(self) -> None:
                 )
                 self.__is_oom.set()
                 self.__oom_wait_time = message.out_of_memory.wait_time
+                with self.__acks_lock:
+                    self.__unresolved_objs.difference_update(message.out_of_memory.uuids)
+                    self.__unresolved_refs.difference_update(message.out_of_memory.beacons)
                 with self.__objs_cache_lock:
                     self.__batch_objects.prepend(
                         [
@@ -454,14 +500,58 @@ def __recv(self) -> None:
         if self.__is_shutting_down.is_set():
             self.__reconnect()
             logger.info("Restarting batch recv after shutdown...")
-            return self.__recv()
+            return True
         elif self.__is_renewing_stream.is_set():
             # restart the stream if we are renewing it (GCP connections have a max lifetime)
             logger.info("Restarting batch recv after renewing stream...")
-            return self.__recv()
+            return True
 
         logger.info("Server closed the stream from its side, shutting down batch")
         self.__shutdown_loop.set()
+        return False
+
+    def __wait_for_reconnect(self, delay: float) -> None:
+        deadline = time.monotonic() + delay
+        while True:
+            if not self.__connection.is_connected():
+                raise WeaviateClosedClientError()
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                return
+            time.sleep(min(remaining, 0.1))
+
+    def __wait_for_channel(self) -> None:
+        self.__wait_for_reconnect(0)
+        if self.__consistency_level == ConsistencyLevel.ALL or self.__number_of_nodes == 1:
+            # check that all nodes are available before reconnecting
+            deadline = time.monotonic() + self.__oom_wait_time
+            up_nodes = self.__cluster.get_nodes_status()
+            while len(up_nodes) != self.__number_of_nodes or any(
+                node["status"] != "HEALTHY" for node in up_nodes
+            ):
+                logger.info(
+                    "Waiting for all nodes to be HEALTHY before reconnecting to batch stream..."
+                )
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise WeaviateBatchFailedToReestablishStreamError(
+                        "Batch stream was not re-established because cluster nodes did not become HEALTHY."
+                    )
+                self.__wait_for_reconnect(min(5, remaining))
+                up_nodes = self.__cluster.get_nodes_status()
+        # gRPC reconnects the existing channel. Closing the client here would
+        # cancel every other stream and query sharing that connection.
+        if self.__connection._skip_init_checks:
+            return
+        for retry in range(6):
+            self.__wait_for_reconnect(0)
+            try:
+                result(self.__connection._ping_grpc("sync"))
+                return
+            except WeaviateGRPCUnavailableError:
+                if retry == 5:
+                    raise
+                self.__wait_for_reconnect(2**retry)
 
     def __reconnect(self, retry: int = 0) -> None:
         if self.__consistency_level == ConsistencyLevel.ALL or self.__number_of_nodes == 1:
@@ -499,35 +589,60 @@ def loop_wrapper() -> None:
                 self.__bg_exception = e
 
         def recv_wrapper() -> None:
-            try:
-                self.__recv()
-                logger.info("exited batch receive thread")
-            except Exception as e:
-                if isinstance(e, WeaviateBatchStreamError) and (
-                    "Socket closed" in e.message
-                    or "context canceled" in e.message
-                    or "Connection reset" in e.message
-                    or "Received RST_STREAM with error code 2" in e.message
-                ):
-                    logger.error(f"Socket hang up detected in batch receive thread: {e.message}")
+            while True:
+                try:
+                    if self.__recv():
+                        continue
+                    logger.info("exited batch receive thread")
+                    return
+                except Exception as e:
+                    if (
+                        not isinstance(e, WeaviateBatchStreamError)
+                        or not (
+                            isinstance(e.__cause__, RpcError)
+                            and cast(Call, e.__cause__).code() == StatusCode.UNAVAILABLE
+                            or "Socket closed" in e.message
+                            or "context canceled" in e.message
+                            or "Connection reset" in e.message
+                            or "Received RST_STREAM with error code 2" in e.message
+                        )
+                        or self.__consecutive_retries >= 5
+                    ):
+                        logger.error(e)
+                        self.__bg_exception = e
+                        return
                     self.__is_hungup.set()
-                else:
-                    logger.error(e)
-                    logger.error(type(e))
-                    self.__bg_exception = e
-            if self.__is_hungup.is_set():
-                # this happens during ungraceful shutdown of the coordinator
-                # lets restart the stream and add the cached objects again
-                logger.warning("Stream closed unexpectedly, restarting...")
-                self.__reconnect()
-                with self.__objs_cache_lock:
-                    self.__batch_objects.prepend(list(self.__objs_cache.values()))
-                with self.__refs_cache_lock:
-                    self.__batch_references.prepend(list(self.__refs_cache.values()))
-                self.__inflight_objs.clear()
-                self.__inflight_refs.clear()
-                # start a new stream with a newly reconnected channel
-                return recv_wrapper()
+                    try:
+                        if self.__consecutive_retries:
+                            self.__wait_for_reconnect(2 ** (self.__consecutive_retries - 1))
+                        self.__consecutive_retries += 1
+                        self.__wait_for_channel()
+                    except Exception as reconnect_error:
+                        self.__bg_exception = reconnect_error
+                        return
+                    # Replay only requests taken by this RPC, not objects still
+                    # waiting in the producer's queues. The producer may already
+                    # have stopped, so send these directly on the replacement RPC.
+                    with self.__objs_cache_lock:
+                        objects = [
+                            obj
+                            for uuid in self.__unresolved_objs
+                            if (obj := self.__objs_cache.get(uuid)) is not None
+                        ]
+                    with self.__refs_cache_lock:
+                        references = [
+                            ref
+                            for beacon in self.__unresolved_refs
+                            if (ref := self.__refs_cache.get(beacon)) is not None
+                        ]
+                    self.__replay_requests.extendleft(
+                        reversed(list(self.__generate_stream_requests(objects, references)))
+                    )
+                    with self.__acks_lock:
+                        self.__inflight_objs.clear()
+                        self.__inflight_refs.clear()
+                        self.__unresolved_objs.clear()
+                        self.__unresolved_refs.clear()
 
         self.__bg_threads = _BgThreads(
             loop=threading.Thread(
diff --git a/weaviate/connect/v4.py b/weaviate/connect/v4.py
index 5dcadfaeb..1e8cc3d67 100644
--- a/weaviate/connect/v4.py
+++ b/weaviate/connect/v4.py
@@ -1088,7 +1088,7 @@ def grpc_batch_stream(
                 raise InsufficientPermissionsError(error)
             if error.code() == StatusCode.ABORTED:
                 raise _BatchStreamShutdownError()
-            raise WeaviateBatchStreamError(f"{error.code()}({error.details()})")
+            raise WeaviateBatchStreamError(f"{error.code()}({error.details()})") from e
 
     def grpc_batch_delete(
         self, request: batch_delete_pb2.BatchDeleteRequest