From 9b2368016c5e10307d01b1ccb3a0230c30f43304 Mon Sep 17 00:00:00 2001
From: feisiyicl <64510805+feisiyicl@users.noreply.github.com>
Date: Wed, 29 Apr 2020 20:01:52 +0800
Subject: [PATCH] #1946 fix load index file CPU2GPU fail during searching
 (#2182)

Signed-off-by: feisiyicl <7764126@qq.com>
---
 CHANGELOG.md                                  |  1 +
 .../helpers/FaissGpuResourceMgr.cpp           | 36 ++++++++++---------
 .../helpers/FaissGpuResourceMgr.h             |  3 +-
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c5f412adec..a77a09f636 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ Please mark all change in change log and use the issue from GitHub
 ## Bug
 -   \#1705 Limit the insert data batch size
 -   \#1929 Skip MySQL meta schema field width check
+-   \#1946 Fix load index file CPU2GPU fail during searching
 -   \#1997 Index file missed after compact
 -   \#2073 Fix CheckDBConfigBackendUrl error message
 -   \#2076 CheckMetricConfigAddress error message
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
index 051c5496d4..b315f538c9 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
+++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
@@ -51,29 +51,31 @@ FaissGpuResourceMgr::InitDevice(int64_t device_id, int64_t pin_mem_size, int64_t
 
 void
 FaissGpuResourceMgr::InitResource() {
-    if (is_init)
-        return;
+    if (!initialized_) {
+        std::lock_guard<std::mutex> lock(init_mutex_);
 
-    is_init = true;
+        if (!initialized_) {
+            for (auto& device : devices_params_) {
+                auto& device_id = device.first;
 
-    for (auto& device : devices_params_) {
-        auto& device_id = device.first;
+                mutex_cache_.emplace(device_id, std::make_unique<std::mutex>());
 
-        mutex_cache_.emplace(device_id, std::make_unique<std::mutex>());
+                auto& device_param = device.second;
+                auto& bq = idle_map_[device_id];
 
-        auto& device_param = device.second;
-        auto& bq = idle_map_[device_id];
+                for (int64_t i = 0; i < device_param.resource_num; ++i) {
+                    auto raw_resource = std::make_shared<faiss::gpu::StandardGpuResources>();
 
-        for (int64_t i = 0; i < device_param.resource_num; ++i) {
-            auto raw_resource = std::make_shared<faiss::gpu::StandardGpuResources>();
+                    // TODO(linxj): enable set pinned memory
+                    auto res_wrapper = std::make_shared<Resource>(raw_resource);
+                    AllocateTempMem(res_wrapper, device_id, 0);
 
-            // TODO(linxj): enable set pinned memory
-            auto res_wrapper = std::make_shared<Resource>(raw_resource);
-            AllocateTempMem(res_wrapper, device_id, 0);
-
-            bq.Put(res_wrapper);
+                    bq.Put(res_wrapper);
+                }
+                LOG_KNOWHERE_DEBUG_ << "DEVICEID " << device_id << ", resource count " << bq.Size();
+            }
+            initialized_ = true;
         }
-        LOG_KNOWHERE_DEBUG_ << "DEVICEID " << device_id << ", resource count " << bq.Size();
     }
 }
 
@@ -115,7 +117,7 @@ FaissGpuResourceMgr::Free() {
             bq.Take();
         }
     }
-    is_init = false;
+    initialized_ = false;
 }
 
 void
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
index e97b2ab05b..5c3ee2aa83 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
+++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
@@ -77,7 +77,8 @@ class FaissGpuResourceMgr {
     Dump();
 
  protected:
-    bool is_init = false;
+    bool initialized_ = false;
+    std::mutex init_mutex_;
 
     std::map<int64_t, std::unique_ptr<std::mutex>> mutex_cache_;
     std::map<int64_t, DeviceParams> devices_params_;