Release API

epsilla-cloud · Feb 21, 2024 · c9c1f46 · c9c1f46
1 parent d3aa661
commit c9c1f46
Show file tree

Hide file tree

Showing 9 changed files with 60 additions and 0 deletions.
diff --git a/engine/db/db_mvp.cpp b/engine/db/db_mvp.cpp
@@ -105,5 +105,19 @@ Status DBMVP::SwapExecutors() {
   return Status::OK();
 }
 
+Status DBMVP::Release() {
+  // Loop through all tables and release
+  for (int64_t i = 0; i < tables_.size(); ++i) {
+    std::shared_ptr<TableMVP> table = tables_[i];
+    if (table != nullptr) {
+      auto status = table->Release();
+      if (!status.ok()) {
+        std::cout << "Release table " << table->table_schema_.name_ << " failed." << std::endl;
+      }
+    }
+  }
+  return Status::OK();
+}
+
 }  // namespace engine
 }  // namespace vectordb
diff --git a/engine/db/db_mvp.hpp b/engine/db/db_mvp.hpp
@@ -30,6 +30,7 @@ class DBMVP {
   std::shared_ptr<TableMVP> GetTable(const std::string& table_name);
   Status Rebuild();
   Status SwapExecutors();
+  Status Release();
 
   void SetWALEnabled(bool enabled) {
     for (auto table : tables_) {

diff --git a/engine/db/db_server.cpp b/engine/db/db_server.cpp
@@ -68,6 +68,16 @@ Status DBServer::UnloadDB(const std::string& db_name) {
   return Status::OK();
 }
 
+Status DBServer::ReleaseDB(const std::string& db_name) {
+  // Release database from memory.
+  auto it = db_name_to_id_map_.find(db_name);
+  if (it == db_name_to_id_map_.end()) {
+    return Status(DB_UNEXPECTED_ERROR, "DB not found: " + db_name);
+  }
+  dbs_[it->second]->Release();
+  return Status::OK();
+}
+
 Status DBServer::GetStatistics(const std::string& db_name,
                                vectordb::Json& response) {
   auto db = GetDB(db_name);

diff --git a/engine/db/db_server.hpp b/engine/db/db_server.hpp
@@ -27,6 +27,7 @@ class DBServer {
 
   Status LoadDB(const std::string& db_name, const std::string& db_catalog_path, int64_t init_table_scale, bool wal_enabled, std::unordered_map<std::string, std::string> &headers);
   Status UnloadDB(const std::string& db_name);
+  Status ReleaseDB(const std::string& db_name);
   Status GetStatistics(const std::string& db_name, vectordb::Json& response);
   Status CreateTable(const std::string& db_name, meta::TableSchema& table_schema, size_t& table_id);
   Status CreateTable(const std::string& db_name, const std::string& table_schema_json, size_t& table_id);

diff --git a/engine/db/table_mvp.cpp b/engine/db/table_mvp.cpp
@@ -267,6 +267,11 @@ Status TableMVP::SwapExecutors() {
   return Status::OK();
 }
 
+Status TableMVP::Release() {
+  // Release the table segment.
+  return table_segment_->Release();
+}
+
 Status TableMVP::Insert(vectordb::Json &record, std::unordered_map<std::string, std::string> &headers, bool upsert) {
   int64_t wal_id =
       wal_->WriteEntry(upsert ? LogEntryType::UPSERT : LogEntryType::INSERT, record.DumpToString());

diff --git a/engine/db/table_mvp.hpp b/engine/db/table_mvp.hpp
@@ -41,6 +41,8 @@ class TableMVP {
   // Swap executors during config change.
   Status SwapExecutors();
 
+  Status Release();
+
   Status Insert(vectordb::Json &records, std::unordered_map<std::string, std::string> &headers, bool upsert = false);
 
   Status InsertPrepare(vectordb::Json &pks, vectordb::Json &result);

diff --git a/engine/db/table_segment_mvp.cpp b/engine/db/table_segment_mvp.cpp
@@ -1113,18 +1113,26 @@ void TableSegmentMVP::Debug(meta::TableSchema& table_schema) {
 }
 
 TableSegmentMVP::~TableSegmentMVP() {
+  Release();
+}
+
+Status TableSegmentMVP::Release() {
   if (attribute_table_ != nullptr) {
     delete[] attribute_table_;
+    attribute_table_ = nullptr;
   }
   if (vector_tables_ != nullptr) {
     for (auto i = 0; i < dense_vector_num_; ++i) {
       delete[] vector_tables_[i];
     }
     delete[] vector_tables_;
+    vector_tables_ = nullptr;
   }
   if (deleted_ != nullptr) {
     delete deleted_;
+    deleted_ = nullptr;
   }
+  return Status::OK();
 }
 
 }  // namespace engine

diff --git a/engine/db/table_segment_mvp.hpp b/engine/db/table_segment_mvp.hpp
@@ -57,6 +57,7 @@ class TableSegmentMVP {
 
   void Debug(meta::TableSchema& table_schema);
 
+  Status Release();
   ~TableSegmentMVP();
 
   std::atomic<bool> skip_sync_disk_;                                              // For default DB, skip sync to disk.

diff --git a/engine/server/web_server/web_controller.hpp b/engine/server/web_server/web_controller.hpp
@@ -160,6 +160,24 @@ class WebController : public oatpp::web::server::api::ApiController {
     return createDtoResponse(Status::CODE_200, dto);
   }
 
+  // Release DB memory for dbfactory.
+  ADD_CORS(ReleaseDB)
+
+  ENDPOINT("POST", "api/{db_name}/release", ReleaseDB, PATH(String, db_name, "db_name")) {
+    vectordb::Status status = db_server->ReleaseDB(db_name);
+
+    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 = "Release " + db_name + " successfully.";
+    return createDtoResponse(Status::CODE_200, dto);
+  }
+
   ADD_CORS(DropDB)
 
   ENDPOINT("DELETE", "api/{db_name}/drop", DropDB, PATH(String, db_name, "db_name")) {