lustre-executor 0.2.0

//! Task abstraction with unique IDs and waker integration.

use crate::core::id::{IdGenerator, TaskId};
use futures::future::Future;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll as TaskPoll, Wake, Waker};

/// A task wrapping an async future with a unique ID.
pub struct Task {
    id: TaskId<u64>,
    future: std::pin::Pin<Box<dyn Future<Output = ()> + 'static>>,
}

impl Task {
    /// Creates a new task from an async future.
    pub fn new(future: impl Future<Output = ()> + 'static) -> Self {
        Self {
            id: TaskId::new(crate::core::id::AtomicU64Generator::generate()),
            future: Box::pin(future),
        }
    }

    pub(crate) fn id(&self) -> TaskId<u64> {
        self.id
    }

    pub(crate) fn poll(&mut self, cx: &mut Context) -> TaskPoll<()> {
        self.future.as_mut().poll(cx)
    }
}

/// Custom waker that enqueues tasks for polling.
pub(crate) struct TaskWaker {
    task_id: TaskId<u64>,
    task_queue: Arc<Mutex<VecDeque<TaskId<u64>>>>,
}

impl TaskWaker {
    pub(crate) fn new(
        task_id: TaskId<u64>,
        task_queue: Arc<Mutex<VecDeque<TaskId<u64>>>>,
    ) -> Waker {
        Waker::from(Arc::new(Self {
            task_id,
            task_queue,
        }))
    }

    fn wake_task(&self) {
        self.task_queue.lock().unwrap().push_front(self.task_id);
    }
}

impl Wake for TaskWaker {
    fn wake(self: Arc<Self>) {
        self.wake_task();
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.wake_task();
    }
}