[ENH] Clean up System instances when clients are destroyed

Author: jairad26

chromadb/api/shared_system_client.py

@@ -9,6 +9,7 @@
 
 class SharedSystemClient:
     _identifier_to_system: ClassVar[Dict[str, System]] = {}
+    _identifier_to_refcount: ClassVar[Dict[str, int]] = {}
     _identifier: str
 
     def __init__(
@@ -17,6 +18,9 @@ def __init__(
     ) -> None:
         self._identifier = SharedSystemClient._get_identifier_from_settings(settings)
         SharedSystemClient._create_system_if_not_exists(self._identifier, settings)
+        SharedSystemClient._identifier_to_refcount[self._identifier] = (
+            SharedSystemClient._identifier_to_refcount.get(self._identifier, 0) + 1
+        )
 
     @classmethod
     def _create_system_if_not_exists(
@@ -25,6 +29,7 @@ def _create_system_if_not_exists(
         if identifier not in cls._identifier_to_system:
             new_system = System(settings)
             cls._identifier_to_system[identifier] = new_system
+            cls._identifier_to_refcount[identifier] = 0
 
             new_system.instance(ProductTelemetryClient)
             new_system.instance(ServerAPI)
@@ -83,9 +88,35 @@ def from_system(cls, system: System) -> "SharedSystemClient":
         instance = cls(system.settings)
         return instance
 
+    def __del__(self) -> None:
+        """Clean up System when this client is destroyed."""
+        if hasattr(self, "_identifier"):
+            SharedSystemClient._decrement_refcount(self._identifier)
+
+    @classmethod
+    def _decrement_refcount(cls, identifier: str) -> None:
+        """Decrement reference count for a System and cleanup if no clients remain."""
+        if identifier not in cls._identifier_to_refcount:
+            return
+
+        cls._identifier_to_refcount[identifier] -= 1
+
+        if cls._identifier_to_refcount[identifier] <= 0:
+            # since no more client using this system, can stop it and remove from cache
+            if identifier in cls._identifier_to_system:
+                system = cls._identifier_to_system[identifier]
+                system.stop()
+                del cls._identifier_to_system[identifier]
+            del cls._identifier_to_refcount[identifier]
+
     @staticmethod
     def clear_system_cache() -> None:
+        """Clear the system cache so that new systems can be created for an existing path.
+        This should only be used for testing purposes."""
+        for system in SharedSystemClient._identifier_to_system.values():
+            system.stop()
         SharedSystemClient._identifier_to_system = {}
+        SharedSystemClient._identifier_to_refcount = {}
 
     @property
     def _system(self) -> System: