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`].
|
/// 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]`.
|
||||||
|
Loading…
Reference in New Issue
Block a user