Skip to content

Commit

Permalink
reduce memory and consolidate shared vs transient cells (#8263)
Browse files Browse the repository at this point in the history 
### Description

This change reworks the structure to drop an option such that data which
is colocated with its type information doesn't need to store it twice.

### Testing Instructions

Existing tests should suffice.

---------

Co-authored-by: Tobias Koppers <[email protected]>
  • Loading branch information
@arlyon @sokra
arlyon and sokra authored Jul 25, 2024
1 parent f32ebe9 commit b953c89
Show file tree
Hide file tree
Showing 13 changed files with 231 additions and 162 deletions.
6 changes: 5 additions & 1 deletion crates/turbo-tasks-memory/src/cell.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,14 +151,18 @@ impl Cell {
CellState::Empty | CellState::Computing { .. } | CellState::TrackedValueless => {
CellContent(None)
}
CellState::Value { content } => content.clone(),
CellState::Value { content } => content.to_owned(),
}
}

/// Assigns a new content to the cell. Will notify dependent tasks if the
/// content has changed.
/// If clean = true, the task inputs weren't changes since the last
/// execution and can be assumed to produce the same content again.
///
/// Safety: This funtion does not check if the type of the content is the
/// same as the type of the cell. It is the caller's responsibility to
/// ensure that the content is of the correct type.
pub fn assign(
&mut self,
content: CellContent,
Expand Down
24 changes: 15 additions & 9 deletions crates/turbo-tasks-memory/src/memory_backend.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use turbo_prehash::{BuildHasherExt, PassThroughHash, PreHashed};
use turbo_tasks::{
backend::{
Backend, BackendJobId, CellContent, PersistentTaskType, TaskCollectiblesMap,
TaskExecutionSpec, TransientTaskType,
TaskExecutionSpec, TransientTaskType, TypedCellContent,
},
event::EventListener,
util::{IdFactoryWithReuse, NoMoveVec},
Expand Down Expand Up @@ -404,11 +404,13 @@ impl Backend for MemoryBackend {
index: CellId,
reader: TaskId,
turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
) -> Result<Result<CellContent, EventListener>> {
) -> Result<Result<TypedCellContent, EventListener>> {
if task_id == reader {
Ok(Ok(self.with_task(task_id, |task| {
task.with_cell(index, |cell| cell.read_own_content_untracked())
})))
Ok(Ok(self
.with_task(task_id, |task| {
task.with_cell(index, |cell| cell.read_own_content_untracked())
})
.into_typed(index.type_id)))
} else {
Task::add_dependency_to_current(TaskEdge::Cell(task_id, index));
self.with_task(task_id, |task| {
Expand All @@ -420,7 +422,7 @@ impl Backend for MemoryBackend {
self,
turbo_tasks,
) {
Ok(content) => Ok(Ok(content)),
Ok(content) => Ok(Ok(content.into_typed(index.type_id))),
Err(ReadCellError::Recomputing(listener)) => Ok(Err(listener)),
Err(ReadCellError::CellRemoved) => Err(anyhow!("Cell doesn't exist")),
}
Expand All @@ -433,9 +435,10 @@ impl Backend for MemoryBackend {
current_task: TaskId,
index: CellId,
_turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
) -> Result<CellContent> {
) -> Result<TypedCellContent> {
Ok(self.with_task(current_task, |task| {
task.with_cell(index, |cell| cell.read_own_content_untracked())
.into_typed(index.type_id)
}))
}

Expand All @@ -444,7 +447,7 @@ impl Backend for MemoryBackend {
task_id: TaskId,
index: CellId,
turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
) -> Result<Result<CellContent, EventListener>> {
) -> Result<Result<TypedCellContent, EventListener>> {
self.with_task(task_id, |task| {
match task.read_cell(
index,
Expand All @@ -454,7 +457,7 @@ impl Backend for MemoryBackend {
self,
turbo_tasks,
) {
Ok(content) => Ok(Ok(content)),
Ok(content) => Ok(Ok(content.into_typed(index.type_id))),
Err(ReadCellError::Recomputing(listener)) => Ok(Err(listener)),
Err(ReadCellError::CellRemoved) => Err(anyhow!("Cell doesn't exist")),
}
Expand Down Expand Up @@ -497,6 +500,9 @@ impl Backend for MemoryBackend {
});
}

/// SAFETY: This function does not validate that the data in `content` is of
/// the same type as in `index`. It is the caller's responsibility to ensure
/// that the content is of the correct type.
fn update_task_cell(
&self,
task: TaskId,
Expand Down
31 changes: 17 additions & 14 deletions crates/turbo-tasks-testing/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use std::{
use anyhow::{anyhow, Result};
use futures::FutureExt;
use turbo_tasks::{
backend::{CellContent, TaskCollectiblesMap},
backend::{CellContent, TaskCollectiblesMap, TypedCellContent},
event::{Event, EventListener},
registry,
test_helpers::with_turbo_tasks_for_testing,
Expand Down Expand Up @@ -189,33 +189,35 @@ impl TurboTasksApi for VcStorage {
&self,
task: TaskId,
index: CellId,
) -> Result<Result<CellContent, EventListener>> {
) -> Result<Result<TypedCellContent, EventListener>> {
let map = self.cells.lock().unwrap();
if let Some(cell) = map.get(&(task, index)) {
Ok(Ok(cell.clone()))
Ok(Ok(if let Some(cell) = map.get(&(task, index)) {
cell.clone()
} else {
Ok(Ok(CellContent::default()))
Default::default()
}
.into_typed(index.type_id)))
}

fn try_read_task_cell_untracked(
&self,
task: TaskId,
index: CellId,
) -> Result<Result<CellContent, EventListener>> {
) -> Result<Result<TypedCellContent, EventListener>> {
let map = self.cells.lock().unwrap();
if let Some(cell) = map.get(&(task, index)) {
Ok(Ok(cell.clone()))
Ok(Ok(if let Some(cell) = map.get(&(task, index)) {
cell.to_owned()
} else {
Ok(Ok(CellContent::default()))
Default::default()
}
.into_typed(index.type_id)))
}

fn try_read_own_task_cell_untracked(
&self,
current_task: TaskId,
index: CellId,
) -> Result<CellContent> {
) -> Result<TypedCellContent> {
self.read_own_task_cell(current_task, index)
}

Expand Down Expand Up @@ -244,13 +246,14 @@ impl TurboTasksApi for VcStorage {
unimplemented!()
}

fn read_own_task_cell(&self, task: TaskId, index: CellId) -> Result<CellContent> {
fn read_own_task_cell(&self, task: TaskId, index: CellId) -> Result<TypedCellContent> {
let map = self.cells.lock().unwrap();
if let Some(cell) = map.get(&(task, index)) {
Ok(cell.clone())
Ok(if let Some(cell) = map.get(&(task, index)) {
cell.to_owned()
} else {
Ok(CellContent::default())
Default::default()
}
.into_typed(index.type_id))
}

fn update_own_task_cell(&self, task: TaskId, index: CellId, content: CellContent) {
Expand Down
60 changes: 31 additions & 29 deletions crates/turbo-tasks/src/backend.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,9 @@ use std::{
time::Duration,
};

use anyhow::{anyhow, bail, Result};
use anyhow::{anyhow, Result};
use auto_hash_map::AutoMap;
use rustc_hash::FxHasher;
use serde::{Deserialize, Serialize};
use tracing::Span;

pub use crate::id::{BackendJobId, ExecutionId};
Expand Down Expand Up @@ -333,10 +332,10 @@ pub struct TaskExecutionSpec<'a> {
pub span: Span,
}

// TODO technically CellContent is already indexed by the ValueTypeId, so we
// don't need to store it here
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Hash, Default)]
pub struct CellContent(pub Option<SharedReference>);
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TypedCellContent(pub ValueTypeId, pub CellContent);

impl Display for CellContent {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
Expand All @@ -347,9 +346,9 @@ impl Display for CellContent {
}
}

impl CellContent {
impl TypedCellContent {
pub fn cast<T: Any + VcValueType>(self) -> Result<ReadRef<T>> {
let data = self.0.ok_or_else(|| anyhow!("Cell is empty"))?;
let data = self.1 .0.ok_or_else(|| anyhow!("Cell is empty"))?;
let data = data
.downcast()
.map_err(|_err| anyhow!("Unexpected type in cell"))?;
Expand All @@ -358,24 +357,35 @@ impl CellContent {

/// # Safety
///
/// The caller must ensure that the CellContent contains a vc that
/// implements T.
/// The caller must ensure that the TypedCellContent contains a vc
/// that implements T.
pub fn cast_trait<T>(self) -> Result<TraitRef<T>>
where
T: VcValueTrait + ?Sized,
{
let shared_reference = self.0.ok_or_else(|| anyhow!("Cell is empty"))?;
if shared_reference.0.is_none() {
bail!("Cell content is untyped");
}
let shared_reference = self
.1
.0
.ok_or_else(|| anyhow!("Cell is empty"))?
.typed(self.0);
Ok(
// Safety: We just checked that the content is typed.
// Safety: It is a TypedSharedReference
TraitRef::new(shared_reference),
)
}

pub fn try_cast<T: Any + VcValueType>(self) -> Option<ReadRef<T>> {
Some(ReadRef::new_arc(self.0?.downcast().ok()?))
Some(ReadRef::new_arc(self.1 .0?.downcast().ok()?))
}

pub fn into_untyped(self) -> CellContent {
self.1
}
}

impl CellContent {
pub fn into_typed(self, type_id: ValueTypeId) -> TypedCellContent {
TypedCellContent(type_id, self)
}
}

Expand Down Expand Up @@ -463,7 +473,7 @@ pub trait Backend: Sync + Send {
index: CellId,
reader: TaskId,
turbo_tasks: &dyn TurboTasksBackendApi<Self>,
) -> Result<Result<CellContent, EventListener>>;
) -> Result<Result<TypedCellContent, EventListener>>;

/// INVALIDATION: Be careful with this, it will not track dependencies, so
/// using it could break cache invalidation.
Expand All @@ -472,7 +482,7 @@ pub trait Backend: Sync + Send {
task: TaskId,
index: CellId,
turbo_tasks: &dyn TurboTasksBackendApi<Self>,
) -> Result<Result<CellContent, EventListener>>;
) -> Result<Result<TypedCellContent, EventListener>>;

/// INVALIDATION: Be careful with this, it will not track dependencies, so
/// using it could break cache invalidation.
Expand All @@ -481,10 +491,10 @@ pub trait Backend: Sync + Send {
current_task: TaskId,
index: CellId,
turbo_tasks: &dyn TurboTasksBackendApi<Self>,
) -> Result<CellContent> {
) -> Result<TypedCellContent> {
match self.try_read_task_cell_untracked(current_task, index, turbo_tasks)? {
Ok(content) => Ok(content),
Err(_) => Ok(CellContent(None)),
Err(_) => Ok(TypedCellContent(index.type_id, CellContent(None))),
}
}

Expand Down Expand Up @@ -575,10 +585,7 @@ impl PersistentTaskType {
name: Cow<'static, str>,
this: RawVc,
) -> Result<FunctionId> {
let CellContent(Some(SharedReference(Some(value_type), _))) = this.into_read().await?
else {
bail!("Cell is empty or untyped");
};
let TypedCellContent(value_type, _) = this.into_read().await?;
Self::resolve_trait_method_from_value(trait_type, value_type, name)
}

Expand All @@ -590,12 +597,7 @@ impl PersistentTaskType {
turbo_tasks: Arc<dyn TurboTasksBackendApi<B>>,
) -> Result<RawVc> {
let this = this.resolve().await?;
let CellContent(Some(SharedReference(this_ty, _))) = this.into_read().await? else {
bail!("Cell is empty");
};
let Some(this_ty) = this_ty else {
bail!("Cell is untyped");
};
let TypedCellContent(this_ty, _) = this.into_read().await?;

let native_fn = Self::resolve_trait_method_from_value(trait_type, this_ty, name)?;
let arg = registry::get_function(native_fn)
Expand Down
Loading

0 comments on commit b953c89

Please sign in to comment.