Add test & fine grain locking for collection manager

This commit is contained in:
ozanarmagan 2023-10-09 13:14:10 +03:00
parent 73ae433ab5
commit 0dad1a369f
2 changed files with 79 additions and 3 deletions

View File

@ -534,12 +534,12 @@ Option<nlohmann::json> CollectionManager::drop_collection(const std::string& col
const auto& embedding_fields = collection->get_embedding_fields();
u_lock.unlock();
for(const auto& embedding_field : embedding_fields) {
const auto& model_name = embedding_field.embed[fields::model_config]["model_name"].get<std::string>();
process_embedding_field_delete(model_name);
}
u_lock.unlock();
// don't hold any collection manager locks here, since this can take some time
delete collection;
@ -1565,8 +1565,7 @@ Option<Collection*> CollectionManager::clone_collection(const string& existing_n
}
void CollectionManager::process_embedding_field_delete(const std::string& model_name) {
// can't have a shared lock here
// because we will be already acquiring a lock on collection manager if we are deleting a collection
std::shared_lock lock(mutex);
bool found = false;
for(const auto& collection: collections) {

View File

@ -2398,6 +2398,83 @@ TEST_F(CollectionVectorTest, TestUnloadingModelsOnDrop) {
drop_op = coll->alter(drop_schema);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
}
TEST_F(CollectionVectorTest, TestUnloadModelsCollectionHaveTwoEmbeddingField) {
nlohmann::json actual_schema = R"({
"name": "test",
"fields": [
{
"name": "title",
"type": "string"
},
{
"name": "title_vec",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
},
{
"name": "title_vec2",
"type": "float[]",
"embed": {
"from": [
"title"
],
"model_config": {
"model_name": "ts/e5-small"
}
}
}
]
})"_json;
TextEmbedderManager::set_model_dir("/tmp/typesense_test/models");
auto schema = actual_schema;
auto collection_create_op = collectionManager.create_collection(schema);
ASSERT_TRUE(collection_create_op.ok());
auto coll = collection_create_op.get();
auto text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
nlohmann::json drop_schema = R"({
"fields": [
{
"name": "title_vec",
"drop": true
}
]
})"_json;
auto drop_op = coll->alter(drop_schema);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(1, text_embedders.size());
drop_schema = R"({
"fields": [
{
"name": "title_vec2",
"drop": true
}
]
})"_json;
drop_op = coll->alter(drop_schema);
ASSERT_TRUE(drop_op.ok());
text_embedders = TextEmbedderManager::get_instance()._get_text_embedders();
ASSERT_EQ(0, text_embedders.size());
}