MS-614 Preload table at startup

Former-commit-id: 15ea4eb6d6c6f3437979a8b741a07bee5eaada5b

cpp/CHANGELOG.md

@@ -28,6 +28,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-609 - Update task construct function
 - MS-611 - Add resources validity check in ResourceMgr
 - MS-619 - Add optimizer class in scheduler
+- MS-614 - Preload table at startup
 
 ## New Feature
 

cpp/conf/server_config.template

@@ -18,6 +18,9 @@ db_config:
                                     # sum of insert_buffer_size and cpu_cache_capacity cannot exceed total memory
   build_index_gpu: 0                # gpu id used for building index
 
+  preload_table:                    # preload data at startup, '*' means load all tables, empty value means no preload
+                                    # you can specify preload tables like this: table1,table2,table3
+
 metric_config:
   enable_monitor: false             # enable monitoring or not
   collector: prometheus             # prometheus

cpp/src/db/DBImpl.cpp

@@ -206,7 +206,7 @@ DBImpl::PreloadTable(const std::string& table_id) {
 
             size += engine->PhysicalSize();
             if (size > available_size) {
-                break;
+                return Status(SERVER_CACHE_FULL, "Cache is full");
             } else {
                 try {
                     // step 1: load index
@@ -639,8 +639,9 @@ DBImpl::MergeFiles(const std::string& table_id, const meta::DateT& date, const m
         ENGINE_LOG_DEBUG << "Merging file " << file_schema.file_id_;
         index_size = index->Size();
 
-        if (index_size >= file_schema.index_file_size_)
+        if (index_size >= file_schema.index_file_size_) {
             break;
+        }
     }
 
     // step 3: serialize to disk

cpp/src/server/Config.cpp

@@ -770,6 +770,16 @@ Config::GetDBConfigStrBuildIndexGPU() {
     return value;
 }
 
+std::string
+Config::GetDBConfigStrPreloadTable() {
+    std::string value;
+    if (!GetConfigValueInMem(CONFIG_DB, CONFIG_DB_PRELOAD_TABLE, value).ok()) {
+        value = GetConfigNode(CONFIG_DB).GetValue(CONFIG_DB_PRELOAD_TABLE);
+        SetConfigValueInMem(CONFIG_DB, CONFIG_DB_PRELOAD_TABLE, value);
+    }
+    return value;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 /* metric config */
 std::string
@@ -988,6 +998,12 @@ Config::GetDBConfigBuildIndexGPU(int32_t& value) {
     return Status::OK();
 }
 
+Status
+Config::GetDBConfigPreloadTable(std::string& value) {
+    value = GetDBConfigStrPreloadTable();
+    return Status::OK();
+}
+
 Status
 Config::GetMetricConfigEnableMonitor(bool& value) {
     std::string str = GetMetricConfigStrEnableMonitor();

cpp/src/server/Config.h

@@ -56,6 +56,7 @@ static const char* CONFIG_DB_INSERT_BUFFER_SIZE = "insert_buffer_size";
 static const char* CONFIG_DB_INSERT_BUFFER_SIZE_DEFAULT = "4";
 static const char* CONFIG_DB_BUILD_INDEX_GPU = "build_index_gpu";
 static const char* CONFIG_DB_BUILD_INDEX_GPU_DEFAULT = "0";
+static const char* CONFIG_DB_PRELOAD_TABLE = "preload_table";
 
 /* cache config */
 static const char* CONFIG_CACHE = "cache_config";
@@ -204,6 +205,8 @@ class Config {
     GetDBConfigStrInsertBufferSize();
     std::string
     GetDBConfigStrBuildIndexGPU();
+    std::string
+    GetDBConfigStrPreloadTable();
 
     /* metric config */
     std::string
@@ -261,6 +264,8 @@ class Config {
     GetDBConfigInsertBufferSize(int32_t& value);
     Status
     GetDBConfigBuildIndexGPU(int32_t& value);
+    Status
+    GetDBConfigPreloadTable(std::string& value);
 
     /* metric config */
     Status

cpp/src/server/DBWrapper.cpp

@@ -155,6 +155,20 @@ DBWrapper::StartService() {
 
     db_->Start();
 
+    // preload table
+    std::string preload_tables;
+    s = config.GetDBConfigPreloadTable(preload_tables);
+    if (!s.ok()) {
+        return s;
+    }
+
+    s = PreloadTables(preload_tables);
+    if (!s.ok()) {
+        std::cerr << "ERROR! Failed to preload tables: " << preload_tables << std::endl;
+        std::cerr << s.ToString() << std::endl;
+        kill(0, SIGUSR1);
+    }
+
     return Status::OK();
 }
 
@@ -167,5 +181,34 @@ DBWrapper::StopService() {
     return Status::OK();
 }
 
+Status
+DBWrapper::PreloadTables(const std::string& preload_tables) {
+    if(preload_tables.empty()) {
+        //do nothing
+    } else if(preload_tables == "*") {
+        //load all tables
+        std::vector<engine::meta::TableSchema> table_schema_array;
+        db_->AllTables(table_schema_array);
+
+        for(auto& schema : table_schema_array) {
+            auto status = db_->PreloadTable(schema.table_id_);
+            if (!status.ok()) {
+                return status;
+            }
+        }
+    } else {
+        std::vector<std::string> table_names;
+        StringHelpFunctions::SplitStringByDelimeter(preload_tables, ",", table_names);
+        for(auto& name : table_names) {
+            auto status = db_->PreloadTable(name);
+            if (!status.ok()) {
+                return status;
+            }
+        }
+    }
+
+    return Status::OK();
+}
+
 }  // namespace server
 }  // namespace milvus

cpp/src/server/DBWrapper.h

@@ -52,6 +52,10 @@ class DBWrapper {
         return db_;
     }
 
+ private:
+    Status
+    PreloadTables(const std::string& preload_tables);
+
  private:
     engine::DBPtr db_;
 };

cpp/src/server/grpc_impl/GrpcRequestScheduler.cpp

@@ -36,7 +36,6 @@ ErrorMap(ErrorCode code) {
         {SERVER_INVALID_ARGUMENT, ::milvus::grpc::ErrorCode::ILLEGAL_ARGUMENT},
         {SERVER_FILE_NOT_FOUND, ::milvus::grpc::ErrorCode::FILE_NOT_FOUND},
         {SERVER_NOT_IMPLEMENT, ::milvus::grpc::ErrorCode::UNEXPECTED_ERROR},
-        {SERVER_BLOCKING_QUEUE_EMPTY, ::milvus::grpc::ErrorCode::UNEXPECTED_ERROR},
         {SERVER_CANNOT_CREATE_FOLDER, ::milvus::grpc::ErrorCode::CANNOT_CREATE_FOLDER},
         {SERVER_CANNOT_CREATE_FILE, ::milvus::grpc::ErrorCode::CANNOT_CREATE_FILE},
         {SERVER_CANNOT_DELETE_FOLDER, ::milvus::grpc::ErrorCode::CANNOT_DELETE_FOLDER},
@@ -57,7 +56,7 @@ ErrorMap(ErrorCode code) {
         {SERVER_INVALID_INDEX_FILE_SIZE, ::milvus::grpc::ErrorCode::ILLEGAL_ARGUMENT},
         {SERVER_ILLEGAL_VECTOR_ID, ::milvus::grpc::ErrorCode::ILLEGAL_VECTOR_ID},
         {SERVER_ILLEGAL_SEARCH_RESULT, ::milvus::grpc::ErrorCode::ILLEGAL_SEARCH_RESULT},
-        {SERVER_CACHE_ERROR, ::milvus::grpc::ErrorCode::CACHE_FAILED},
+        {SERVER_CACHE_FULL, ::milvus::grpc::ErrorCode::CACHE_FAILED},
         {DB_META_TRANSACTION_FAILED, ::milvus::grpc::ErrorCode::META_FAILED},
         {SERVER_BUILD_INDEX_ERROR, ::milvus::grpc::ErrorCode::BUILD_INDEX_ERROR},
         {SERVER_OUT_OF_MEMORY, ::milvus::grpc::ErrorCode::OUT_OF_MEMORY},

cpp/src/utils/Error.h

@@ -56,7 +56,6 @@ constexpr ErrorCode SERVER_NULL_POINTER = ToServerErrorCode(3);
 constexpr ErrorCode SERVER_INVALID_ARGUMENT = ToServerErrorCode(4);
 constexpr ErrorCode SERVER_FILE_NOT_FOUND = ToServerErrorCode(5);
 constexpr ErrorCode SERVER_NOT_IMPLEMENT = ToServerErrorCode(6);
-constexpr ErrorCode SERVER_BLOCKING_QUEUE_EMPTY = ToServerErrorCode(7);
 constexpr ErrorCode SERVER_CANNOT_CREATE_FOLDER = ToServerErrorCode(8);
 constexpr ErrorCode SERVER_CANNOT_CREATE_FILE = ToServerErrorCode(9);
 constexpr ErrorCode SERVER_CANNOT_DELETE_FOLDER = ToServerErrorCode(10);
@@ -74,7 +73,7 @@ constexpr ErrorCode SERVER_INVALID_ROWRECORD_ARRAY = ToServerErrorCode(107);
 constexpr ErrorCode SERVER_INVALID_TOPK = ToServerErrorCode(108);
 constexpr ErrorCode SERVER_ILLEGAL_VECTOR_ID = ToServerErrorCode(109);
 constexpr ErrorCode SERVER_ILLEGAL_SEARCH_RESULT = ToServerErrorCode(110);
-constexpr ErrorCode SERVER_CACHE_ERROR = ToServerErrorCode(111);
+constexpr ErrorCode SERVER_CACHE_FULL = ToServerErrorCode(111);
 constexpr ErrorCode SERVER_WRITE_ERROR = ToServerErrorCode(112);
 constexpr ErrorCode SERVER_INVALID_NPROBE = ToServerErrorCode(113);
 constexpr ErrorCode SERVER_INVALID_INDEX_NLIST = ToServerErrorCode(114);