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