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, 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
33const META_KEY_OPERATIONS: u32 = 0;
34const META_KEY_NEXT_FREE_TASK_ID: u32 = 1;
35
36struct IntKey([u8; 4]);
37
38impl IntKey {
39 fn new(value: u32) -> Self {
40 Self(value.to_le_bytes())
41 }
42}
43
44impl AsRef<[u8]> for IntKey {
45 fn as_ref(&self) -> &[u8] {
46 &self.0
47 }
48}
49
50fn as_u32(bytes: impl Borrow<[u8]>) -> Result<u32> {
51 let n = u32::from_le_bytes(bytes.borrow().try_into()?);
52 Ok(n)
53}
54
55fn should_invalidate_on_panic() -> bool {
64 fn env_is_falsy(key: &str) -> bool {
65 env::var_os(key)
66 .is_none_or(|value| ["".as_ref(), "0".as_ref(), "false".as_ref()].contains(&&*value))
67 }
68 static SHOULD_INVALIDATE: LazyLock<bool> = LazyLock::new(|| {
69 env_is_falsy("TURBO_ENGINE_SKIP_INVALIDATE_ON_PANIC") && env_is_falsy("__NEXT_TEST_MODE")
70 });
71 *SHOULD_INVALIDATE
72}
73
74struct TurboBackingStorageInner {
75 database: TurboKeyValueDatabase,
76 base_path: Option<PathBuf>,
79 invalidated: Mutex<bool>,
81 _panic_hook_guard: Option<PanicHookGuard>,
84}
85
86pub struct TurboBackingStorage {
94 inner: Arc<TurboBackingStorageInner>,
96}
97
98impl TurboBackingStorage {
99 pub(crate) fn new_in_memory(database: TurboKeyValueDatabase) -> Self {
100 Self {
101 inner: Arc::new(TurboBackingStorageInner {
102 database,
103 base_path: None,
104 invalidated: Mutex::new(false),
105 _panic_hook_guard: None,
106 }),
107 }
108 }
109
110 pub(crate) fn open_versioned_on_disk(
122 base_path: PathBuf,
123 version_info: &GitVersionInfo,
124 is_ci: bool,
125 database: impl FnOnce(PathBuf) -> Result<TurboKeyValueDatabase>,
126 ) -> Result<(Self, StartupCacheState)> {
127 let startup_cache_state = check_db_invalidation_and_cleanup(&base_path)
128 .context("Failed to check database invalidation and cleanup")?;
129 let versioned_path = handle_db_versioning(&base_path, version_info, is_ci)
130 .context("Failed to handle database versioning")?;
131 let database = (database)(versioned_path).context("Failed to open database")?;
132 let backing_storage = Self {
133 inner: Arc::new_cyclic(move |weak_inner: &Weak<TurboBackingStorageInner>| {
134 let panic_hook_guard = if should_invalidate_on_panic() {
135 let weak_inner = weak_inner.clone();
136 Some(register_panic_hook(Box::new(move |_| {
137 let Some(inner) = weak_inner.upgrade() else {
138 return;
139 };
140 let _ = inner.invalidate(invalidation_reasons::PANIC);
145 })))
146 } else {
147 None
148 };
149 TurboBackingStorageInner {
150 database,
151 base_path: Some(base_path),
152 invalidated: Mutex::new(false),
153 _panic_hook_guard: panic_hook_guard,
154 }
155 }),
156 };
157 Ok((backing_storage, startup_cache_state))
158 }
159}
160
161impl TurboBackingStorageInner {
162 fn invalidate(&self, reason_code: &str) -> Result<()> {
163 if let Some(base_path) = &self.base_path {
165 let mut invalidated_guard = self
168 .invalidated
169 .lock()
170 .unwrap_or_else(PoisonError::into_inner);
171 if *invalidated_guard {
172 return Ok(());
173 }
174 invalidate_db(base_path, reason_code)?;
178 self.database.prevent_writes();
179 *invalidated_guard = true;
181 }
182 Ok(())
183 }
184
185 fn get_infra_u32(&self, key: u32) -> Result<Option<u32>> {
187 self.database
188 .get(KeySpace::Infra, IntKey::new(key).as_ref())?
189 .map(as_u32)
190 .transpose()
191 }
192}
193
194impl TurboBackingStorage {
195 pub(crate) fn invalidate(&self, reason_code: &str) -> Result<()> {
203 self.inner.invalidate(reason_code)
204 }
205
206 pub(crate) fn next_free_task_id(&self) -> Result<TaskId> {
207 Ok(self
208 .inner
209 .get_infra_u32(META_KEY_NEXT_FREE_TASK_ID)
210 .context("Unable to read next free task id from database")?
211 .map_or(Ok(TaskId::MIN), TaskId::try_from)?)
212 }
213
214 pub(crate) fn uncompleted_operations(&self) -> Result<Vec<AnyOperation>> {
215 fn get(database: &TurboKeyValueDatabase) -> Result<Vec<AnyOperation>> {
216 let Some(operations) =
217 database.get(KeySpace::Infra, IntKey::new(META_KEY_OPERATIONS).as_ref())?
218 else {
219 return Ok(Vec::new());
220 };
221 let operations = turbo_bincode_decode(operations.borrow())?;
222 Ok(operations)
223 }
224 get(&self.inner.database).context("Unable to read uncompleted operations from database")
225 }
226
227 pub(crate) fn save_snapshot<I>(
228 &self,
229 operations: Vec<Arc<AnyOperation>>,
230 snapshots: Vec<I>,
231 ) -> Result<SnapshotMeta>
232 where
233 I: IntoIterator<Item = SnapshotItem> + Send + Sync,
234 {
235 let _span = tracing::info_span!("save snapshot", operations = operations.len()).entered();
236 let batch = self.inner.database.write_batch()?;
237
238 {
239 let span = tracing::trace_span!("update task data");
240 let mut snapshot_meta =
241 parallel::map_collect_owned::<_, _, Result<Vec<_>>>(snapshots, |shard: I| {
242 let _span = span.clone().entered();
243 let mut max_new_task_id = 0;
244 let mut data_items = 0;
245 let mut meta_items = 0;
246 let mut task_cache_items = 0;
247 for item in shard {
248 match item {
249 SnapshotItem::Put {
250 task_id,
251 meta,
252 data,
253 task_type_hash,
254 } => {
255 let key = IntKey::new(*task_id);
256 let key = key.as_ref();
257 if let Some(meta) = meta {
258 batch.put(
259 KeySpace::TaskMeta,
260 WriteBuffer::Borrowed(key),
261 WriteBuffer::SmallVec(meta),
262 )?;
263 meta_items += 1;
264 }
265 if let Some(data) = data {
266 batch.put(
267 KeySpace::TaskData,
268 WriteBuffer::Borrowed(key),
269 WriteBuffer::SmallVec(data),
270 )?;
271 data_items += 1;
272 }
273 if let Some(task_type_hash) = task_type_hash {
275 batch.put(
276 KeySpace::TaskCache,
277 WriteBuffer::Borrowed(&task_type_hash),
278 WriteBuffer::Borrowed(key),
279 )?;
280 task_cache_items += 1;
281 max_new_task_id = max_new_task_id.max(*task_id);
282 }
283 }
284 SnapshotItem::Delete {
285 task_id,
286 task_type_hash,
287 } => {
288 let key = IntKey::new(*task_id);
289 let key = key.as_ref();
290 batch.delete(KeySpace::TaskMeta, WriteBuffer::Borrowed(key))?;
291 batch.delete(KeySpace::TaskData, WriteBuffer::Borrowed(key))?;
292 batch.delete_value(
294 KeySpace::TaskCache,
295 WriteBuffer::Borrowed(&task_type_hash[..]),
296 WriteBuffer::Borrowed(key),
297 )?;
298 }
299 }
300 }
301 Ok(SnapshotMeta {
302 data_items,
303 meta_items,
304 task_cache_items,
305 bytes_written: 0,
308 bytes_deleted: 0,
309 max_next_task_id: max_new_task_id,
310 })
311 })?
312 .into_iter()
313 .reduce(|t1, t2| t1.merge(t2))
314 .unwrap_or_default();
315
316 let span = tracing::trace_span!("flush task data");
317 parallel::try_for_each(
318 &[KeySpace::TaskMeta, KeySpace::TaskData, KeySpace::TaskCache],
319 |&key_space| {
320 let _span = span.clone().entered();
321 unsafe { batch.flush(key_space) }
324 },
325 )?;
326
327 let mut next_task_id = get_next_free_task_id(&batch)?;
328 next_task_id = next_task_id.max(snapshot_meta.max_next_task_id + 1);
329
330 save_infra(&batch, next_task_id, operations)?;
331 {
332 let _span = tracing::trace_span!("commit").entered();
333 let stats = batch.commit().context("Unable to commit snapshot")?;
336 snapshot_meta.bytes_written = stats.bytes_written;
337 snapshot_meta.bytes_deleted = stats.bytes_deleted;
338 }
339 Ok(snapshot_meta)
340 }
341 }
342
343 pub(crate) fn lookup_task_candidates(
344 &self,
345 native_fn: &'static NativeFunction,
346 this: Option<RawVc>,
347 arg: &dyn DynTaskInputs,
348 ) -> Result<SmallVec<[TaskId; 1]>> {
349 let inner = &*self.inner;
350 if inner.database.is_empty() {
351 return Ok(SmallVec::new());
354 }
355 let hash = compute_task_type_hash_from_components(native_fn, this, arg);
356 let buffers = inner
357 .database
358 .get_multiple(KeySpace::TaskCache, &hash)
359 .with_context(|| {
360 format!("Looking up task id for {native_fn:?}(this={this:?}) from database failed")
361 })?;
362
363 let mut task_ids = SmallVec::with_capacity(buffers.len());
364 for bytes in buffers {
365 let bytes = Borrow::<[u8]>::borrow(&bytes).try_into()?;
366 let id = TaskId::try_from(u32::from_le_bytes(bytes)).unwrap();
367 task_ids.push(id);
368 }
369 Ok(task_ids)
370 }
371
372 pub(crate) fn lookup_data(
378 &self,
379 task_id: TaskId,
380 category: SpecificTaskDataCategory,
381 ) -> Result<Option<TaskStorage>> {
382 let inner = &*self.inner;
383 let Some(bytes) = inner
384 .database
385 .get(category.key_space(), IntKey::new(*task_id).as_ref())
386 .with_context(|| {
387 format!("Looking up task storage for {task_id} from database failed")
388 })?
389 else {
390 return Ok(None);
391 };
392 let mut storage = TaskStorage::default();
393 let mut decoder = new_turbo_bincode_decoder(bytes.borrow());
394 storage
395 .decode(category, &mut decoder)
396 .with_context(|| format!("Failed to decode {category:?}"))?;
397 Ok(Some(storage))
398 }
399
400 pub(crate) fn batch_lookup_data(
401 &self,
402 task_ids: &[TaskId],
403 category: SpecificTaskDataCategory,
404 ) -> Result<Vec<TaskStorage>> {
405 let inner = &*self.inner;
406 let int_keys: Vec<_> = task_ids.iter().map(|&id| IntKey::new(*id)).collect();
407 let keys = int_keys.iter().map(|k| k.as_ref()).collect::<Vec<_>>();
408 let bytes = inner
409 .database
410 .batch_get(category.key_space(), &keys)
411 .with_context(|| {
412 format!(
413 "Looking up typed data for {} tasks from database failed",
414 task_ids.len()
415 )
416 })?;
417 bytes
418 .into_iter()
419 .map(|opt_bytes| {
420 let mut storage = TaskStorage::new();
421 if let Some(bytes) = opt_bytes {
422 let mut decoder = new_turbo_bincode_decoder(bytes.borrow());
423 storage
424 .decode(category, &mut decoder)
425 .map_err(|e| anyhow::anyhow!("Failed to decode {category:?}: {e:?}"))?;
426 }
427 Ok(storage)
428 })
429 .collect::<Result<Vec<_>>>()
430 }
431
432 pub(crate) fn compact(&self) -> Result<Option<CommitStats>> {
433 self.inner.database.compact()
434 }
435
436 pub(crate) fn shutdown(&self) -> Result<()> {
437 self.inner.database.shutdown()
438 }
439
440 pub(crate) fn has_unrecoverable_write_error(&self) -> bool {
441 self.inner.database.has_unrecoverable_write_error()
442 }
443}
444
445fn get_next_free_task_id(batch: &TurboWriteBatch<'_>) -> Result<u32, anyhow::Error> {
446 Ok(
447 match batch.get(
448 KeySpace::Infra,
449 IntKey::new(META_KEY_NEXT_FREE_TASK_ID).as_ref(),
450 )? {
451 Some(bytes) => u32::from_le_bytes(Borrow::<[u8]>::borrow(&bytes).try_into()?),
452 None => 1,
453 },
454 )
455}
456
457fn save_infra(
458 batch: &TurboWriteBatch<'_>,
459 next_task_id: u32,
460 operations: Vec<Arc<AnyOperation>>,
461) -> Result<(), anyhow::Error> {
462 batch
463 .put(
464 KeySpace::Infra,
465 WriteBuffer::Borrowed(IntKey::new(META_KEY_NEXT_FREE_TASK_ID).as_ref()),
466 WriteBuffer::Borrowed(&next_task_id.to_le_bytes()),
467 )
468 .context("Unable to write next free task id")?;
469 {
470 let _span =
471 tracing::trace_span!("update operations", operations = operations.len()).entered();
472 let operations =
473 turbo_bincode_encode(&operations).context("Unable to serialize operations")?;
474 batch
475 .put(
476 KeySpace::Infra,
477 WriteBuffer::Borrowed(IntKey::new(META_KEY_OPERATIONS).as_ref()),
478 WriteBuffer::SmallVec(operations),
479 )
480 .context("Unable to write operations")?;
481 }
482 unsafe { batch.flush(KeySpace::Infra)? };
484 Ok(())
485}
486
487#[cfg(test)]
488mod tests {
489 use std::borrow::Borrow;
490
491 use turbo_tasks::TaskId;
492
493 use super::*;
494 use crate::database::{turbo::TurboKeyValueDatabase, write_batch::WriteBuffer};
495
496 fn write_task_cache_entry(
498 db: &TurboKeyValueDatabase,
499 hash: u64,
500 task_id: TaskId,
501 ) -> Result<()> {
502 let batch = db.write_batch()?;
503 batch.put(
504 KeySpace::TaskCache,
505 WriteBuffer::Borrowed(&hash.to_le_bytes()),
506 WriteBuffer::Borrowed(&(*task_id).to_le_bytes()),
507 )?;
508 batch.commit()?;
509 Ok(())
510 }
511
512 fn task_cache_ids(db: &TurboKeyValueDatabase, hash: u64) -> Result<Vec<TaskId>> {
514 let mut ids: Vec<TaskId> = db
515 .get_multiple(KeySpace::TaskCache, &hash.to_le_bytes())?
516 .iter()
517 .map(|bytes| {
518 let bytes: [u8; 4] = Borrow::<[u8]>::borrow(bytes).try_into().unwrap();
519 TaskId::try_from(u32::from_le_bytes(bytes)).unwrap()
520 })
521 .collect();
522 ids.sort_by_key(|id| **id);
523 Ok(ids)
524 }
525
526 #[tokio::test(flavor = "multi_thread")]
532 async fn test_hash_collision_returns_multiple_candidates() -> Result<()> {
533 let tempdir = tempfile::tempdir()?;
534 let path = tempdir.path();
535
536 let db = TurboKeyValueDatabase::new(path.to_path_buf(), false, true, false)?;
539
540 let collision_hash: u64 = 0xDEADBEEF;
542 let task_id_1 = TaskId::try_from(100u32).unwrap();
543 let task_id_2 = TaskId::try_from(200u32).unwrap();
544 let task_id_3 = TaskId::try_from(300u32).unwrap();
545
546 write_task_cache_entry(&db, collision_hash, task_id_1)?;
549 write_task_cache_entry(&db, collision_hash, task_id_2)?;
550 write_task_cache_entry(&db, collision_hash, task_id_3)?;
551
552 assert_eq!(
554 task_cache_ids(&db, collision_hash)?,
555 vec![task_id_1, task_id_2, task_id_3],
556 "Should return all 3 task IDs for the colliding hash"
557 );
558
559 db.shutdown()?;
560 Ok(())
561 }
562
563 #[tokio::test(flavor = "multi_thread")]
566 async fn test_batch_write_with_flush_and_reopen() -> Result<()> {
567 let tempdir = tempfile::tempdir()?;
568 let path = tempdir.path();
569
570 let n = 100_000;
571 let hashes: Vec<u64> = (0..n).map(|i| 0x1000 + i as u64).collect();
572 let task_ids: Vec<TaskId> = (1..=n as u32)
573 .map(|i| TaskId::try_from(i).unwrap())
574 .collect();
575
576 {
578 let db = TurboKeyValueDatabase::new(path.to_path_buf(), false, true, false)?;
579 let batch = db.write_batch()?;
580
581 for (hash, task_id) in hashes.iter().zip(task_ids.iter()) {
582 batch.put(
583 KeySpace::TaskCache,
584 WriteBuffer::Borrowed(&hash.to_le_bytes()),
585 WriteBuffer::Borrowed(&(**task_id).to_le_bytes()),
586 )?;
587 }
588 unsafe { batch.flush(KeySpace::TaskCache) }?;
590 batch.commit()?;
591
592 db.shutdown()?;
593 }
594
595 {
597 let db = TurboKeyValueDatabase::new(path.to_path_buf(), false, true, false)?;
598 let mut found = 0;
599 let mut missing = 0;
600 for (hash, expected_id) in hashes.iter().zip(task_ids.iter()) {
601 let results = db.get_multiple(KeySpace::TaskCache, &hash.to_le_bytes())?;
602 if results.is_empty() {
603 missing += 1;
604 } else {
605 found += 1;
606 let bytes: [u8; 4] = Borrow::<[u8]>::borrow(&results[0]).try_into().unwrap();
607 let id = TaskId::try_from(u32::from_le_bytes(bytes)).unwrap();
608 assert_eq!(id, *expected_id, "Task ID mismatch for hash {hash:#x}");
609 }
610 }
611 assert_eq!(missing, 0, "Found {found}/{n} entries, missing {missing}");
612 db.shutdown()?;
613 }
614
615 Ok(())
616 }
617
618 #[tokio::test(flavor = "multi_thread")]
626 async fn test_save_snapshot_delete_tombstones_task() -> Result<()> {
627 let tempdir = tempfile::tempdir()?;
628 let path = tempdir.path();
629
630 let collision_hash: u64 = 0xC0FFEE;
631 let deleted_id = TaskId::try_from(111u32).unwrap();
632 let survivor_id = TaskId::try_from(222u32).unwrap();
633 let deleted_key = (*deleted_id).to_le_bytes();
634
635 let db = TurboKeyValueDatabase::new(path.to_path_buf(), false, true, false)?;
636
637 write_task_cache_entry(&db, collision_hash, deleted_id)?;
640 write_task_cache_entry(&db, collision_hash, survivor_id)?;
641 let batch = db.write_batch()?;
642 batch.put(
643 KeySpace::TaskMeta,
644 WriteBuffer::Borrowed(&deleted_key),
645 WriteBuffer::Borrowed(b"meta-bytes"),
646 )?;
647 batch.put(
648 KeySpace::TaskData,
649 WriteBuffer::Borrowed(&deleted_key),
650 WriteBuffer::Borrowed(b"data-bytes"),
651 )?;
652 batch.commit()?;
653
654 assert!(db.get(KeySpace::TaskMeta, &deleted_key)?.is_some());
656 assert!(db.get(KeySpace::TaskData, &deleted_key)?.is_some());
657 assert_eq!(
658 task_cache_ids(&db, collision_hash)?,
659 vec![deleted_id, survivor_id],
660 );
661
662 let storage = TurboBackingStorage::new_in_memory(db);
663
664 storage.save_snapshot(
666 Vec::new(),
667 vec![vec![SnapshotItem::Delete {
668 task_id: deleted_id,
669 task_type_hash: collision_hash.to_le_bytes(),
670 }]],
671 )?;
672
673 let db = &storage.inner.database;
674 assert!(
675 db.get(KeySpace::TaskMeta, &deleted_key)?.is_none(),
676 "TaskMeta should be tombstoned"
677 );
678 assert!(
679 db.get(KeySpace::TaskData, &deleted_key)?.is_none(),
680 "TaskData should be tombstoned"
681 );
682 assert_eq!(
683 task_cache_ids(db, collision_hash)?,
684 vec![survivor_id],
685 "save_snapshot should delete only the named id from the bucket"
686 );
687
688 db.shutdown()?;
689 Ok(())
690 }
691}