refine code (#3089)

Signed-off-by: Nyhmo <yihua.mo@zilliz.com>

core/src/cache/CpuCacheMgr.cpp

@@ -32,18 +32,23 @@ CpuCacheMgr::~CpuCacheMgr() {
     ConfigMgr::GetInstance().Detach("cache.cache_size", this);
 }
 
-CpuCacheMgr*
+CpuCacheMgr&
 CpuCacheMgr::GetInstance() {
     static CpuCacheMgr s_mgr;
-    return &s_mgr;
+    return s_mgr;
 }
 
 DataObjPtr
-CpuCacheMgr::GetIndex(const std::string& key) {
+CpuCacheMgr::GetDataObj(const std::string& key) {
     DataObjPtr obj = GetItem(key);
     return obj;
 }
 
+void
+CpuCacheMgr::SetDataObj(const std::string& key, const milvus::cache::DataObjPtr& data) {
+    CacheMgr<DataObjPtr>::InsertItem(key, data);
+}
+
 void
 CpuCacheMgr::ConfigUpdate(const std::string& name) {
     SetCapacity(config.cache.cache_size());

core/src/cache/CpuCacheMgr.h

@@ -28,12 +28,14 @@ class CpuCacheMgr : public CacheMgr<DataObjPtr>, public ConfigObserver {
     ~CpuCacheMgr();
 
  public:
-    // TODO(myh): use smart pointer instead
-    static CpuCacheMgr*
+    static CpuCacheMgr&
     GetInstance();
 
     DataObjPtr
-    GetIndex(const std::string& key);
+    GetDataObj(const std::string& key);
+
+    void
+    SetDataObj(const std::string& key, const DataObjPtr& data);
 
  public:
     void

core/src/cache/GpuCacheMgr.cpp

@@ -40,17 +40,26 @@ GpuCacheMgr::~GpuCacheMgr() {
     ConfigMgr::GetInstance().Detach("gpu.cache_threshold", this);
 }
 
+GpuCacheMgrPtr
+GpuCacheMgr::GetInstance(int64_t gpu_id) {
+    if (instance_.find(gpu_id) == instance_.end()) {
+        std::lock_guard<std::mutex> lock(global_mutex_);
+        if (instance_.find(gpu_id) == instance_.end()) {
+            instance_[gpu_id] = std::make_shared<GpuCacheMgr>(gpu_id);
+        }
+    }
+    return instance_[gpu_id];
+}
+
 DataObjPtr
-GpuCacheMgr::GetIndex(const std::string& key) {
+GpuCacheMgr::GetDataObj(const std::string& key) {
     DataObjPtr obj = GetItem(key);
     return obj;
 }
 
 void
-GpuCacheMgr::InsertItem(const std::string& key, const milvus::cache::DataObjPtr& data) {
-    if (gpu_enable_) {
-        CacheMgr<DataObjPtr>::InsertItem(key, data);
-    }
+GpuCacheMgr::SetDataObj(const std::string& key, const milvus::cache::DataObjPtr& data) {
+    CacheMgr<DataObjPtr>::InsertItem(key, data);
 }
 
 bool
@@ -58,20 +67,12 @@ GpuCacheMgr::Reserve(const int64_t size) {
     return CacheMgr<DataObjPtr>::Reserve(size);
 }
 
-GpuCacheMgrPtr
-GpuCacheMgr::GetInstance(int64_t gpu_id) {
-    if (instance_.find(gpu_id) == instance_.end()) {
-        std::lock_guard<std::mutex> lock(global_mutex_);
-        if (instance_.find(gpu_id) == instance_.end()) {
-            instance_[gpu_id] = std::make_shared<GpuCacheMgr>(gpu_id);
-        }
-    }
-    return instance_[gpu_id];
-}
-
 void
 GpuCacheMgr::ConfigUpdate(const std::string& name) {
-    for (auto& it : instance_) it.second->SetCapacity(config.gpu.cache_size());
+    std::lock_guard<std::mutex> lock(global_mutex_);
+    for (auto& it : instance_) {
+        it.second->SetCapacity(config.gpu.cache_size());
+    }
 }
 
 #endif

core/src/cache/GpuCacheMgr.h

@@ -33,24 +33,23 @@ class GpuCacheMgr : public CacheMgr<DataObjPtr>, public ConfigObserver {
 
     ~GpuCacheMgr();
 
+    static GpuCacheMgrPtr
+    GetInstance(int64_t gpu_id);
+
     DataObjPtr
-    GetIndex(const std::string& key);
+    GetDataObj(const std::string& key);
 
     void
-    InsertItem(const std::string& key, const DataObjPtr& data);
+    SetDataObj(const std::string& key, const DataObjPtr& data);
 
     bool
     Reserve(const int64_t size);
 
-    static GpuCacheMgrPtr
-    GetInstance(int64_t gpu_id);
-
  public:
     void
     ConfigUpdate(const std::string& name) override;
 
  private:
-    bool gpu_enable_ = true;
     int64_t gpu_id_;
     static std::mutex global_mutex_;
     static std::unordered_map<int64_t, GpuCacheMgrPtr> instance_;

core/src/db/DBImpl.cpp

@@ -857,8 +857,8 @@ DBImpl::TimingFlushThread() {
 void
 DBImpl::StartMetricTask() {
     server::Metrics::GetInstance().KeepingAliveCounterIncrement(BACKGROUND_METRIC_INTERVAL);
-    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
-    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
+    int64_t cache_usage = cache::CpuCacheMgr::GetInstance().CacheUsage();
+    int64_t cache_total = cache::CpuCacheMgr::GetInstance().CacheCapacity();
     fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);
 
     if (cache_total > 0) {

core/unittest/db/test_db.cpp

@@ -392,6 +392,9 @@ TEST_F(DBTest, InsertTest) {
     auto status = CreateCollection2(db_, collection_name, 0);
     ASSERT_TRUE(status.ok());
 
+    status = db_->Flush();
+    ASSERT_TRUE(status.ok());
+
     const uint64_t entity_count = 100;
     milvus::engine::DataChunkPtr data_chunk;
     BuildEntities(entity_count, 0, data_chunk);
@@ -535,6 +538,7 @@ TEST_F(DBTest, IndexTest) {
 
         milvus::engine::CollectionIndex index_get;
         status = db_->DescribeIndex(collection_name, VECTOR_FIELD_NAME, index_get);
+        ASSERT_TRUE(status.ok());
         ASSERT_TRUE(index_get.index_name_.empty());
     }
 
@@ -544,6 +548,7 @@ TEST_F(DBTest, IndexTest) {
 
         milvus::engine::CollectionIndex index_get;
         status = db_->DescribeIndex(collection_name, "field_0", index_get);
+        ASSERT_TRUE(status.ok());
         ASSERT_TRUE(index_get.index_name_.empty());
     }
 }