From b504d7b9a014abdd5045ae6dd085bacef46f63e0 Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Sun, 4 Oct 2026 23:52:04 +0100 Subject: [PATCH] fix(batch): release stream references that touch a rejected object When the server rejects an object in a `stream()` batch, the error branch dropped it from the object cache but left its uuid in the lookup that holds back references until both ends are processed. Only the success branch cleared it. A reference to or from the rejected object therefore stayed queued forever: the batch loop never reached its stop sentinel, the sync `with` block took twice the shutdown timeout to exit, the reference was neither sent nor reported in `failed_references`, `flush()` never returned, and the async client raised "Background batch tasks did not terminate after forced shutdown." Clear the uuid in the error branch too, as the success branch and the client-side batching already do. The reference is then sent and its outcome comes back from the server like any other reference. --- mock_tests/test_batch_stream_references.py | 159 +++++++++++++++++++++ weaviate/collections/batch/async_.py | 2 + weaviate/collections/batch/sync.py | 2 + 3 files changed, 163 insertions(+) create mode 100644 mock_tests/test_batch_stream_references.py diff --git a/mock_tests/test_batch_stream_references.py b/mock_tests/test_batch_stream_references.py new file mode 100644 index 000000000..e62bcc591 --- /dev/null +++ b/mock_tests/test_batch_stream_references.py @@ -0,0 +1,159 @@ +import threading +import time +import uuid +from typing import AsyncGenerator, Generator, List, Set + +import grpc +import pytest +import pytest_asyncio +import weaviate +from weaviate.classes.init import AdditionalConfig, Timeout +from weaviate.proto.v1 import batch_pb2, weaviate_pb2_grpc +from .conftest import MOCK_IP, MOCK_PORT, MOCK_PORT_GRPC, mock_class, HTTPServer + +FAILING_UUID = str(uuid.UUID(int=1)) +SUCCEEDING_UUID = str(uuid.UUID(int=2)) + +# A stuck stream is only abandoned after the shutdown timeout (insert + 5s), so a short insert +# timeout keeps a regression from blocking the suite while still failing the exit-time check. +INSERT_TIMEOUT = 1 +MAX_EXIT_SECONDS = 5 + + +class MockRejectingWeaviateService(weaviate_pb2_grpc.WeaviateServicer): + """Rejects the object FAILING_UUID and every reference that touches it.""" + + def __init__(self) -> None: + self.references: List[batch_pb2.BatchReference] = [] + + def BatchStream( + self, + request_iterator: Generator[batch_pb2.BatchStreamRequest, None, None], + context: grpc.ServicerContext, + ) -> Generator[batch_pb2.BatchStreamReply, None, None]: + yield batch_pb2.BatchStreamReply(started=batch_pb2.BatchStreamReply.Started()) + for request in request_iterator: + if request.HasField("data"): + uuids: List[str] = [] + beacons: List[str] = [] + errors: List[batch_pb2.BatchStreamReply.Results.Error] = [] + successes: List[batch_pb2.BatchStreamReply.Results.Success] = [] + for obj in request.data.objects.values: + uuids.append(obj.uuid) + if obj.uuid == FAILING_UUID: + errors.append( + batch_pb2.BatchStreamReply.Results.Error( + uuid=obj.uuid, error="mock object failure" + ) + ) + else: + successes.append(batch_pb2.BatchStreamReply.Results.Success(uuid=obj.uuid)) + for ref in request.data.references.values: + self.references.append(ref) + beacon = ( + f"weaviate://localhost/{ref.from_collection}/{ref.from_uuid}/{ref.name}" + ) + beacons.append(beacon) + if FAILING_UUID in (ref.from_uuid, ref.to_uuid): + errors.append( + batch_pb2.BatchStreamReply.Results.Error( + beacon=beacon, error="mock reference failure" + ) + ) + else: + successes.append(batch_pb2.BatchStreamReply.Results.Success(beacon=beacon)) + yield batch_pb2.BatchStreamReply( + acks=batch_pb2.BatchStreamReply.Acks(uuids=uuids, beacons=beacons) + ) + yield batch_pb2.BatchStreamReply( + results=batch_pb2.BatchStreamReply.Results(errors=errors, successes=successes) + ) + if request.HasField("stop"): + return + + +def _alive_batch_loop_threads() -> Set[threading.Thread]: + return {t for t in threading.enumerate() if t.name == "BgBatchLoop" and t.is_alive()} + + +@pytest.fixture(scope="function") +def rejecting_service(start_grpc_server: grpc.Server) -> MockRejectingWeaviateService: + service = MockRejectingWeaviateService() + weaviate_pb2_grpc.add_WeaviateServicer_to_server(service, start_grpc_server) + return service + + +@pytest.fixture(scope="function") +def short_insert_timeout_client( + weaviate_mock: HTTPServer, start_grpc_server: grpc.Server +) -> Generator[weaviate.WeaviateClient, None, None]: + weaviate_mock.expect_request(f"/v1/schema/{mock_class['class']}").respond_with_json(mock_class) + client = weaviate.connect_to_local( + port=MOCK_PORT, + host=MOCK_IP, + grpc_port=MOCK_PORT_GRPC, + additional_config=AdditionalConfig(timeout=Timeout(insert=INSERT_TIMEOUT)), + ) + yield client + client.close() + + +@pytest_asyncio.fixture +async def short_insert_timeout_client_async( + weaviate_mock: HTTPServer, start_grpc_server: grpc.Server +) -> AsyncGenerator[weaviate.WeaviateAsyncClient, None]: + weaviate_mock.expect_request(f"/v1/schema/{mock_class['class']}").respond_with_json(mock_class) + client = weaviate.use_async_with_local( + port=MOCK_PORT, + host=MOCK_IP, + grpc_port=MOCK_PORT_GRPC, + additional_config=AdditionalConfig(timeout=Timeout(insert=INSERT_TIMEOUT)), + ) + await client.connect() + yield client + await client.close() + + +@pytest.mark.timeout(60) +def test_ssb_stream_sends_reference_to_failed_object( + short_insert_timeout_client: weaviate.WeaviateClient, + rejecting_service: MockRejectingWeaviateService, +) -> None: + collection = short_insert_timeout_client.collections.use(mock_class["class"]) + loops_before = _alive_batch_loop_threads() + + start = time.time() + with collection.batch.stream() as batch: + batch.add_object({"name": "fails"}, uuid=FAILING_UUID) + batch.add_object({"name": "succeeds"}, uuid=SUCCEEDING_UUID) + batch.add_reference(from_uuid=SUCCEEDING_UUID, from_property="ref", to=FAILING_UUID) + elapsed = time.time() - start + + assert elapsed < MAX_EXIT_SECONDS + assert _alive_batch_loop_threads() - loops_before == set() + failed_refs = collection.batch.failed_references + assert [ref.to_uuid for ref in rejecting_service.references] == [FAILING_UUID] + assert [str(err.object_.uuid) for err in collection.batch.failed_objects] == [FAILING_UUID] + assert [str(err.reference.to_object_uuid) for err in failed_refs] == [FAILING_UUID] + + +@pytest.mark.timeout(60) +@pytest.mark.asyncio +async def test_ssb_stream_sends_reference_to_failed_object_async( + short_insert_timeout_client_async: weaviate.WeaviateAsyncClient, + rejecting_service: MockRejectingWeaviateService, +) -> None: + collection = short_insert_timeout_client_async.collections.use(mock_class["class"]) + + start = time.time() + async with collection.batch.stream() as batch: + await batch.add_object({"name": "fails"}, uuid=FAILING_UUID) + await batch.add_object({"name": "succeeds"}, uuid=SUCCEEDING_UUID) + await batch.add_reference(from_uuid=SUCCEEDING_UUID, from_property="ref", to=FAILING_UUID) + elapsed = time.time() - start + + assert elapsed < MAX_EXIT_SECONDS + failed_refs = collection.batch.failed_references + assert [ref.to_uuid for ref in rejecting_service.references] == [FAILING_UUID] + assert [str(err.object_.uuid) for err in collection.batch.failed_objects] == [FAILING_UUID] + assert [str(err.reference.to_object_uuid) for err in failed_refs] == [FAILING_UUID] diff --git a/weaviate/collections/batch/async_.py b/weaviate/collections/batch/async_.py index 07083fe9b..004cc890a 100644 --- a/weaviate/collections/batch/async_.py +++ b/weaviate/collections/batch/async_.py @@ -407,6 +407,8 @@ async def __recv(self) -> None: try: async with self.__objs_cache_lock: cached = self.__objs_cache.pop(error.uuid) + async with self.__uuid_lookup_lock: + self.__uuid_lookup.discard(error.uuid) except KeyError: continue err = ErrorObject( diff --git a/weaviate/collections/batch/sync.py b/weaviate/collections/batch/sync.py index 6cf8c1edc..34ede023e 100644 --- a/weaviate/collections/batch/sync.py +++ b/weaviate/collections/batch/sync.py @@ -356,6 +356,8 @@ def __recv(self) -> None: try: with self.__objs_cache_lock: cached = self.__objs_cache.pop(error.uuid) + with self.__uuid_lookup_lock: + self.__uuid_lookup.discard(error.uuid) except KeyError: continue err = ErrorObject(