#[cfg(any(feature = "std", feature = "alloc"))]
use core::future::Future;
#[cfg(any(feature = "std", feature = "alloc"))]
use core::marker::PhantomData;
#[cfg(any(feature = "std", feature = "alloc"))]
use core::task::{RawWaker, RawWakerVTable, Waker};
#[cfg(any(feature = "std", feature = "alloc"))]
use crate::valtron::{Ctx, QueueReadiness, TaskIterator, TaskStatus};
#[cfg(any(feature = "std", feature = "alloc"))]
use concurrent_queue::ConcurrentQueue;
#[cfg(any(feature = "std", feature = "alloc"))]
use core::pin::Pin;
#[cfg(any(feature = "std", feature = "alloc"))]
use core::task::{Context, Poll};
#[cfg(feature = "std")]
use std::boxed::Box;
#[cfg(feature = "std")]
use std::sync::Arc;
#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::boxed::Box;
#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::sync::Arc;
#[cfg(all(feature = "std", feature = "multi"))]
use crate::valtron::executors::{
local::{NotificationItem, NotifyQueue, NotifyRecvIterator},
DEFAULT_NOTIFY_QUEUE_WAIT_TIMEOUT,
};
#[cfg(all(feature = "std", feature = "multi"))]
use concurrent_queue::PushError;
use crate::valtron::EventReadinessPtr;
#[cfg(any(feature = "std", feature = "alloc"))]
fn create_noop_waker() -> Waker {
const VTABLE: RawWakerVTable = RawWakerVTable::new(
|_| RAW_WAKER, |_| {}, |_| {}, |_| {}, );
const RAW_WAKER: RawWaker = RawWaker::new(core::ptr::null(), &VTABLE);
unsafe { Waker::from_raw(RAW_WAKER) }
}
#[cfg(all(feature = "std", not(feature = "multi")))]
fn get_noop_waker() -> Waker {
thread_local! {
static NOOP_WAKER: Waker = create_noop_waker();
}
NOOP_WAKER.with(std::clone::Clone::clone)
}
#[cfg(all(feature = "alloc", not(feature = "std"), not(feature = "multi")))]
fn get_noop_waker() -> Waker {
create_noop_waker()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WakeToken {
pub id: u64,
}
#[cfg(any(feature = "std", feature = "alloc"))]
struct WakerState {
queue: Arc<ConcurrentQueue<WakeToken>>,
id: u64,
}
#[cfg(any(feature = "std", feature = "alloc"))]
static QUEUE_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
queue_waker_clone,
queue_waker_wake,
queue_waker_wake_by_ref,
queue_waker_drop,
);
#[cfg(any(feature = "std", feature = "alloc"))]
unsafe fn queue_waker_clone(data: *const ()) -> RawWaker {
Arc::increment_strong_count(data.cast::<WakerState>());
RawWaker::new(data, &QUEUE_WAKER_VTABLE)
}
#[cfg(any(feature = "std", feature = "alloc"))]
unsafe fn queue_waker_wake(data: *const ()) {
let state = Arc::from_raw(data.cast::<WakerState>());
let _ = state.queue.push(WakeToken { id: state.id });
}
#[cfg(any(feature = "std", feature = "alloc"))]
unsafe fn queue_waker_wake_by_ref(data: *const ()) {
let state = core::mem::ManuallyDrop::new(Arc::from_raw(data.cast::<WakerState>()));
let _ = state.queue.push(WakeToken { id: state.id });
}
#[cfg(any(feature = "std", feature = "alloc"))]
unsafe fn queue_waker_drop(data: *const ()) {
drop(Arc::from_raw(data.cast::<WakerState>()));
}
#[cfg(any(feature = "std", feature = "alloc"))]
#[must_use]
pub fn queue_waker(queue: Arc<ConcurrentQueue<WakeToken>>, id: u64) -> Waker {
let state = Arc::new(WakerState { queue, id });
let raw = RawWaker::new(Arc::into_raw(state).cast::<()>(), &QUEUE_WAKER_VTABLE);
unsafe { Waker::from_raw(raw) }
}
#[cfg(any(feature = "std", feature = "alloc"))]
pub struct FutureTask<F>
where
F: Future,
{
future: Pin<Box<F>>,
completed: bool,
wake_queue: Arc<ConcurrentQueue<WakeToken>>,
}
#[cfg(any(feature = "std", feature = "alloc"))]
impl<F> FutureTask<F>
where
F: Future,
{
pub fn new(future: F) -> Self {
Self {
future: Box::pin(future),
completed: false,
wake_queue: Arc::new(ConcurrentQueue::unbounded()),
}
}
#[must_use]
pub fn from_pinned(future: Pin<Box<F>>) -> Self {
Self {
future,
completed: false,
wake_queue: Arc::new(ConcurrentQueue::unbounded()),
}
}
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
impl<F> TaskIterator for FutureTask<F>
where
F: Future + 'static,
F::Output: 'static,
{
type Ready = F::Output;
type Pending = ();
fn next_status(&mut self, _ctx: Ctx) -> Option<TaskStatus<Self::Ready, Self::Pending>> {
if self.completed {
return None;
}
while self.wake_queue.pop().is_ok() {}
let waker = queue_waker(self.wake_queue.clone(), 0);
let mut cx = Context::from_waker(&waker);
match self.future.as_mut().poll(&mut cx) {
Poll::Ready(output) => {
self.completed = true;
Some(TaskStatus::Ready(output))
}
Poll::Pending => {
if self.wake_queue.is_empty() {
Some(TaskStatus::Depends(Arc::new(QueueReadiness::new(
self.wake_queue.clone(),
))))
} else {
Some(TaskStatus::Ignore)
}
}
}
}
}
#[cfg(feature = "multi")]
impl<F> TaskIterator for FutureTask<F>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
type Ready = F::Output;
type Pending = ();
fn next_status(&mut self, _ctx: Ctx) -> Option<TaskStatus<Self::Ready, Self::Pending>> {
if self.completed {
return None;
}
while self.wake_queue.pop().is_ok() {}
let waker = queue_waker(self.wake_queue.clone(), 0);
let mut cx = Context::from_waker(&waker);
match self.future.as_mut().poll(&mut cx) {
Poll::Ready(output) => {
self.completed = true;
Some(TaskStatus::Ready(output))
}
Poll::Pending => {
if self.wake_queue.is_empty() {
Some(TaskStatus::Depends(Arc::new(QueueReadiness::new(
self.wake_queue.clone(),
))))
} else {
Some(TaskStatus::Ignore)
}
}
}
}
}
#[cfg(any(feature = "std", feature = "alloc"))]
pub struct StreamTask<S>
where
S: futures_core::Stream,
{
stream: Pin<Box<S>>,
exhausted: bool,
wake_queue: Arc<ConcurrentQueue<WakeToken>>,
}
#[cfg(any(feature = "std", feature = "alloc"))]
impl<S> StreamTask<S>
where
S: futures_core::Stream,
{
pub fn new(stream: S) -> Self {
Self {
stream: Box::pin(stream),
exhausted: false,
wake_queue: Arc::new(ConcurrentQueue::unbounded()),
}
}
}
#[cfg(feature = "multi")]
impl<S> TaskIterator for StreamTask<S>
where
S: futures_core::Stream + Send + 'static,
S::Item: Send + 'static,
{
type Ready = Option<S::Item>;
type Pending = ();
fn next_status(&mut self, _ctx: Ctx) -> Option<TaskStatus<Self::Ready, Self::Pending>> {
if self.exhausted {
return None;
}
while self.wake_queue.pop().is_ok() {}
let waker = queue_waker(self.wake_queue.clone(), 0);
let mut cx = Context::from_waker(&waker);
match self.stream.as_mut().poll_next(&mut cx) {
Poll::Ready(Some(item)) => Some(TaskStatus::Ready(Some(item))),
Poll::Ready(None) => {
self.exhausted = true;
Some(TaskStatus::Ready(None))
}
Poll::Pending => {
if self.wake_queue.is_empty() {
Some(TaskStatus::Depends(Arc::new(QueueReadiness::new(
self.wake_queue.clone(),
))))
} else {
Some(TaskStatus::Ignore)
}
}
}
}
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
impl<S> TaskIterator for StreamTask<S>
where
S: futures_core::Stream + 'static,
S::Item: 'static,
{
type Ready = Option<S::Item>;
type Pending = ();
fn next_status(&mut self, _ctx: Ctx) -> Option<TaskStatus<Self::Ready, Self::Pending>> {
if self.exhausted {
return None;
}
while self.wake_queue.pop().is_ok() {}
let waker = queue_waker(self.wake_queue.clone(), 0);
let mut cx = Context::from_waker(&waker);
match self.stream.as_mut().poll_next(&mut cx) {
Poll::Ready(Some(item)) => Some(TaskStatus::Ready(Some(item))),
Poll::Ready(None) => {
self.exhausted = true;
Some(TaskStatus::Ready(None))
}
Poll::Pending => {
if self.wake_queue.is_empty() {
Some(TaskStatus::Depends(Arc::new(QueueReadiness::new(
self.wake_queue.clone(),
))))
} else {
Some(TaskStatus::Ignore)
}
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CancelOutcome {
Cancelled,
}
#[cfg(any(feature = "std", feature = "alloc"))]
struct CancelOrReadiness {
inner: EventReadinessPtr,
cancel: Arc<core::sync::atomic::AtomicBool>,
}
#[cfg(any(feature = "std", feature = "alloc"))]
impl crate::valtron::EventReadiness for CancelOrReadiness {
fn is_ready(&self, dur: Option<core::time::Duration>) -> bool {
self.cancel.load(core::sync::atomic::Ordering::Acquire) || self.inner.is_ready(dur)
}
}
#[cfg(any(feature = "std", feature = "alloc"))]
pub struct CancellableFutureTask<F>
where
F: Future,
{
inner: FutureTask<F>,
cancel: Arc<core::sync::atomic::AtomicBool>,
}
#[cfg(any(feature = "std", feature = "alloc"))]
impl<F: Future> CancellableFutureTask<F> {
pub fn new(future: F, cancel: Arc<core::sync::atomic::AtomicBool>) -> Self {
Self {
inner: FutureTask::new(future),
cancel,
}
}
pub fn request_cancel(&self) {
self.cancel
.store(true, core::sync::atomic::Ordering::Release);
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.cancel.load(core::sync::atomic::Ordering::Acquire)
}
#[must_use]
pub fn cancel_signal(&self) -> Arc<core::sync::atomic::AtomicBool> {
self.cancel.clone()
}
}
#[cfg(feature = "multi")]
impl<F> TaskIterator for CancellableFutureTask<F>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
type Ready = Result<F::Output, CancelOutcome>;
type Pending = ();
fn next_status(&mut self, ctx: Ctx) -> Option<TaskStatus<Self::Ready, Self::Pending>> {
if self.is_cancelled() {
return Some(TaskStatus::Ready(Err(CancelOutcome::Cancelled)));
}
match self.inner.next_status(ctx) {
Some(TaskStatus::Ready(v)) => Some(TaskStatus::Ready(Ok(v))),
Some(TaskStatus::Pending(p)) => Some(TaskStatus::Pending(p)),
Some(TaskStatus::Init) => Some(TaskStatus::Init),
Some(TaskStatus::Ignore) => Some(TaskStatus::Ignore),
Some(TaskStatus::Wait) => Some(TaskStatus::Wait),
Some(TaskStatus::Delayed(d)) => Some(TaskStatus::Delayed(d)),
Some(TaskStatus::Spawn(s)) => Some(TaskStatus::Spawn(s)),
Some(TaskStatus::Depends(r)) => {
Some(TaskStatus::Depends(Arc::new(CancelOrReadiness {
inner: r,
cancel: self.cancel.clone(),
})))
}
Some(TaskStatus::Spread(items)) => {
use crate::valtron::TaskSpread;
Some(TaskStatus::Spread(
items
.into_iter()
.map(|item| match item {
TaskSpread::Ready(v) => TaskSpread::Ready(Ok(v)),
TaskSpread::Pending(p) => TaskSpread::Pending(p),
})
.collect(),
))
}
None => None,
}
}
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
impl<F> TaskIterator for CancellableFutureTask<F>
where
F: Future + 'static,
F::Output: 'static,
{
type Ready = Result<F::Output, CancelOutcome>;
type Pending = ();
fn next_status(&mut self, ctx: Ctx) -> Option<TaskStatus<Self::Ready, Self::Pending>> {
if self.is_cancelled() {
return Some(TaskStatus::Ready(Err(CancelOutcome::Cancelled)));
}
match self.inner.next_status(ctx) {
Some(TaskStatus::Ready(v)) => Some(TaskStatus::Ready(Ok(v))),
Some(TaskStatus::Pending(p)) => Some(TaskStatus::Pending(p)),
Some(TaskStatus::Init) => Some(TaskStatus::Init),
Some(TaskStatus::Ignore) => Some(TaskStatus::Ignore),
Some(TaskStatus::Wait) => Some(TaskStatus::Wait),
Some(TaskStatus::Delayed(d)) => Some(TaskStatus::Delayed(d)),
Some(TaskStatus::Spawn(s)) => Some(TaskStatus::Spawn(s)),
Some(TaskStatus::Depends(r)) => {
Some(TaskStatus::Depends(Arc::new(CancelOrReadiness {
inner: r,
cancel: self.cancel.clone(),
})))
}
Some(TaskStatus::Spread(items)) => {
use crate::valtron::TaskSpread;
Some(TaskStatus::Spread(
items
.into_iter()
.map(|item| match item {
TaskSpread::Ready(v) => TaskSpread::Ready(Ok(v)),
TaskSpread::Pending(p) => TaskSpread::Pending(p),
})
.collect(),
))
}
None => None,
}
}
}
#[derive(Debug)]
pub enum ThreadedValue<T, E> {
Value(Result<T, E>),
Waiting,
}
#[cfg(feature = "multi")]
pub struct FutureIterator<T, E> {
iter: NotifyRecvIterator<ThreadedValue<T, E>>,
}
#[cfg(feature = "multi")]
impl<T, E> Iterator for FutureIterator<T, E> {
type Item = ThreadedValue<T, E>;
fn next(&mut self) -> Option<Self::Item> {
match self.iter.next() {
Some(NotificationItem::Ready(value)) => Some(value),
Some(NotificationItem::None) => Some(ThreadedValue::Waiting),
None => None,
}
}
}
#[cfg(feature = "multi")]
pub struct ThreadedIterFuture<F, Fut, I, T, E>
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<I, E>> + 'static,
I: Iterator<Item = Result<T, E>> + 'static,
T: Send + 'static,
E: Send + 'static,
{
future_fn: F,
queue_size: usize,
backpressure_sleep: Option<std::time::Duration>,
_phantom: PhantomData<(Fut, I, T, E)>,
}
#[cfg(feature = "multi")]
impl<F, Fut, I, T, E> ThreadedIterFuture<F, Fut, I, T, E>
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<I, E>> + 'static,
I: Iterator<Item = Result<T, E>> + 'static,
T: Send + 'static,
E: Send + 'static,
{
pub fn new(future_fn: F) -> Self {
Self {
future_fn,
queue_size: 16,
backpressure_sleep: Some(std::time::Duration::from_millis(10)),
_phantom: PhantomData,
}
}
pub fn with_queue_size(future_fn: F, queue_size: usize) -> Self {
Self {
future_fn,
queue_size: queue_size.max(2),
backpressure_sleep: Some(std::time::Duration::from_millis(10)),
_phantom: PhantomData,
}
}
pub fn execute(
self,
) -> crate::valtron::GenericResult<impl Iterator<Item = ThreadedValue<T, E>>> {
let queue = Arc::new(NotifyQueue::bounded(self.queue_size));
let worker_queue = queue.clone();
let backpressure_sleep = self.backpressure_sleep;
super::run_background_job(move || {
let waker = create_noop_waker();
let mut cx = Context::from_waker(&waker);
let mut future = Box::pin((self.future_fn)());
let iterator = loop {
match future.as_mut().poll(&mut cx) {
Poll::Ready(Ok(iter)) => break Some(iter),
Poll::Ready(Err(e)) => {
let _ = worker_queue.push(ThreadedValue::Value(Err(e)));
break None;
}
Poll::Pending => std::thread::yield_now(),
}
};
if let Some(iter) = iterator {
for result in iter {
let mut value = ThreadedValue::Value(result);
loop {
match worker_queue.push(value) {
Ok(()) => break,
Err(PushError::Full(v)) => {
value = v;
if let Some(sleep_dur) = backpressure_sleep {
std::thread::sleep(sleep_dur);
} else {
std::hint::spin_loop();
}
}
Err(PushError::Closed(_)) => return,
}
}
}
}
let _ = worker_queue.close();
})?;
Ok(FutureIterator {
iter: NotifyRecvIterator::from_notify_queue(
queue,
DEFAULT_NOTIFY_QUEUE_WAIT_TIMEOUT,
),
})
}
pub fn with_backpressure_sleep(
future_fn: F,
queue_size: usize,
backpressure_sleep: Option<std::time::Duration>,
) -> Self {
Self {
future_fn,
queue_size: queue_size.max(2),
backpressure_sleep,
_phantom: PhantomData,
}
}
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
pub struct FutureIterator<F, Fut, I, T, E>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<I, E>>,
I: Iterator<Item = Result<T, E>>,
{
future: Option<F>,
inner_iter: Option<I>,
_phantom: PhantomData<(Fut, T, E)>,
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
impl<F, Fut, I, T, E> Iterator for FutureIterator<F, Fut, I, T, E>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<I, E>>,
I: Iterator<Item = Result<T, E>>,
{
type Item = ThreadedValue<T, E>;
fn next(&mut self) -> Option<Self::Item> {
if let Some(iter) = &mut self.inner_iter {
return iter.next().map(ThreadedValue::Value);
}
let mut future = {
let f = self.future.take()?;
Box::pin((f)())
};
let waker = get_noop_waker();
let mut cx = Context::from_waker(&waker);
match future.as_mut().poll(&mut cx) {
Poll::Ready(Ok(iter)) => {
self.inner_iter = Some(iter);
if let Some(iter) = &mut self.inner_iter {
return iter.next().map(ThreadedValue::Value);
}
None
}
Poll::Ready(Err(e)) => Some(ThreadedValue::Value(Err(e))),
Poll::Pending => Some(ThreadedValue::Waiting),
}
}
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
pub struct ThreadedIterFuture<F, Fut, I, T, E>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<I, E>>,
I: Iterator<Item = Result<T, E>>,
{
future_fn: F,
_phantom: PhantomData<(Fut, I, T, E)>,
}
#[cfg(all(not(feature = "multi"), any(feature = "std", feature = "alloc")))]
impl<F, Fut, I, T, E> ThreadedIterFuture<F, Fut, I, T, E>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<I, E>>,
I: Iterator<Item = Result<T, E>>,
{
pub fn new(future_fn: F) -> Self {
Self {
future_fn,
_phantom: PhantomData,
}
}
pub fn with_queue_size(future_fn: F, _queue_size: usize) -> Self {
Self {
future_fn,
_phantom: PhantomData,
}
}
pub fn execute(self) -> impl Iterator<Item = ThreadedValue<T, E>> {
FutureIterator {
future: Some(self.future_fn),
inner_iter: None,
_phantom: PhantomData,
}
}
}