Make AdapterMgr.GetAdapter thread-safe (#11674)

Signed-off-by: dragondriver <jiquan.long@zilliz.com>
pull/11555/head
dragondriver 2021-11-11 20:42:44 +08:00 committed by GitHub
parent 72cf285f3a
commit daaeb27ee1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 44 additions and 5 deletions

View File

@ -19,9 +19,8 @@ namespace knowhere {
ConfAdapterPtr
AdapterMgr::GetAdapter(const IndexType type) {
if (!init_) {
RegisterAdapter();
}
auto register_wrapper = [&, this]() { RegisterAdapter(); };
std::call_once(once_, register_wrapper);
try {
return collection_.at(type)();
@ -34,8 +33,6 @@ AdapterMgr::GetAdapter(const IndexType type) {
void
AdapterMgr::RegisterAdapter() {
init_ = true;
REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_FAISS_IDMAP, idmap_adapter);
REGISTER_CONF_ADAPTER(IVFConfAdapter, IndexEnum::INDEX_FAISS_IVFFLAT, ivf_adapter);
REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexEnum::INDEX_FAISS_IVFPQ, ivfpq_adapter);
@ -56,6 +53,9 @@ AdapterMgr::RegisterAdapter() {
REGISTER_CONF_ADAPTER(RHNSWSQConfAdapter, IndexEnum::INDEX_RHNSWSQ, rhnswsq_adapter);
REGISTER_CONF_ADAPTER(NGTPANNGConfAdapter, IndexEnum::INDEX_NGTPANNG, ngtpanng_adapter);
REGISTER_CONF_ADAPTER(NGTONNGConfAdapter, IndexEnum::INDEX_NGTONNG, ngtonng_adapter);
// though init_ won't be used later, it's better to set `init_` to true after registration was done.
init_ = true;
}
} // namespace knowhere

View File

@ -14,6 +14,7 @@
#include <functional>
#include <memory>
#include <unordered_map>
#include <mutex>
#include "knowhere/index/IndexType.h"
#include "knowhere/index/vector_index/ConfAdapter.h"
@ -45,6 +46,7 @@ class AdapterMgr {
protected:
bool init_ = false;
std::unordered_map<IndexType, std::function<ConfAdapterPtr()>> collection_;
std::once_flag once_;
};
} // namespace knowhere

View File

@ -35,6 +35,7 @@ set(MILVUS_TEST_FILES
test_span.cpp
test_timestamp_index.cpp
test_reduce_c.cpp
test_conf_adapter_mgr.cpp
)
add_executable(all_tests

View File

@ -0,0 +1,36 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License
#include <vector>
#include <thread>
#include <gtest/gtest.h>
#include "index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h"
TEST(AdapterMgr, MultiThread) {
auto run_case = [&]() {
auto& ins = milvus::knowhere::AdapterMgr::GetInstance();
auto adapter = ins.GetAdapter(milvus::knowhere::IndexEnum::INDEX_HNSW);
ASSERT_TRUE(adapter != nullptr);
ASSERT_ANY_THROW(ins.GetAdapter("not supported now!"));
};
size_t num = 4;
std::vector<std::thread> threads;
for (auto i = 0; i < num; i++) {
threads.emplace_back(std::move(std::thread(run_case)));
}
for (auto i = 0; i < num; i++) {
threads[i].join();
}
}