From 13ccb51418761cb4293734cbdc0ea862b0934ad0 Mon Sep 17 00:00:00 2001
From: Tobias Koppers <tobias.koppers@googlemail.com>
Date: Mon, 3 Jun 2024 18:41:04 +0200
Subject: [PATCH] use prehash to avoid rehashing the key in the task cache
 (vercel/turbo#8174)

### Description

This avoids hashing the task type multiple times, e. g. for inserting
after lookup or on rehashing of the hash map

### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->
---
 crates/turbo-tasks-memory/Cargo.toml           |  1 +
 .../turbo-tasks-memory/src/memory_backend.rs   | 14 +++++++++++---
 crates/turbo-tasks-memory/src/task.rs          | 18 ++++++++++++------
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/crates/turbo-tasks-memory/Cargo.toml b/crates/turbo-tasks-memory/Cargo.toml
index 0e2c722c3c951..a7408fb16c723 100644
--- a/crates/turbo-tasks-memory/Cargo.toml
+++ b/crates/turbo-tasks-memory/Cargo.toml
@@ -28,6 +28,7 @@ rustc-hash = { workspace = true }
 smallvec = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
+turbo-prehash = { workspace = true }
 turbo-tasks = { workspace = true }
 turbo-tasks-hash = { workspace = true }
 turbo-tasks-malloc = { workspace = true, default-features = false }
diff --git a/crates/turbo-tasks-memory/src/memory_backend.rs b/crates/turbo-tasks-memory/src/memory_backend.rs
index a2822316d400c..e83dd0bb5a6d1 100644
--- a/crates/turbo-tasks-memory/src/memory_backend.rs
+++ b/crates/turbo-tasks-memory/src/memory_backend.rs
@@ -17,6 +17,7 @@ use dashmap::{mapref::entry::Entry, DashMap};
 use rustc_hash::FxHasher;
 use tokio::task::futures::TaskLocalFuture;
 use tracing::trace_span;
+use turbo_prehash::{BuildHasherExt, PassThroughHash, PreHashed};
 use turbo_tasks::{
     backend::{
         Backend, BackendJobId, CellContent, PersistentTaskType, TaskExecutionSpec,
@@ -34,11 +35,16 @@ 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<PassThroughHash>>,
     memory_limit: usize,
     gc_queue: Option<GcQueue>,
     idle_gc_active: AtomicBool,
@@ -513,10 +519,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 +532,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 (task_type_hash, mut task_type) = 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_hash, task_type));
             // 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
index cd9e3cd4fca69..20232fec49d9f 100644
--- a/crates/turbo-tasks-memory/src/task.rs
+++ b/crates/turbo-tasks-memory/src/task.rs
@@ -18,6 +18,7 @@ use rustc_hash::FxHasher;
 use smallvec::SmallVec;
 use tokio::task_local;
 use tracing::Span;
+use turbo_prehash::PreHashed;
 use turbo_tasks::{
     backend::{PersistentTaskType, TaskExecutionSpec},
     event::{Event, EventListener},
@@ -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 {