use std::fmt::Debug;
use std::future::Future;
use std::hash::Hash;
use std::marker::Unpin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::Poll;
use futures::{Sink, SinkExt, Stream, StreamExt};
use futures::task::AtomicWaker;
use parking_lot::Mutex;
use parking_lot::RwLock;
#[cfg(feature = "rate")]
use update_rate::{DiscreteRateCounter, RateCounter};
use queue_ext::{Action, QueueExt, Reply};
use super::{
assert_future, close::Close, Counter, Error, ErrorType, flush::Flush,
GroupTaskQueue, IndexSet, local_builder::SyncSender, LocalSpawner, PendingOnce,
};
type DashMap<K, V> = dashmap::DashMap<K, V, ahash::RandomState>;
type GroupChannels<G> = Arc<DashMap<G, Arc<Mutex<GroupTaskQueue<LocalTaskType>>>>>;
pub type LocalTaskType = Box<dyn std::future::Future<Output=()> + 'static + Unpin>;
pub struct LocalExecutor<Tx = SyncSender, G = (), D = ()> {
pub(crate) tx: Tx,
workers: usize,
queue_max: isize,
active_count: Counter,
pub(crate) waiting_count: Counter,
completed_count: Counter,
#[cfg(feature = "rate")]
rate_counter: Arc<RwLock<DiscreteRateCounter>>,
flush_waker: Arc<AtomicWaker>,
is_flushing: Arc<AtomicBool>,
is_closed: Arc<AtomicBool>,
group_channels: GroupChannels<G>,
_d: std::marker::PhantomData<D>,
}
impl<Tx, G, D> Clone for LocalExecutor<Tx, G, D>
where
Tx: Clone,
{
#[inline]
fn clone(&self) -> Self {
Self {
tx: self.tx.clone(),
workers: self.workers,
queue_max: self.queue_max,
active_count: self.active_count.clone(),
waiting_count: self.waiting_count.clone(),
completed_count: self.completed_count.clone(),
#[cfg(feature = "rate")]
rate_counter: self.rate_counter.clone(),
flush_waker: self.flush_waker.clone(),
is_flushing: self.is_flushing.clone(),
is_closed: self.is_closed.clone(),
group_channels: self.group_channels.clone(),
_d: std::marker::PhantomData,
}
}
}
impl<Tx, G, D> LocalExecutor<Tx, G, D>
where
Tx: Clone + Sink<(D, LocalTaskType)> + Unpin + Sync + 'static,
G: Hash + Eq + Clone + Debug + Sync + 'static,
{
#[inline]
pub(crate) fn with_channel<Rx>(
workers: usize,
queue_max: usize,
tx: Tx,
rx: Rx,
) -> (Self, impl Future<Output=()>)
where
Rx: Stream<Item=(D, LocalTaskType)> + Unpin,
{
let exec = Self {
tx,
workers,
queue_max: queue_max as isize,
active_count: Counter::new(),
waiting_count: Counter::new(),
completed_count: Counter::new(),
#[cfg(feature = "rate")]
rate_counter: Arc::new(RwLock::new(DiscreteRateCounter::new(100))),
flush_waker: Arc::new(AtomicWaker::new()),
is_flushing: Arc::new(AtomicBool::new(false)),
is_closed: Arc::new(AtomicBool::new(false)),
group_channels: Arc::new(DashMap::default()),
_d: std::marker::PhantomData,
};
let runner = exec.clone().run(rx);
(exec, runner)
}
#[inline]
pub fn spawn_with<T>(&mut self, msg: T, name: D) -> LocalSpawner<'_, T, Tx, G, D>
where
D: Clone,
T: Future + 'static,
T::Output: 'static,
{
let fut = LocalSpawner::new(self, msg, name);
assert_future::<Result<(), _>, _>(fut)
}
#[inline]
pub fn flush(&self) -> Flush<Tx, D> {
self.is_flushing.store(true, Ordering::SeqCst);
Flush::new(
self.tx.clone(),
self.waiting_count.clone(),
self.active_count.clone(),
self.is_flushing.clone(),
self.flush_waker.clone(),
)
}
#[inline]
pub fn close(&self) -> Close<Tx, D> {
self.is_flushing.store(true, Ordering::SeqCst);
self.is_closed.store(true, Ordering::SeqCst);
Close::new(
self.tx.clone(),
self.waiting_count.clone(),
self.active_count.clone(),
self.is_flushing.clone(),
self.flush_waker.clone(),
)
}
#[inline]
pub fn workers(&self) -> usize {
self.workers
}
#[inline]
pub fn active_count(&self) -> isize {
self.active_count.value()
}
#[inline]
pub fn waiting_count(&self) -> isize {
self.waiting_count.value()
}
#[inline]
pub fn completed_count(&self) -> isize {
self.completed_count.value()
}
#[inline]
#[cfg(feature = "rate")]
pub fn rate(&self) -> f64 {
self.rate_counter.read().rate()
}
#[inline]
pub fn is_full(&self) -> bool {
self.waiting_count() >= self.queue_max
}
#[inline]
pub fn is_closed(&self) -> bool {
self.is_closed.load(Ordering::SeqCst)
}
#[inline]
pub fn is_flushing(&self) -> bool {
self.is_flushing.load(Ordering::SeqCst)
}
async fn run<Rx>(self, mut task_rx: Rx)
where
Rx: Stream<Item=(D, LocalTaskType)> + Unpin,
{
let exec = self;
let idle_waker = Arc::new(AtomicWaker::new());
let channel = || {
let rx = OneValue::new().queue_stream(|s, _| match s.take() {
None => Poll::Pending,
Some(m) => Poll::Ready(Some(m)),
});
let tx = rx.clone().queue_sender(|s, action| match action {
Action::Send(item) => Reply::Send(s.set(item)),
Action::IsFull => Reply::IsFull(s.is_full()),
Action::IsEmpty => Reply::IsEmpty(s.is_empty()),
});
(tx, rx)
};
let idle_idxs = IndexSet::new();
let mut txs = Vec::new();
let mut rxs = Vec::new();
for i in 0..exec.workers {
let (tx, mut rx) = channel();
let idle_waker = idle_waker.clone();
let idle_idxs = idle_idxs.clone();
idle_idxs.insert(i);
let exec = exec.clone();
let rx_fut = async move {
loop {
match rx.next().await {
Some(task) => {
exec.active_count.inc();
task.await;
exec.completed_count.inc();
exec.active_count.dec();
#[cfg(feature = "rate")]
exec.rate_counter.write().update();
}
None => break,
}
if !rx.is_full() {
idle_idxs.insert(i);
idle_waker.wake();
}
if exec.is_flushing() && rx.is_empty() {
exec.flush_waker.wake();
}
}
};
txs.push(tx);
rxs.push(rx_fut);
}
let tasks_bus = async move {
while let Some((_, task)) = task_rx.next().await {
loop {
if idle_idxs.is_empty() {
PendingOnce::new(idle_waker.clone()).await;
} else if let Some(idx) = idle_idxs.pop() {
if let Some(tx) = txs.get_mut(idx) {
if let Err(_t) = tx.send(task).await {
log::error!("send error ...");
}
}
break;
};
}
}
};
futures::future::join(tasks_bus, futures::future::join_all(rxs)).await;
log::info!("exit local task executor");
}
}
impl<Tx, G> LocalExecutor<Tx, G, ()>
where
Tx: Clone + Sink<((), LocalTaskType)> + Unpin + Sync + 'static,
G: Hash + Eq + Clone + Debug + Sync + 'static,
{
#[inline]
pub fn spawn<T>(&mut self, msg: T) -> LocalSpawner<'_, T, Tx, G, ()>
where
T: Future + 'static,
T::Output: 'static,
{
let fut = LocalSpawner::new(self, msg, ());
assert_future::<Result<(), _>, _>(fut)
}
#[inline]
pub(crate) async fn group_send(
&self,
name: G,
task: LocalTaskType,
) -> Result<(), Error<LocalTaskType>> {
if self.is_closed() {
return Err(Error::SendError(ErrorType::Closed(Some(task))));
}
let gt_queue = self
.group_channels
.entry(name.clone())
.or_insert_with(|| Arc::new(Mutex::new(GroupTaskQueue::new())))
.value()
.clone();
let exec = self.clone();
let group_channels = self.group_channels.clone();
let runner_task = {
let mut task_tx = gt_queue.lock();
if task_tx.is_running() {
task_tx.push(task);
drop(task_tx);
drop(gt_queue);
None
} else {
task_tx.set_running(true);
drop(task_tx);
let task_rx = gt_queue; let runner_task = async move {
exec.active_count.inc();
task.await;
exec.active_count.dec();
loop {
let task: Option<LocalTaskType> = task_rx.lock().pop();
if let Some(task) = task {
exec.active_count.inc();
task.await;
exec.completed_count.inc();
exec.active_count.dec();
} else {
group_channels.remove(&name);
break;
}
}
};
Some(runner_task)
}
};
if let Some(runner_task) = runner_task {
if (self
.tx
.clone()
.send(((), Box::new(Box::pin(runner_task))))
.await)
.is_err()
{
Err(Error::SendError(ErrorType::Closed(None)))
} else {
Ok(())
}
} else {
Ok(())
}
}
}
#[derive(Clone)]
struct OneValue(Arc<RwLock<Option<LocalTaskType>>>);
unsafe impl Sync for OneValue {}
unsafe impl Send for OneValue {}
impl OneValue {
#[inline]
fn new() -> Self {
Self(Arc::new(RwLock::new(None)))
}
#[inline]
fn set(&self, val: LocalTaskType) -> Option<LocalTaskType> {
self.0.write().replace(val)
}
#[inline]
fn take(&self) -> Option<LocalTaskType> {
self.0.write().take()
}
#[inline]
fn is_full(&self) -> bool {
self.0.read().is_some()
}
fn is_empty(&self) -> bool {
self.0.read().is_none()
}
}