Turbopack: run tasks with priority based on distance to leaf (#87662)

### What?

* Compute the approximate distance to leaf in the dependency graph
* Prioritize tasks re-computing by distance to leaf, running leaf tasks first

This ensures that it's less likely that tasks are re-executed multiple times due to bad ordering.
This adresses the "diamond" problem where a task is invalidated multiple times due to multiple paths of invalidation to that task.
This commit is contained in:
Tobias Koppers
2026-01-15 20:18:03 +01:00
committed by GitHub
parent 3933d47b5e
commit 125b7d89f4
18 changed files with 1375 additions and 293 deletions
@@ -21,6 +21,7 @@ verify_aggregation_graph = []
verify_immutable = []
verify_determinism = []
trace_aggregation_update = []
trace_leaf_distance_update = []
trace_find_and_schedule = []
trace_task_completion = []
trace_task_dirty = []
@@ -27,8 +27,8 @@ use tokio::time::{Duration, Instant};
use tracing::{Span, trace_span};
use turbo_tasks::{
CellId, FxDashMap, KeyValuePair, RawVc, ReadCellOptions, ReadCellTracking, ReadConsistency,
ReadOutputOptions, ReadTracking, TRANSIENT_TASK_BIT, TaskExecutionReason, TaskId, TraitTypeId,
TurboTasksBackendApi, ValueTypeId,
ReadOutputOptions, ReadTracking, TRANSIENT_TASK_BIT, TaskExecutionReason, TaskId, TaskPriority,
TraitTypeId, TurboTasksBackendApi, ValueTypeId,
backend::{
Backend, CachedTaskType, CellContent, TaskExecutionSpec, TransientTaskRoot,
TransientTaskType, TurboTasksExecutionError, TypedCellContent, VerificationMode,
@@ -51,8 +51,8 @@ use crate::{
operation::{
AggregationUpdateJob, AggregationUpdateQueue, ChildExecuteContext,
CleanupOldEdgesOperation, ComputeDirtyAndCleanUpdate, ConnectChildOperation,
ExecuteContext, ExecuteContextImpl, Operation, OutdatedEdge, TaskGuard,
connect_children, get_aggregation_number, get_uppers, is_root_node,
ExecuteContext, ExecuteContextImpl, LeafDistanceUpdateQueue, Operation, OutdatedEdge,
TaskGuard, connect_children, get_aggregation_number, get_uppers, is_root_node,
make_task_dirty_internal, prepare_new_children,
},
storage::{
@@ -575,7 +575,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
// Check the dirty count of the root node
let has_dirty_containers = task.has_dirty_containers();
if has_dirty_containers || is_dirty {
if has_dirty_containers || is_dirty.is_some() {
let activeness = get_mut!(task, Activeness);
let mut task_ids_to_schedule: Vec<_> = Vec::new();
// When there are dirty task, subscribe to the all_clean_event
@@ -641,7 +641,11 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
let has_dirty_containers = task.has_dirty_containers();
let task_description = ctx.get_task_description(task_id);
let is_dirty_label = if is_dirty { ", dirty" } else { "" };
let is_dirty_label = if let Some(parent_priority) = is_dirty {
format!(", dirty({parent_priority})")
} else {
String::new()
};
let has_dirty_containers_label = if has_dirty_containers {
", dirty containers"
} else {
@@ -733,10 +737,25 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
dependent_task = ?reader
)
.entered();
let _ = task.add(CachedDataItem::OutputDependent {
task: reader.unwrap(),
let mut queue = LeafDistanceUpdateQueue::new();
let reader = reader.unwrap();
if task.add(CachedDataItem::OutputDependent {
task: reader,
value: (),
});
}) {
// Ensure that dependent leaf distance is strictly monotonic increasing
let leaf_distance = get!(task, LeafDistance).copied().unwrap_or_default();
let reader_leaf_distance =
get!(reader_task, LeafDistance).copied().unwrap_or_default();
if reader_leaf_distance.distance <= leaf_distance.distance {
queue.push(
reader,
leaf_distance.distance,
leaf_distance.max_distance_in_buffer,
);
}
}
drop(task);
// Note: We use `task_pair` earlier to lock the task and its reader at the same
@@ -753,6 +772,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
value: (),
});
}
drop(reader_task);
queue.execute(&mut ctx);
}
return result;
@@ -779,7 +801,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
// It's not possible that the task is InProgress at this point. If it is InProgress {
// done: true } it must have Output and would early return.
task.add_new(item);
ctx.schedule_task(task);
ctx.schedule_task(task, TaskPriority::Initial);
Ok(Err(listener))
}
@@ -805,10 +827,11 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
if let Some(mut reader_task) = reader_task
&& (!task.is_immutable() || cfg!(feature = "verify_immutable"))
{
let reader = reader.unwrap();
let _ = task.add(CachedDataItem::CellDependent {
cell,
key,
task: reader.unwrap(),
task: reader,
value: (),
});
drop(task);
@@ -832,6 +855,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
value: (),
});
}
drop(reader_task);
}
}
@@ -932,7 +956,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
TaskExecutionReason::CellNotAvailable,
|| self.get_task_desc_fn(task_id),
));
ctx.schedule_task(task);
ctx.schedule_task(task, TaskPriority::Initial);
Ok(Err(listener))
}
@@ -1632,6 +1656,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
fn try_start_task_execution(
&self,
task_id: TaskId,
priority: TaskPriority,
turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
) -> Option<TaskExecutionSpec<'_>> {
enum TaskType {
@@ -1753,7 +1778,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
arg,
} = &*task_type;
(
native_fn.span(task_id.persistence(), execution_reason),
native_fn.span(task_id.persistence(), execution_reason, priority),
native_fn.execute(*this, &**arg),
)
}
@@ -2439,7 +2464,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
let old_dirtyness = get!(task, Dirty).cloned();
let (old_self_dirty, old_current_session_self_clean) = match old_dirtyness {
None => (false, false),
Some(Dirtyness::Dirty) => (true, false),
Some(Dirtyness::Dirty(_)) => (true, false),
Some(Dirtyness::SessionDependent) => {
let clean_in_current_session = get!(task, CurrentSessionClean).is_some();
(true, clean_in_current_session)
@@ -2944,7 +2969,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
let mut task = ctx.task(task_id, TaskDataCategory::All);
let is_dirty = task.is_dirty();
let has_dirty_containers = task.has_dirty_containers();
if is_dirty || has_dirty_containers {
if is_dirty.is_some() || has_dirty_containers {
if let Some(activeness_state) = get_mut!(task, Activeness) {
// We will finish the task, but it would be removed after the task is done
activeness_state.unset_root_type();
@@ -3283,9 +3308,11 @@ impl<B: BackingStorage> Backend for TurboTasksBackend<B> {
fn try_start_task_execution(
&self,
task_id: TaskId,
priority: TaskPriority,
turbo_tasks: &dyn TurboTasksBackendApi<Self>,
) -> Option<TaskExecutionSpec<'_>> {
self.0.try_start_task_execution(task_id, turbo_tasks)
self.0
.try_start_task_execution(task_id, priority, turbo_tasks)
}
fn task_execution_completed(
@@ -883,8 +883,8 @@ pub struct AggregationUpdateQueue {
#[bincode(with = "encode_jobs")]
jobs: VecDeque<AggregationUpdateJobItem>,
#[bincode(with = "turbo_bincode::indexmap")]
number_updates: FxIndexMap<TaskId, AggregationNumberUpdate>,
done_number_updates: FxHashMap<TaskId, AggregationNumberUpdate>,
aggregation_number_updates: FxIndexMap<TaskId, AggregationNumberUpdate>,
done_aggregation_number_updates: FxHashMap<TaskId, AggregationNumberUpdate>,
#[bincode(with = "turbo_bincode::ringset")]
find_and_schedule: FxRingSet<FindAndScheduleJob>,
#[bincode(with = "turbo_bincode::ringset")]
@@ -898,8 +898,8 @@ impl AggregationUpdateQueue {
pub fn new() -> Self {
Self {
jobs: VecDeque::with_capacity(0),
number_updates: FxIndexMap::default(),
done_number_updates: FxHashMap::default(),
aggregation_number_updates: FxIndexMap::default(),
done_aggregation_number_updates: FxHashMap::default(),
find_and_schedule: FxRingSet::default(),
balance_queue: FxRingSet::default(),
optimize_queue: FxRingSet::default(),
@@ -910,14 +910,14 @@ impl AggregationUpdateQueue {
pub fn is_empty(&self) -> bool {
let Self {
jobs,
number_updates,
aggregation_number_updates,
find_and_schedule,
balance_queue,
optimize_queue,
done_number_updates: _,
done_aggregation_number_updates: _,
} = self;
jobs.is_empty()
&& number_updates.is_empty()
&& aggregation_number_updates.is_empty()
&& find_and_schedule.is_empty()
&& balance_queue.is_empty()
&& optimize_queue.is_empty()
@@ -931,7 +931,7 @@ impl AggregationUpdateQueue {
base_aggregation_number,
distance,
} => {
match self.number_updates.entry(task_id) {
match self.aggregation_number_updates.entry(task_id) {
Entry::Occupied(mut entry) => {
let update = entry.get_mut();
update.base_aggregation_number =
@@ -945,7 +945,7 @@ impl AggregationUpdateQueue {
}
}
Entry::Vacant(entry) => {
match self.done_number_updates.entry(task_id) {
match self.done_aggregation_number_updates.entry(task_id) {
HashMapEntry::Occupied(mut entry) => {
let update = entry.get_mut();
let change =
@@ -1207,7 +1207,7 @@ impl AggregationUpdateQueue {
}
}
false
} else if !self.number_updates.is_empty() {
} else if !self.aggregation_number_updates.is_empty() {
let mut remaining = MAX_COUNT_BEFORE_YIELD;
while remaining > 0 {
if let Some((
@@ -1218,11 +1218,11 @@ impl AggregationUpdateQueue {
#[cfg(feature = "trace_aggregation_update")]
span,
},
)) = self.number_updates.pop()
)) = self.aggregation_number_updates.pop()
{
#[cfg(feature = "trace_aggregation_update")]
let _guard = span.map(|s| s.entered());
self.done_number_updates.insert(
self.done_aggregation_number_updates.insert(
task_id,
AggregationNumberUpdate {
base_aggregation_number,
@@ -1498,10 +1498,13 @@ impl AggregationUpdateQueue {
) {
// Task need to be scheduled if it's dirty or doesn't have output
let dirty = task.is_dirty();
let should_schedule = if dirty {
Some(TaskExecutionReason::ActivateDirty)
let should_schedule = if let Some(parent_priority) = dirty {
Some((TaskExecutionReason::ActivateDirty, parent_priority))
} else if !task.has_key(&CachedDataItemKey::Output {}) {
Some(TaskExecutionReason::ActivateInitial)
Some((
TaskExecutionReason::ActivateInitial,
ctx.get_current_task_priority(),
))
} else {
None
};
@@ -1512,7 +1515,7 @@ impl AggregationUpdateQueue {
if !is_active_until_clean {
let mut dirty_containers = task.dirty_containers().peekable();
let is_empty = dirty_containers.peek().is_none();
if !is_empty || dirty {
if !is_empty || dirty.is_some() {
self.extend_find_and_schedule_dirty(dirty_containers);
let activeness_state =
@@ -1520,11 +1523,11 @@ impl AggregationUpdateQueue {
activeness_state.set_active_until_clean();
}
}
if let Some(reason) = should_schedule {
if let Some((reason, parent_priority)) = should_schedule {
let description = || ctx.get_task_desc_fn(task_id);
if task.add(CachedDataItem::new_scheduled(reason, description)) {
drop(task);
ctx.schedule(task_id);
ctx.schedule(task_id, parent_priority);
}
}
}
@@ -34,7 +34,7 @@ impl ConnectChildOperation {
new_children, ..
})) = get_mut!(parent_task, InProgress)
else {
panic!("Task is not in progress while calling another task");
panic!("Task is not in progress while calling another task: {parent_task:?}");
};
// Quick skip if the child was already connected before
@@ -74,7 +74,7 @@ impl ConnectChildOperation {
|| ctx.get_task_desc_fn(child_task_id),
))
{
ctx.schedule_task(child_task);
ctx.schedule_task(child_task, ctx.get_current_task_priority());
}
}
@@ -13,10 +13,7 @@ use crate::{
},
storage::{get, get_mut, remove},
},
data::{
CachedDataItem, CachedDataItemKey, CachedDataItemValue, Dirtyness, InProgressState,
InProgressStateInner,
},
data::{CachedDataItem, CachedDataItemKey, Dirtyness, InProgressState, InProgressStateInner},
};
#[derive(Encode, Decode, Clone, Default)]
@@ -247,13 +244,9 @@ pub fn make_task_dirty_internal(
.entered();
*stale = true;
}
let old = task.insert(CachedDataItem::Dirty {
value: Dirtyness::Dirty,
});
let (old_self_dirty, old_current_session_self_clean) = match old {
Some(CachedDataItemValue::Dirty {
value: Dirtyness::Dirty,
}) => {
let current = get!(task, Dirty);
let (old_self_dirty, old_current_session_self_clean, parent_priority) = match current {
Some(Dirtyness::Dirty(current_priority)) => {
#[cfg(feature = "trace_task_dirty")]
let _span = tracing::trace_span!(
"task already dirty",
@@ -263,17 +256,26 @@ pub fn make_task_dirty_internal(
)
.entered();
// already dirty
let parent_priority = ctx.get_current_task_priority();
if current_priority >= &parent_priority {
// Update the priority to be the lower one
task.insert(CachedDataItem::Dirty {
value: Dirtyness::Dirty(parent_priority),
});
}
return;
}
Some(CachedDataItemValue::Dirty {
value: Dirtyness::SessionDependent,
}) => {
Some(Dirtyness::SessionDependent) => {
let parent_priority = ctx.get_current_task_priority();
task.insert(CachedDataItem::Dirty {
value: Dirtyness::Dirty(parent_priority),
});
// It was a session-dependent dirty before, so we need to remove that clean count
let was_current_session_clean = remove!(task, CurrentSessionClean).is_some();
if was_current_session_clean {
// There was a clean count for a session. If it was the current session, we need to
// propagate that change.
(true, true)
(true, true, parent_priority)
} else {
#[cfg(feature = "trace_task_dirty")]
let _span = tracing::trace_span!(
@@ -287,10 +289,13 @@ pub fn make_task_dirty_internal(
}
}
None => {
let parent_priority = ctx.get_current_task_priority();
task.insert(CachedDataItem::Dirty {
value: Dirtyness::Dirty(parent_priority),
});
// It was clean before, so we need to increase the dirty count
(false, false)
(false, false, parent_priority)
}
_ => unreachable!(),
};
let new_self_dirty = true;
@@ -343,7 +348,7 @@ pub fn make_task_dirty_internal(
)) {
drop(task);
let task = ctx.task(task_id, TaskDataCategory::All);
ctx.schedule_task(task);
ctx.schedule_task(task, parent_priority);
}
}
}
@@ -0,0 +1,198 @@
use std::{
cmp::Reverse,
collections::{BinaryHeap, hash_map::Entry},
};
use bincode::{Decode, Encode};
use rustc_hash::FxHashMap;
#[cfg(feature = "trace_leaf_distance_update")]
use tracing::{span::Span, trace_span};
use turbo_tasks::TaskId;
use crate::{
backend::{
TaskDataCategory,
operation::{ExecuteContext, Operation, TaskGuard},
storage::{get, iter_many},
},
data::CachedDataItem,
};
/// The maximum number of leaf distance updates to process before yielding back to the executor.
/// This prevents long blocking operations and allows to interrupt the processing for persistent
/// caching.
const MAX_COUNT_BEFORE_YIELD: usize = 1000;
/// We avoid incrementing the leaf distance by 1 each time to avoid frequent updates.
/// Instead we use a buffer zone that shrinks as the leaf distance increases.
/// This constant defines the size of that buffer zone at leaf distance 0.
const BASE_LEAF_DISTANCE_BUFFER: u32 = 128;
/// An leaf distance update job that is enqueued.
#[derive(Encode, Decode, Clone)]
struct LeafDistanceUpdate {
dependencies_distance: u32,
dependencies_max_distance_in_buffer: u32,
done: bool,
#[cfg(feature = "trace_leaf_distance_update")]
#[bincode(skip, default)]
span: Option<Span>,
}
impl LeafDistanceUpdate {
fn add(&mut self, dependency_distance: u32, dependency_max_distance_in_buffer: u32) {
self.dependencies_distance = self.dependencies_distance.max(dependency_distance);
self.dependencies_max_distance_in_buffer = self
.dependencies_max_distance_in_buffer
.max(dependency_max_distance_in_buffer);
}
}
/// A queue of leaf distance update jobs.
/// It will execute these jobs in order of their minimum dependency leaf distance.
/// This ensures that we never have to re-process a task.
#[derive(Default, Encode, Decode, Clone)]
pub struct LeafDistanceUpdateQueue {
queue: BinaryHeap<(Reverse<u32>, TaskId)>,
leaf_distance_updates: FxHashMap<TaskId, LeafDistanceUpdate>,
}
impl LeafDistanceUpdateQueue {
pub fn new() -> Self {
Self::default()
}
pub fn is_empty(&self) -> bool {
self.queue.is_empty()
}
pub fn push(
&mut self,
task_id: TaskId,
dependency_distance: u32,
dependency_max_distance_in_buffer: u32,
) {
match self.leaf_distance_updates.entry(task_id) {
Entry::Occupied(mut entry) => {
let update = entry.get_mut();
if update.done && update.dependencies_distance < dependency_distance {
update.done = false;
self.queue.push((Reverse(dependency_distance), task_id));
}
update.add(dependency_distance, dependency_max_distance_in_buffer);
}
Entry::Vacant(entry) => {
entry.insert(LeafDistanceUpdate {
dependencies_distance: dependency_distance,
dependencies_max_distance_in_buffer: dependency_max_distance_in_buffer,
done: false,
});
self.queue.push((Reverse(dependency_distance), task_id));
}
};
}
/// Executes a single step of the queue. Returns true, when the queue is empty.
pub fn process(&mut self, ctx: &mut impl ExecuteContext) -> bool {
let mut remaining = MAX_COUNT_BEFORE_YIELD;
while remaining > 0 {
if let Some((Reverse(queue_dependencies_distance), task_id)) = self.queue.pop() {
let &mut LeafDistanceUpdate {
dependencies_distance,
dependencies_max_distance_in_buffer,
ref mut done,
#[cfg(feature = "trace_leaf_distance_update")]
span,
} = self.leaf_distance_updates.get_mut(&task_id).unwrap();
if queue_dependencies_distance != dependencies_distance {
// Stale entry in queue
// Re-enqueue to keep the ordering correct
self.queue.push((Reverse(dependencies_distance), task_id));
continue;
}
#[cfg(feature = "trace_leaf_distance_update")]
let _guard = span.map(|s| s.entered());
*done = true;
self.update_leaf_distance(
ctx,
task_id,
dependencies_distance,
dependencies_max_distance_in_buffer,
);
remaining -= 1;
} else {
return true;
}
}
false
}
fn update_leaf_distance(
&mut self,
ctx: &mut impl ExecuteContext,
task_id: TaskId,
dependencies_distance: u32,
dependencies_max_distance_in_buffer: u32,
) {
#[cfg(feature = "trace_leaf_distance_update")]
let _span = trace_span!(
"update leaf distance",
dependencies_distance,
dependencies_max_distance_in_buffer
)
.entered();
let mut task = ctx.task(
task_id,
// For performance reasons this should stay `Data` and not `All`
TaskDataCategory::Data,
);
debug_assert!(dependencies_max_distance_in_buffer < u32::MAX / 2);
let mut leaf_distance = get!(task, LeafDistance).copied().unwrap_or_default();
if leaf_distance.distance > dependencies_distance {
// It is strictly monotonic. No need to update.
return;
}
// It's not strictly monotonic, we need to update
if leaf_distance.max_distance_in_buffer <= dependencies_distance {
// We overshoot the buffer zone.
let old_value = leaf_distance.distance;
leaf_distance.distance = dependencies_max_distance_in_buffer + 1;
let buffer_size = BASE_LEAF_DISTANCE_BUFFER
- BASE_LEAF_DISTANCE_BUFFER.saturating_mul(old_value) / leaf_distance.distance;
leaf_distance.max_distance_in_buffer = leaf_distance.distance + buffer_size;
} else {
// We are within the buffer zone, keep the max as is
leaf_distance.distance = dependencies_distance + 1;
}
let dependents = iter_many!(task, OutputDependent { task } => task)
// TODO Technically this is also needed, but there are cycles in the CellDependent graph
// So we need to handle that properly first
// When enabling this, make sure to also call the leaf update queue when adding CellDependents
// .chain(iter_many!(task, CellDependent { task, .. } => task))
;
for dependent_id in dependents {
self.push(
dependent_id,
leaf_distance.distance,
leaf_distance.max_distance_in_buffer,
);
}
task.insert(CachedDataItem::LeafDistance {
value: leaf_distance,
});
}
}
impl Operation for LeafDistanceUpdateQueue {
fn execute(mut self, ctx: &mut impl ExecuteContext) {
if self.is_empty() {
return;
}
loop {
ctx.operation_suspend_point(&self);
if self.process(ctx) {
return;
}
}
}
}
@@ -3,6 +3,7 @@ mod cleanup_old_edges;
mod connect_child;
mod connect_children;
mod invalidate;
mod leaf_distance_update;
mod prepare_new_children;
mod update_cell;
mod update_collectible;
@@ -15,7 +16,8 @@ use std::{
use bincode::{Decode, Encode};
use turbo_tasks::{
CellId, FxIndexMap, KeyValuePair, TaskId, TurboTasksBackendApi, TypedSharedReference,
CellId, FxIndexMap, KeyValuePair, TaskId, TaskPriority, TurboTasksBackendApi,
TypedSharedReference,
};
use crate::{
@@ -77,8 +79,9 @@ pub trait ExecuteContext<'e>: Sized {
task_id2: TaskId,
category: TaskDataCategory,
) -> (Self::TaskGuardImpl, Self::TaskGuardImpl);
fn schedule(&mut self, task_id: TaskId);
fn schedule_task(&self, task: Self::TaskGuardImpl);
fn schedule(&mut self, task_id: TaskId, parent_priority: TaskPriority);
fn schedule_task(&self, task: Self::TaskGuardImpl, parent_priority: TaskPriority);
fn get_current_task_priority(&self) -> TaskPriority;
fn operation_suspend_point<T>(&mut self, op: &T)
where
T: Clone + Into<AnyOperation>;
@@ -557,13 +560,28 @@ where
)
}
fn schedule(&mut self, task_id: TaskId) {
fn schedule(&mut self, task_id: TaskId, parent_priority: TaskPriority) {
let task = self.task(task_id, TaskDataCategory::All);
self.schedule_task(task);
self.schedule_task(task, parent_priority);
}
fn schedule_task(&self, task: Self::TaskGuardImpl) {
self.turbo_tasks.schedule(task.id());
fn schedule_task(&self, task: Self::TaskGuardImpl, parent_priority: TaskPriority) {
let priority = if get!(task, Output).is_some() {
TaskPriority::invalidation(
get!(task, LeafDistance)
.copied()
.unwrap_or_default()
.distance,
)
} else {
TaskPriority::initial()
};
self.turbo_tasks
.schedule(task.id(), priority.in_parent(parent_priority));
}
fn get_current_task_priority(&self) -> TaskPriority {
self.turbo_tasks.get_current_task_priority()
}
fn operation_suspend_point<T: Clone + Into<AnyOperation>>(&mut self, op: &T) {
@@ -667,15 +685,21 @@ pub trait TaskGuard: Debug {
fn is_immutable(&self) -> bool {
self.has_key(&CachedDataItemKey::Immutable {})
}
fn is_dirty(&self) -> bool {
get!(self, Dirty).is_some_and(|dirtyness| match dirtyness {
Dirtyness::Dirty => true,
Dirtyness::SessionDependent => get!(self, CurrentSessionClean).is_none(),
fn is_dirty(&self) -> Option<TaskPriority> {
get!(self, Dirty).and_then(|dirtyness| match dirtyness {
Dirtyness::Dirty(priority) => Some(*priority),
Dirtyness::SessionDependent => {
if get!(self, CurrentSessionClean).is_none() {
Some(TaskPriority::leaf())
} else {
None
}
}
})
}
fn dirtyness_and_session(&self) -> Option<(Dirtyness, bool)> {
match get!(self, Dirty)? {
Dirtyness::Dirty => Some((Dirtyness::Dirty, false)),
Dirtyness::Dirty(priority) => Some((Dirtyness::Dirty(*priority), false)),
Dirtyness::SessionDependent => Some((
Dirtyness::SessionDependent,
get!(self, CurrentSessionClean).is_some(),
@@ -686,7 +710,7 @@ pub trait TaskGuard: Debug {
fn dirty(&self) -> (bool, bool) {
match get!(self, Dirty) {
None => (false, false),
Some(Dirtyness::Dirty) => (true, false),
Some(Dirtyness::Dirty(_)) => (true, false),
Some(Dirtyness::SessionDependent) => (true, get!(self, CurrentSessionClean).is_some()),
}
}
@@ -1043,6 +1067,7 @@ pub enum AnyOperation {
UpdateCell(update_cell::UpdateCellOperation),
CleanupOldEdges(cleanup_old_edges::CleanupOldEdgesOperation),
AggregationUpdate(aggregation_update::AggregationUpdateQueue),
LeafDistanceUpdate(leaf_distance_update::LeafDistanceUpdateQueue),
Nested(Vec<AnyOperation>),
}
@@ -1054,6 +1079,7 @@ impl AnyOperation {
AnyOperation::UpdateCell(op) => op.execute(ctx),
AnyOperation::CleanupOldEdges(op) => op.execute(ctx),
AnyOperation::AggregationUpdate(op) => op.execute(ctx),
AnyOperation::LeafDistanceUpdate(op) => op.execute(ctx),
AnyOperation::Nested(ops) => {
for op in ops {
op.execute(ctx);
@@ -1068,6 +1094,7 @@ impl_operation!(Invalidate invalidate::InvalidateOperation);
impl_operation!(UpdateCell update_cell::UpdateCellOperation);
impl_operation!(CleanupOldEdges cleanup_old_edges::CleanupOldEdgesOperation);
impl_operation!(AggregationUpdate aggregation_update::AggregationUpdateQueue);
impl_operation!(LeafDistanceUpdate leaf_distance_update::LeafDistanceUpdateQueue);
#[cfg(feature = "trace_task_dirty")]
pub use self::invalidate::TaskDirtyCause;
@@ -12,7 +12,8 @@ use crate::{
backend::dynamic_storage::DynamicStorage,
data::{
AggregationNumber, CachedDataItem, CachedDataItemKey, CachedDataItemType,
CachedDataItemValue, CachedDataItemValueRef, CachedDataItemValueRefMut, OutputValue,
CachedDataItemValue, CachedDataItemValueRef, CachedDataItemValueRefMut, LeafDistance,
OutputValue,
},
data_storage::{AutoMapStorage, DefaultStorage, OptionStorage},
utils::{
@@ -138,6 +139,7 @@ impl InnerStorageState {
}
pub struct InnerStorageSnapshot {
leaf_distance: DefaultStorage<LeafDistance>,
aggregation_number: DefaultStorage<AggregationNumber>,
output_dependent: AutoMapStorage<TaskId, ()>,
output: OptionStorage<OutputValue>,
@@ -150,6 +152,7 @@ pub struct InnerStorageSnapshot {
impl From<&InnerStorage> for InnerStorageSnapshot {
fn from(inner: &InnerStorage) -> Self {
Self {
leaf_distance: inner.leaf_distance.clone(),
aggregation_number: inner.aggregation_number.clone(),
output_dependent: inner.output_dependent.clone(),
output: inner.output.clone(),
@@ -168,6 +171,12 @@ impl InnerStorageSnapshot {
use crate::data_storage::Storage;
self.dynamic
.iter_all()
.chain(self.leaf_distance.iter().map(|(_, value)| {
(
CachedDataItemKey::LeafDistance {},
CachedDataItemValueRef::LeafDistance { value },
)
}))
.chain(self.aggregation_number.iter().map(|(_, value)| {
(
CachedDataItemKey::AggregationNumber {},
@@ -197,6 +206,7 @@ impl InnerStorageSnapshot {
pub fn len(&self) -> usize {
use crate::data_storage::Storage;
self.dynamic.len()
+ self.leaf_distance.len()
+ self.aggregation_number.len()
+ self.output.len()
+ self.upper.len()
@@ -206,6 +216,7 @@ impl InnerStorageSnapshot {
#[derive(Debug, Clone)]
pub struct InnerStorage {
leaf_distance: DefaultStorage<LeafDistance>,
aggregation_number: DefaultStorage<AggregationNumber>,
output_dependent: AutoMapStorage<TaskId, ()>,
output: OptionStorage<OutputValue>,
@@ -217,6 +228,7 @@ pub struct InnerStorage {
impl InnerStorage {
fn new() -> Self {
Self {
leaf_distance: Default::default(),
aggregation_number: Default::default(),
output_dependent: Default::default(),
output: Default::default(),
@@ -553,32 +565,36 @@ macro_rules! generate_inner_storage {
}
generate_inner_storage!(
LeafDistance => leaf_distance,
AggregationNumber => aggregation_number,
OutputDependent task => output_dependent,
Output => output,
Upper task => upper,
);
enum InnerStorageIter<A, B, C, D, E> {
AggregationNumber(A),
OutputDependent(B),
Output(C),
Upper(D),
Dynamic(E),
enum InnerStorageIter<A, B, C, D, E, F> {
LeafDistance(A),
AggregationNumber(B),
OutputDependent(C),
Output(D),
Upper(E),
Dynamic(F),
}
impl<T, A, B, C, D, E> Iterator for InnerStorageIter<A, B, C, D, E>
impl<T, A, B, C, D, E, F> Iterator for InnerStorageIter<A, B, C, D, E, F>
where
A: Iterator<Item = T>,
B: Iterator<Item = T>,
C: Iterator<Item = T>,
D: Iterator<Item = T>,
E: Iterator<Item = T>,
F: Iterator<Item = T>,
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
match self {
InnerStorageIter::LeafDistance(iter) => iter.next(),
InnerStorageIter::AggregationNumber(iter) => iter.next(),
InnerStorageIter::OutputDependent(iter) => iter.next(),
InnerStorageIter::Output(iter) => iter.next(),
@@ -595,6 +611,12 @@ impl InnerStorage {
use crate::data_storage::Storage;
self.dynamic
.iter_all()
.chain(self.leaf_distance.iter().map(|(_, value)| {
(
CachedDataItemKey::LeafDistance {},
CachedDataItemValueRef::LeafDistance { value },
)
}))
.chain(self.aggregation_number.iter().map(|(_, value)| {
(
CachedDataItemKey::AggregationNumber {},
@@ -624,6 +646,7 @@ impl InnerStorage {
pub fn len(&self) -> usize {
use crate::data_storage::Storage;
self.dynamic.len()
+ self.leaf_distance.len()
+ self.aggregation_number.len()
+ self.output.len()
+ self.upper.len()
@@ -1231,12 +1254,12 @@ mod tests {
// We track this to catch unexpected bloat from type changes.
assert_eq!(
std::mem::size_of::<InnerStorage>(),
128,
136,
"InnerStorage size changed - please review if this is intentional"
);
assert_eq!(
std::mem::size_of::<InnerStorageSnapshot>(),
128,
136,
"InnerStorageSnapshot size changed - please review if this is intentional"
);
}
@@ -453,8 +453,8 @@ mod tests {
// Lazy direct fields
assert!(storage.get_dirty().is_none());
storage.set_dirty(Dirtyness::Dirty);
assert_eq!(storage.get_dirty(), Some(&Dirtyness::Dirty));
storage.set_dirty(Dirtyness::SessionDependent);
assert_eq!(storage.get_dirty(), Some(&Dirtyness::SessionDependent));
// Lazy collection fields (None until accessed via _mut)
assert!(storage.output_dependencies().is_none());
@@ -596,7 +596,7 @@ mod tests {
original.set_output(OutputValue::Output(TaskId::new(42).unwrap()));
original.upper_mut().insert(TaskId::new(100).unwrap(), 7);
original.upper_mut().insert(TaskId::new(200).unwrap(), 3);
original.set_dirty(Dirtyness::Dirty);
original.set_dirty(Dirtyness::SessionDependent);
original.set_aggregated_dirty_container_count(5);
original
.aggregated_dirty_containers_mut()
@@ -1,7 +1,7 @@
use bincode::{Decode, Encode};
use rustc_hash::FxHashSet;
use turbo_tasks::{
CellId, KeyValuePair, SharedReference, TaskExecutionReason, TaskId, TraitTypeId,
CellId, KeyValuePair, SharedReference, TaskExecutionReason, TaskId, TaskPriority, TraitTypeId,
TypedSharedReference, ValueTypeId,
backend::TurboTasksExecutionError,
event::{Event, EventListener},
@@ -167,7 +167,7 @@ impl Eq for ActivenessState {}
#[derive(Debug, Clone, Copy, Encode, Decode, PartialEq, Eq)]
pub enum Dirtyness {
Dirty,
Dirty(TaskPriority),
SessionDependent,
}
@@ -237,6 +237,29 @@ pub struct AggregationNumber {
pub effective: u32,
}
/// Monotonic increasing distance range to leaf nodes when following "dependencies" edges.
/// It is a range and ranges might overlap. There is a strictly monotonic increasing `distance`
/// value. `max_distance_in_buffer` value might not be monotonic. The `max_distance_in_buffer` value
/// is used as buffer zone to avoid too many updates to dependent nodes when the leaf distance
/// increases slightly. When the leaf distance is increased it tries to keep the
/// `max_distance_in_buffer` value equal. When increasing there are three cases:
/// - `distance` >= `distance` of the dependency + 1: no change.
/// - `distance` <= `max_distance_in_buffer`: only `distance` is increased to the smallest possible
/// value.
/// - `distance` > `max_distance_in_buffer`: `distance` is increased to the `max_distance_in_buffer`
/// value of the dependency + 1 and `max_distance_in_buffer` is increased to `distance` + buffer
/// zone.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Encode, Decode)]
pub struct LeafDistance {
/// This is the strictly monotonic increasing minimum leaf distance.
pub distance: u32,
/// A buffer zone value in which is usually safe to increase the leaf distance without causing
/// too many updates to dependent nodes.
/// Newly added dependents might be added within this buffer zone to avoid propagating updates,
/// therefore one can't rely on this being safe. It's only "often safe".
pub max_distance_in_buffer: u32,
}
#[derive(Debug, Clone, KeyValuePair, Encode, Decode)]
pub enum CachedDataItem {
// Output
@@ -313,6 +336,11 @@ pub enum CachedDataItem {
value: (),
},
// Priority
LeafDistance {
value: LeafDistance,
},
// Aggregation Graph
AggregationNumber {
value: AggregationNumber,
@@ -443,6 +471,7 @@ impl CachedDataItem {
CachedDataItem::AggregationNumber { .. } => true,
CachedDataItem::Follower { task, .. } => !task.is_transient(),
CachedDataItem::Upper { task, .. } => !task.is_transient(),
CachedDataItem::LeafDistance { .. } => true,
CachedDataItem::AggregatedDirtyContainer { task, .. } => !task.is_transient(),
CachedDataItem::AggregatedCurrentSessionCleanContainer { .. } => false,
CachedDataItem::AggregatedCollectible { collectible, .. } => {
@@ -508,7 +537,8 @@ impl CachedDataItem {
| Self::CellDependency { .. }
| Self::CollectiblesDependency { .. }
| Self::OutputDependent { .. }
| Self::CellDependent { .. } => TaskDataCategory::Data,
| Self::CellDependent { .. }
| Self::LeafDistance { .. } => TaskDataCategory::Data,
Self::Collectible { .. }
| Self::Output { .. }
@@ -573,6 +603,7 @@ impl CachedDataItemKey {
CachedDataItemKey::AggregationNumber { .. } => true,
CachedDataItemKey::Follower { task, .. } => !task.is_transient(),
CachedDataItemKey::Upper { task, .. } => !task.is_transient(),
CachedDataItemKey::LeafDistance { .. } => true,
CachedDataItemKey::AggregatedDirtyContainer { task, .. } => !task.is_transient(),
CachedDataItemKey::AggregatedCurrentSessionCleanContainer { .. } => false,
CachedDataItemKey::AggregatedCollectible { collectible, .. } => {
@@ -606,7 +637,8 @@ impl CachedDataItemType {
| Self::CellDependency { .. }
| Self::CollectiblesDependency { .. }
| Self::OutputDependent { .. }
| Self::CellDependent { .. } => TaskDataCategory::Data,
| Self::CellDependent { .. }
| Self::LeafDistance { .. } => TaskDataCategory::Data,
Self::Collectible { .. }
| Self::Output { .. }
@@ -653,6 +685,7 @@ impl CachedDataItemType {
| Self::AggregationNumber
| Self::Follower
| Self::Upper
| Self::LeafDistance
| Self::AggregatedDirtyContainer
| Self::AggregatedCollectible
| Self::AggregatedDirtyContainerCount
@@ -73,15 +73,14 @@ async fn spawns_detached(
notify: TransientInstance<NotifyTaskInput>,
sender: TransientInstance<WatchSenderTaskInput<Option<Vc<u32>>>>,
) -> Vc<()> {
tokio::spawn(turbo_tasks().detached_for_testing(Box::pin(async move {
turbo_tasks().spawn_detached_for_testing(Box::pin(async move {
println!("spawns_detached: waiting for notify");
notify.0.notified().await;
println!("spawns_detached: notified, sending value");
// creating cells after the normal lifetime of the task should be okay, as the parent task
// is waiting on us before exiting!
sender.0.send(Some(Vc::cell(42))).unwrap();
Ok(())
})));
}));
Vc::cell(())
}
@@ -145,10 +144,10 @@ async fn spawns_detached_changing(
changing_input_outer: Vc<ChangingInput>,
) -> Vc<u32> {
let tt = turbo_tasks();
tokio::spawn(tt.clone().detached_for_testing(Box::pin(async move {
tt.clone().spawn_detached_for_testing(Box::pin(async move {
sleep(Duration::from_millis(100)).await;
// nested detached_for_testing calls should work
tokio::spawn(tt.clone().detached_for_testing(Box::pin(async move {
// nested spawn_detached_for_testing calls should work
tt.clone().spawn_detached_for_testing(Box::pin(async move {
sleep(Duration::from_millis(100)).await;
// creating cells after the normal lifetime of the task should be okay, as the parent
// task is waiting on us before exiting!
@@ -158,10 +157,8 @@ async fn spawns_detached_changing(
*read_changing_input(changing_input_detached).await.unwrap(),
)))
.unwrap();
Ok(())
})));
Ok(())
})));
}));
}));
Vc::cell(*read_changing_input(changing_input_outer).await.unwrap())
}
@@ -292,10 +292,10 @@ impl TurboTasksApi for VcStorage {
// no-op
}
fn detached_for_testing(
fn spawn_detached_for_testing(
&self,
_f: std::pin::Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>,
) -> std::pin::Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>> {
_f: std::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
) {
unimplemented!()
}
+2 -1
View File
@@ -26,7 +26,7 @@ use turbo_rcstr::RcStr;
use crate::{
RawVc, ReadCellOptions, ReadOutputOptions, ReadRef, SharedReference, TaskId, TaskIdSet,
TraitRef, TraitTypeId, TurboTasksPanic, ValueTypeId, VcValueTrait, VcValueType,
TaskPriority, TraitRef, TraitTypeId, TurboTasksPanic, ValueTypeId, VcValueTrait, VcValueType,
event::EventListener, macro_helpers::NativeFunction, magic_any::MagicAny,
manager::TurboTasksBackendApi, raw_vc::CellId, registry,
task::shared_reference::TypedSharedReference, task_statistics::TaskStatisticsApi,
@@ -414,6 +414,7 @@ pub trait Backend: Sync + Send {
fn try_start_task_execution<'a>(
&'a self,
task: TaskId,
priority: TaskPriority,
turbo_tasks: &dyn TurboTasksBackendApi<Self>,
) -> Option<TaskExecutionSpec<'a>>;
+4 -2
View File
@@ -38,6 +38,7 @@
#![feature(ptr_metadata)]
#![feature(sync_unsafe_cell)]
#![feature(async_fn_traits)]
#![feature(impl_trait_in_assoc_type)]
pub mod backend;
mod capture_future;
@@ -68,6 +69,7 @@ mod output;
pub mod panic_hooks;
pub mod parallel;
pub mod primitives;
mod priority_runner;
mod raw_vc;
mod read_options;
mod read_ref;
@@ -112,8 +114,8 @@ pub use crate::{
magic_any::MagicAny,
manager::{
CurrentCellRef, ReadCellTracking, ReadConsistency, ReadTracking, TaskPersistence,
TurboTasks, TurboTasksApi, TurboTasksBackendApi, TurboTasksCallApi, Unused, UpdateInfo,
dynamic_call, emit, get_serialization_invalidator, mark_finished, mark_root,
TaskPriority, TurboTasks, TurboTasksApi, TurboTasksBackendApi, TurboTasksCallApi, Unused,
UpdateInfo, dynamic_call, emit, get_serialization_invalidator, mark_finished, mark_root,
mark_session_dependent, prevent_gc, run, run_once, run_once_with_reason, trait_call,
turbo_tasks, turbo_tasks_scope, turbo_tasks_weak, with_turbo_tasks,
},
+312 -185
View File
@@ -1,5 +1,6 @@
use std::{
fmt::Display,
cmp::Reverse,
fmt::{Debug, Display},
future::Future,
hash::{BuildHasher, BuildHasherDefault},
mem::take,
@@ -14,12 +15,13 @@ use std::{
use anyhow::{Result, anyhow};
use auto_hash_map::AutoMap;
use bincode::{Decode, Encode};
use either::Either;
use futures::stream::FuturesUnordered;
use rustc_hash::{FxBuildHasher, FxHasher};
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use tokio::{select, sync::mpsc::Receiver, task_local};
use tokio_util::task::TaskTracker;
use tracing::{Instrument, instrument};
use tracing::{Instrument, Span, instrument};
use crate::{
Completion, InvalidationReason, InvalidationReasonSet, OutputContent, ReadCellOptions,
@@ -37,6 +39,7 @@ use crate::{
macro_helpers::NativeFunction,
magic_any::MagicAny,
message_queue::{CompilationEvent, CompilationEventQueue},
priority_runner::{Executor, JoinHandle, PriorityRunner},
raw_vc::{CellId, RawVc},
registry,
serialization_invalidation::SerializationInvalidator,
@@ -177,10 +180,7 @@ pub trait TurboTasksApi: TurboTasksCallApi + Sync + Send {
///
/// Beware: this method is not safe to use in production code. It is only intended for use in
/// tests and for debugging purposes.
fn detached_for_testing(
&self,
f: Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>,
) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>;
fn spawn_detached_for_testing(&self, f: Pin<Box<dyn Future<Output = ()> + Send + 'static>>);
fn task_statistics(&self) -> &TaskStatisticsApi;
@@ -242,7 +242,10 @@ pub trait TurboTasksBackendApi<B: Backend + 'static>: TurboTasksCallApi + Sync +
unsafe fn reuse_transient_task_id(&self, id: Unused<TaskId>);
/// Schedule a task for execution.
fn schedule(&self, task: TaskId);
fn schedule(&self, task: TaskId, priority: TaskPriority);
/// Returns the priority of the current task.
fn get_current_task_priority(&self) -> TaskPriority;
/// Schedule a foreground backend job for execution.
fn schedule_backend_foreground_job(&self, job: B::BackendJob);
@@ -399,6 +402,75 @@ impl Display for ReadTracking {
}
}
#[derive(Encode, Decode, Default, Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub enum TaskPriority {
#[default]
Initial,
Invalidation {
priority: Reverse<u32>,
},
}
impl TaskPriority {
pub fn invalidation(priority: u32) -> Self {
Self::Invalidation {
priority: Reverse(priority),
}
}
pub fn initial() -> Self {
Self::Initial
}
pub fn leaf() -> Self {
Self::Invalidation {
priority: Reverse(0),
}
}
pub fn in_parent(&self, parent_priority: TaskPriority) -> Self {
match self {
TaskPriority::Initial => parent_priority,
TaskPriority::Invalidation { priority } => {
if let TaskPriority::Invalidation {
priority: parent_priority,
} = parent_priority
&& priority.0 < parent_priority.0
{
Self::Invalidation {
priority: Reverse(parent_priority.0.saturating_add(1)),
}
} else {
*self
}
}
}
}
}
impl Display for TaskPriority {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
TaskPriority::Initial => write!(f, "initial"),
TaskPriority::Invalidation { priority } => write!(f, "invalidation({})", priority.0),
}
}
}
enum ScheduledTask {
Task {
task_id: TaskId,
span: Span,
},
LocalTask {
ty: LocalTaskSpec,
persistence: TaskPersistence,
local_task_id: LocalTaskId,
global_task_state: Arc<RwLock<CurrentTaskState>>,
span: Span,
},
}
pub struct TurboTasks<B: Backend + 'static> {
this: Weak<Self>,
backend: B,
@@ -409,6 +481,8 @@ pub struct TurboTasks<B: Backend + 'static> {
currently_scheduled_foreground_jobs: AtomicUsize,
currently_scheduled_background_jobs: AtomicUsize,
scheduled_tasks: AtomicUsize,
priority_runner:
Arc<PriorityRunner<TurboTasks<B>, ScheduledTask, TaskPriority, TurboTasksExecutor>>,
start: Mutex<Option<Instant>>,
aggregated_update: Mutex<(Option<(Duration, usize)>, InvalidationReasonSet)>,
/// Event that is triggered when currently_scheduled_foreground_jobs becomes non-zero
@@ -422,6 +496,10 @@ pub struct TurboTasks<B: Backend + 'static> {
compilation_events: CompilationEventQueue,
}
type LocalTaskTracker = Option<
FuturesUnordered<Either<JoinHandle, Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>>>,
>;
/// Information about a non-local task. A non-local task can contain multiple "local" tasks, which
/// all share the same non-local task state.
///
@@ -433,6 +511,7 @@ pub struct TurboTasks<B: Backend + 'static> {
struct CurrentTaskState {
task_id: Option<TaskId>,
execution_id: ExecutionId,
priority: TaskPriority,
/// True if the current task uses an external invalidator
has_invalidator: bool,
@@ -447,30 +526,32 @@ struct CurrentTaskState {
local_tasks: Vec<LocalTask>,
/// Tracks currently running local tasks, and defers cleanup of the global task until those
/// complete. Also used by `detached_for_testing`.
local_task_tracker: TaskTracker,
/// complete. Also used by `spawn_detached_for_testing`.
local_task_tracker: LocalTaskTracker,
}
impl CurrentTaskState {
fn new(task_id: TaskId, execution_id: ExecutionId) -> Self {
fn new(task_id: TaskId, execution_id: ExecutionId, priority: TaskPriority) -> Self {
Self {
task_id: Some(task_id),
execution_id,
priority,
has_invalidator: false,
cell_counters: Some(AutoMap::default()),
local_tasks: Vec::new(),
local_task_tracker: TaskTracker::new(),
local_task_tracker: None,
}
}
fn new_temporary(execution_id: ExecutionId) -> Self {
fn new_temporary(execution_id: ExecutionId, priority: TaskPriority) -> Self {
Self {
task_id: None,
execution_id,
priority,
has_invalidator: false,
cell_counters: None,
local_tasks: Vec::new(),
local_task_tracker: TaskTracker::new(),
local_task_tracker: None,
}
}
@@ -535,6 +616,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
currently_scheduled_foreground_jobs: AtomicUsize::new(0),
currently_scheduled_background_jobs: AtomicUsize::new(0),
scheduled_tasks: AtomicUsize::new(0),
priority_runner: Arc::new(PriorityRunner::new(TurboTasksExecutor)),
start: Default::default(),
aggregated_update: Default::default(),
event_foreground_done: Event::new(|| {
@@ -574,7 +656,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
})),
self,
);
self.schedule(id);
self.schedule(id, TaskPriority::initial());
id
}
@@ -598,7 +680,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
})),
self,
);
self.schedule(id);
self.schedule(id, TaskPriority::initial());
}
pub async fn run_once<T: TraceRawVcs + Send + 'static>(
@@ -624,8 +706,10 @@ impl<B: Backend + 'static> TurboTasks<B> {
self.begin_foreground_job();
// it's okay for execution ids to overflow and wrap, they're just used for an assert
let execution_id = self.execution_id_factory.wrapping_get();
let current_task_state =
Arc::new(RwLock::new(CurrentTaskState::new_temporary(execution_id)));
let current_task_state = Arc::new(RwLock::new(CurrentTaskState::new_temporary(
execution_id,
TaskPriority::initial(),
)));
let result = TURBO_TASKS
.scope(
@@ -634,10 +718,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
let result = CaptureFuture::new(future).await;
// wait for all spawned local tasks using `local` to finish
let ltt =
CURRENT_TASK_STATE.with(|ts| ts.read().unwrap().local_task_tracker.clone());
ltt.close();
ltt.wait().await;
wait_for_local_tasks().await;
match result {
Ok(Ok(raw_vc)) => Ok(raw_vc),
@@ -745,80 +826,18 @@ impl<B: Backend + 'static> TurboTasks<B> {
}
#[track_caller]
pub(crate) fn schedule(&self, task_id: TaskId) {
pub(crate) fn schedule(&self, task_id: TaskId, priority: TaskPriority) {
self.begin_foreground_job();
self.scheduled_tasks.fetch_add(1, Ordering::AcqRel);
let this = self.pin();
let future = async move {
let mut schedule_again = true;
while schedule_again {
// it's okay for execution ids to overflow and wrap, they're just used for an assert
let execution_id = this.execution_id_factory.wrapping_get();
let current_task_state =
Arc::new(RwLock::new(CurrentTaskState::new(task_id, execution_id)));
let single_execution_future = async {
if this.stopped.load(Ordering::Acquire) {
this.backend.task_execution_canceled(task_id, &*this);
return false;
}
let Some(TaskExecutionSpec { future, span }) =
this.backend.try_start_task_execution(task_id, &*this)
else {
return false;
};
async {
let result = CaptureFuture::new(future).await;
// wait for all spawned local tasks using `local` to finish
let ltt = CURRENT_TASK_STATE
.with(|ts| ts.read().unwrap().local_task_tracker.clone());
ltt.close();
ltt.wait().await;
let result = match result {
Ok(Ok(raw_vc)) => Ok(raw_vc),
Ok(Err(err)) => Err(err.into()),
Err(err) => Err(TurboTasksExecutionError::Panic(Arc::new(err))),
};
let FinishedTaskState { has_invalidator } =
this.finish_current_task_state();
let cell_counters = CURRENT_TASK_STATE
.with(|ts| ts.write().unwrap().cell_counters.take().unwrap());
this.backend.task_execution_completed(
task_id,
result,
&cell_counters,
has_invalidator,
&*this,
)
}
.instrument(span)
.await
};
schedule_again = CURRENT_TASK_STATE
.scope(current_task_state, single_execution_future)
.await;
}
this.finish_foreground_job();
anyhow::Ok(())
};
let future = TURBO_TASKS.scope(self.pin(), future).in_current_span();
#[cfg(feature = "tokio_tracing")]
{
let description = self.backend.get_task_description(task_id);
tokio::task::Builder::new()
.name(&description)
.spawn(future)
.unwrap();
}
#[cfg(not(feature = "tokio_tracing"))]
tokio::task::spawn(future);
self.priority_runner.schedule(
&self.pin(),
ScheduledTask::Task {
task_id,
span: Span::current(),
},
priority,
);
}
fn schedule_local_task(
@@ -828,8 +847,8 @@ impl<B: Backend + 'static> TurboTasks<B> {
persistence: TaskPersistence,
) -> RawVc {
let task_type = ty.task_type;
let (global_task_state, parent_task_id, execution_id, local_task_id) = CURRENT_TASK_STATE
.with(|gts| {
let (global_task_state, execution_id, priority, local_task_id) =
CURRENT_TASK_STATE.with(|gts| {
let mut gts_write = gts.write().unwrap();
let local_task_id = gts_write.create_local_task(LocalTask::Scheduled {
done_event: Event::new(move || {
@@ -838,89 +857,29 @@ impl<B: Backend + 'static> TurboTasks<B> {
});
(
Arc::clone(gts),
gts_write.task_id,
gts_write.execution_id,
gts_write.priority,
local_task_id,
)
});
#[cfg(feature = "tokio_tracing")]
let description = format!(
"[local] (parent: {:?}) {}",
parent_task_id.map(|id| self.backend.get_task_description(id)),
ty.task_type,
let future = self.priority_runner.schedule_with_join_handle(
&self.pin(),
ScheduledTask::LocalTask {
ty,
persistence,
local_task_id,
global_task_state: global_task_state.clone(),
span: Span::current(),
},
priority,
);
#[cfg(not(feature = "tokio_tracing"))]
let _ = parent_task_id; // suppress unused variable warning
let this = self.pin();
let future = async move {
let span = match &ty.task_type {
LocalTaskType::ResolveNative { native_fn } => native_fn.resolve_span(),
LocalTaskType::ResolveTrait { trait_method } => trait_method.resolve_span(),
};
async move {
let result = match ty.task_type {
LocalTaskType::ResolveNative { native_fn } => {
LocalTaskType::run_resolve_native(
native_fn,
ty.this,
&*ty.arg,
persistence,
this,
)
.await
}
LocalTaskType::ResolveTrait { trait_method } => {
LocalTaskType::run_resolve_trait(
trait_method,
ty.this.unwrap(),
&*ty.arg,
persistence,
this,
)
.await
}
};
let output = match result {
Ok(raw_vc) => OutputContent::Link(raw_vc),
Err(err) => OutputContent::Error(
TurboTasksExecutionError::from(err).with_task_context(task_type, None),
),
};
let local_task = LocalTask::Done { output };
let done_event = CURRENT_TASK_STATE.with(move |gts| {
let mut gts_write = gts.write().unwrap();
let scheduled_task =
std::mem::replace(gts_write.get_mut_local_task(local_task_id), local_task);
let LocalTask::Scheduled { done_event } = scheduled_task else {
panic!("local task finished, but was not in the scheduled state?");
};
done_event
});
done_event.notify(usize::MAX)
}
.instrument(span)
.await
};
let future = global_task_state
.read()
global_task_state
.write()
.unwrap()
.local_task_tracker
.track_future(future);
let future = CURRENT_TASK_STATE.scope(global_task_state, future);
let future = TURBO_TASKS.scope(self.pin(), future).in_current_span();
#[cfg(feature = "tokio_tracing")]
tokio::task::Builder::new()
.name(&description)
.spawn(future)
.unwrap();
#[cfg(not(feature = "tokio_tracing"))]
tokio::task::spawn(future);
.get_or_insert_default()
.push(Either::Left(future));
RawVc::LocalOutput(execution_id, local_task_id, persistence)
}
@@ -1193,6 +1152,158 @@ impl<B: Backend + 'static> TurboTasks<B> {
}
}
struct TurboTasksExecutor;
impl<B: Backend> Executor<TurboTasks<B>, ScheduledTask, TaskPriority> for TurboTasksExecutor {
type Future = impl Future<Output = ()> + Send + 'static;
fn execute(
&self,
this: &Arc<TurboTasks<B>>,
scheduled_task: ScheduledTask,
priority: TaskPriority,
) -> Self::Future {
match scheduled_task {
ScheduledTask::Task { task_id, span } => {
let this2 = this.clone();
let this = this.clone();
let future = async move {
let mut schedule_again = true;
while schedule_again {
// it's okay for execution ids to overflow and wrap, they're just used for
// an assert
let execution_id = this.execution_id_factory.wrapping_get();
let current_task_state = Arc::new(RwLock::new(CurrentTaskState::new(
task_id,
execution_id,
priority,
)));
let single_execution_future = async {
if this.stopped.load(Ordering::Acquire) {
this.backend.task_execution_canceled(task_id, &*this);
return false;
}
let Some(TaskExecutionSpec { future, span }) = this
.backend
.try_start_task_execution(task_id, priority, &*this)
else {
return false;
};
async {
let result = CaptureFuture::new(future).await;
// wait for all spawned local tasks using `local` to finish
wait_for_local_tasks().await;
let result = match result {
Ok(Ok(raw_vc)) => Ok(raw_vc),
Ok(Err(err)) => Err(err.into()),
Err(err) => Err(TurboTasksExecutionError::Panic(Arc::new(err))),
};
let FinishedTaskState { has_invalidator } =
this.finish_current_task_state();
let cell_counters = CURRENT_TASK_STATE
.with(|ts| ts.write().unwrap().cell_counters.take().unwrap());
this.backend.task_execution_completed(
task_id,
result,
&cell_counters,
has_invalidator,
&*this,
)
}
.instrument(span)
.await
};
schedule_again = CURRENT_TASK_STATE
.scope(current_task_state, single_execution_future)
.await;
}
this.finish_foreground_job();
};
Either::Left(TURBO_TASKS.scope(this2, future).instrument(span))
}
ScheduledTask::LocalTask {
ty,
persistence,
local_task_id,
global_task_state,
span,
} => {
let this2 = this.clone();
let this = this.clone();
let task_type = ty.task_type;
let future = async move {
let span = match &ty.task_type {
LocalTaskType::ResolveNative { native_fn } => {
native_fn.resolve_span(priority)
}
LocalTaskType::ResolveTrait { trait_method } => {
trait_method.resolve_span(priority)
}
};
async move {
let result = match ty.task_type {
LocalTaskType::ResolveNative { native_fn } => {
LocalTaskType::run_resolve_native(
native_fn,
ty.this,
&*ty.arg,
persistence,
this,
)
.await
}
LocalTaskType::ResolveTrait { trait_method } => {
LocalTaskType::run_resolve_trait(
trait_method,
ty.this.unwrap(),
&*ty.arg,
persistence,
this,
)
.await
}
};
let output = match result {
Ok(raw_vc) => OutputContent::Link(raw_vc),
Err(err) => OutputContent::Error(
TurboTasksExecutionError::from(err)
.with_task_context(task_type, None),
),
};
let local_task = LocalTask::Done { output };
let done_event = CURRENT_TASK_STATE.with(move |gts| {
let mut gts_write = gts.write().unwrap();
let scheduled_task = std::mem::replace(
gts_write.get_mut_local_task(local_task_id),
local_task,
);
let LocalTask::Scheduled { done_event } = scheduled_task else {
panic!("local task finished, but was not in the scheduled state?");
};
done_event
});
done_event.notify(usize::MAX)
}
.instrument(span)
.await
};
let future = CURRENT_TASK_STATE.scope(global_task_state, future);
Either::Right(TURBO_TASKS.scope(this2, future).instrument(span))
}
}
}
}
struct FinishedTaskState {
/// True if the task uses an external invalidator
has_invalidator: bool,
@@ -1434,21 +1545,21 @@ impl<B: Backend + 'static> TurboTasksApi for TurboTasks<B> {
/// Creates a future that inherits the current task id and task state. The current global task
/// will wait for this future to be dropped before exiting.
fn detached_for_testing(
&self,
fut: Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>,
) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>> {
fn spawn_detached_for_testing(&self, fut: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) {
// this is similar to what happens for a local task, except that we keep the local task's
// state as well.
let global_task_state = CURRENT_TASK_STATE.with(|ts| ts.clone());
let tracked_fut = {
let ts = global_task_state.read().unwrap();
ts.local_task_tracker.track_future(fut)
};
Box::pin(TURBO_TASKS.scope(
let fut = tokio::spawn(TURBO_TASKS.scope(
turbo_tasks(),
CURRENT_TASK_STATE.scope(global_task_state, tracked_fut),
))
CURRENT_TASK_STATE.scope(global_task_state.clone(), fut),
));
let fut = Box::pin(async move {
fut.await.unwrap();
});
let mut ts = global_task_state.write().unwrap();
ts.local_task_tracker
.get_or_insert_default()
.push(Either::Right(fut));
}
fn task_statistics(&self) -> &TaskStatisticsApi {
@@ -1505,8 +1616,14 @@ impl<B: Backend + 'static> TurboTasksBackendApi<B> for TurboTasks<B> {
}
#[track_caller]
fn schedule(&self, task: TaskId) {
self.schedule(task)
fn schedule(&self, task: TaskId, priority: TaskPriority) {
self.schedule(task, priority)
}
fn get_current_task_priority(&self) -> TaskPriority {
CURRENT_TASK_STATE
.try_with(|task_state| task_state.read().unwrap().priority)
.unwrap_or(TaskPriority::initial())
}
fn program_duration_until(&self, instant: Instant) -> Duration {
@@ -1538,6 +1655,15 @@ impl<B: Backend + 'static> TurboTasksBackendApi<B> for TurboTasks<B> {
}
}
async fn wait_for_local_tasks() {
while let Some(mut ltt) =
CURRENT_TASK_STATE.with(|ts| ts.write().unwrap().local_task_tracker.take())
{
use futures::StreamExt;
while ltt.next().await.is_some() {}
}
}
pub(crate) fn current_task_if_available(from: &str) -> Option<TaskId> {
match CURRENT_TASK_STATE.try_with(|ts| ts.read().unwrap().task_id) {
Ok(id) => id,
@@ -1671,6 +1797,7 @@ pub fn with_turbo_tasks_for_testing<T>(
Arc::new(RwLock::new(CurrentTaskState::new(
current_task,
execution_id,
TaskPriority::initial(),
))),
f,
),
@@ -1681,8 +1808,8 @@ pub fn with_turbo_tasks_for_testing<T>(
///
/// Beware: this method is not safe to use in production code. It is only
/// intended for use in tests and for debugging purposes.
pub fn spawn_detached_for_testing(f: impl Future<Output = Result<()>> + Send + 'static) {
tokio::spawn(turbo_tasks().detached_for_testing(Box::pin(f.in_current_span())));
pub fn spawn_detached_for_testing(f: impl Future<Output = ()> + Send + 'static) {
turbo_tasks().spawn_detached_for_testing(Box::pin(f));
}
pub fn current_task_for_testing() -> Option<TaskId> {
@@ -8,7 +8,7 @@ use tracing::Span;
use turbo_bincode::{AnyDecodeFn, AnyEncodeFn};
use crate::{
RawVc, TaskExecutionReason, TaskInput, TaskPersistence,
RawVc, TaskExecutionReason, TaskInput, TaskPersistence, TaskPriority,
magic_any::{MagicAny, any_as_encode},
task::{
IntoTaskFn, TaskFn,
@@ -238,7 +238,12 @@ impl NativeFunction {
}
}
pub fn span(&'static self, persistence: TaskPersistence, reason: TaskExecutionReason) -> Span {
pub fn span(
&'static self,
persistence: TaskPersistence,
reason: TaskExecutionReason,
priority: TaskPriority,
) -> Span {
let flags = match persistence {
TaskPersistence::Persistent => "",
TaskPersistence::Transient => "transient",
@@ -246,13 +251,14 @@ impl NativeFunction {
tracing::trace_span!(
"turbo_tasks::function",
name = self.name,
priority = %priority,
flags = flags,
reason = reason.as_str()
)
}
pub fn resolve_span(&'static self) -> Span {
tracing::trace_span!("turbo_tasks::resolve_call", name = self.name)
pub fn resolve_span(&'static self, priority: TaskPriority) -> Span {
tracing::trace_span!("turbo_tasks::resolve_call", name = self.name, priority = %priority)
}
}
impl PartialEq for NativeFunction {
@@ -0,0 +1,631 @@
use std::{
collections::BinaryHeap,
fmt::Debug,
future::Future,
pin::Pin,
ptr::drop_in_place,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
task::{Context, Poll},
time::{Duration, Instant},
};
use parking_lot::Mutex;
use pin_project_lite::pin_project;
use tokio::sync::oneshot::{Receiver, Sender};
pub trait Executor<C, T, P>: Send + Sync {
type Future: Future<Output = ()> + Send;
fn execute(&self, execute_context: &Arc<C>, task: T, priority: P) -> Self::Future;
}
impl<C, T, P, F, Fut> Executor<C, T, P> for F
where
F: Fn(&Arc<C>, T, P) -> Fut + Send + Sync,
Fut: Future<Output = ()> + Send,
{
type Future = Fut;
fn execute(&self, execute_context: &Arc<C>, task: T, priority: P) -> Self::Future {
(self)(execute_context, task, priority)
}
}
struct HeapItem<P, T> {
priority: P,
task: T,
tx: Option<Sender<()>>,
}
impl<P: Eq, T> PartialEq for HeapItem<P, T> {
fn eq(&self, other: &Self) -> bool {
self.priority == other.priority
}
}
impl<P: Eq, T> Eq for HeapItem<P, T> {}
impl<P: Ord, T> Ord for HeapItem<P, T> {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.priority.cmp(&other.priority)
}
}
impl<P: Ord, T> PartialOrd for HeapItem<P, T> {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
pub struct PriorityRunner<
C: Send + Sync + 'static,
T: Send + 'static,
P: Ord + Send + 'static,
E: Executor<C, T, P> + 'static,
> {
executor: E,
/// The target number of workers to spawn.
target_workers: usize,
/// The queue of tasks to execute. These tasks are not scheduled yet.
queue: Mutex<BinaryHeap<HeapItem<P, T>>>,
/// The number of active workers currently polling tasks.
/// Workers that responded with Poll::Pending are not counted until they are polled again.
active_workers: AtomicUsize,
phantom: std::marker::PhantomData<C>,
}
impl<
C: Send + Sync + 'static,
T: Send + 'static,
P: Debug + Ord + Send + 'static,
E: Executor<C, T, P> + 'static,
> PriorityRunner<C, T, P, E>
{
pub fn new(executor: E) -> Self {
Self {
executor,
target_workers: tokio::runtime::Handle::current().metrics().num_workers(),
queue: Mutex::new(BinaryHeap::new()),
active_workers: AtomicUsize::new(0),
phantom: std::marker::PhantomData,
}
}
pub fn schedule(self: &Arc<Self>, execute_context: &Arc<C>, task: T, priority: P) {
self.schedule_internal(execute_context, task, priority, None);
}
pub fn schedule_with_join_handle(
self: &Arc<Self>,
execute_context: &Arc<C>,
task: T,
priority: P,
) -> JoinHandle {
let (tx, rx) = tokio::sync::oneshot::channel();
self.schedule_internal(execute_context, task, priority, Some(tx));
JoinHandle { receiver: rx }
}
fn schedule_internal(
self: &Arc<Self>,
execute_context: &Arc<C>,
task: T,
priority: P,
tx: Option<Sender<()>>,
) {
let mut queue = self.queue.lock();
if !queue.is_empty() {
// If there is already work in the queue, we don't have any
// free capacity so we can just push the task to the queue.
// It will be picked up by existing workers.
queue.push(HeapItem { priority, task, tx });
return;
}
// The queue is empty, so we might have free capacity to spawn a new worker.
let active_workers = self.active_workers.fetch_add(1, Ordering::Relaxed);
if active_workers < self.target_workers {
// We have free capacity, spawn a new worker to execute this task immediately.
drop(queue);
let future = self.executor.execute(execute_context, task, priority);
WorkerFuture::spawn(future, tx, execute_context.clone(), self.clone());
} else {
// No free capacity, push the task to the queue.
queue.push(HeapItem { priority, task, tx });
drop(queue);
// Undo the added active worker since we didn't spawn a new worker.
self.decrease_active_workers(execute_context);
}
}
/// Tries to decrease the active worker count by 1.
/// If there is work available in the queue, a new worker is spawned instead.
fn reuse_or_decrease_active_workers(self: &Arc<Self>, execute_context: &Arc<C>) {
let active_workers = self.active_workers.load(Ordering::Relaxed) - 1;
if active_workers >= self.target_workers
|| !self.spawn_worker_if_work_available(execute_context, true)
{
// Undo the added active worker since we didn't spawn a new worker.
// Beware the race condition here:
// If the active workers became lower in the meantime we might have free
// capacity now, so we try to spawn a new worker if
// there is work available.
self.decrease_active_workers(execute_context);
}
}
/// Tries to decrease the active worker count by 1.
/// If there is work available in the queue, a new worker is spawned instead.
fn decrease_active_workers(self: &Arc<Self>, execute_context: &Arc<C>) {
// If the active workers became lower we might have free
// capacity now, so we try to spawn a new worker if
// there is work available.
let active_workers = self.active_workers.fetch_sub(1, Ordering::Relaxed) - 1;
if active_workers < self.target_workers {
self.spawn_worker_if_work_available(execute_context, false);
}
}
fn pop_future_from_worker(
&self,
execute_context: &Arc<C>,
) -> Option<(E::Future, Option<Sender<()>>)> {
let mut queue = self.queue.lock();
if let Some(heap_item) = queue.pop() {
if queue.len() * 2 + 16 < queue.capacity() {
queue.shrink_to_fit();
}
drop(queue);
let tx = heap_item.tx;
Some((
self.executor
.execute(execute_context, heap_item.task, heap_item.priority),
tx,
))
} else {
None
}
}
fn spawn_worker_if_work_available(
self: &Arc<Self>,
execute_context: &Arc<C>,
unused_active_count: bool,
) -> bool {
let mut queue = self.queue.lock();
if let Some(heap_item) = queue.pop() {
if queue.len() * 2 + 16 < queue.capacity() {
queue.shrink_to_fit();
}
drop(queue);
let tx = heap_item.tx;
let new_future =
self.executor
.execute(execute_context, heap_item.task, heap_item.priority);
if !unused_active_count {
self.active_workers.fetch_add(1, Ordering::Relaxed);
}
WorkerFuture::spawn(new_future, tx, execute_context.clone(), self.clone());
true
} else {
false
}
}
}
#[derive(Debug)]
enum WorkerState {
UnfinishedFuture,
PendingFuture,
Done,
Closed,
}
pin_project! {
struct WorkerFuture<C, T, P, E>
where
// pin_project doesn't support bounds with +
C: Send,
C: Sync,
C: 'static,
T: Send,
T: 'static,
P: Ord,
P: Send,
P: 'static,
E: Executor<C, T, P>,
E: 'static,
{
#[pin]
future: E::Future,
tx: Option<Sender<()>>,
execute_context: Arc<C>,
runner: Arc<PriorityRunner<C, T, P, E>>,
state: WorkerState,
}
}
impl<
C: Send + Sync + 'static,
T: Send + 'static,
P: Debug + Ord + Send + 'static,
E: Executor<C, T, P> + 'static,
> WorkerFuture<C, T, P, E>
{
fn spawn(
future: E::Future,
tx: Option<Sender<()>>,
execute_context: Arc<C>,
runner: Arc<PriorityRunner<C, T, P, E>>,
) {
tokio::task::spawn(Self {
future,
tx,
execute_context,
runner,
state: WorkerState::UnfinishedFuture,
});
}
}
impl<
C: Send + Sync + 'static,
T: Send + 'static,
P: Debug + Ord + Send + 'static,
E: Executor<C, T, P> + 'static,
> Future for WorkerFuture<C, T, P, E>
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();
if matches!(this.state, WorkerState::PendingFuture) {
// When the worker is not active (it previously returned Poll::Pending),
// we need to mark it as active again since it is being polled now.
this.runner.active_workers.fetch_add(1, Ordering::Relaxed);
*this.state = WorkerState::UnfinishedFuture;
}
let last_yield = Instant::now();
loop {
match this.state {
WorkerState::Closed => return Poll::Ready(()),
WorkerState::PendingFuture => unreachable!(),
WorkerState::UnfinishedFuture => {
match this.future.as_mut().poll(cx) {
Poll::Ready(()) => {
// Notify that the task is done
if let Some(tx) = this.tx.take() {
let _ = tx.send(());
}
*this.state = WorkerState::Done;
if last_yield.elapsed() > Duration::from_millis(5) {
cx.waker().wake_by_ref();
return Poll::Pending;
}
}
Poll::Pending => {
// The current future is still pending, we need to suspend this worker.
// But we if there are free capacity we can spawn a new worker to pick
// up other tasks in the queue.
this.runner
.reuse_or_decrease_active_workers(this.execute_context);
*this.state = WorkerState::PendingFuture;
return Poll::Pending;
}
}
}
WorkerState::Done => {
let active_workers = this.runner.active_workers.load(Ordering::Relaxed);
if active_workers > this.runner.target_workers {
// There are more active workers than target, so we should end this
// worker.
this.runner.decrease_active_workers(this.execute_context);
*this.state = WorkerState::Closed;
return Poll::Ready(());
}
// This future is done, we need to check the queue for more tasks,
// so we can continue working on a new future in this worker.
if let Some((new_future, new_tx)) =
this.runner.pop_future_from_worker(this.execute_context)
{
// We are replacing the future with a new one, but the current future is
// pinned. So we need to drop the future in place
// and replace it with the new future, which becomes
// pinned in that place.
// SAFETY: The pinned future is dropped in place
unsafe {
let future_slot = this.future.as_mut().get_unchecked_mut();
let future_slot: *mut E::Future = future_slot;
drop_in_place(future_slot);
future_slot.write(new_future);
}
*this.tx = new_tx;
*this.state = WorkerState::UnfinishedFuture;
} else {
// No more tasks to execute
// This worker ends here
this.runner.decrease_active_workers(this.execute_context);
*this.state = WorkerState::Closed;
return Poll::Ready(());
}
}
}
}
}
}
pub struct JoinHandle {
receiver: Receiver<()>,
}
impl Future for JoinHandle {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
match Pin::new(&mut this.receiver).poll(cx) {
Poll::Ready(result) => {
let _ = result;
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, thread::sleep, time::Duration};
use super::*;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_cpu_bound_tasks() {
struct ExecutorImpl;
impl Executor<Mutex<Vec<u32>>, u32, u32> for ExecutorImpl {
type Future = Pin<Box<dyn Future<Output = ()> + Send>>;
fn execute(
&self,
execute_context: &Arc<Mutex<Vec<u32>>>,
task: u32,
_priority: u32,
) -> Self::Future {
let execute_context = execute_context.clone();
Box::pin(async move {
println!("Executing task {}...", task);
sleep(Duration::from_millis((task as u64 + 1) * 10));
execute_context.lock().push(task);
println!("Finished task {}.", task);
})
}
}
let executor = ExecutorImpl;
let runner: Arc<PriorityRunner<Mutex<Vec<u32>>, u32, u32, _>> =
Arc::new(PriorityRunner::new(executor));
let results = Arc::new(Mutex::new(Vec::new()));
for i in 0..10 {
let results = results.clone();
println!("Scheduling task {}...", i);
runner.schedule(&results, i, i);
}
while results.lock().len() < 10 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
let results = results.lock();
println!("Results: {:?}", *results);
// The first two tasks are directly spawned without queuing
assert_eq!(&results[0..2], &[0, 1]);
// All tasks after that are queued and therefore prioritized
// This means the highest priority tasks are executed next
assert!(results[2..4].contains(&9));
assert!(results[2..4].contains(&8));
// The last tasks are the tasks with the lowest priority
assert!(results[8..10].contains(&2));
assert!(results[8..10].contains(&3));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_cpu_bound_with_yield_tasks() {
struct ExecutorImpl;
impl Executor<Mutex<Vec<u32>>, u32, u32> for ExecutorImpl {
type Future = Pin<Box<dyn Future<Output = ()> + Send>>;
fn execute(
&self,
execute_context: &Arc<Mutex<Vec<u32>>>,
task: u32,
_priority: u32,
) -> Self::Future {
let execute_context = execute_context.clone();
Box::pin(async move {
println!("Executing task {}...", task);
sleep(Duration::from_millis((task as u64 + 1) * 10));
execute_context.lock().push(task);
println!("Finished task {}.", task);
tokio::task::yield_now().await;
})
}
}
let executor = ExecutorImpl;
let runner: Arc<PriorityRunner<Mutex<Vec<u32>>, u32, u32, _>> =
Arc::new(PriorityRunner::new(executor));
let results = Arc::new(Mutex::new(Vec::new()));
for i in 0..10 {
let results = results.clone();
println!("Scheduling task {}...", i);
runner.schedule(&results, i, i);
}
while results.lock().len() < 10 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
let results = results.lock();
println!("Results: {:?}", *results);
// The first two tasks are directly spawned without queuing
assert_eq!(&results[0..2], &[0, 1]);
// All tasks after that are queued and therefore prioritized
// This means the highest priority tasks are executed next
assert!(results[2..4].contains(&9));
assert!(results[2..4].contains(&8));
// The last tasks are the tasks with the lowest priority
assert!(results[8..10].contains(&2));
assert!(results[8..10].contains(&3));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_waiting_tasks() {
struct ExecutorImpl;
impl Executor<Mutex<Vec<u32>>, u32, u32> for ExecutorImpl {
type Future = Pin<Box<dyn Future<Output = ()> + Send>>;
fn execute(
&self,
execute_context: &Arc<Mutex<Vec<u32>>>,
task: u32,
_priority: u32,
) -> Self::Future {
let execute_context = execute_context.clone();
Box::pin(async move {
println!("Executing task {}...", task);
tokio::time::sleep(Duration::from_millis((task as u64 + 1) * 10)).await;
execute_context.lock().push(task);
println!("Finished task {}.", task);
})
}
}
let executor = ExecutorImpl;
let runner: Arc<PriorityRunner<Mutex<Vec<u32>>, u32, u32, _>> =
Arc::new(PriorityRunner::new(executor));
let results = Arc::new(Mutex::new(Vec::new()));
for i in 0..10 {
let results = results.clone();
println!("Scheduling task {}...", i);
runner.schedule(&results, i, i);
}
while results.lock().len() < 10 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
let results = results.lock();
println!("Results: {:?}", *results);
assert_eq!(*results, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
}
#[test]
fn test_mixed_cpu_bound_and_waiting_tasks() {
tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.event_interval(1)
.global_queue_interval(1)
.disable_lifo_slot()
.enable_all()
.build()
.unwrap()
.block_on(test_mixed_cpu_bound_and_waiting_tasks_impl());
}
async fn test_mixed_cpu_bound_and_waiting_tasks_impl() {
struct ExecutorImpl;
impl Executor<Mutex<Vec<u32>>, u32, u32> for ExecutorImpl {
type Future = Pin<Box<dyn Future<Output = ()> + Send>>;
fn execute(
&self,
execute_context: &Arc<Mutex<Vec<u32>>>,
task: u32,
_priority: u32,
) -> Self::Future {
let execute_context = execute_context.clone();
println!("Created task {}", task);
Box::pin(async move {
let cpu_bound = task < 10;
if cpu_bound {
println!("Executing cpu-bound task {}...", task);
// CPU bound task
sleep(Duration::from_millis((task as u64 + 1) * 10));
} else {
println!("Executing waiting task {}...", task);
// Waiting task
tokio::time::sleep(Duration::from_millis((task as u64 + 1) * 10)).await;
}
execute_context.lock().push(task);
if cpu_bound {
println!("Finished cpu-bound task {}.", task);
} else {
println!("Finished waiting task {}.", task);
}
})
}
}
let executor = ExecutorImpl;
let runner: Arc<PriorityRunner<Mutex<Vec<u32>>, u32, u32, _>> =
Arc::new(PriorityRunner::new(executor));
let results = Arc::new(Mutex::new(Vec::new()));
for i in 0..20 {
let results = results.clone();
println!("Scheduling task {}...", i);
runner.schedule(&results, i, i);
}
while results.lock().len() < 20 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
let results = results.lock();
println!("Results: {:?}", *results);
// The first two tasks are directly spawned without queuing
assert_eq!(&results[0..2], &[0, 1]);
// All tasks after that are queued and therefore prioritized
// The waiting tasks are just waiting, so all of them are executed.
// And the two highest priority cpu-bound tasks are executed too.
// Since we only have 2 workers, the waiting tasks ain't polled until the cpu-bound tasks
// are done.
assert!(results[2..4].contains(&9));
assert!(results[2..4].contains(&8));
let waiting_task_pos = results
.iter()
.position(|&x| x >= 10)
.expect("Waiting task should be executed");
// Waiting tasks should be interleaved with cpu-bound tasks
assert!(waiting_task_pos < 8);
let cpu_bound_results = results
.iter()
.copied()
.filter(|&x| x < 10)
.collect::<Vec<_>>();
// The last tasks are the tasks with the lowest priority
assert!(cpu_bound_results[8..10].contains(&2));
assert!(cpu_bound_results[8..10].contains(&3));
}
}
@@ -9,9 +9,9 @@ use tracing::Span;
use turbo_bincode::{AnyDecodeFn, AnyEncodeFn};
use crate::{
RawVc, SharedReference, VcValueType, id::TraitTypeId, macro_helpers::NativeFunction,
magic_any::any_as_encode, registry, task::shared_reference::TypedSharedReference,
vc::VcCellMode,
RawVc, SharedReference, TaskPriority, VcValueType, id::TraitTypeId,
macro_helpers::NativeFunction, magic_any::any_as_encode, registry,
task::shared_reference::TypedSharedReference, vc::VcCellMode,
};
type RawCellFactoryFn = fn(TypedSharedReference) -> RawVc;
@@ -227,10 +227,11 @@ impl Debug for TraitMethod {
}
}
impl TraitMethod {
pub(crate) fn resolve_span(&self) -> Span {
pub(crate) fn resolve_span(&self, priority: TaskPriority) -> Span {
tracing::trace_span!(
"turbo_tasks::resolve_trait_call",
name = format_args!("{}::{}", &self.trait_name, &self.method_name),
priority = %priority,
)
}
}