Use async select rather than flume::Selector
This commit is contained in:
parent
8fc9100b41
commit
2884d854ff
45
src/lib.rs
45
src/lib.rs
|
@ -1,11 +1,10 @@
|
|||
mod selector;
|
||||
|
||||
use std::{
|
||||
num::{NonZeroU16, NonZeroUsize},
|
||||
sync::{
|
||||
atomic::{AtomicBool, AtomicU64},
|
||||
Arc, Mutex,
|
||||
},
|
||||
sync::{atomic::AtomicU64, Arc, Mutex},
|
||||
thread::JoinHandle,
|
||||
time::{Duration, Instant},
|
||||
time::Instant,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
|
@ -178,9 +177,7 @@ impl CpuPool {
|
|||
let mut threads = state.take_threads();
|
||||
|
||||
for thread in &mut threads {
|
||||
thread
|
||||
.signal
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
thread.signal.take();
|
||||
}
|
||||
|
||||
for thread in &mut threads {
|
||||
|
@ -373,9 +370,7 @@ impl ThreadVec {
|
|||
impl Drop for ThreadVec {
|
||||
fn drop(&mut self) {
|
||||
for thread in &mut self.threads {
|
||||
thread
|
||||
.signal
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
thread.signal.take();
|
||||
}
|
||||
|
||||
for thread in &mut self.threads {
|
||||
|
@ -388,14 +383,13 @@ impl Drop for ThreadVec {
|
|||
|
||||
struct Thread {
|
||||
handle: Option<JoinHandle<()>>,
|
||||
signal: Arc<AtomicBool>,
|
||||
signal: Option<SendOnDrop>,
|
||||
closed: flume::Receiver<()>,
|
||||
}
|
||||
|
||||
impl Thread {
|
||||
async fn reap(mut self) {
|
||||
self.signal
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
self.signal.take();
|
||||
|
||||
let _ = self.closed.recv_async().await;
|
||||
|
||||
|
@ -417,19 +411,19 @@ impl Drop for SendOnDrop {
|
|||
|
||||
fn spawn(name: &'static str, id: u64, receiver: flume::Receiver<SendFn>) -> Thread {
|
||||
let (closed_tx, closed) = flume::bounded(1);
|
||||
let (signal, signal_rx) = flume::bounded(1);
|
||||
|
||||
let signal = Arc::new(AtomicBool::new(false));
|
||||
let signal = SendOnDrop { sender: signal };
|
||||
let closed_tx = SendOnDrop { sender: closed_tx };
|
||||
|
||||
let signal2 = signal.clone();
|
||||
let handle = std::thread::Builder::new()
|
||||
.name(format!("{name}-{id}"))
|
||||
.spawn(move || run(name, id, receiver, signal2, closed_tx))
|
||||
.spawn(move || run(name, id, receiver, signal_rx, closed_tx))
|
||||
.expect("Failed to spawn new thread");
|
||||
|
||||
Thread {
|
||||
handle: Some(handle),
|
||||
signal,
|
||||
signal: Some(signal),
|
||||
closed,
|
||||
}
|
||||
}
|
||||
|
@ -471,22 +465,15 @@ fn run(
|
|||
name: &'static str,
|
||||
id: u64,
|
||||
receiver: flume::Receiver<SendFn>,
|
||||
signal: Arc<AtomicBool>,
|
||||
signal: flume::Receiver<()>,
|
||||
closed_tx: SendOnDrop,
|
||||
) {
|
||||
let guard = MetricsGuard::guard(name, id);
|
||||
|
||||
loop {
|
||||
match receiver.recv_timeout(Duration::from_millis(50)) {
|
||||
Ok(send_fn) => {
|
||||
invoke_send_fn(name, id, send_fn);
|
||||
}
|
||||
Err(flume::RecvTimeoutError::Disconnected) => break,
|
||||
Err(flume::RecvTimeoutError::Timeout) => {
|
||||
if signal.load(std::sync::atomic::Ordering::Acquire) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
match selector::blocking_select(signal.recv_async(), receiver.recv_async()) {
|
||||
selector::Either::Left(_) | selector::Either::Right(Err(_)) => break,
|
||||
selector::Either::Right(Ok(send_fn)) => invoke_send_fn(name, id, send_fn),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
146
src/selector.rs
Normal file
146
src/selector.rs
Normal file
|
@ -0,0 +1,146 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
pin::Pin,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
},
|
||||
task::{Context, Poll, Wake, Waker},
|
||||
};
|
||||
|
||||
struct ThreadWaker {
|
||||
thread: std::thread::Thread,
|
||||
}
|
||||
|
||||
impl Wake for ThreadWaker {
|
||||
fn wake(self: Arc<Self>) {
|
||||
self.thread.unpark();
|
||||
}
|
||||
|
||||
fn wake_by_ref(self: &Arc<Self>) {
|
||||
self.thread.unpark();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) enum Either<L, R> {
|
||||
Left(L),
|
||||
Right(R),
|
||||
}
|
||||
|
||||
struct Select<F1, F2> {
|
||||
left: F1,
|
||||
left_woken: Arc<AtomicBool>,
|
||||
|
||||
right: F2,
|
||||
right_woken: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
struct SelectWaker {
|
||||
inner: Waker,
|
||||
flag: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl Wake for SelectWaker {
|
||||
fn wake_by_ref(self: &Arc<Self>) {
|
||||
self.flag.store(true, Ordering::Release);
|
||||
|
||||
self.inner.wake_by_ref();
|
||||
}
|
||||
|
||||
fn wake(self: Arc<Self>) {
|
||||
self.flag.store(true, Ordering::Release);
|
||||
|
||||
match Arc::try_unwrap(self) {
|
||||
Ok(this) => this.inner.wake(),
|
||||
Err(this) => this.inner.wake_by_ref(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<F1, F2> Future for Select<F1, F2>
|
||||
where
|
||||
F1: Future + Unpin,
|
||||
F2: Future + Unpin,
|
||||
{
|
||||
type Output = Either<F1::Output, F2::Output>;
|
||||
|
||||
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
let left_waker = Arc::new(SelectWaker {
|
||||
inner: cx.waker().clone(),
|
||||
flag: self.left_woken.clone(),
|
||||
})
|
||||
.into();
|
||||
|
||||
let mut left_ctx = Context::from_waker(&left_waker);
|
||||
|
||||
if let Poll::Ready(left_out) = Pin::new(&mut self.left).poll(&mut left_ctx) {
|
||||
return Poll::Ready(Either::Left(left_out));
|
||||
}
|
||||
|
||||
let right_waker = Arc::new(SelectWaker {
|
||||
inner: cx.waker().clone(),
|
||||
flag: self.right_woken.clone(),
|
||||
})
|
||||
.into();
|
||||
|
||||
let mut right_ctx = Context::from_waker(&right_waker);
|
||||
|
||||
if let Poll::Ready(right_out) = Pin::new(&mut self.right).poll(&mut right_ctx) {
|
||||
return Poll::Ready(Either::Right(right_out));
|
||||
}
|
||||
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn blocking_select<Left, Right>(
|
||||
left: Left,
|
||||
right: Right,
|
||||
) -> Either<Left::Output, Right::Output>
|
||||
where
|
||||
Left: Future,
|
||||
Right: Future,
|
||||
{
|
||||
block_on(select(left, right))
|
||||
}
|
||||
|
||||
fn block_on<F>(fut: F) -> F::Output
|
||||
where
|
||||
F: Future,
|
||||
{
|
||||
let thread_waker = Arc::new(ThreadWaker {
|
||||
thread: std::thread::current(),
|
||||
})
|
||||
.into();
|
||||
|
||||
let mut ctx = Context::from_waker(&thread_waker);
|
||||
|
||||
let mut fut = std::pin::pin!(fut);
|
||||
|
||||
loop {
|
||||
if let Poll::Ready(out) = fut.as_mut().poll(&mut ctx) {
|
||||
return out;
|
||||
}
|
||||
|
||||
// doesn't race - unpark followed by park will result in park returning immediately
|
||||
std::thread::park();
|
||||
}
|
||||
}
|
||||
|
||||
async fn select<Left, Right>(left: Left, right: Right) -> Either<Left::Output, Right::Output>
|
||||
where
|
||||
Left: Future,
|
||||
Right: Future,
|
||||
{
|
||||
let left = std::pin::pin!(left);
|
||||
let right = std::pin::pin!(right);
|
||||
|
||||
Select {
|
||||
left,
|
||||
left_woken: Arc::new(AtomicBool::new(true)),
|
||||
|
||||
right,
|
||||
right_woken: Arc::new(AtomicBool::new(true)),
|
||||
}
|
||||
.await
|
||||
}
|
Loading…
Reference in a new issue