Skip to content

Commit 0967bee

Browse files
committed
handle recomputation of cells on read
1 parent c45b56e commit 0967bee

File tree

2 files changed

+96
-47
lines changed

2 files changed

+96
-47
lines changed

crates/turbo-tasks-memory/src/memory_backend.rs

Lines changed: 23 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,10 @@ use turbo_tasks::{
3030
};
3131

3232
use crate::{
33-
cell::RecomputingCell,
3433
edges_set::{TaskEdge, TaskEdgesSet},
3534
gc::{GcQueue, PERCENTAGE_IDLE_TARGET_MEMORY, PERCENTAGE_TARGET_MEMORY},
3635
output::Output,
37-
task::{Task, DEPENDENCIES_TO_TRACK},
36+
task::{ReadCellError, Task, DEPENDENCIES_TO_TRACK},
3837
task_statistics::TaskStatisticsApi,
3938
};
4039

@@ -412,23 +411,17 @@ impl Backend for MemoryBackend {
412411
} else {
413412
Task::add_dependency_to_current(TaskEdge::Cell(task_id, index));
414413
self.with_task(task_id, |task| {
415-
match task.access_cell_for_read(index, self.gc_queue.as_ref(), |cell| {
416-
cell.map(|cell| {
417-
cell.read_content(
418-
reader,
419-
move || format!("{task_id} {index}"),
420-
move || format!("reading {} {} from {}", task_id, index, reader),
421-
)
422-
})
423-
}) {
424-
None => Err(anyhow!("Cell doesn't exist (anymore)")),
425-
Some(Ok(content)) => Ok(Ok(content)),
426-
Some(Err(RecomputingCell { listener, schedule })) => {
427-
if schedule {
428-
task.recompute(self, turbo_tasks);
429-
}
430-
Ok(Err(listener))
431-
}
414+
match task.read_cell(
415+
index,
416+
self.gc_queue.as_ref(),
417+
move || format!("reading {} {} from {}", task_id, index, reader),
418+
Some(reader),
419+
self,
420+
turbo_tasks,
421+
) {
422+
Ok(content) => Ok(Ok(content)),
423+
Err(ReadCellError::Recomputing(listener)) => Ok(Err(listener)),
424+
Err(ReadCellError::CellRemoved) => Err(anyhow!("Cell {index:?} doesn't exist")),
432425
}
433426
})
434427
}
@@ -452,22 +445,17 @@ impl Backend for MemoryBackend {
452445
turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
453446
) -> Result<Result<CellContent, EventListener>> {
454447
self.with_task(task_id, |task| {
455-
match task.access_cell_for_read(index, self.gc_queue.as_ref(), |cell| {
456-
cell.map(|cell| {
457-
cell.read_content_untracked(
458-
move || format!("{task_id} {index}"),
459-
move || format!("reading {} {} untracked", task_id, index),
460-
)
461-
})
462-
}) {
463-
None => Err(anyhow!("Cell doesn't exist (anymore)")),
464-
Some(Ok(content)) => Ok(Ok(content)),
465-
Some(Err(RecomputingCell { listener, schedule })) => {
466-
if schedule {
467-
task.recompute(self, turbo_tasks);
468-
}
469-
Ok(Err(listener))
470-
}
448+
match task.read_cell(
449+
index,
450+
self.gc_queue.as_ref(),
451+
move || format!("reading {} {} untracked", task_id, index),
452+
None,
453+
self,
454+
turbo_tasks,
455+
) {
456+
Ok(content) => Ok(Ok(content)),
457+
Err(ReadCellError::Recomputing(listener)) => Ok(Err(listener)),
458+
Err(ReadCellError::CellRemoved) => Err(anyhow!("Cell {index:?} doesn't exist")),
471459
}
472460
})
473461
}

crates/turbo-tasks-memory/src/task.rs

Lines changed: 73 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use tokio::task_local;
2121
use tracing::Span;
2222
use turbo_prehash::PreHashed;
2323
use turbo_tasks::{
24-
backend::{PersistentTaskType, TaskExecutionSpec},
24+
backend::{CellContent, PersistentTaskType, TaskExecutionSpec},
2525
event::{Event, EventListener},
2626
get_invalidator, registry, CellId, Invalidator, RawVc, TaskId, TaskIdSet, TraitTypeId,
2727
TurboTasksBackendApi, ValueTypeId,
@@ -32,7 +32,7 @@ use crate::{
3232
aggregation_data, handle_new_edge, prepare_aggregation_data, query_root_info,
3333
AggregationDataGuard, PreparedOperation,
3434
},
35-
cell::Cell,
35+
cell::{Cell, RecomputingCell},
3636
edges_set::{TaskEdge, TaskEdgesList, TaskEdgesSet},
3737
gc::{GcQueue, GcTaskState},
3838
output::{Output, OutputContent},
@@ -448,6 +448,11 @@ pub enum GcResult {
448448
Unloaded,
449449
}
450450

451+
pub enum ReadCellError {
452+
CellRemoved,
453+
Recomputing(EventListener),
454+
}
455+
451456
impl Task {
452457
pub(crate) fn new_persistent(
453458
id: TaskId,
@@ -1286,26 +1291,82 @@ impl Task {
12861291
}
12871292
}
12881293

1289-
/// Access to a cell.
1290-
pub(crate) fn access_cell_for_read<T>(
1294+
/// Read a cell.
1295+
pub(crate) fn read_cell(
12911296
&self,
12921297
index: CellId,
12931298
gc_queue: Option<&GcQueue>,
1294-
func: impl FnOnce(Option<&mut Cell>) -> T,
1295-
) -> T {
1299+
note: impl Fn() -> String + Sync + Send + 'static,
1300+
reader: Option<TaskId>,
1301+
backend: &MemoryBackend,
1302+
turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
1303+
) -> Result<CellContent, ReadCellError> {
1304+
let task_id = self.id;
12961305
let mut state = self.full_state_mut();
12971306
if let Some(gc_queue) = gc_queue {
12981307
let generation = gc_queue.generation();
12991308
if state.gc.on_read(generation) {
13001309
let _ = gc_queue.task_accessed(self.id);
13011310
}
13021311
}
1303-
let list = state.cells.entry(index.type_id).or_default();
1304-
let i = index.index as usize;
1305-
if list.len() <= i {
1306-
func(None)
1307-
} else {
1308-
func(Some(&mut list[i]))
1312+
match state.state_type {
1313+
Done { .. } | InProgress(..) | InProgressDirty { .. } => {
1314+
let is_done = matches!(state.state_type, Done { .. });
1315+
let list = state.cells.entry(index.type_id).or_default();
1316+
let i = index.index as usize;
1317+
if list.len() <= i {
1318+
if is_done {
1319+
return Err(ReadCellError::CellRemoved);
1320+
} else {
1321+
list.resize_with(i + 1, Default::default);
1322+
}
1323+
}
1324+
let cell = &mut list[i];
1325+
let description = move || format!("{task_id} {index}");
1326+
let read_result = if let Some(reader) = reader {
1327+
cell.read_content(reader, description, note)
1328+
} else {
1329+
cell.read_content_untracked(description, note)
1330+
};
1331+
drop(state);
1332+
match read_result {
1333+
Ok(content) => Ok(content),
1334+
Err(RecomputingCell { listener, schedule }) => {
1335+
if schedule {
1336+
self.recompute(backend, turbo_tasks);
1337+
}
1338+
Err(ReadCellError::Recomputing(listener))
1339+
}
1340+
}
1341+
}
1342+
Dirty {
1343+
ref mut outdated_edges,
1344+
} => {
1345+
let mut aggregation_context = TaskAggregationContext::new(turbo_tasks, backend);
1346+
let description = self.get_event_description();
1347+
let event = Event::new(move || format!("TaskState({})::event", description()));
1348+
let listener = event.listen_with_note(note);
1349+
state.state_type = Scheduled {
1350+
event,
1351+
outdated_edges: take(outdated_edges),
1352+
clean: false,
1353+
};
1354+
let change_job = state.aggregation_node.apply_change(
1355+
&aggregation_context,
1356+
TaskChange {
1357+
dirty_tasks_update: vec![(self.id, -1)],
1358+
..Default::default()
1359+
},
1360+
);
1361+
drop(state);
1362+
turbo_tasks.schedule(self.id);
1363+
change_job.apply(&aggregation_context);
1364+
aggregation_context.apply_queued_updates();
1365+
Err(ReadCellError::Recomputing(listener))
1366+
}
1367+
Scheduled { ref event, .. } => {
1368+
Err(ReadCellError::Recomputing(event.listen_with_note(note)))
1369+
}
13091370
}
13101371
}
13111372

0 commit comments

Comments
 (0)