use presence of task reference to roll back state on drop

This commit is contained in:
Hailey Somerville 2023-09-14 15:26:38 +10:00
parent 63d54c6a62
commit 12450ffeb4

View File

@ -189,7 +189,7 @@ impl<F: Future + 'static> TaskStorage<F> {
/// An uninitialized [`TaskStorage`]. /// An uninitialized [`TaskStorage`].
pub struct AvailableTask<F: Future + 'static> { pub struct AvailableTask<F: Future + 'static> {
task: &'static TaskStorage<F>, task: Option<&'static TaskStorage<F>>,
} }
impl<F: Future + 'static> AvailableTask<F> { impl<F: Future + 'static> AvailableTask<F> {
@ -201,21 +201,23 @@ impl<F: Future + 'static> AvailableTask<F> {
.state .state
.compare_exchange(STATE_ELIGIBLE, STATE_CLAIMED, Ordering::AcqRel, Ordering::Acquire) .compare_exchange(STATE_ELIGIBLE, STATE_CLAIMED, Ordering::AcqRel, Ordering::Acquire)
.ok() .ok()
.map(|_| Self { task }) .map(|_| Self { task: Some(task) })
} }
fn initialize_impl<S>(self, future: impl FnOnce() -> F) -> SpawnToken<S> { fn initialize_impl<S>(mut self, future: impl FnOnce() -> F) -> SpawnToken<S> {
unsafe { unsafe {
self.task.raw.poll_fn.set(Some(TaskStorage::<F>::poll)); let task = self.task.take().unwrap();
self.task.future.write_in_place(future);
self.task // initialize task body
.raw task.raw.poll_fn.set(Some(TaskStorage::<F>::poll));
task.future.write_in_place(future);
// set state now that task is valid
task.raw
.state .state
.store(STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::Release); .store(STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::Release);
let task = TaskRef::new(self.task); let task = TaskRef::new(task);
SpawnToken::new(task) SpawnToken::new(task)
} }
} }
@ -262,6 +264,18 @@ impl<F: Future + 'static> AvailableTask<F> {
} }
} }
impl<F: Future + 'static> Drop for AvailableTask<F> {
fn drop(&mut self) {
if let Some(task) = self.task {
// if AvailableTask is dropped before a task is spawned on it,
// set task back to STATE_ELIGIBLE so we can re-use the slot.
// This also ensures that we can re-use this slot even if the
// user callback in initialize_impl panics.
task.raw.state.store(STATE_ELIGIBLE, Ordering::Release);
}
}
}
/// Raw storage that can hold up to N tasks of the same type. /// Raw storage that can hold up to N tasks of the same type.
/// ///
/// This is essentially a `[TaskStorage<F>; N]`. /// This is essentially a `[TaskStorage<F>; N]`.