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::{
|
use std::{
|
||||||
num::{NonZeroU16, NonZeroUsize},
|
num::{NonZeroU16, NonZeroUsize},
|
||||||
sync::{
|
sync::{atomic::AtomicU64, Arc, Mutex},
|
||||||
atomic::{AtomicBool, AtomicU64},
|
|
||||||
Arc, Mutex,
|
|
||||||
},
|
|
||||||
thread::JoinHandle,
|
thread::JoinHandle,
|
||||||
time::{Duration, Instant},
|
time::Instant,
|
||||||
};
|
};
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
|
@ -178,9 +177,7 @@ impl CpuPool {
|
||||||
let mut threads = state.take_threads();
|
let mut threads = state.take_threads();
|
||||||
|
|
||||||
for thread in &mut threads {
|
for thread in &mut threads {
|
||||||
thread
|
thread.signal.take();
|
||||||
.signal
|
|
||||||
.store(true, std::sync::atomic::Ordering::Release);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
for thread in &mut threads {
|
for thread in &mut threads {
|
||||||
|
@ -373,9 +370,7 @@ impl ThreadVec {
|
||||||
impl Drop for ThreadVec {
|
impl Drop for ThreadVec {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
for thread in &mut self.threads {
|
for thread in &mut self.threads {
|
||||||
thread
|
thread.signal.take();
|
||||||
.signal
|
|
||||||
.store(true, std::sync::atomic::Ordering::Release);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
for thread in &mut self.threads {
|
for thread in &mut self.threads {
|
||||||
|
@ -388,14 +383,13 @@ impl Drop for ThreadVec {
|
||||||
|
|
||||||
struct Thread {
|
struct Thread {
|
||||||
handle: Option<JoinHandle<()>>,
|
handle: Option<JoinHandle<()>>,
|
||||||
signal: Arc<AtomicBool>,
|
signal: Option<SendOnDrop>,
|
||||||
closed: flume::Receiver<()>,
|
closed: flume::Receiver<()>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Thread {
|
impl Thread {
|
||||||
async fn reap(mut self) {
|
async fn reap(mut self) {
|
||||||
self.signal
|
self.signal.take();
|
||||||
.store(true, std::sync::atomic::Ordering::Release);
|
|
||||||
|
|
||||||
let _ = self.closed.recv_async().await;
|
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 {
|
fn spawn(name: &'static str, id: u64, receiver: flume::Receiver<SendFn>) -> Thread {
|
||||||
let (closed_tx, closed) = flume::bounded(1);
|
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 closed_tx = SendOnDrop { sender: closed_tx };
|
||||||
|
|
||||||
let signal2 = signal.clone();
|
|
||||||
let handle = std::thread::Builder::new()
|
let handle = std::thread::Builder::new()
|
||||||
.name(format!("{name}-{id}"))
|
.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");
|
.expect("Failed to spawn new thread");
|
||||||
|
|
||||||
Thread {
|
Thread {
|
||||||
handle: Some(handle),
|
handle: Some(handle),
|
||||||
signal,
|
signal: Some(signal),
|
||||||
closed,
|
closed,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -471,22 +465,15 @@ fn run(
|
||||||
name: &'static str,
|
name: &'static str,
|
||||||
id: u64,
|
id: u64,
|
||||||
receiver: flume::Receiver<SendFn>,
|
receiver: flume::Receiver<SendFn>,
|
||||||
signal: Arc<AtomicBool>,
|
signal: flume::Receiver<()>,
|
||||||
closed_tx: SendOnDrop,
|
closed_tx: SendOnDrop,
|
||||||
) {
|
) {
|
||||||
let guard = MetricsGuard::guard(name, id);
|
let guard = MetricsGuard::guard(name, id);
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
match receiver.recv_timeout(Duration::from_millis(50)) {
|
match selector::blocking_select(signal.recv_async(), receiver.recv_async()) {
|
||||||
Ok(send_fn) => {
|
selector::Either::Left(_) | selector::Either::Right(Err(_)) => break,
|
||||||
invoke_send_fn(name, id, send_fn);
|
selector::Either::Right(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;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
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