339 lines
9.4 KiB
C++
339 lines
9.4 KiB
C++
//
|
|
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2022
|
|
//
|
|
// Distributed under the Boost Software License, Version 1.0. (See accompanying
|
|
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
|
|
//
|
|
#pragma once
|
|
|
|
#include "td/utils/common.h"
|
|
#include "td/utils/logging.h"
|
|
#include "td/utils/port/sleep.h"
|
|
#include "td/utils/port/thread.h"
|
|
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
#include <mutex>
|
|
|
|
namespace td {
|
|
|
|
class MpmcEagerWaiter {
|
|
public:
|
|
struct Slot {
|
|
private:
|
|
friend class MpmcEagerWaiter;
|
|
int yields;
|
|
uint32 worker_id;
|
|
};
|
|
static void init_slot(Slot &slot, uint32 worker_id) {
|
|
slot.yields = 0;
|
|
slot.worker_id = worker_id;
|
|
}
|
|
|
|
void wait(Slot &slot) {
|
|
if (slot.yields < RoundsTillSleepy) {
|
|
yield();
|
|
slot.yields++;
|
|
} else if (slot.yields == RoundsTillSleepy) {
|
|
auto state = state_.load(std::memory_order_relaxed);
|
|
if (!State::has_worker(state)) {
|
|
auto new_state = State::with_worker(state, slot.worker_id);
|
|
if (state_.compare_exchange_strong(state, new_state, std::memory_order_acq_rel)) {
|
|
yield();
|
|
slot.yields++;
|
|
return;
|
|
}
|
|
if (state == State::awake()) {
|
|
slot.yields = 0;
|
|
return;
|
|
}
|
|
}
|
|
yield();
|
|
slot.yields = 0;
|
|
} else if (slot.yields < RoundsTillAsleep) {
|
|
auto state = state_.load(std::memory_order_acquire);
|
|
if (State::still_sleepy(state, slot.worker_id)) {
|
|
yield();
|
|
slot.yields++;
|
|
return;
|
|
}
|
|
slot.yields = 0;
|
|
} else {
|
|
auto state = state_.load(std::memory_order_acquire);
|
|
if (State::still_sleepy(state, slot.worker_id)) {
|
|
std::unique_lock<std::mutex> lock(mutex_);
|
|
if (state_.compare_exchange_strong(state, State::asleep(), std::memory_order_acq_rel)) {
|
|
condition_variable_.wait(lock);
|
|
}
|
|
}
|
|
slot.yields = 0;
|
|
}
|
|
}
|
|
|
|
void stop_wait(Slot &slot) {
|
|
if (slot.yields > RoundsTillSleepy) {
|
|
notify_cold();
|
|
}
|
|
slot.yields = 0;
|
|
}
|
|
|
|
void close() {
|
|
}
|
|
|
|
void notify() {
|
|
std::atomic_thread_fence(std::memory_order_seq_cst);
|
|
if (state_.load(std::memory_order_acquire) == State::awake()) {
|
|
return;
|
|
}
|
|
notify_cold();
|
|
}
|
|
|
|
private:
|
|
struct State {
|
|
static constexpr uint32 awake() {
|
|
return 0;
|
|
}
|
|
static constexpr uint32 asleep() {
|
|
return 1;
|
|
}
|
|
static bool is_asleep(uint32 state) {
|
|
return (state & 1) != 0;
|
|
}
|
|
static bool has_worker(uint32 state) {
|
|
return (state >> 1) != 0;
|
|
}
|
|
static int32 with_worker(uint32 state, uint32 worker) {
|
|
return state | ((worker + 1) << 1);
|
|
}
|
|
static bool still_sleepy(uint32 state, uint32 worker) {
|
|
return (state >> 1) == (worker + 1);
|
|
}
|
|
};
|
|
enum { RoundsTillSleepy = 32, RoundsTillAsleep = 64 };
|
|
// enum { RoundsTillSleepy = 1, RoundsTillAsleep = 2 };
|
|
std::atomic<uint32> state_{State::awake()};
|
|
std::mutex mutex_;
|
|
std::condition_variable condition_variable_;
|
|
|
|
void notify_cold() {
|
|
auto old_state = state_.exchange(State::awake(), std::memory_order_release);
|
|
if (State::is_asleep(old_state)) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
condition_variable_.notify_all();
|
|
}
|
|
}
|
|
static void yield() {
|
|
// whatever, this is better than sched_yield
|
|
usleep_for(1);
|
|
}
|
|
};
|
|
|
|
class MpmcSleepyWaiter {
|
|
public:
|
|
struct Slot {
|
|
private:
|
|
friend class MpmcSleepyWaiter;
|
|
|
|
enum State { Search, Work, Sleep } state_{Work};
|
|
|
|
void park() {
|
|
std::unique_lock<std::mutex> guard(mutex_);
|
|
condition_variable_.wait(guard, [&] { return unpark_flag_; });
|
|
unpark_flag_ = false;
|
|
}
|
|
|
|
bool cancel_park() {
|
|
auto res = unpark_flag_;
|
|
unpark_flag_ = false;
|
|
return res;
|
|
}
|
|
|
|
void unpark() {
|
|
//TODO: try to unlock guard before notify_all
|
|
std::unique_lock<std::mutex> guard(mutex_);
|
|
unpark_flag_ = true;
|
|
condition_variable_.notify_all();
|
|
}
|
|
|
|
std::mutex mutex_;
|
|
std::condition_variable condition_variable_;
|
|
bool unpark_flag_{false}; // TODO: move out of lock
|
|
int yield_cnt{0};
|
|
int32 worker_id{0};
|
|
|
|
public:
|
|
char padding[TD_CONCURRENCY_PAD];
|
|
};
|
|
|
|
// There are a lot of workers
|
|
// Each has a slot
|
|
//
|
|
// States of a worker:
|
|
// - searching for work | Search
|
|
// - processing work | Work
|
|
// - sleeping | Sleep
|
|
//
|
|
// When somebody adds a work it calls notify
|
|
//
|
|
// notify
|
|
// if there are workers in search phase do nothing.
|
|
// if all workers are awake do nothing
|
|
// otherwise wake some random worker
|
|
//
|
|
// Initially all workers are in Search mode.
|
|
//
|
|
// When worker found nothing it may try to call wait.
|
|
// This may put it in a Sleep for some time.
|
|
// After wait returns worker will be in Search state again.
|
|
//
|
|
// Suppose worker found a work and ready to process it.
|
|
// Then it may call stop_wait. This will cause transition from
|
|
// Search to Work state.
|
|
//
|
|
// Main invariant:
|
|
// After notify is called there should be at least on worker in Search or Work state.
|
|
// If possible - in Search state
|
|
//
|
|
|
|
static void init_slot(Slot &slot, int32 worker_id) {
|
|
slot.state_ = Slot::State::Work;
|
|
slot.unpark_flag_ = false;
|
|
slot.worker_id = worker_id;
|
|
VLOG(waiter) << "Init slot " << worker_id;
|
|
}
|
|
|
|
static constexpr int VERBOSITY_NAME(waiter) = VERBOSITY_NAME(DEBUG) + 10;
|
|
void wait(Slot &slot) {
|
|
if (slot.state_ == Slot::State::Work) {
|
|
VLOG(waiter) << "Work -> Search";
|
|
state_++;
|
|
slot.state_ = Slot::State::Search;
|
|
slot.yield_cnt = 0;
|
|
return;
|
|
}
|
|
if (slot.state_ == Slot::Search) {
|
|
if (slot.yield_cnt++ < 10 && false) {
|
|
// TODO some sleep backoff is possible
|
|
return;
|
|
}
|
|
|
|
slot.state_ = Slot::State::Sleep;
|
|
std::unique_lock<std::mutex> guard(sleepers_mutex_);
|
|
auto state_view = StateView(state_.fetch_add((1 << PARKING_SHIFT) - 1));
|
|
CHECK(state_view.searching_count != 0);
|
|
bool should_search = state_view.searching_count == 1;
|
|
if (closed_) {
|
|
return;
|
|
}
|
|
sleepers_.push_back(&slot);
|
|
LOG_CHECK(slot.unpark_flag_ == false) << slot.worker_id;
|
|
VLOG(waiter) << "Add to sleepers " << slot.worker_id;
|
|
//guard.unlock();
|
|
if (should_search) {
|
|
VLOG(waiter) << "Search -> Search once, then Sleep ";
|
|
return;
|
|
}
|
|
VLOG(waiter) << "Search -> Sleep " << state_view.searching_count << " " << state_view.parked_count;
|
|
}
|
|
|
|
CHECK(slot.state_ == Slot::State::Sleep);
|
|
VLOG(waiter) << "Park " << slot.worker_id;
|
|
slot.park();
|
|
VLOG(waiter) << "Resume " << slot.worker_id;
|
|
slot.state_ = Slot::State::Search;
|
|
slot.yield_cnt = 0;
|
|
}
|
|
|
|
void stop_wait(Slot &slot) {
|
|
if (slot.state_ == Slot::State::Work) {
|
|
return;
|
|
}
|
|
if (slot.state_ == Slot::State::Sleep) {
|
|
VLOG(waiter) << "Search once, then Sleep -> Work/Search " << slot.worker_id;
|
|
slot.state_ = Slot::State::Work;
|
|
std::unique_lock<std::mutex> guard(sleepers_mutex_);
|
|
auto it = std::find(sleepers_.begin(), sleepers_.end(), &slot);
|
|
if (it != sleepers_.end()) {
|
|
sleepers_.erase(it);
|
|
VLOG(waiter) << "Remove from sleepers " << slot.worker_id;
|
|
state_.fetch_sub((1 << PARKING_SHIFT) - 1);
|
|
guard.unlock();
|
|
} else {
|
|
guard.unlock();
|
|
VLOG(waiter) << "Not in sleepers" << slot.worker_id;
|
|
CHECK(slot.cancel_park());
|
|
}
|
|
}
|
|
VLOG(waiter) << "Search once, then Sleep -> Work " << slot.worker_id;
|
|
slot.state_ = Slot::State::Search;
|
|
auto state_view = StateView(state_.fetch_sub(1));
|
|
CHECK(state_view.searching_count != 0);
|
|
CHECK(state_view.searching_count < 1000);
|
|
bool should_notify = state_view.searching_count == 1;
|
|
if (should_notify) {
|
|
VLOG(waiter) << "Notify others";
|
|
notify();
|
|
}
|
|
VLOG(waiter) << "Search -> Work ";
|
|
slot.state_ = Slot::State::Work;
|
|
}
|
|
|
|
void notify() {
|
|
auto view = StateView(state_.load());
|
|
//LOG(ERROR) << view.parked_count;
|
|
if (view.searching_count > 0 || view.parked_count == 0) {
|
|
VLOG(waiter) << "Ingore notify: " << view.searching_count << " " << view.parked_count;
|
|
return;
|
|
}
|
|
|
|
VLOG(waiter) << "Notify: " << view.searching_count << " " << view.parked_count;
|
|
std::unique_lock<std::mutex> guard(sleepers_mutex_);
|
|
|
|
view = StateView(state_.load());
|
|
if (view.searching_count > 0) {
|
|
VLOG(waiter) << "Skip notify: got searching";
|
|
return;
|
|
}
|
|
|
|
CHECK(view.parked_count == static_cast<int>(sleepers_.size()));
|
|
if (sleepers_.empty()) {
|
|
VLOG(waiter) << "Skip notify: no sleepers";
|
|
return;
|
|
}
|
|
|
|
auto sleeper = sleepers_.back();
|
|
sleepers_.pop_back();
|
|
state_.fetch_sub((1 << PARKING_SHIFT) - 1);
|
|
VLOG(waiter) << "Unpark " << sleeper->worker_id;
|
|
sleeper->unpark();
|
|
}
|
|
|
|
void close() {
|
|
StateView state(state_.load());
|
|
LOG_CHECK(state.parked_count == 0) << state.parked_count;
|
|
LOG_CHECK(state.searching_count == 0) << state.searching_count;
|
|
}
|
|
|
|
private:
|
|
static constexpr int32 PARKING_SHIFT = 16;
|
|
struct StateView {
|
|
int32 parked_count;
|
|
int32 searching_count;
|
|
explicit StateView(int32 x) {
|
|
parked_count = x >> PARKING_SHIFT;
|
|
searching_count = x & ((1 << PARKING_SHIFT) - 1);
|
|
}
|
|
};
|
|
std::atomic<int32> state_{0};
|
|
|
|
std::mutex sleepers_mutex_;
|
|
vector<Slot *> sleepers_;
|
|
|
|
bool closed_ = false;
|
|
};
|
|
|
|
using MpmcWaiter = MpmcSleepyWaiter;
|
|
|
|
} // namespace td
|