use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::sync::{Arc, Weak};
use kovan_queue::seg_queue::SegQueue;
use super::runtime::Inner;
use crate::sync::atomic_waker::AtomicWaker;
use crate::sync::internal::{AtomicBool, AtomicUsize, Ordering, fence};
pub(crate) fn time() -> std::time::Duration {
super::runtime::io_knobs::current().linger
}
pub(crate) type Work = Pin<Box<dyn Future<Output = ()> + Send>>;
pub(crate) fn hand<F>(work: F)
where
F: Future<Output = ()> + Send + 'static,
{
let work: Work = Box::pin(work);
match super::runtime::running_inner() {
Some(inner) => inner.take_linger(work),
None => alone(work),
}
}
pub(crate) struct Lingering {
list: SegQueue<Work>,
active: AtomicUsize,
waker: AtomicWaker,
closing: AtomicBool,
over: AtomicBool,
}
impl Lingering {
pub(crate) fn new() -> Self {
Self {
list: SegQueue::new(),
active: AtomicUsize::new(0),
waker: AtomicWaker::new(),
closing: AtomicBool::new(false),
over: AtomicBool::new(false),
}
}
pub(crate) fn active(&self) -> usize {
self.active.load(Ordering::SeqCst)
}
fn ended(&self) {
self.active.fetch_sub(1, Ordering::SeqCst);
self.waker.wake();
}
pub(crate) fn nudge(&self) {
self.waker.wake();
}
pub(crate) fn close(&self) {
self.closing.store(true, Ordering::SeqCst);
}
pub(crate) fn is_closing(&self) -> bool {
self.closing.load(Ordering::SeqCst)
}
pub(crate) fn is_over(&self) -> bool {
self.over.load(Ordering::SeqCst)
}
pub(crate) fn push(&self, work: Work) {
self.list.push(work);
self.waker.wake();
fence(Ordering::SeqCst);
if self.over.load(Ordering::SeqCst) {
self.run_rest_alone();
}
}
pub(crate) fn end(&self) {
self.over.store(true, Ordering::SeqCst);
fence(Ordering::SeqCst);
self.run_rest_alone();
}
fn run_rest_alone(&self) {
while let Ok(work) = self.list.try_pop() {
self.ended();
alone(work);
}
}
}
pub(crate) struct RunAll<'a> {
inner: &'a Inner,
running: Vec<Work>,
}
impl Future for RunAll<'_> {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.get_mut();
let lingering = this.inner.lingering();
lingering.waker.register(cx.waker());
loop {
while let Ok(work) = lingering.list.try_pop() {
this.running.push(work);
}
this.running.retain_mut(|work| {
let ended = work.as_mut().poll(cx).is_ready();
if ended {
lingering.ended();
}
!ended
});
if lingering.list.is_empty() {
break;
}
}
if this.inner.lingers_done() {
Poll::Ready(())
} else {
Poll::Pending
}
}
}
impl Drop for RunAll<'_> {
fn drop(&mut self) {
let lingering = self.inner.lingering();
for work in self.running.drain(..) {
lingering.ended();
alone(work);
}
}
}
impl Inner {
pub(crate) fn lingers_done(&self) -> bool {
self.workers.shards().all(|shard| shard.settled()) && self.lingering().active() == 0
}
pub(crate) fn run_lingers(&self) -> RunAll<'_> {
RunAll {
inner: self,
running: Vec::new(),
}
}
pub(crate) fn take_linger(self: &Arc<Self>, work: Work) {
self.lingering().active.fetch_add(1, Ordering::SeqCst);
self.route_linger(work, true);
}
fn route_linger(self: &Arc<Self>, work: Work, here: bool) {
let lingering = self.lingering();
if !lingering.is_closing() {
let guarded = Guarded {
work: Some(work),
owner: Arc::downgrade(self),
};
match super::runtime::spawn_linger(self, guarded, here) {
Ok(()) => return,
Err(guarded) => return guarded.into_list(),
}
}
lingering.push(work);
}
}
pub(crate) struct Guarded {
work: Option<Work>,
owner: Weak<Inner>,
}
impl Guarded {
fn into_list(mut self) {
let Some(work) = self.work.take() else {
return;
};
match self.owner.upgrade() {
Some(inner) => inner.lingering().push(work),
None => alone(work),
}
}
}
impl Future for Guarded {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.get_mut();
let Some(work) = this.work.as_mut() else {
return Poll::Ready(());
};
let poll = work.as_mut().poll(cx);
if poll.is_ready() {
this.work = None;
if let Some(inner) = this.owner.upgrade() {
inner.lingering().ended();
}
}
poll
}
}
impl Drop for Guarded {
fn drop(&mut self) {
let Some(work) = self.work.take() else {
return;
};
match self.owner.upgrade() {
Some(inner) => inner.route_linger(work, false),
None => alone(work),
}
}
}
pub(crate) fn alone(work: Work) {
let (give, take) = crate::sync::oneshot::channel::<Work>();
let started = std::thread::Builder::new()
.name("nuclei-linger".into())
.spawn(move || {
if let Ok(work) = take.blocking_recv() {
run_with_own_runtime(work);
}
});
let unsent = match started {
Ok(_) => give.send(work).err(),
Err(err) => {
log::warn!(
"no thread could start to deliver the bytes a dropped handle left ({err}): they are delivered on the thread that dropped it"
);
Some(work)
}
};
if let Some(work) = unsent {
run_with_own_runtime(work);
}
}
fn run_with_own_runtime(work: Work) {
match crate::Runtime::builder().worker_threads(0).build() {
Ok(rt) => rt.block_on(work),
Err(err) => log::error!(
"the bytes a dropped handle left were not delivered: no runtime could start for them ({err})"
),
}
}