From 3651baadd523b7eb1d806845e6e29afd6caefdf5 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Tue, 16 Apr 2024 20:34:42 +0300
Subject: [PATCH] test: Added invariants for fd in multithreaded tests

---
 chromadb/test/property/invariants.py | 18 ++++++++++++++++++
 chromadb/test/test_multithreaded.py  |  4 +++-
 requirements_dev.txt                 |  1 +
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/chromadb/test/property/invariants.py b/chromadb/test/property/invariants.py
index 803cac2ac075..ca0845ed3cea 100644
--- a/chromadb/test/property/invariants.py
+++ b/chromadb/test/property/invariants.py
@@ -1,4 +1,8 @@
+import gc
 import math
+
+import psutil
+
 from chromadb.test.property.strategies import NormalizedRecordSet, RecordSet
 from typing import Callable, Optional, Tuple, Union, List, TypeVar, cast
 from typing_extensions import Literal
@@ -169,6 +173,20 @@ def is_metadata_valid(normalized_record_set: NormalizedRecordSet) -> bool:
     return not any([len(m) == 0 for m in normalized_record_set["metadatas"]])
 
 
+def fd_not_exceeding_threadpool_size(threadpool_size: int) -> None:
+    """
+    Checks that the open file descriptors are not exceeding the threadpool size
+    works only for SegmentAPI
+    """
+    current_process = psutil.Process()
+    open_files = current_process.open_files()
+    if len([p.path for p in open_files if "sqlite3" in p.path]) - 1 <= threadpool_size:
+        gc.collect()  # GC to collect the orphaned TLS objects
+    assert (
+        len([p.path for p in open_files if "sqlite3" in p.path]) - 1 <= threadpool_size
+    )
+
+
 def ann_accuracy(
     collection: Collection,
     record_set: RecordSet,
diff --git a/chromadb/test/test_multithreaded.py b/chromadb/test/test_multithreaded.py
index c0b05e883243..a35687daaa15 100644
--- a/chromadb/test/test_multithreaded.py
+++ b/chromadb/test/test_multithreaded.py
@@ -7,6 +7,7 @@
 
 from chromadb.api import ServerAPI
 import chromadb.test.property.invariants as invariants
+from chromadb.api.segment import SegmentAPI
 from chromadb.test.property.strategies import RecordSet
 from chromadb.test.property.strategies import test_hnsw_config
 from chromadb.types import Metadata
@@ -191,7 +192,8 @@ def perform_operation(
         exception = future.exception()
         if exception is not None:
             raise exception
-
+    if isinstance(api, SegmentAPI):  # we can't check invariants for FastAPI
+        invariants.fd_not_exceeding_threadpool_size(num_workers)
     # Check that invariants hold
     invariants.count(coll, records_set)
     invariants.ids_match(coll, records_set)
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 4df73fb8c564..eae609e18b19 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -6,6 +6,7 @@ hypothesis<=6.98.9 # > Than this version has API changes we don't currently supp
 hypothesis[numpy]<=6.98.9
 mypy-protobuf
 pre-commit
+psutil
 pytest
 pytest-asyncio
 setuptools_scm