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(