aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/lib/core/blocking_counter.h
blob: f141be2c76cd0425fd9a276e6e6995d983adfda2 (plain)
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 TENSORFLOW_LIB_CORE_BLOCKING_COUNTER_H_
#define TENSORFLOW_LIB_CORE_BLOCKING_COUNTER_H_

#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/port.h"

namespace tensorflow {

class BlockingCounter {
 public:
  BlockingCounter(int initial_count) : count_(initial_count) {
    CHECK_GE(count_, 0);
  }

  ~BlockingCounter() {}

  inline void DecrementCount() {
    mutex_lock l(mu_);
    --count_;
    CHECK(count_ >= 0);
    if (count_ == 0) {
      cond_var_.notify_all();
    }
  }

  inline void Wait() {
    mutex_lock l(mu_);
    while (count_ > 0) {
      cond_var_.wait(l);
    }
  }

 private:
  int count_;
  mutex mu_;
  condition_variable cond_var_;
};

}  // namespace tensorflow

#endif  // TENSORFLOW_LIB_CORE_BLOCKING_COUNTER_H_