use prehash to avoid rehashing the key in the task cache

vercel · May 17, 2024 · 6cca3f2 · 6cca3f2
1 parent a4ad586
commit 6cca3f2
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 10 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/turbo-tasks-memory/Cargo.toml b/crates/turbo-tasks-memory/Cargo.toml
@@ -22,6 +22,7 @@ nohash-hasher = { workspace = true }
 num_cpus = "1.13.1"
 once_cell = { workspace = true }
 parking_lot = { workspace = true }
+prehash = "1.0.0"
 priority-queue = "1.3.0"
 ref-cast = "1.0.20"
 rustc-hash = { workspace = true }

diff --git a/crates/turbo-tasks-memory/src/memory_backend.rs b/crates/turbo-tasks-memory/src/memory_backend.rs
@@ -14,6 +14,7 @@ use std::{
 use anyhow::{bail, Result};
 use auto_hash_map::AutoMap;
 use dashmap::{mapref::entry::Entry, DashMap};
+use prehash::{Passthru, Prehashed, Prehasher};
 use rustc_hash::FxHasher;
 use tokio::task::futures::TaskLocalFuture;
 use tracing::trace_span;
@@ -34,11 +35,15 @@ use crate::{
     task::{Task, TaskDependency, TaskDependencySet, DEPENDENCIES_TO_TRACK},
 };
 
+fn prehash_task_type(task_type: PersistentTaskType) -> Prehashed<PersistentTaskType> {
+    BuildHasherDefault::<FxHasher>::prehash(&Default::default(), task_type)
+}
+
 pub struct MemoryBackend {
     memory_tasks: NoMoveVec<Task, 13>,
     backend_jobs: NoMoveVec<Job>,
     backend_job_id_factory: IdFactory<BackendJobId>,
-    task_cache: DashMap<Arc<PersistentTaskType>, TaskId, BuildHasherDefault<FxHasher>>,
+    task_cache: DashMap<Arc<Prehashed<PersistentTaskType>>, TaskId, BuildHasherDefault<Passthru>>,
     memory_limit: usize,
     gc_queue: Option<GcQueue>,
     idle_gc_active: AtomicBool,
@@ -513,10 +518,11 @@ impl Backend for MemoryBackend {
 
     fn get_or_create_persistent_task(
         &self,
-        mut task_type: PersistentTaskType,
+        task_type: PersistentTaskType,
         parent_task: TaskId,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
     ) -> TaskId {
+        let task_type = prehash_task_type(task_type);
         if let Some(task) =
             self.lookup_and_connect_task(parent_task, &self.task_cache, &task_type, turbo_tasks)
         {
@@ -525,8 +531,9 @@ impl Backend for MemoryBackend {
         } else {
             // It's important to avoid overallocating memory as this will go into the task
             // cache and stay there forever. We can to be as small as possible.
+            let (mut task_type, task_type_hash) = Prehashed::into_parts(task_type);
             task_type.shrink_to_fit();
-            let task_type = Arc::new(task_type);
+            let task_type = Arc::new(Prehashed::new(task_type, task_type_hash));
             // slow pass with key lock
             let id = turbo_tasks.get_fresh_task_id();
             let task = Task::new_persistent(

diff --git a/crates/turbo-tasks-memory/src/task.rs b/crates/turbo-tasks-memory/src/task.rs
@@ -14,6 +14,7 @@ use anyhow::Result;
 use auto_hash_map::{AutoMap, AutoSet};
 use nohash_hasher::BuildNoHashHasher;
 use parking_lot::{Mutex, RwLock};
+use prehash::Prehashed;
 use rustc_hash::FxHasher;
 use smallvec::SmallVec;
 use tokio::task_local;
@@ -76,13 +77,15 @@ enum TaskType {
     Once(Box<OnceTaskFn>),
 
     /// A normal persistent task
-    Persistent { ty: Arc<PersistentTaskType> },
+    Persistent {
+        ty: Arc<Prehashed<PersistentTaskType>>,
+    },
 }
 
 enum TaskTypeForDescription {
     Root,
     Once,
-    Persistent(Arc<PersistentTaskType>),
+    Persistent(Arc<Prehashed<PersistentTaskType>>),
 }
 
 impl TaskTypeForDescription {
@@ -406,7 +409,10 @@ use self::{
 };
 
 impl Task {
-    pub(crate) fn new_persistent(id: TaskId, task_type: Arc<PersistentTaskType>) -> Self {
+    pub(crate) fn new_persistent(
+        id: TaskId,
+        task_type: Arc<Prehashed<PersistentTaskType>>,
+    ) -> Self {
         let ty = TaskType::Persistent { ty: task_type };
         let description = Self::get_event_description_static(id, &ty);
         Self {
@@ -516,7 +522,7 @@ impl Task {
 
     pub(crate) fn get_function_name(&self) -> Option<Cow<'static, str>> {
         if let TaskType::Persistent { ty, .. } = &self.ty {
-            match &**ty {
+            match &***ty {
                 PersistentTaskType::Native(native_fn, _)
                 | PersistentTaskType::ResolveNative(native_fn, _) => {
                     return Some(Cow::Borrowed(&registry::get_function(*native_fn).name));
@@ -539,7 +545,7 @@ impl Task {
         match ty {
             TaskTypeForDescription::Root => format!("[{}] root", id),
             TaskTypeForDescription::Once => format!("[{}] once", id),
-            TaskTypeForDescription::Persistent(ty) => match &**ty {
+            TaskTypeForDescription::Persistent(ty) => match &***ty {
                 PersistentTaskType::Native(native_fn, _) => {
                     format!("[{}] {}", id, registry::get_function(*native_fn).name)
                 }
@@ -742,7 +748,7 @@ impl Task {
                     tracing::trace_span!("turbo_tasks::once_task"),
                 )
             }
-            TaskType::Persistent { ty, .. } => match &**ty {
+            TaskType::Persistent { ty, .. } => match &***ty {
                 PersistentTaskType::Native(native_fn, inputs) => {
                     let result =
                         if let PrepareTaskType::Native(func, bound_fn) = &state.prepared_type {