From c8a59b273c90f31012bac02c7df6b131f8015488 Mon Sep 17 00:00:00 2001
From: groot <yhmo@zeronedata.com>
Date: Fri, 10 Apr 2020 22:51:01 +0800
Subject: [PATCH] #1276 SQLite throw exception when create 50000+ partitions in
 a table (#1879)

* #1276 SQLite throw exception when create 50000+ partitions in a table

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

* #1276 SQLite throw exception when create 50000+ partitions in a table

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

* remove some log

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

* remove some log

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

* typo

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

* typo

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

* fix unittest

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

Co-authored-by: JinHai-CN <hai.jin@zilliz.com>
---
 CHANGELOG.md                                           |  3 ++-
 core/src/db/DBImpl.cpp                                 | 10 ++++++----
 .../server/delivery/request/CreatePartitionRequest.cpp | 10 ++++++++++
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9c9dd3df81..bc13527646 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,9 +5,10 @@ Please mark all change in change log and use the issue from GitHub
 # Milvus 0.8.0 (TBD)
 
 ## Bug
+-   \#1276 SQLite throw exception after create 50000+ partitions in a table
 -   \#1762 Server is not forbidden to create new partition which tag is `_default`
 -   \#1873 Fix index file serialize to incorrect path
--   \#1881 Fix Annoy index search fail
+-   \#1881 Fix Annoy index search failure
 
 ## Feature
 -   \#261  Integrate ANNOY into Milvus
diff --git a/core/src/db/DBImpl.cpp b/core/src/db/DBImpl.cpp
index e361ce3024..e2e7427743 100644
--- a/core/src/db/DBImpl.cpp
+++ b/core/src/db/DBImpl.cpp
@@ -171,6 +171,8 @@ DBImpl::Stop() {
             bg_flush_thread_.join();
         }
 
+        WaitMergeFileFinish();
+
         swn_index_.Notify();
         bg_index_thread_.join();
 
@@ -1262,22 +1264,22 @@ DBImpl::BackgroundIndexThread() {
 
 void
 DBImpl::WaitMergeFileFinish() {
-    ENGINE_LOG_DEBUG << "Begin WaitMergeFileFinish";
+    //    ENGINE_LOG_DEBUG << "Begin WaitMergeFileFinish";
     std::lock_guard<std::mutex> lck(merge_result_mutex_);
     for (auto& iter : merge_thread_results_) {
         iter.wait();
     }
-    ENGINE_LOG_DEBUG << "End WaitMergeFileFinish";
+    //    ENGINE_LOG_DEBUG << "End WaitMergeFileFinish";
 }
 
 void
 DBImpl::WaitBuildIndexFinish() {
-    ENGINE_LOG_DEBUG << "Begin WaitBuildIndexFinish";
+    //    ENGINE_LOG_DEBUG << "Begin WaitBuildIndexFinish";
     std::lock_guard<std::mutex> lck(index_result_mutex_);
     for (auto& iter : index_thread_results_) {
         iter.wait();
     }
-    ENGINE_LOG_DEBUG << "End WaitBuildIndexFinish";
+    //    ENGINE_LOG_DEBUG << "End WaitBuildIndexFinish";
 }
 
 void
diff --git a/core/src/server/delivery/request/CreatePartitionRequest.cpp b/core/src/server/delivery/request/CreatePartitionRequest.cpp
index 048440a31d..e33c71a8d5 100644
--- a/core/src/server/delivery/request/CreatePartitionRequest.cpp
+++ b/core/src/server/delivery/request/CreatePartitionRequest.cpp
@@ -18,10 +18,13 @@
 #include <fiu-local.h>
 #include <memory>
 #include <string>
+#include <vector>
 
 namespace milvus {
 namespace server {
 
+constexpr uint64_t MAX_PARTITION_LIMIT = 5000;
+
 CreatePartitionRequest::CreatePartitionRequest(const std::shared_ptr<milvus::server::Context>& context,
                                                const std::string& collection_name, const std::string& tag)
     : BaseRequest(context, BaseRequest::kCreatePartition), collection_name_(collection_name), tag_(tag) {
@@ -76,6 +79,13 @@ CreatePartitionRequest::OnExecute() {
             }
         }
 
+        // check partition total count
+        std::vector<engine::meta::CollectionSchema> schema_array;
+        status = DBWrapper::DB()->ShowPartitions(collection_name_, schema_array);
+        if (schema_array.size() >= MAX_PARTITION_LIMIT) {
+            return Status(SERVER_UNSUPPORTED_ERROR, "The number of partitions exceeds the upper limit(5000)");
+        }
+
         rc.RecordSection("check validation");
 
         // step 2: create partition