From d3cbc1ccad981e06f1f70c0698080f0932c7d86a Mon Sep 17 00:00:00 2001
From: Xu Peng <xupeng3112@163.com>
Date: Sun, 5 May 2019 17:26:50 +0800
Subject: [PATCH] fix(db): db desctruction bug fix

Former-commit-id: e456a0d76c2b1401107d47600c34d4622ca8cd8a
---
 cpp/src/db/DBImpl.cpp |  1 +
 cpp/src/db/Env.cpp    | 35 ++++++++++++++++++++++++++++++++---
 cpp/src/db/Env.h      |  4 ++++
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/cpp/src/db/DBImpl.cpp b/cpp/src/db/DBImpl.cpp
index e4144cd50f..f85e018781 100644
--- a/cpp/src/db/DBImpl.cpp
+++ b/cpp/src/db/DBImpl.cpp
@@ -400,6 +400,7 @@ DBImpl<EngineT>::~DBImpl() {
     }
     std::vector<std::string> ids;
     _pMemMgr->serialize(ids);
+    _env->Stop();
 }
 
 /*
diff --git a/cpp/src/db/Env.cpp b/cpp/src/db/Env.cpp
index 4d4613f43e..f3a4b64810 100644
--- a/cpp/src/db/Env.cpp
+++ b/cpp/src/db/Env.cpp
@@ -3,18 +3,23 @@
  * Unauthorized copying of this file, via any medium is strictly prohibited.
  * Proprietary and confidential.
  ******************************************************************************/
+#include <easylogging++.h>
 #include <assert.h>
+#include <atomic>
 #include "Env.h"
 
 namespace zilliz {
 namespace vecwise {
 namespace engine {
 
-Env::Env() : _bg_work_started(false) {
+Env::Env()
+    : _bg_work_started(false),
+      _shutting_down(false) {
 }
 
 void Env::schedule(void (*function_)(void* arg_), void* arg_) {
     std::unique_lock<std::mutex> lock(_bg_work_mutex);
+    if (_shutting_down) return;
 
     if (!_bg_work_started) {
         _bg_work_started = true;
@@ -30,12 +35,14 @@ void Env::schedule(void (*function_)(void* arg_), void* arg_) {
 }
 
 void Env::backgroud_thread_main() {
-    while (true) {
+    while (!_shutting_down) {
         std::unique_lock<std::mutex> lock(_bg_work_mutex);
-        while (_bg_work_queue.empty()) {
+        while (_bg_work_queue.empty() && !_shutting_down) {
             _bg_work_cv.wait(lock);
         }
 
+        if (_shutting_down) break;
+
         assert(!_bg_work_queue.empty());
         auto bg_function = _bg_work_queue.front()._function;
         void* bg_arg = _bg_work_queue.front()._arg;
@@ -44,6 +51,28 @@ void Env::backgroud_thread_main() {
         lock.unlock();
         bg_function(bg_arg);
     }
+
+    std::unique_lock<std::mutex> lock(_bg_work_mutex);
+    _bg_work_started = false;
+    _bg_work_cv.notify_all();
+}
+
+void Env::Stop() {
+    {
+        std::unique_lock<std::mutex> lock(_bg_work_mutex);
+        if (_shutting_down || !_bg_work_started) return;
+    }
+    _shutting_down = true;
+    {
+        std::unique_lock<std::mutex> lock(_bg_work_mutex);
+        if (_bg_work_queue.empty()) {
+            _bg_work_cv.notify_one();
+        }
+        while (_bg_work_started) {
+            _bg_work_cv.wait(lock);
+        }
+    }
+    _shutting_down = false;
 }
 
 Env::~Env() {}
diff --git a/cpp/src/db/Env.h b/cpp/src/db/Env.h
index 648644a948..9eed68cddb 100644
--- a/cpp/src/db/Env.h
+++ b/cpp/src/db/Env.h
@@ -9,6 +9,7 @@
 #include <thread>
 #include <mutex>
 #include <queue>
+#include <atomic>
 
 namespace zilliz {
 namespace vecwise {
@@ -23,6 +24,8 @@ public:
 
     void schedule(void (*function_)(void* arg_), void* arg_);
 
+    virtual void Stop();
+
     virtual ~Env();
 
     static Env* Default();
@@ -45,6 +48,7 @@ protected:
     std::condition_variable _bg_work_cv;
     std::queue<BGWork> _bg_work_queue;
     bool _bg_work_started;
+    std::atomic<bool> _shutting_down;
 
 }; // Env