This is an automated email from the ASF dual-hosted git repository. wenjin272 pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 4e01a73e9f55459dc9ba4e93fb94edf8a52d7f7f Author: WenjinXie <[email protected]> AuthorDate: Thu Aug 6 22:46:13 2026 +0800 [python] Synchronize Chroma client initialization Generated-by: Codex CLI 0.144.5 --- .../vector_stores/chroma/chroma_vector_store.py | 64 +++++++++++++--------- .../chroma/tests/test_chroma_vector_store.py | 41 +++++++++++++- 2 files changed, 77 insertions(+), 28 deletions(-) diff --git a/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py b/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py index 29dacf90..cab7ec3d 100644 --- a/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py +++ b/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################ import uuid +from threading import Lock from typing import Any, Dict, Generator, List import chromadb @@ -37,6 +38,11 @@ DEFAULT_COLLECTION = "flink_agents_chroma_collection" # anything larger must be split. _MAX_CHUNK_SIZE = 41665 +# Chroma's embedded clients share a process-wide System cache whose initial +# startup is not synchronized. Serialize client construction so concurrent +# first access cannot observe a RustBindingsAPI before its bindings are ready. +_CHROMA_CLIENT_CREATION_LOCK = Lock() + class ChromaVectorStore(CollectionManageableVectorStore): """ChromaDB vector store that handles connection and semantic search. @@ -167,33 +173,37 @@ class ChromaVectorStore(CollectionManageableVectorStore): if self.__client is not None: return self.__client - if self.api_key is not None: - self.__client = CloudClient( - tenant=self.tenant, - database=self.database, - api_key=self.api_key, - ) - elif self.host is not None: - self.__client = chromadb.HttpClient( - host=self.host, - port=self.port, - settings=self.client_settings, - tenant=self.tenant, - database=self.database, - ) - elif self.persist_directory is not None: - self.__client = chromadb.PersistentClient( - path=self.persist_directory, - settings=self.client_settings, - tenant=self.tenant, - database=self.database, - ) - else: - self.__client = chromadb.EphemeralClient( - settings=self.client_settings, - tenant=self.tenant, - database=self.database, - ) + with _CHROMA_CLIENT_CREATION_LOCK: + if self.__client is not None: + return self.__client + + if self.api_key is not None: + self.__client = CloudClient( + tenant=self.tenant, + database=self.database, + api_key=self.api_key, + ) + elif self.host is not None: + self.__client = chromadb.HttpClient( + host=self.host, + port=self.port, + settings=self.client_settings, + tenant=self.tenant, + database=self.database, + ) + elif self.persist_directory is not None: + self.__client = chromadb.PersistentClient( + path=self.persist_directory, + settings=self.client_settings, + tenant=self.tenant, + database=self.database, + ) + else: + self.__client = chromadb.EphemeralClient( + settings=self.client_settings, + tenant=self.tenant, + database=self.database, + ) return self.__client @property diff --git a/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py b/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py index 135ded8b..7b67ba51 100644 --- a/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py +++ b/python/flink_agents/integrations/vector_stores/chroma/tests/test_chroma_vector_store.py @@ -16,6 +16,9 @@ # limitations under the License. ################################################################################ import os +import time +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier, Lock from typing import Any, Dict, List from unittest.mock import MagicMock @@ -25,7 +28,7 @@ from chromadb.errors import NotFoundError from flink_agents.api.resource_context import ResourceContext try: - import chromadb # noqa: F401 + import chromadb chromadb_available = True except ImportError: @@ -89,6 +92,42 @@ def _populate_test_data( return documents +def test_client_is_initialized_once_under_concurrent_access( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Concurrent first access must not create multiple Chroma systems.""" + expected_client = MagicMock() + start_barrier = Barrier(2) + call_count = 0 + call_count_lock = Lock() + + def create_client(**kwargs: Any) -> MagicMock: + nonlocal call_count + with call_count_lock: + call_count += 1 + # Release the GIL long enough for the other accessor to enter the + # client factory when lazy initialization is not synchronized. + time.sleep(0.1) + return expected_client + + monkeypatch.setattr(chromadb, "PersistentClient", create_client) + vector_store = ChromaVectorStore( + name="chroma_vector_store", + embedding_model="mock_embeddings", + persist_directory="/tmp/chroma-concurrent-client-test", + ) + + def get_client() -> Any: + start_barrier.wait() + return vector_store.client + + with ThreadPoolExecutor(max_workers=2) as executor: + clients = list(executor.map(lambda _: get_client(), range(2))) + + assert all(client is expected_client for client in clients) + assert call_count == 1 + + @pytest.mark.skipif(not chromadb_available, reason="ChromaDB is not available") def test_local_chroma_vector_store() -> None: """Test ChromaDB vector store with embedding model integration."""
