Move more responsibilities into the KeyValueDatabaseBackingStorage constructor

Author: bgw

turbopack/crates/turbo-tasks-backend/src/database/db_invalidation.rs

@@ -4,19 +4,21 @@ use std::{
     path::Path,
 };
 
+use anyhow::Context;
+
 const INVALIDATION_MARKER: &str = "__turbo_tasks_invalidated_db";
 
 /// Atomically write an invalidation marker.
 ///
 /// Because attempting to delete currently open database files could cause issues, actual deletion
-/// of files is deferred until the next start-up (in [`check_db_invalidation`]).
+/// of files is deferred until the next start-up (in [`check_db_invalidation_and_cleanup`]).
 ///
 /// In the case that no database is currently open (e.g. via a separate CLI subcommand), you should
-/// call [`cleanup_db`] after this to eagerly remove the database files.
+/// call [`cleanup_db`] *after* this to eagerly remove the database files.
 ///
 /// This should be run with the base (non-versioned) path, as that likely aligns closest with user
 /// expectations (e.g. if they're clearing the cache for disk space reasons).
-pub fn invalidate_db(base_path: &Path) -> io::Result<()> {
+pub fn invalidate_db(base_path: &Path) -> anyhow::Result<()> {
     fs::write(base_path.join(INVALIDATION_MARKER), [0u8; 0])?;
     Ok(())
 }
@@ -25,7 +27,7 @@ pub fn invalidate_db(base_path: &Path) -> io::Result<()> {
 /// delete any invalidated database files.
 ///
 /// This should be run with the base (non-versioned) path.
-pub fn check_db_invalidation_and_cleanup(base_path: &Path) -> io::Result<()> {
+pub fn check_db_invalidation_and_cleanup(base_path: &Path) -> anyhow::Result<()> {
     if fs::exists(base_path.join(INVALIDATION_MARKER))? {
         // if this cleanup fails, we might try to open an invalid database later, so it's best to
         // just propagate the error here.
@@ -34,8 +36,21 @@ pub fn check_db_invalidation_and_cleanup(base_path: &Path) -> io::Result<()> {
     Ok(())
 }
 
-/// Helper for `invalidate_db` and `check_db_invalidation`.
-pub fn cleanup_db(base_path: &Path) -> io::Result<()> {
+/// Helper for [`check_db_invalidation_and_cleanup`]. You can call this to explicitly clean up a
+/// database after running [`invalidate_db`] when turbo-tasks is not running.
+///
+/// You should not run this if the database has not yet been invalidated, as this operation is not
+/// atomic and could result in a partially-deleted and corrupted database.
+pub fn cleanup_db(base_path: &Path) -> anyhow::Result<()> {
+    cleanup_db_inner(base_path).with_context(|| {
+        format!(
+            "Unable to remove invalid database. If this issue persists you can work around by \
+             deleting {base_path:?}."
+        )
+    })
+}
+
+fn cleanup_db_inner(base_path: &Path) -> io::Result<()> {
     let Ok(contents) = read_dir(base_path) else {
         return Ok(());
     };
@@ -44,7 +59,7 @@ pub fn cleanup_db(base_path: &Path) -> io::Result<()> {
     for entry in contents {
         let entry = entry?;
         if entry.file_name() != INVALIDATION_MARKER {
-            let _ = fs::remove_dir_all(entry.path());
+            fs::remove_dir_all(entry.path())?;
         }
     }
 

turbopack/crates/turbo-tasks-backend/src/database/key_value_database.rs

@@ -60,7 +60,8 @@ pub trait KeyValueDatabase {
     /// during [`KeyValueDatabase::shutdown`].
     ///
     /// This is a best-effort optimization hint, and the database may choose to ignore this and
-    /// continue file writes.
+    /// continue file writes. This happens after the database is invalidated, so it is valid for
+    /// this to leave the database in a half-updated and corrupted state.
     fn prevent_writes(&self) {
         // this is an optional performance hint to the database
     }

turbopack/crates/turbo-tasks-backend/src/database/turbo.rs

@@ -24,7 +24,7 @@ pub struct TurboKeyValueDatabase {
 
 impl TurboKeyValueDatabase {
     pub fn new(versioned_path: PathBuf) -> Result<Self> {
-        let db = Arc::new(TurboPersistence::open(versioned_path.to_path_buf())?);
+        let db = Arc::new(TurboPersistence::open(versioned_path)?);
         let mut this = Self {
             db: db.clone(),
             compact_join_handle: Mutex::new(None),

turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs

@@ -8,11 +8,13 @@ use tracing::Span;
 use turbo_tasks::{SessionId, TaskId, backend::CachedTaskType, turbo_tasks_scope};
 
 use crate::{
+    GitVersionInfo,
     backend::{AnyOperation, TaskDataCategory},
     backing_storage::BackingStorage,
     data::CachedDataItem,
     database::{
-        db_invalidation::invalidate_db,
+        db_invalidation::{check_db_invalidation_and_cleanup, invalidate_db},
+        db_versioning::handle_db_versioning,
         key_value_database::{KeySpace, KeyValueDatabase},
         write_batch::{
             BaseWriteBatch, ConcurrentWriteBatch, SerialWriteBatch, WriteBatch, WriteBatchRef,
@@ -89,13 +91,26 @@ pub struct KeyValueDatabaseBackingStorage<T: KeyValueDatabase> {
 }
 
 impl<T: KeyValueDatabase> KeyValueDatabaseBackingStorage<T> {
-    pub fn new(database: T, base_path: Option<PathBuf>) -> Self {
+    pub fn new_in_memory(database: T) -> Self {
         Self {
             database,
-            base_path,
+            base_path: None,
         }
     }
 
+    pub fn open_versioned_on_disk(
+        base_path: PathBuf,
+        version_info: &GitVersionInfo,
+        database: impl FnOnce(PathBuf) -> Result<T>,
+    ) -> Result<Self> {
+        check_db_invalidation_and_cleanup(&base_path)?;
+        let versioned_path = handle_db_versioning(&base_path, version_info)?;
+        Ok(Self {
+            database: (database)(versioned_path)?,
+            base_path: Some(base_path),
+        })
+    }
+
     fn with_tx<R>(
         &self,
         tx: Option<&T::ReadTransaction<'_>>,

turbopack/crates/turbo-tasks-backend/src/lib.rs

@@ -14,11 +14,8 @@ mod utils;
 use std::path::Path;
 
 use anyhow::Result;
-use database::db_invalidation::check_db_invalidation_and_cleanup;
 
-use crate::database::{
-    db_versioning::handle_db_versioning, noop_kv::NoopKvDb, turbo::TurboKeyValueDatabase,
-};
+use crate::database::{noop_kv::NoopKvDb, turbo::TurboKeyValueDatabase};
 pub use crate::{
     backend::{BackendOptions, StorageMode, TurboTasksBackend},
     database::db_versioning::GitVersionInfo,
@@ -47,18 +44,18 @@ pub fn lmdb_backing_storage(
         startup_cache::StartupCacheLayer,
     };
 
-    check_db_invalidation_and_cleanup(base_path)?;
-    let versioned_path = handle_db_versioning(base_path, version_info)?;
-    let fresh_db = is_fresh(&versioned_path);
-    let database = crate::database::lmdb::LmbdKeyValueDatabase::new(&versioned_path)?;
-    let database = FreshDbOptimization::new(database, fresh_db);
-    let database =
-        StartupCacheLayer::new(database, versioned_path.join("startup.cache"), fresh_db)?;
-    let database = ReadTransactionCache::new(database);
-    Ok(KeyValueDatabaseBackingStorage::new(
-        database,
-        Some(base_path.to_owned()),
-    ))
+    KeyValueDatabaseBackingStorage::open_versioned_on_disk(
+        base_path.to_owned(),
+        version_info,
+        |versioned_path| {
+            let fresh_db = is_fresh(&versioned_path);
+            let database = crate::database::lmdb::LmbdKeyValueDatabase::new(&versioned_path)?;
+            let database = FreshDbOptimization::new(database, fresh_db);
+            let database =
+                StartupCacheLayer::new(database, versioned_path.join("startup.cache"), fresh_db)?;
+            Ok(ReadTransactionCache::new(database))
+        },
+    )
 }
 
 pub type TurboBackingStorage = KeyValueDatabaseBackingStorage<TurboKeyValueDatabase>;
@@ -67,19 +64,17 @@ pub fn turbo_backing_storage(
     base_path: &Path,
     version_info: &GitVersionInfo,
 ) -> Result<TurboBackingStorage> {
-    check_db_invalidation_and_cleanup(base_path)?;
-    let versioned_path = handle_db_versioning(base_path, version_info)?;
-    let database = TurboKeyValueDatabase::new(versioned_path)?;
-    Ok(KeyValueDatabaseBackingStorage::new(
-        database,
-        Some(base_path.to_owned()),
-    ))
+    KeyValueDatabaseBackingStorage::open_versioned_on_disk(
+        base_path.to_owned(),
+        version_info,
+        TurboKeyValueDatabase::new,
+    )
 }
 
 pub type NoopBackingStorage = KeyValueDatabaseBackingStorage<NoopKvDb>;
 
 pub fn noop_backing_storage() -> NoopBackingStorage {
-    KeyValueDatabaseBackingStorage::new(NoopKvDb, None)
+    KeyValueDatabaseBackingStorage::new_in_memory(NoopKvDb)
 }
 
 #[cfg(feature = "lmdb")]