latch.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

#ifndef UTILS_TIMER_H
#define UTILS_TIMER_H

#include <condition_variable>
#include <mutex>

#include <cassert>

/// latch
///
/// A simple one-shot latch based on a mutex and condition variable.
///
/// The latch is initialized with a number N of threads to wait for. Different
/// threads can call arrive_and_wait() on the latch. Those threads enter a wait
/// state until the Nth thread calls arrive_and_wait(), which then releases all
/// threads at the same time.
struct latch {
  explicit latch(std::size_t cnt) : m_cnt{cnt} {}
  ~latch() {
    assert(m_cnt == 0);
  }

  void arrive_and_wait() {
    std::unique_lock<std::mutex> lk(m_mtx);
    m_cnt -= 1;

    if (m_cnt != 0) {
      m_cv.wait(lk, [this]() { return m_cnt == 0; });
    } else {
      m_cv.notify_all();
    }
  }

 private:
  std::size_t m_cnt;

  std::mutex m_mtx;
  std::condition_variable m_cv;
};

#endif