use presence of task reference to roll back state on drop
This commit is contained in:
parent
63d54c6a62
commit
12450ffeb4
@ -189,7 +189,7 @@ impl<F: Future + 'static> TaskStorage<F> {
|
||||
|
||||
/// An uninitialized [`TaskStorage`].
|
||||
pub struct AvailableTask<F: Future + 'static> {
|
||||
task: &'static TaskStorage<F>,
|
||||
task: Option<&'static TaskStorage<F>>,
|
||||
}
|
||||
|
||||
impl<F: Future + 'static> AvailableTask<F> {
|
||||
@ -201,21 +201,23 @@ impl<F: Future + 'static> AvailableTask<F> {
|
||||
.state
|
||||
.compare_exchange(STATE_ELIGIBLE, STATE_CLAIMED, Ordering::AcqRel, Ordering::Acquire)
|
||||
.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 {
|
||||
self.task.raw.poll_fn.set(Some(TaskStorage::<F>::poll));
|
||||
self.task.future.write_in_place(future);
|
||||
let task = self.task.take().unwrap();
|
||||
|
||||
self.task
|
||||
.raw
|
||||
// initialize task body
|
||||
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
|
||||
.store(STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::Release);
|
||||
|
||||
let task = TaskRef::new(self.task);
|
||||
|
||||
let task = TaskRef::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.
|
||||
///
|
||||
/// This is essentially a `[TaskStorage<F>; N]`.
|
||||
|
Loading…
Reference in New Issue
Block a user