From daaeb27ee1a2ff1f435527f10222cf93e94d4e01 Mon Sep 17 00:00:00 2001
From: dragondriver <jiquan.long@zilliz.com>
Date: Thu, 11 Nov 2021 20:42:44 +0800
Subject: [PATCH] Make AdapterMgr.GetAdapter thread-safe (#11674)

Signed-off-by: dragondriver <jiquan.long@zilliz.com>
---
 .../index/vector_index/ConfAdapterMgr.cpp     | 10 +++---
 .../index/vector_index/ConfAdapterMgr.h       |  2 ++
 internal/core/unittest/CMakeLists.txt         |  1 +
 .../core/unittest/test_conf_adapter_mgr.cpp   | 36 +++++++++++++++++++
 4 files changed, 44 insertions(+), 5 deletions(-)
 create mode 100644 internal/core/unittest/test_conf_adapter_mgr.cpp

diff --git a/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp b/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp
index b83384942d..2f21acb2dc 100644
--- a/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp
+++ b/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp
@@ -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
diff --git a/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h b/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h
index 83b9d0f584..5ebc31b38d 100644
--- a/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h
+++ b/internal/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h
@@ -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
diff --git a/internal/core/unittest/CMakeLists.txt b/internal/core/unittest/CMakeLists.txt
index 418878b275..b6a228c1a4 100644
--- a/internal/core/unittest/CMakeLists.txt
+++ b/internal/core/unittest/CMakeLists.txt
@@ -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
diff --git a/internal/core/unittest/test_conf_adapter_mgr.cpp b/internal/core/unittest/test_conf_adapter_mgr.cpp
new file mode 100644
index 0000000000..0a9d1d4d61
--- /dev/null
+++ b/internal/core/unittest/test_conf_adapter_mgr.cpp
@@ -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();
+    }
+}