Support auto rebuild and on-demand rebuild (#32)

* Support auto rebuild and on-demand rebuild

* Change back rebuild internal

engine/cmd/main.cpp

@@ -31,6 +31,7 @@ int main(int argc, char *argv[]) {
   static struct option long_options[] = {{"conf_file", required_argument, nullptr, 'c'},
                                          {"help", no_argument, nullptr, 'h'},
                                          {"port", required_argument, 0, 'p'},
+                                         {"rebuild", required_argument, 0, 'r'},
                                          {nullptr, 0, nullptr, 0}};
 
   int option_index = 0;
@@ -48,7 +49,8 @@ int main(int argc, char *argv[]) {
 
   int value;
   uint16_t port = 8888;
-  while ((value = getopt_long(argc, argv, "c:p:h", long_options, &option_index)) != -1) {
+  bool rebuild = true;
+  while ((value = getopt_long(argc, argv, "c:p:r:h", long_options, &option_index)) != -1) {
     switch (value) {
       case 'c': {
         char *config_filename_ptr = strdup(optarg);
@@ -62,6 +64,11 @@ int main(int argc, char *argv[]) {
         port = (uint16_t)(stoi(server_port));
         break;
       }
+      case 'r': {
+        std::string start_rebuild = optarg;
+        rebuild = start_rebuild[0] == 't' || start_rebuild[0] == 'T';
+        break;
+      }
       case 'h':
         print_help(app_name);
         return EXIT_SUCCESS;
@@ -73,7 +80,7 @@ int main(int argc, char *argv[]) {
 
   server.Init(config_filename);
 
-  status = server.Start(port);
+  status = server.Start(port, rebuild);
   if (status.ok()) {
     std::cout << "Epsilla Vector Database server started successfully!" << std::endl;
     std::cout << "Server running on http://localhost:" << port << std::endl;

engine/db/db_server.cpp

@@ -8,9 +8,6 @@ namespace engine {
 DBServer::DBServer() {
   // Initialize the meta database
   meta_ = std::make_shared<meta::BasicMetaImpl>();
-
-  // Start the thread to periodically call Rebuild
-  rebuild_thread_ = std::thread(&DBServer::RebuildPeriodically, this);
 }
 
 DBServer::~DBServer() {

engine/db/db_server.hpp

@@ -27,7 +27,6 @@ class DBServer {
   Status CreateTable(const std::string& db_name, meta::TableSchema& table_schema);
   Status DropTable(const std::string& db_name, const std::string& table_name);
   std::shared_ptr<DBMVP> GetDB(const std::string& db_name);
-  Status Rebuild();
   Status Insert(const std::string& db_name, const std::string& table_name, vectordb::Json& records);
   Status Search(
     const std::string& db_name,
@@ -48,13 +47,33 @@ class DBServer {
     vectordb::Json& result
   );
 
+  void StartRebuild() {
+    if (rebuild_started_) {
+      return;
+    }
+    rebuild_started_ = true;
+    // Start the thread to periodically call Rebuild
+    rebuild_thread_ = std::thread(&DBServer::RebuildPeriodically, this);
+  }
+
+  Status RebuildOndemand() {
+    if (rebuild_started_) {
+      return Status(DB_UNEXPECTED_ERROR, "Auto rebuild is enabled. Cannot conduct on-demand rebuild.");
+    }
+    rebuild_started_ = true;
+    Rebuild();
+    rebuild_started_ = false;
+    return Status::OK();
+  }
+
  private:
   std::shared_ptr<meta::Meta> meta_;                           // The db meta.
   // TODO: change to concurrent version.
   std::unordered_map<std::string, size_t> db_name_to_id_map_;  // The db name to db index map.
   std::vector<std::shared_ptr<DBMVP>> dbs_;                    // The dbs.
   std::thread rebuild_thread_;
   bool stop_rebuild_thread_ = false;
+  bool rebuild_started_ = false;
 
   // periodically in a separate thread
   void RebuildPeriodically() {
@@ -67,6 +86,8 @@ class DBServer {
       std::this_thread::sleep_for(rebuild_interval);
     }
   };
+
+  Status Rebuild();
 };
 
 

engine/server/server.cpp

@@ -18,9 +18,10 @@ void Server::Init(const std::string& config_filename) {
   config_filename_ = config_filename;
 }
 
-Status Server::Start(uint16_t port) {
+Status Server::Start(uint16_t port, bool rebuild) {
   try {
     web::WebServer::GetInstance().SetPort(port);
+    web::WebServer::GetInstance().SetRebuild(rebuild);
     return StartService();
   } catch (std::exception& ex) {
     std::string str = "Epsilla VectorDB server encounter exception: " + std::string(ex.what());

engine/server/server.hpp

@@ -12,7 +12,7 @@ class Server {
 
   void Init(const std::string& config_filename);
 
-  Status Start(uint16_t port);
+  Status Start(uint16_t port, bool rebuild);
   void Stop();
 
  private:

engine/server/web_server/web_controller.hpp

@@ -531,6 +531,23 @@ class WebController : public oatpp::web::server::api::ApiController {
         return createDtoResponse(Status::CODE_200, res_dto);
     }
 
+    ADD_CORS(Rebuild)
+
+    ENDPOINT("POST", "/api/rebuild", Rebuild) {
+        vectordb::Status status = db_server->RebuildOndemand();
+
+        auto dto = StatusDto::createShared();
+        if (!status.ok()) {
+            dto->statusCode = Status::CODE_500.code;
+            dto->message = status.message();
+            return createDtoResponse(Status::CODE_500, dto);
+        }
+
+        dto->statusCode = Status::CODE_200.code;
+        dto->message = "Rebuild finished!";
+        return createDtoResponse(Status::CODE_200, dto);
+    }
+
 /**
  *  Finish ENDPOINTs generation ('ApiController' codegen)
  */

engine/server/web_server/web_server.cpp

@@ -35,6 +35,10 @@ WebServer::StartService() {
 
         /* create ApiControllers and add endpoints to router */
         auto user_controller = WebController::createShared();
+        // Start rebuild thread
+        if (rebuild_) {
+            user_controller->db_server->StartRebuild();
+        }
         auto router = components.http_router.getObject();
         router->addController(user_controller);
 

engine/server/web_server/web_server.hpp

@@ -17,6 +17,7 @@ class WebServer {
   std::atomic_bool try_stop_;
   std::shared_ptr<std::thread> thread_ptr_;
   uint16_t port_;
+  bool rebuild_;
 
  private:
   WebServer() {
@@ -25,10 +26,8 @@ class WebServer {
 
   ~WebServer() = default;
 
-  Status
-  StartService();
-  Status
-  StopService();
+  Status StartService();
+  Status StopService();
 
  public:
   static WebServer& GetInstance() {
@@ -43,6 +42,10 @@ class WebServer {
   void SetPort(uint16_t port) {
     port_ = port;
   }
+
+  void SetRebuild(bool rebuild) {
+    rebuild_ = rebuild;
+  }
 };
 }  // namespace web
 }  // namespace server