turbo_tasks_backend/
kv_backing_storage.rs

1use std::{
2    borrow::Borrow,
3    env,
4    path::PathBuf,
5    sync::{Arc, LazyLock, Mutex, PoisonError, Weak},
6};
7
8use anyhow::{Context, Result};
9use smallvec::SmallVec;
10use turbo_bincode::{new_turbo_bincode_decoder, turbo_bincode_decode, turbo_bincode_encode};
11use turbo_persistence::CommitStats;
12use turbo_tasks::{
13    DynTaskInputs, RawVc, TaskId,
14    macro_helpers::NativeFunction,
15    panic_hooks::{PanicHookGuard, register_panic_hook},
16    parallel,
17};
18
19use crate::{
20    GitVersionInfo,
21    backend::{AnyOperation, SpecificTaskDataCategory, TtlCounter, storage_schema::TaskStorage},
22    backing_storage::{SnapshotItem, SnapshotMeta, compute_task_type_hash_from_components},
23    database::{
24        db_invalidation::{StartupCacheState, check_db_invalidation_and_cleanup, invalidate_db},
25        db_versioning::handle_db_versioning,
26        key_value_database::KeySpace,
27        turbo::{TurboKeyValueDatabase, TurboWriteBatch},
28        write_batch::WriteBuffer,
29    },
30    db_invalidation::invalidation_reasons,
31};
32
33/// The fixed keys in the [`KeySpace::Infra`] keyspace.
34#[derive(Clone, Copy)]
35#[repr(u8)]
36enum InfraKey {
37    Operations = 0,
38    NextFreeTaskId = 1,
39    GcRoots = 2,
40}
41
42impl InfraKey {
43    fn key(self) -> ByteKey {
44        ByteKey::new(self as u8)
45    }
46}
47
48struct ByteKey([u8; 1]);
49
50impl ByteKey {
51    fn new(value: u8) -> Self {
52        Self([value])
53    }
54}
55
56impl AsRef<[u8]> for ByteKey {
57    fn as_ref(&self) -> &[u8] {
58        &self.0
59    }
60}
61
62struct IntKey([u8; 4]);
63
64impl IntKey {
65    fn new(value: u32) -> Self {
66        Self(value.to_le_bytes())
67    }
68}
69
70impl AsRef<[u8]> for IntKey {
71    fn as_ref(&self) -> &[u8] {
72        &self.0
73    }
74}
75
76fn as_u32(bytes: impl Borrow<[u8]>) -> Result<u32> {
77    let n = u32::from_le_bytes(bytes.borrow().try_into()?);
78    Ok(n)
79}
80
81// We want to invalidate the cache on panic for most users, but this is a band-aid to underlying
82// problems in turbo-tasks.
83//
84// If we invalidate the cache upon panic and it "fixes" the issue upon restart, users typically
85// won't report bugs to us, and we'll never find root-causes for these problems.
86//
87// These overrides let us avoid the cache invalidation / error suppression within Vercel so that we
88// feel these pain points and fix the root causes of bugs.
89fn should_invalidate_on_panic() -> bool {
90    fn env_is_falsy(key: &str) -> bool {
91        env::var_os(key)
92            .is_none_or(|value| ["".as_ref(), "0".as_ref(), "false".as_ref()].contains(&&*value))
93    }
94    static SHOULD_INVALIDATE: LazyLock<bool> = LazyLock::new(|| {
95        env_is_falsy("TURBO_ENGINE_SKIP_INVALIDATE_ON_PANIC") && env_is_falsy("__NEXT_TEST_MODE")
96    });
97    *SHOULD_INVALIDATE
98}
99
100struct TurboBackingStorageInner {
101    database: TurboKeyValueDatabase,
102    /// Used when calling [`TurboBackingStorage::invalidate`]. Can be `None` in the
103    /// memory-only/no-op storage case.
104    base_path: Option<PathBuf>,
105    /// Used to skip calling [`invalidate_db`] when the database has already been invalidated.
106    invalidated: Mutex<bool>,
107    /// We configure a panic hook to invalidate the cache. This guard cleans up our panic hook upon
108    /// drop.
109    _panic_hook_guard: Option<PanicHookGuard>,
110}
111
112/// The higher-level backing storage passed to [`TurboTasksBackend::new`], used by
113/// [`crate::turbo_backing_storage`] and [`crate::noop_backing_storage`].
114///
115/// Wraps a low-level [`TurboKeyValueDatabase`] and adapts it into the persistence operations the
116/// backend needs (snapshots, task-candidate lookups, etc.).
117///
118/// [`TurboTasksBackend::new`]: crate::TurboTasksBackend::new
119pub struct TurboBackingStorage {
120    // wrapped so that `register_panic_hook` can hold a weak reference to `inner`.
121    inner: Arc<TurboBackingStorageInner>,
122}
123
124impl TurboBackingStorage {
125    pub(crate) fn new_in_memory(database: TurboKeyValueDatabase) -> Self {
126        Self {
127            inner: Arc::new(TurboBackingStorageInner {
128                database,
129                base_path: None,
130                invalidated: Mutex::new(false),
131                _panic_hook_guard: None,
132            }),
133        }
134    }
135
136    /// Handles boilerplate logic for an on-disk persisted database with versioning.
137    ///
138    /// - Creates a directory per version, with a maximum number of old versions and performs
139    ///   automatic cleanup of old versions.
140    /// - Checks for a database invalidation marker file, and cleans up the database as needed.
141    /// - [Registers a dynamic panic hook][turbo_tasks::panic_hooks] to invalidate the database upon
142    ///   a panic. This invalidates the database using [`invalidation_reasons::PANIC`].
143    ///
144    /// Along with returning a [`TurboBackingStorage`], this returns a
145    /// [`StartupCacheState`], which can be used by the application for logging information to the
146    /// user or telemetry about the cache.
147    pub(crate) fn open_versioned_on_disk(
148        base_path: PathBuf,
149        version_info: &GitVersionInfo,
150        is_ci: bool,
151        database: impl FnOnce(PathBuf) -> Result<TurboKeyValueDatabase>,
152    ) -> Result<(Self, StartupCacheState)> {
153        let startup_cache_state = check_db_invalidation_and_cleanup(&base_path)
154            .context("Failed to check database invalidation and cleanup")?;
155        let versioned_path = handle_db_versioning(&base_path, version_info, is_ci)
156            .context("Failed to handle database versioning")?;
157        let database = (database)(versioned_path).context("Failed to open database")?;
158        let backing_storage = Self {
159            inner: Arc::new_cyclic(move |weak_inner: &Weak<TurboBackingStorageInner>| {
160                let panic_hook_guard = if should_invalidate_on_panic() {
161                    let weak_inner = weak_inner.clone();
162                    Some(register_panic_hook(Box::new(move |_| {
163                        let Some(inner) = weak_inner.upgrade() else {
164                            return;
165                        };
166                        // If a panic happened that must mean something deep inside of turbopack
167                        // or turbo-tasks failed, and it may be hard to recover. We don't want
168                        // the cache to stick around, as that may persist bugs. Make a
169                        // best-effort attempt to invalidate the database (ignoring failures).
170                        let _ = inner.invalidate(invalidation_reasons::PANIC);
171                    })))
172                } else {
173                    None
174                };
175                TurboBackingStorageInner {
176                    database,
177                    base_path: Some(base_path),
178                    invalidated: Mutex::new(false),
179                    _panic_hook_guard: panic_hook_guard,
180                }
181            }),
182        };
183        Ok((backing_storage, startup_cache_state))
184    }
185}
186
187impl TurboBackingStorageInner {
188    fn invalidate(&self, reason_code: &str) -> Result<()> {
189        // `base_path` is `None` for in-memory backing storage (see `noop_backing_storage`).
190        if let Some(base_path) = &self.base_path {
191            // Invalidation could happen frequently if there's a bunch of panics. We only need to
192            // invalidate once, so grab a lock.
193            let mut invalidated_guard = self
194                .invalidated
195                .lock()
196                .unwrap_or_else(PoisonError::into_inner);
197            if *invalidated_guard {
198                return Ok(());
199            }
200            // Invalidate first, as it's a very fast atomic operation. `prevent_writes` is allowed
201            // to be slower (e.g. wait for a lock) and is allowed to corrupt the database with
202            // partial writes.
203            invalidate_db(base_path, reason_code)?;
204            self.database.prevent_writes();
205            // Avoid redundant invalidations from future panics
206            *invalidated_guard = true;
207        }
208        Ok(())
209    }
210
211    /// Used to read the next free task ID from the database.
212    fn get_infra_u32(&self, key: InfraKey) -> Result<Option<u32>> {
213        self.database
214            .get(KeySpace::Infra, key.key().as_ref())?
215            .map(as_u32)
216            .transpose()
217    }
218}
219
220impl TurboBackingStorage {
221    /// Called when the database should be invalidated upon re-initialization.
222    ///
223    /// This typically means that we'll restart the process or `turbo-tasks` soon with a fresh
224    /// database. If this happens, there's no point in writing anything else to disk, or flushing
225    /// during [`TurboTasksBackend::stop`].
226    ///
227    /// [`TurboTasksBackend::stop`]: turbo_tasks::backend::Backend::stop
228    pub(crate) fn invalidate(&self, reason_code: &str) -> Result<()> {
229        self.inner.invalidate(reason_code)
230    }
231
232    pub(crate) fn next_free_task_id(&self) -> Result<TaskId> {
233        Ok(self
234            .inner
235            .get_infra_u32(InfraKey::NextFreeTaskId)
236            .context("Unable to read next free task id from database")?
237            .map_or(Ok(TaskId::MIN), TaskId::try_from)?)
238    }
239
240    pub(crate) fn uncompleted_operations(&self) -> Result<Vec<AnyOperation>> {
241        fn get(database: &TurboKeyValueDatabase) -> Result<Vec<AnyOperation>> {
242            let Some(operations) =
243                database.get(KeySpace::Infra, InfraKey::Operations.key().as_ref())?
244            else {
245                return Ok(Vec::new());
246            };
247            let operations = turbo_bincode_decode(operations.borrow())?;
248            Ok(operations)
249        }
250        get(&self.inner.database).context("Unable to read uncompleted operations from database")
251    }
252
253    /// Reads the persisted GC roots set (see [`InfraKey::GcRoots`]). Empty on a fresh database.
254    pub(crate) fn roots(&self) -> Result<Vec<(TaskId, TtlCounter)>> {
255        fn get(database: &TurboKeyValueDatabase) -> Result<Vec<(TaskId, TtlCounter)>> {
256            let Some(roots) = database.get(KeySpace::Infra, InfraKey::GcRoots.key().as_ref())?
257            else {
258                return Ok(Vec::new());
259            };
260            let roots = turbo_bincode_decode(roots.borrow())?;
261            Ok(roots)
262        }
263        get(&self.inner.database).context("Unable to read GC roots from database")
264    }
265
266    pub(crate) fn save_snapshot<I>(
267        &self,
268        operations: Vec<Arc<AnyOperation>>,
269        roots: Option<Vec<(TaskId, TtlCounter)>>,
270        snapshots: Vec<I>,
271    ) -> Result<SnapshotMeta>
272    where
273        I: IntoIterator<Item = SnapshotItem> + Send + Sync,
274    {
275        let _span = tracing::info_span!("save snapshot", operations = operations.len()).entered();
276        let batch = self.inner.database.write_batch()?;
277
278        {
279            let span = tracing::trace_span!("update task data");
280            let mut snapshot_meta =
281                parallel::map_collect_owned::<_, _, Result<Vec<_>>>(snapshots, |shard: I| {
282                    let _span = span.clone().entered();
283                    let mut max_new_task_id = 0;
284                    let mut data_items = 0;
285                    let mut meta_items = 0;
286                    let mut task_cache_items = 0;
287                    for item in shard {
288                        match item {
289                            SnapshotItem::Put {
290                                task_id,
291                                meta,
292                                data,
293                                task_type_hash,
294                            } => {
295                                let key = IntKey::new(*task_id);
296                                let key = key.as_ref();
297                                if let Some(meta) = meta {
298                                    batch.put(
299                                        KeySpace::TaskMeta,
300                                        WriteBuffer::Borrowed(key),
301                                        WriteBuffer::SmallVec(meta),
302                                    )?;
303                                    meta_items += 1;
304                                }
305                                if let Some(data) = data {
306                                    batch.put(
307                                        KeySpace::TaskData,
308                                        WriteBuffer::Borrowed(key),
309                                        WriteBuffer::SmallVec(data),
310                                    )?;
311                                    data_items += 1;
312                                }
313                                // Register the task type only for new tasks.
314                                if let Some(task_type_hash) = task_type_hash {
315                                    batch.put(
316                                        KeySpace::TaskCache,
317                                        WriteBuffer::Borrowed(&task_type_hash),
318                                        WriteBuffer::Borrowed(key),
319                                    )?;
320                                    task_cache_items += 1;
321                                    max_new_task_id = max_new_task_id.max(*task_id);
322                                }
323                            }
324                            SnapshotItem::Delete {
325                                task_id,
326                                task_type_hash,
327                            } => {
328                                let key = IntKey::new(*task_id);
329                                let key = key.as_ref();
330                                batch.delete(KeySpace::TaskMeta, WriteBuffer::Borrowed(key))?;
331                                batch.delete(KeySpace::TaskData, WriteBuffer::Borrowed(key))?;
332                                // TaskCache is MultiValue, delete just this id from the bucket.
333                                batch.delete_value(
334                                    KeySpace::TaskCache,
335                                    WriteBuffer::Borrowed(&task_type_hash[..]),
336                                    WriteBuffer::Borrowed(key),
337                                )?;
338                            }
339                        }
340                    }
341                    Ok(SnapshotMeta {
342                        data_items,
343                        meta_items,
344                        task_cache_items,
345                        // The on-disk byte totals aren't known until the batch is committed
346                        // below; they're filled in from `CommitStats` after `batch.commit()`.
347                        bytes_written: 0,
348                        bytes_deleted: 0,
349                        max_next_task_id: max_new_task_id,
350                    })
351                })?
352                .into_iter()
353                .reduce(|t1, t2| t1.merge(t2))
354                .unwrap_or_default();
355
356            let span = tracing::trace_span!("flush task data");
357            parallel::try_for_each(
358                &[KeySpace::TaskMeta, KeySpace::TaskData, KeySpace::TaskCache],
359                |&key_space| {
360                    let _span = span.clone().entered();
361                    // Safety: `map_collect_owned` has returned, so no concurrent `put` or `delete`
362                    // on these key spaces are in-flight.
363                    unsafe { batch.flush(key_space) }
364                },
365            )?;
366
367            let mut next_task_id = get_next_free_task_id(&batch)?;
368            next_task_id = next_task_id.max(snapshot_meta.max_next_task_id + 1);
369
370            save_infra(&batch, next_task_id, operations, roots)?;
371            {
372                let _span = tracing::trace_span!("commit").entered();
373                // Byte totals are the physical on-disk bytes (post-compression, including .sst /
374                // .blob / .meta files) produced and removed by the commit.
375                let stats = batch.commit().context("Unable to commit snapshot")?;
376                snapshot_meta.bytes_written = stats.bytes_written;
377                snapshot_meta.bytes_deleted = stats.bytes_deleted;
378            }
379            Ok(snapshot_meta)
380        }
381    }
382
383    pub(crate) fn lookup_task_candidates(
384        &self,
385        native_fn: &'static NativeFunction,
386        this: Option<RawVc>,
387        arg: &dyn DynTaskInputs,
388    ) -> Result<SmallVec<[TaskId; 1]>> {
389        let inner = &*self.inner;
390        if inner.database.is_empty() {
391            // Checking if the database is empty is a performance optimization
392            // to avoid computing the hash.
393            return Ok(SmallVec::new());
394        }
395        let hash = compute_task_type_hash_from_components(native_fn, this, arg);
396        let buffers = inner
397            .database
398            .get_multiple(KeySpace::TaskCache, &hash)
399            .with_context(|| {
400                format!("Looking up task id for {native_fn:?}(this={this:?}) from database failed")
401            })?;
402
403        let mut task_ids = SmallVec::with_capacity(buffers.len());
404        for bytes in buffers {
405            let bytes = Borrow::<[u8]>::borrow(&bytes).try_into()?;
406            let id = TaskId::try_from(u32::from_le_bytes(bytes)).unwrap();
407            task_ids.push(id);
408        }
409        Ok(task_ids)
410    }
411
412    /// Reads the stored `category` for `task_id`.
413    ///
414    /// `None` means the database had no key for it. That is distinct from `Some` of an empty
415    /// [`TaskStorage`] (a key that decoded to nothing), which is what lets a `MustExist` open tell
416    /// "absent everywhere" from "present but empty".
417    pub(crate) fn lookup_data(
418        &self,
419        task_id: TaskId,
420        category: SpecificTaskDataCategory,
421    ) -> Result<Option<TaskStorage>> {
422        let inner = &*self.inner;
423        let Some(bytes) = inner
424            .database
425            .get(category.key_space(), IntKey::new(*task_id).as_ref())
426            .with_context(|| {
427                format!("Looking up task storage for {task_id} from database failed")
428            })?
429        else {
430            return Ok(None);
431        };
432        let mut storage = TaskStorage::default();
433        let mut decoder = new_turbo_bincode_decoder(bytes.borrow());
434        storage
435            .decode(category, &mut decoder)
436            .with_context(|| format!("Failed to decode {category:?}"))?;
437        Ok(Some(storage))
438    }
439
440    pub(crate) fn batch_lookup_data(
441        &self,
442        task_ids: &[TaskId],
443        category: SpecificTaskDataCategory,
444    ) -> Result<Vec<TaskStorage>> {
445        let inner = &*self.inner;
446        let int_keys: Vec<_> = task_ids.iter().map(|&id| IntKey::new(*id)).collect();
447        let keys = int_keys.iter().map(|k| k.as_ref()).collect::<Vec<_>>();
448        let bytes = inner
449            .database
450            .batch_get(category.key_space(), &keys)
451            .with_context(|| {
452                format!(
453                    "Looking up typed data for {} tasks from database failed",
454                    task_ids.len()
455                )
456            })?;
457        bytes
458            .into_iter()
459            .map(|opt_bytes| {
460                let mut storage = TaskStorage::new();
461                if let Some(bytes) = opt_bytes {
462                    let mut decoder = new_turbo_bincode_decoder(bytes.borrow());
463                    storage
464                        .decode(category, &mut decoder)
465                        .map_err(|e| anyhow::anyhow!("Failed to decode {category:?}: {e:?}"))?;
466                }
467                Ok(storage)
468            })
469            .collect::<Result<Vec<_>>>()
470    }
471
472    pub(crate) fn compact(&self) -> Result<Option<CommitStats>> {
473        self.inner.database.compact()
474    }
475
476    pub(crate) fn shutdown(&self) -> Result<()> {
477        self.inner.database.shutdown()
478    }
479
480    pub(crate) fn has_unrecoverable_write_error(&self) -> bool {
481        self.inner.database.has_unrecoverable_write_error()
482    }
483}
484
485fn get_next_free_task_id(batch: &TurboWriteBatch<'_>) -> Result<u32, anyhow::Error> {
486    Ok(
487        match batch.get(KeySpace::Infra, InfraKey::NextFreeTaskId.key().as_ref())? {
488            Some(bytes) => u32::from_le_bytes(Borrow::<[u8]>::borrow(&bytes).try_into()?),
489            None => 1,
490        },
491    )
492}
493
494fn save_infra(
495    batch: &TurboWriteBatch<'_>,
496    next_task_id: u32,
497    operations: Vec<Arc<AnyOperation>>,
498    roots: Option<Vec<(TaskId, TtlCounter)>>,
499) -> Result<(), anyhow::Error> {
500    batch
501        .put(
502            KeySpace::Infra,
503            WriteBuffer::Borrowed(InfraKey::NextFreeTaskId.key().as_ref()),
504            WriteBuffer::Borrowed(&next_task_id.to_le_bytes()),
505        )
506        .context("Unable to write next free task id")?;
507    {
508        let _span =
509            tracing::trace_span!("update operations", operations = operations.len()).entered();
510        let operations =
511            turbo_bincode_encode(&operations).context("Unable to serialize operations")?;
512        batch
513            .put(
514                KeySpace::Infra,
515                WriteBuffer::Borrowed(InfraKey::Operations.key().as_ref()),
516                WriteBuffer::SmallVec(operations),
517            )
518            .context("Unable to write operations")?;
519    }
520    if let Some(roots) = roots {
521        let _span = tracing::trace_span!("update roots", roots = roots.len()).entered();
522        let roots = turbo_bincode_encode(&roots).context("Unable to serialize GC roots")?;
523        batch
524            .put(
525                KeySpace::Infra,
526                WriteBuffer::Borrowed(InfraKey::GcRoots.key().as_ref()),
527                WriteBuffer::SmallVec(roots),
528            )
529            .context("Unable to write GC roots")?;
530    }
531    // Safety: save_infra is called after all concurrent writes to Infra are done.
532    unsafe { batch.flush(KeySpace::Infra)? };
533    Ok(())
534}
535
536#[cfg(test)]
537mod tests {
538    use std::borrow::Borrow;
539
540    use turbo_tasks::TaskId;
541
542    use super::*;
543    use crate::{
544        BackingStorageOptions,
545        database::{turbo::TurboKeyValueDatabase, write_batch::WriteBuffer},
546    };
547
548    /// Options used by these tests. `is_short_session` disables background compaction, which
549    /// requires a turbo-tasks context that these tests don't set up.
550    const TEST_STORAGE_OPTIONS: BackingStorageOptions = BackingStorageOptions {
551        is_ci: false,
552        is_short_session: true,
553        skip_compaction: false,
554    };
555
556    /// Helper to write to the database using the concurrent batch API.
557    fn write_task_cache_entry(
558        db: &TurboKeyValueDatabase,
559        hash: u64,
560        task_id: TaskId,
561    ) -> Result<()> {
562        let batch = db.write_batch()?;
563        batch.put(
564            KeySpace::TaskCache,
565            WriteBuffer::Borrowed(&hash.to_le_bytes()),
566            WriteBuffer::Borrowed(&(*task_id).to_le_bytes()),
567        )?;
568        batch.commit()?;
569        Ok(())
570    }
571
572    /// Reads the TaskIds stored under `hash` in `TaskCache`, sorted for stable comparison.
573    fn task_cache_ids(db: &TurboKeyValueDatabase, hash: u64) -> Result<Vec<TaskId>> {
574        let mut ids: Vec<TaskId> = db
575            .get_multiple(KeySpace::TaskCache, &hash.to_le_bytes())?
576            .iter()
577            .map(|bytes| {
578                let bytes: [u8; 4] = Borrow::<[u8]>::borrow(bytes).try_into().unwrap();
579                TaskId::try_from(u32::from_le_bytes(bytes)).unwrap()
580            })
581            .collect();
582        ids.sort_by_key(|id| **id);
583        Ok(ids)
584    }
585
586    /// Tests that `get_multiple` correctly returns multiple TaskIds when the same hash key
587    /// is used (simulating a hash collision scenario).
588    ///
589    /// This is a lower-level test that verifies the database layer correctly handles
590    /// the case where multiple task IDs are stored under the same hash key.
591    #[tokio::test(flavor = "multi_thread")]
592    async fn test_hash_collision_returns_multiple_candidates() -> Result<()> {
593        let tempdir = tempfile::tempdir()?;
594        let path = tempdir.path();
595
596        let db = TurboKeyValueDatabase::new(path.to_path_buf(), TEST_STORAGE_OPTIONS)?;
597
598        // Simulate a hash collision by writing multiple TaskIds with the same hash key
599        let collision_hash: u64 = 0xDEADBEEF;
600        let task_id_1 = TaskId::try_from(100u32).unwrap();
601        let task_id_2 = TaskId::try_from(200u32).unwrap();
602        let task_id_3 = TaskId::try_from(300u32).unwrap();
603
604        // Write three task IDs under the same hash key (simulating collision)
605        // Each write creates a new SST file, so all three will be returned by get_multiple
606        write_task_cache_entry(&db, collision_hash, task_id_1)?;
607        write_task_cache_entry(&db, collision_hash, task_id_2)?;
608        write_task_cache_entry(&db, collision_hash, task_id_3)?;
609
610        // Now query using get_multiple - should return all three TaskIds
611        assert_eq!(
612            task_cache_ids(&db, collision_hash)?,
613            vec![task_id_1, task_id_2, task_id_3],
614            "Should return all 3 task IDs for the colliding hash"
615        );
616
617        db.shutdown()?;
618        Ok(())
619    }
620
621    /// Tests that multiple distinct keys written in a single batch with flush can be read back.
622    /// This mirrors the actual save_snapshot pattern: write many TaskCache entries, flush, commit.
623    #[tokio::test(flavor = "multi_thread")]
624    async fn test_batch_write_with_flush_and_reopen() -> Result<()> {
625        let tempdir = tempfile::tempdir()?;
626        let path = tempdir.path();
627
628        let n = 100_000;
629        let hashes: Vec<u64> = (0..n).map(|i| 0x1000 + i as u64).collect();
630        let task_ids: Vec<TaskId> = (1..=n as u32)
631            .map(|i| TaskId::try_from(i).unwrap())
632            .collect();
633
634        // Write all entries in a single batch with flush (like save_snapshot does)
635        {
636            let db = TurboKeyValueDatabase::new(path.to_path_buf(), TEST_STORAGE_OPTIONS)?;
637            let batch = db.write_batch()?;
638
639            for (hash, task_id) in hashes.iter().zip(task_ids.iter()) {
640                batch.put(
641                    KeySpace::TaskCache,
642                    WriteBuffer::Borrowed(&hash.to_le_bytes()),
643                    WriteBuffer::Borrowed(&(**task_id).to_le_bytes()),
644                )?;
645            }
646            // Flush TaskCache (like the new code does)
647            unsafe { batch.flush(KeySpace::TaskCache) }?;
648            batch.commit()?;
649
650            db.shutdown()?;
651        }
652
653        // Reopen and verify all entries are readable
654        {
655            let db = TurboKeyValueDatabase::new(path.to_path_buf(), TEST_STORAGE_OPTIONS)?;
656            let mut found = 0;
657            let mut missing = 0;
658            for (hash, expected_id) in hashes.iter().zip(task_ids.iter()) {
659                let results = db.get_multiple(KeySpace::TaskCache, &hash.to_le_bytes())?;
660                if results.is_empty() {
661                    missing += 1;
662                } else {
663                    found += 1;
664                    let bytes: [u8; 4] = Borrow::<[u8]>::borrow(&results[0]).try_into().unwrap();
665                    let id = TaskId::try_from(u32::from_le_bytes(bytes)).unwrap();
666                    assert_eq!(id, *expected_id, "Task ID mismatch for hash {hash:#x}");
667                }
668            }
669            assert_eq!(missing, 0, "Found {found}/{n} entries, missing {missing}");
670            db.shutdown()?;
671        }
672
673        Ok(())
674    }
675
676    /// `save_snapshot` delete path: a `Delete` item must erase the task's `TaskMeta` and
677    /// `TaskData` (`SingleValue`) entries and remove *only* that id from its `TaskCache`
678    /// (`MultiValue`) bucket.
679    ///
680    /// The colliding survivor is never read or rewritten — the key-value tombstone names the
681    /// single id it deletes, so anything else in the bucket is untouched whether or not this
682    /// commit knows about it.
683    #[tokio::test(flavor = "multi_thread")]
684    async fn test_save_snapshot_delete_tombstones_task() -> Result<()> {
685        let tempdir = tempfile::tempdir()?;
686        let path = tempdir.path();
687
688        let collision_hash: u64 = 0xC0FFEE;
689        let deleted_id = TaskId::try_from(111u32).unwrap();
690        let survivor_id = TaskId::try_from(222u32).unwrap();
691        let deleted_key = (*deleted_id).to_le_bytes();
692
693        let db = TurboKeyValueDatabase::new(
694            path.to_path_buf(),
695            BackingStorageOptions {
696                is_ci: false,
697                is_short_session: true,
698                skip_compaction: false,
699            },
700        )?;
701
702        // Both ids collide in one TaskCache bucket, purely on disk; the deleted task also has
703        // meta and data entries.
704        write_task_cache_entry(&db, collision_hash, deleted_id)?;
705        write_task_cache_entry(&db, collision_hash, survivor_id)?;
706        let batch = db.write_batch()?;
707        batch.put(
708            KeySpace::TaskMeta,
709            WriteBuffer::Borrowed(&deleted_key),
710            WriteBuffer::Borrowed(b"meta-bytes"),
711        )?;
712        batch.put(
713            KeySpace::TaskData,
714            WriteBuffer::Borrowed(&deleted_key),
715            WriteBuffer::Borrowed(b"data-bytes"),
716        )?;
717        batch.commit()?;
718
719        // Sanity: everything is present before the delete.
720        assert!(db.get(KeySpace::TaskMeta, &deleted_key)?.is_some());
721        assert!(db.get(KeySpace::TaskData, &deleted_key)?.is_some());
722        assert_eq!(
723            task_cache_ids(&db, collision_hash)?,
724            vec![deleted_id, survivor_id],
725        );
726
727        let storage = TurboBackingStorage::new_in_memory(db);
728
729        // Snapshot with no task data, just the one deletion.
730        storage.save_snapshot(
731            Vec::new(),
732            None,
733            vec![vec![SnapshotItem::Delete {
734                task_id: deleted_id,
735                task_type_hash: collision_hash.to_le_bytes(),
736            }]],
737        )?;
738
739        let db = &storage.inner.database;
740        assert!(
741            db.get(KeySpace::TaskMeta, &deleted_key)?.is_none(),
742            "TaskMeta should be tombstoned"
743        );
744        assert!(
745            db.get(KeySpace::TaskData, &deleted_key)?.is_none(),
746            "TaskData should be tombstoned"
747        );
748        assert_eq!(
749            task_cache_ids(db, collision_hash)?,
750            vec![survivor_id],
751            "save_snapshot should delete only the named id from the bucket"
752        );
753
754        db.shutdown()?;
755        Ok(())
756    }
757}
turbo_tasks_backend/kv_backing_storage.rs

turbo_tasks_backend/
kv_backing_storage.rs
