luminal/runtime/
task.rs

1//! Task representation and management
2//!
3//! This module defines the core task structures used in the Luminal runtime,
4//! including task identifiers and the task structure itself.
5
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::task::{Context, Poll};
10
11/// Type alias for a boxed future that can be sent across threads
12///
13/// This represents a task's future that has been pinned and boxed
14/// to allow for dynamic dispatch and safe sending across thread boundaries.
15pub type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
16
17/// A unique identifier for tasks within the runtime
18///
19/// Each task is assigned a unique ID when it's created, which is used
20/// for tracking and managing the task throughout its lifetime.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22pub struct TaskId(pub(crate) u64);
23
24impl TaskId {
25    /// Creates a new unique task ID
26    ///
27    /// Uses an atomic counter to ensure uniqueness across threads.
28    /// This provides a reliable way to identify tasks even in a 
29    /// multithreaded environment.
30    ///
31    /// # Returns
32    ///
33    /// A new unique `TaskId`
34    pub fn new() -> Self {
35        static NEXT_ID: AtomicU64 = AtomicU64::new(1);
36        TaskId(NEXT_ID.fetch_add(1, Ordering::Relaxed))
37    }
38    
39    /// Gets the raw ID value
40    ///
41    /// # Returns
42    ///
43    /// The underlying u64 identifier
44    #[inline]
45    pub fn get(&self) -> u64 {
46        self.0
47    }
48}
49
50/// Represents an async task that can be scheduled and executed by the runtime
51///
52/// Contains the task's unique identifier and its underlying future.
53/// This is the core unit of work within the Luminal runtime.
54pub struct Task {
55    /// The unique identifier for this task
56    pub(crate) id: TaskId,
57    
58    /// The actual future that will be executed
59    pub(crate) future: BoxFuture,
60}
61
62impl Task {
63    /// Creates a new task with the given ID and future
64    ///
65    /// # Parameters
66    ///
67    /// * `id` - Unique identifier for the task
68    /// * `future` - The future to be executed as part of this task
69    ///
70    /// # Returns
71    ///
72    /// A new `Task` instance
73    pub fn new(id: TaskId, future: BoxFuture) -> Self {
74        Self { id, future }
75    }
76    
77    /// Polls the task's future, advancing its execution
78    ///
79    /// This is the core method used by the runtime to make progress on tasks.
80    /// It calls the underlying future's poll method with the given context.
81    ///
82    /// # Parameters
83    ///
84    /// * `cx` - The poll context, containing the waker
85    ///
86    /// # Returns
87    ///
88    /// `Poll::Ready(())` when the future completes, or `Poll::Pending` if it's not ready yet.
89    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<()> {
90        self.future.as_mut().poll(cx)
91    }
92}
93
94impl Default for Task {
95    /// Creates a default task with a completed future
96    ///
97    /// This is used as a placeholder when taking ownership of tasks
98    /// to avoid memory alignment issues.
99    fn default() -> Self {
100        Self {
101            id: TaskId::new(),
102            future: Box::pin(async {}),
103        }
104    }
105}