[ENH] Add close() method and context manager support to Client#6373
[ENH] Add close() method and context manager support to Client#6373rescrv merged 6 commits intochroma-core:mainfrom
Conversation
Fixes chroma-core#5868 Added close() method to the Client class that properly releases resources by calling self._system.stop(). This is essential for PersistentClient to avoid SQLite file locking issues when uploading database files to S3 or performing other operations that require exclusive file access. Also implemented context manager support (__enter__ and __exit__) to enable the "with" statement pattern for automatic resource cleanup. Changes: - Added Client.close() method that stops the underlying system - Added __enter__ and __exit__ methods for context manager support - Added comprehensive tests for both close() and context manager usage - Tests cover both PersistentClient and EphemeralClient Example usage: ```python # Explicit close client = chromadb.PersistentClient(path="./chroma_db") # ... use client ... client.close() # Context manager with chromadb.PersistentClient(path="./chroma_db") as client: # ... use client ... # Automatically closed ```
Reviewer ChecklistPlease leverage this checklist to ensure your code review is thorough before approving Testing, Bugs, Errors, Logs, Documentation
System Compatibility
Quality
|
- Add reference counting to SharedSystemClient to track multiple clients using the same System instance - Only stop() the System when the last client calls close() - This fixes the issue where closing one client would break other clients sharing the same System (e.g., multiple PersistentClient instances with the same path) - Add proper type annotations to __exit__ method parameters - Update clear_system_cache() to also clear reference counts This addresses the propel-code-bot feedback on PR chroma-core#6373 which identified that the shared System singleton needs reference counting to avoid invalidating other clients when one closes.
|
The update also introduces thread-safe reference counting for shared systems so that only the final client release actually stops the underlying services, preventing leaks when initialization fails and ensuring tests cache the system handle before those shared entries are cleared. Possible Issues• Refcount bookkeeping relies on every codepath calling This summary was automatically generated by @propel-code-bot |
|
Reference counting implementation verified and approved. This critical fix prevents client shutdown race conditions. Ready for merge. |
Use pop() instead of direct dict access so that calling close() multiple times (common with context managers + explicit close) is a safe no-op instead of raising KeyError.
|
Addressed the idempotent close() feedback from @propel-code-bot. The |
…d safety - Fix off-by-one: AdminClient created in Client.__init__ also increments refcount for the same shared system. close() now decrements the admin client's refcount before decrementing its own, so the system properly stops when the last client closes. - Fix resource leak: if Client.__init__ raises after super().__init__() increments the refcount, the caller never gets the object to call close(). Wrap the init body in try/except to decrement on failure. - Fix race condition: += and -= on dict values are not atomic. Protect _identifier_to_refcount with a class-level threading.Lock. - Add _closed flag to make close() fully idempotent without relying on refcount state.
|
Hi team -- friendly ping on this PR. It's been open for a couple of days now and I've addressed all the automated review feedback (idempotent close(), reference counting logic, etc.). Would love to get a human review when someone has a chance. Thanks! |
Stop and remove the System from cache when init fails and refcount drops to zero, preventing orphaned running systems. Fix tests to save a local reference to the System before close() since close() pops it from the shared cache.
|
Addressed the remaining propel-code-bot feedback in c60539d:
|
Extract the duplicated "decrement refcount + conditional stop" pattern from both Client.close() and the Client.__init__ exception handler into SharedSystemClient._release_system(). This eliminates code duplication and ensures consistent cleanup behavior in both paths.
|
Only failure is because third parties cannot access secrets in the repo. Looks good otherwise. |
Summary
Fixes #5868
Added
close()method and context manager support to theClientclass to properly release resources and close SQLite connections. This resolves file locking issues when usingPersistentClientwith operations that require exclusive file access, such as uploading database files to S3.Changes
Client.close()method that callsself._system.stop()to properly stop the system and release resources__enter__and__exit__methods for context manager supportclose()and context manager usage withPersistentClientandEphemeralClientTest Plan
Added four new test cases in
chromadb/test/test_client.py:test_persistent_client_close()- Tests explicit close() call with PersistentClienttest_persistent_client_context_manager()- Tests context manager with PersistentClienttest_ephemeral_client_close()- Tests explicit close() call with EphemeralClienttest_ephemeral_client_context_manager()- Tests context manager with EphemeralClientAll tests verify that:
Usage Examples
Related Issues
This addresses the core issue reported in #5868 where users experienced SQLite file locking problems when trying to upload ChromaDB files to AWS S3 from Lambda functions. The workaround mentioned in the issue was to use
client._system.stop(), which this PR now exposes as a public API method.