From e738ee1068a63ae5c107a1bb3ecd3f0255562a9c Mon Sep 17 00:00:00 2001
From: Xinlong Chen <xl_chen@zju.edu.cn>
Date: Sat, 21 Oct 2023 09:52:27 +0000
Subject: [PATCH] [refactor] CurveBS: add some common functions

Signed-off-by: Xinlong Chen <xl_chen@zju.edu.cn>
---
 .../concurrent_apply/concurrent_apply.cpp     | 32 ++++++++++++++++---
 .../concurrent_apply/concurrent_apply.h       |  3 +-
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/chunkserver/concurrent_apply/concurrent_apply.cpp b/src/chunkserver/concurrent_apply/concurrent_apply.cpp
index dd7b219502..48d7130681 100644
--- a/src/chunkserver/concurrent_apply/concurrent_apply.cpp
+++ b/src/chunkserver/concurrent_apply/concurrent_apply.cpp
@@ -124,8 +124,11 @@ void ConcurrentApplyModule::Run(ApplyTaskType type, int index) {
 }
 
 void ConcurrentApplyModule::Stop() {
+    if (!start_.exchange(false)) {
+        return;
+    }
+
     LOG(INFO) << "stop ConcurrentApplyModule...";
-    start_ = false;
     auto wakeup = []() {};
     for (auto iter : rapplyMap_) {
         iter.second->tq.Push(wakeup);
@@ -145,15 +148,36 @@ void ConcurrentApplyModule::Stop() {
 }
 
 void ConcurrentApplyModule::Flush() {
+    if (!start_.load(std::memory_order_relaxed)) {
+        return;
+    }
+
     CountDownEvent event(wconcurrentsize_);
-    auto flushtask = [&event]() {
-        event.Signal();
-    };
+    auto flushtask = [&event]() { event.Signal(); };
+
+    for (int i = 0; i < wconcurrentsize_; i++) {
+        wapplyMap_[i]->tq.Push(flushtask);
+    }
+
+    event.Wait();
+}
+
+void ConcurrentApplyModule::FlushAll() {
+    if (!start_.load(std::memory_order_relaxed)) {
+        return;
+    }
+
+    CountDownEvent event(wconcurrentsize_ + rconcurrentsize_);
+    auto flushtask = [&event]() { event.Signal(); };
 
     for (int i = 0; i < wconcurrentsize_; i++) {
         wapplyMap_[i]->tq.Push(flushtask);
     }
 
+    for (int i = 0; i < rconcurrentsize_; i++) {
+        rapplyMap_[i]->tq.Push(flushtask);
+    }
+
     event.Wait();
 }
 
diff --git a/src/chunkserver/concurrent_apply/concurrent_apply.h b/src/chunkserver/concurrent_apply/concurrent_apply.h
index a4e9712348..60b10a3e55 100644
--- a/src/chunkserver/concurrent_apply/concurrent_apply.h
+++ b/src/chunkserver/concurrent_apply/concurrent_apply.h
@@ -107,6 +107,7 @@ class CURVE_CACHELINE_ALIGNMENT ConcurrentApplyModule {
      * Flush: finish all task in write threads
      */
     void Flush();
+    void FlushAll();
 
     void Stop();
 
@@ -128,7 +129,7 @@ class CURVE_CACHELINE_ALIGNMENT ConcurrentApplyModule {
         explicit TaskThread(size_t capacity) : tq(capacity) {}
     };
 
-    bool start_;
+    std::atomic<bool> start_;
     int rconcurrentsize_;
     int rqueuedepth_;
     int wconcurrentsize_;