Skip to content

[ENH] Add close() method and context manager support to Client#6373

Merged
rescrv merged 6 commits intochroma-core:mainfrom
veeceey:fix/issue-5868-add-close-method
Feb 20, 2026
Merged

[ENH] Add close() method and context manager support to Client#6373
rescrv merged 6 commits intochroma-core:mainfrom
veeceey:fix/issue-5868-add-close-method

Conversation

@veeceey
Copy link
Copy Markdown
Contributor

@veeceey veeceey commented Feb 8, 2026

Summary

Fixes #5868

Added close() method and context manager support to the Client class to properly release resources and close SQLite connections. This resolves file locking issues when using PersistentClient with operations that require exclusive file access, such as uploading database files to S3.

Changes

  • Added Client.close() method that calls self._system.stop() to properly stop the system and release resources
  • Implemented __enter__ and __exit__ methods for context manager support
  • Added comprehensive tests for both close() and context manager usage with PersistentClient and EphemeralClient

Test Plan

Added four new test cases in chromadb/test/test_client.py:

  • test_persistent_client_close() - Tests explicit close() call with PersistentClient
  • test_persistent_client_context_manager() - Tests context manager with PersistentClient
  • test_ephemeral_client_close() - Tests explicit close() call with EphemeralClient
  • test_ephemeral_client_context_manager() - Tests context manager with EphemeralClient

All tests verify that:

  1. The system is running before close/context exit
  2. The system is stopped after close/context exit
  3. Data persists correctly (for PersistentClient)

Usage Examples

# Explicit close
client = chromadb.PersistentClient(path="./chroma_db")
collection = client.create_collection("my_collection")
collection.add(ids=["id1"], documents=["doc1"])
client.close()

# Context manager (recommended)
with chromadb.PersistentClient(path="./chroma_db") as client:
    collection = client.create_collection("my_collection")
    collection.add(ids=["id1"], documents=["doc1"])
# Automatically closed, safe to upload files to S3 now

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.

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
```
@github-actions
Copy link
Copy Markdown

github-actions bot commented Feb 8, 2026

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

  • Can you think of any use case in which the code does not behave as intended? Have they been tested?
  • Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  • If appropriate, are there adequate property based tests?
  • If appropriate, are there adequate unit tests?
  • Should any logging, debugging, tracing information be added or removed?
  • Are error messages user-friendly?
  • Have all documentation changes needed been made?
  • Have all non-obvious changes been commented?

System Compatibility

  • Are there any potential impacts on other parts of the system or backward compatibility?
  • Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

  • Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

- 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.
@propel-code-bot
Copy link
Copy Markdown
Contributor

propel-code-bot bot commented Feb 8, 2026

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 SharedSystemClient._release_system(); a forgotten release elsewhere could leave systems running indefinitely.
Client.close() assumes the internal _admin_client was fully initialized; if instantiation failed before _admin_client was set, hasattr guard covers it but corner cases should be watched.
• Lock contention on _refcount_lock is unlikely but worth monitoring in highly parallel workloads.

This summary was automatically generated by @propel-code-bot

@veeceey
Copy link
Copy Markdown
Contributor Author

veeceey commented Feb 8, 2026

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.
@veeceey
Copy link
Copy Markdown
Contributor Author

veeceey commented Feb 8, 2026

Addressed the idempotent close() feedback from @propel-code-bot. The close() method now uses pop() to safely remove the system from the cache before stopping it, so calling close() multiple times (common when using context managers + explicit close) is a safe no-op instead of raising KeyError. Added a dedicated test_client_close_idempotent test to verify this behavior.

…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.
@veeceey
Copy link
Copy Markdown
Contributor Author

veeceey commented Feb 10, 2026

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.
@veeceey
Copy link
Copy Markdown
Contributor Author

veeceey commented Feb 12, 2026

Addressed the remaining propel-code-bot feedback in c60539d:

  1. Resource leak on init failure: If Client.__init__ raises after the refcount was incremented and the refcount drops to 0, the System is now properly stopped and removed from the cache. Previously it would remain running with no owner.

  2. KeyError in tests after close(): Tests now save a local reference to _system before calling close(), since close() pops the system from the shared _identifier_to_system cache. Accessing client._system after close() would raise KeyError since the property looks up the identifier in the class-level dict.

@veeceey veeceey changed the title Add close() method and context manager support to Client [ENH] Add close() method and context manager support to Client Feb 12, 2026
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.
@rescrv
Copy link
Copy Markdown
Contributor

rescrv commented Feb 20, 2026

Only failure is because third parties cannot access secrets in the repo. Looks good otherwise.

@rescrv rescrv merged commit 5435892 into chroma-core:main Feb 20, 2026
63 of 65 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Unable to Close Persistent Client

2 participants