@@ -13,6 +13,15 @@ use crate::raw::RawTask;
1313use crate :: state:: * ;
1414use crate :: Task ;
1515
16+ mod sealed {
17+ use super :: * ;
18+ pub trait Sealed < M > { }
19+
20+ impl < M , F > Sealed < M > for F where F : Fn ( Runnable < M > ) { }
21+
22+ impl < M , F > Sealed < M > for WithInfo < F > where F : Fn ( Runnable < M > , ScheduleInfo ) { }
23+ }
24+
1625/// A builder that creates a new task.
1726#[ derive( Debug ) ]
1827pub struct Builder < M > {
@@ -30,6 +39,135 @@ impl<M: Default> Default for Builder<M> {
3039 }
3140}
3241
42+ /// Extra scheduling information that can be passed to the scheduling function.
43+ ///
44+ /// The data source of this struct is directly from the actual implementation
45+ /// of the crate itself, different from [`Runnable`]'s metadata, which is
46+ /// managed by the caller.
47+ ///
48+ /// # Examples
49+ ///
50+ /// ```
51+ /// use async_task::{Runnable, ScheduleInfo, WithInfo};
52+ /// use std::sync::{Arc, Mutex};
53+ ///
54+ /// // The future inside the task.
55+ /// let future = async {
56+ /// println!("Hello, world!");
57+ /// };
58+ ///
59+ /// // If the task gets woken up while running, it will be sent into this channel.
60+ /// let (s, r) = flume::unbounded();
61+ /// // Otherwise, it will be placed into this slot.
62+ /// let lifo_slot = Arc::new(Mutex::new(None));
63+ /// let schedule = move |runnable: Runnable, info: ScheduleInfo| {
64+ /// if info.woken_while_running {
65+ /// s.send(runnable).unwrap()
66+ /// } else {
67+ /// let last = lifo_slot.lock().unwrap().replace(runnable);
68+ /// if let Some(last) = last {
69+ /// s.send(last).unwrap()
70+ /// }
71+ /// }
72+ /// };
73+ ///
74+ /// // Create the actual scheduler to be spawned with some future.
75+ /// let scheduler = WithInfo(schedule);
76+ /// // Create a task with the future and the scheduler.
77+ /// let (runnable, task) = async_task::spawn(future, scheduler);
78+ /// ```
79+ #[ derive( Debug , Copy , Clone ) ]
80+ #[ non_exhaustive]
81+ pub struct ScheduleInfo {
82+ /// Indicates whether the task gets woken up while running.
83+ ///
84+ /// It is set to true usually because the task has yielded itself to the
85+ /// scheduler.
86+ pub woken_while_running : bool ,
87+ }
88+
89+ impl ScheduleInfo {
90+ pub ( crate ) fn new ( woken_while_running : bool ) -> Self {
91+ ScheduleInfo {
92+ woken_while_running,
93+ }
94+ }
95+ }
96+
97+ /// The trait for scheduling functions.
98+ pub trait Schedule < M = ( ) > : sealed:: Sealed < M > {
99+ /// The actual scheduling procedure.
100+ fn schedule ( & self , runnable : Runnable < M > , info : ScheduleInfo ) ;
101+ }
102+
103+ impl < M , F > Schedule < M > for F
104+ where
105+ F : Fn ( Runnable < M > ) ,
106+ {
107+ fn schedule ( & self , runnable : Runnable < M > , _: ScheduleInfo ) {
108+ self ( runnable)
109+ }
110+ }
111+
112+ /// Pass a scheduling function with more scheduling information - a.k.a.
113+ /// [`ScheduleInfo`].
114+ ///
115+ /// Sometimes, it's useful to pass the runnable's state directly to the
116+ /// scheduling function, such as whether it's woken up while running. The
117+ /// scheduler can thus use the information to determine its scheduling
118+ /// strategy.
119+ ///
120+ /// The data source of [`ScheduleInfo`] is directly from the actual
121+ /// implementation of the crate itself, different from [`Runnable`]'s metadata,
122+ /// which is managed by the caller.
123+ ///
124+ /// # Examples
125+ ///
126+ /// ```
127+ /// use async_task::{ScheduleInfo, WithInfo};
128+ /// use std::sync::{Arc, Mutex};
129+ ///
130+ /// // The future inside the task.
131+ /// let future = async {
132+ /// println!("Hello, world!");
133+ /// };
134+ ///
135+ /// // If the task gets woken up while running, it will be sent into this channel.
136+ /// let (s, r) = flume::unbounded();
137+ /// // Otherwise, it will be placed into this slot.
138+ /// let lifo_slot = Arc::new(Mutex::new(None));
139+ /// let schedule = move |runnable, info: ScheduleInfo| {
140+ /// if info.woken_while_running {
141+ /// s.send(runnable).unwrap()
142+ /// } else {
143+ /// let last = lifo_slot.lock().unwrap().replace(runnable);
144+ /// if let Some(last) = last {
145+ /// s.send(last).unwrap()
146+ /// }
147+ /// }
148+ /// };
149+ ///
150+ /// // Create a task with the future and the schedule function.
151+ /// let (runnable, task) = async_task::spawn(future, WithInfo(schedule));
152+ /// ```
153+ #[ derive( Debug ) ]
154+ pub struct WithInfo < F > ( pub F ) ;
155+
156+ impl < F > From < F > for WithInfo < F > {
157+ fn from ( value : F ) -> Self {
158+ WithInfo ( value)
159+ }
160+ }
161+
162+ impl < M , F > Schedule < M > for WithInfo < F >
163+ where
164+ F : Fn ( Runnable < M > , ScheduleInfo ) ,
165+ {
166+ fn schedule ( & self , runnable : Runnable < M > , info : ScheduleInfo ) {
167+ ( self . 0 ) ( runnable, info)
168+ }
169+ }
170+
33171impl Builder < ( ) > {
34172 /// Creates a new task builder.
35173///
@@ -226,7 +364,7 @@ impl<M> Builder<M> {
226364 F : FnOnce ( & M ) -> Fut ,
227365 Fut : Future + Send + ' static ,
228366 Fut :: Output : Send + ' static ,
229- S : Fn ( Runnable < M > ) + Send + Sync + ' static ,
367+ S : Schedule < M > + Send + Sync + ' static ,
230368 {
231369 unsafe { self . spawn_unchecked ( future, schedule) }
232370 }
@@ -273,7 +411,7 @@ impl<M> Builder<M> {
273411 F : FnOnce ( & M ) -> Fut ,
274412 Fut : Future + ' static ,
275413 Fut :: Output : ' static ,
276- S : Fn ( Runnable < M > ) + Send + Sync + ' static ,
414+ S : Schedule < M > + Send + Sync + ' static ,
277415 {
278416 use std:: mem:: ManuallyDrop ;
279417 use std:: pin:: Pin ;
@@ -370,7 +508,7 @@ impl<M> Builder<M> {
370508 where
371509 F : FnOnce ( & ' a M ) -> Fut ,
372510 Fut : Future + ' a ,
373- S : Fn ( Runnable < M > ) ,
511+ S : Schedule < M > ,
374512 M : ' a ,
375513 {
376514 // Allocate large futures on the heap.
@@ -432,7 +570,7 @@ pub fn spawn<F, S>(future: F, schedule: S) -> (Runnable, Task<F::Output>)
432570where
433571 F : Future + Send + ' static ,
434572 F :: Output : Send + ' static ,
435- S : Fn ( Runnable ) + Send + Sync + ' static ,
573+ S : Schedule + Send + Sync + ' static ,
436574{
437575 unsafe { spawn_unchecked ( future, schedule) }
438576}
@@ -474,7 +612,7 @@ pub fn spawn_local<F, S>(future: F, schedule: S) -> (Runnable, Task<F::Output>)
474612where
475613 F : Future + ' static ,
476614 F :: Output : ' static ,
477- S : Fn ( Runnable ) + Send + Sync + ' static ,
615+ S : Schedule + Send + Sync + ' static ,
478616{
479617 Builder :: new ( ) . spawn_local ( move |( ) | future, schedule)
480618}
@@ -511,7 +649,7 @@ where
511649pub unsafe fn spawn_unchecked < F , S > ( future : F , schedule : S ) -> ( Runnable , Task < F :: Output > )
512650where
513651 F : Future ,
514- S : Fn ( Runnable ) ,
652+ S : Schedule ,
515653{
516654 Builder :: new ( ) . spawn_unchecked ( move |( ) | future, schedule)
517655}
@@ -604,7 +742,7 @@ impl<M> Runnable<M> {
604742 mem:: forget ( self ) ;
605743
606744 unsafe {
607- ( ( * header) . vtable . schedule ) ( ptr) ;
745+ ( ( * header) . vtable . schedule ) ( ptr, ScheduleInfo :: new ( false ) ) ;
608746 }
609747 }
610748
0 commit comments