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::~DBImpl() { } std::vector 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 #include +#include #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 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 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 lock(_bg_work_mutex); + _bg_work_started = false; + _bg_work_cv.notify_all(); +} + +void Env::Stop() { + { + std::unique_lock lock(_bg_work_mutex); + if (_shutting_down || !_bg_work_started) return; + } + _shutting_down = true; + { + std::unique_lock 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 #include #include +#include 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 _bg_work_queue; bool _bg_work_started; + std::atomic _shutting_down; }; // Env