71 lines
1.4 KiB
C++
71 lines
1.4 KiB
C++
//
|
|
// Created by zzy on 2022/9/6.
|
|
//
|
|
|
|
#pragma once
|
|
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <condition_variable>
|
|
#include <cstdint>
|
|
#include <mutex>
|
|
|
|
class THD;
|
|
|
|
namespace polarx_rpc {
|
|
|
|
class CflowControl {
|
|
public:
|
|
inline void set_thd(THD *thd) { thd_ = thd; }
|
|
|
|
inline bool flow_consume(int32_t token) {
|
|
auto before = flow_counter_.fetch_sub(token);
|
|
auto after = before - token;
|
|
return after > 0;
|
|
}
|
|
|
|
inline bool flow_wait() {
|
|
bool exit;
|
|
wait_begin();
|
|
{
|
|
std::unique_lock<std::mutex> lck(mutex_);
|
|
while (!(exit = exit_.load(std::memory_order_acquire)) &&
|
|
flow_counter_.load(std::memory_order_acquire) <= 0)
|
|
/// use 1s timeout to prevent missing notify
|
|
cv_.wait_for(lck, std::chrono::seconds(1));
|
|
}
|
|
wait_end();
|
|
return !exit;
|
|
}
|
|
|
|
inline void flow_reset(int32_t token) {
|
|
flow_counter_.store(token, std::memory_order_release);
|
|
}
|
|
|
|
inline void flow_offer(int32_t token) {
|
|
std::lock_guard<std::mutex> lck(mutex_);
|
|
flow_reset(token);
|
|
cv_.notify_one(); /// only one wait thread
|
|
}
|
|
|
|
inline void exit() {
|
|
std::lock_guard<std::mutex> lck(mutex_);
|
|
exit_.store(true, std::memory_order_release);
|
|
cv_.notify_all();
|
|
}
|
|
|
|
private:
|
|
THD *thd_{nullptr};
|
|
|
|
std::atomic<bool> exit_{false};
|
|
|
|
std::mutex mutex_;
|
|
std::condition_variable cv_;
|
|
std::atomic<int64_t> flow_counter_{0};
|
|
|
|
void wait_begin();
|
|
void wait_end();
|
|
};
|
|
|
|
} // namespace polarx_rpc
|