添加 CountDownLatch

添加 CountDownLatch, 批量等待多个Coroutine结束
yyzybb537 · Sep 1, 2021 · 4737dcd · 4737dcd
1 parent 4780002
commit 4737dcd
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/libgo/coroutine.h b/libgo/coroutine.h
@@ -1,5 +1,6 @@
 #pragma once
 #define __const__
+#include <atomic>
 #include "common/config.h"
 #include "common/pp.h"
 #include "common/syntax_helper.h"
@@ -71,3 +72,29 @@ typedef ::co::CoTimer::TimerId co_timer_id;
 #define co_last_defer() ::co::GetLastDefer()
 #define co_defer_scope co_defer [&]
 
+class CountDownLatch {
+public:
+    explicit CountDownLatch(size_t n = 1) : mFlyingCount(n) {}
+
+    void Done() {
+        --mFlyingCount;
+    }
+    void Wait() {
+        uint64_t usec = 100 * 1000;
+        while (mFlyingCount) {
+            if (co_sched.TaskCount() <= co_sched.ProcessCount()) {
+                usleep(usec);
+                usec <<= 1;
+            } else {
+                co_yield;
+            }
+        }
+    }
+private:
+    std::atomic<size_t> mFlyingCount;
+
+    CountDownLatch(CountDownLatch const &) = delete;
+    CountDownLatch(CountDownLatch &&) = delete;
+    CountDownLatch& operator=(CountDownLatch const &) = delete;
+    CountDownLatch& operator=(CountDownLatch &&) = delete;
+};
diff --git a/libgo/scheduler/scheduler.h b/libgo/scheduler/scheduler.h
@@ -56,6 +56,9 @@ class Scheduler
     // 使用独立的定时器线程
     void UseAloneTimerThread();
 
+    // 当前调度器中的线程数量
+    size_t ProcessCount() { return processers_.size(); };
+
     // 当前调度器中的协程数量
     uint32_t TaskCount();