Skip to main content

turbo_tasks_backend/
backing_storage.rs

1use std::{any::type_name, sync::Arc};
2
3use anyhow::Result;
4use either::Either;
5use smallvec::SmallVec;
6use turbo_bincode::TurboBincodeBuffer;
7use turbo_tasks::{TaskId, backend::CachedTaskType};
8
9use crate::{
10    backend::{AnyOperation, SpecificTaskDataCategory, storage_schema::TaskStorage},
11    utils::chunked_vec::ChunkedVec,
12};
13
14pub struct SnapshotItem {
15    pub task_id: TaskId,
16    pub data: Option<TurboBincodeBuffer>,
17    pub meta: Option<TurboBincodeBuffer>,
18}
19
20impl SnapshotItem {
21    pub fn is_empty(&self) -> bool {
22        self.meta.is_none() && self.data.is_none()
23    }
24}
25
26/// Represents types accepted by [`TurboTasksBackend::new`]. Typically this is the value returned by
27/// [`default_backing_storage`] or [`noop_backing_storage`].
28///
29/// This trait is [sealed]. External crates are not allowed to implement it.
30///
31/// [`default_backing_storage`]: crate::default_backing_storage
32/// [`noop_backing_storage`]: crate::noop_backing_storage
33/// [`TurboTasksBackend::new`]: crate::TurboTasksBackend::new
34/// [sealed]: https://predr.ag/blog/definitive-guide-to-sealed-traits-in-rust/
35pub trait BackingStorage: BackingStorageSealed {
36    /// Called when the database should be invalidated upon re-initialization.
37    ///
38    /// This typically means that we'll restart the process or `turbo-tasks` soon with a fresh
39    /// database. If this happens, there's no point in writing anything else to disk, or flushing
40    /// during [`KeyValueDatabase::shutdown`].
41    ///
42    /// This can be implemented by calling [`invalidate_db`] with
43    /// the database's non-versioned base path.
44    ///
45    /// [`KeyValueDatabase::shutdown`]: crate::database::key_value_database::KeyValueDatabase::shutdown
46    /// [`invalidate_db`]: crate::database::db_invalidation::invalidate_db
47    fn invalidate(&self, reason_code: &str) -> Result<()>;
48}
49
50/// Private methods used by [`BackingStorage`]. This trait is `pub` (because of the sealed-trait
51/// pattern), but should not be exported outside of the crate.
52///
53/// [`BackingStorage`] is exported for documentation reasons and to expose the public
54/// [`BackingStorage::invalidate`] method.
55pub trait BackingStorageSealed: 'static + Send + Sync {
56    type ReadTransaction<'l>;
57    fn next_free_task_id(&self) -> Result<TaskId>;
58    fn uncompleted_operations(&self) -> Result<Vec<AnyOperation>>;
59
60    fn save_snapshot<I>(
61        &self,
62        operations: Vec<Arc<AnyOperation>>,
63        task_cache_updates: Vec<ChunkedVec<(Arc<CachedTaskType>, TaskId)>>,
64        snapshots: Vec<I>,
65    ) -> Result<()>
66    where
67        I: Iterator<Item = SnapshotItem> + Send + Sync;
68    fn start_read_transaction(&self) -> Option<Self::ReadTransaction<'_>>;
69    /// Returns all task IDs that match the given task type (hash collision candidates).
70    ///
71    /// Since TaskCache uses hash-based keys, multiple task types may (rarely) hash to the same key.
72    /// The caller must verify each returned TaskId by comparing the stored task type which will
73    /// require a second database read
74    ///
75    /// # Safety
76    ///
77    /// `tx` must be a transaction from this BackingStorage instance.
78    unsafe fn lookup_task_candidates(
79        &self,
80        tx: Option<&Self::ReadTransaction<'_>>,
81        key: &CachedTaskType,
82    ) -> Result<SmallVec<[TaskId; 1]>>;
83    /// # Safety
84    ///
85    /// `tx` must be a transaction from this BackingStorage instance.
86    unsafe fn lookup_data(
87        &self,
88        tx: Option<&Self::ReadTransaction<'_>>,
89        task_id: TaskId,
90        category: SpecificTaskDataCategory,
91        storage: &mut TaskStorage,
92    ) -> Result<()>;
93
94    /// Batch lookup and decode data for multiple tasks directly into TypedStorage instances.
95    /// Returns a vector of TypedStorage, one for each task_id in the input slice.
96    /// # Safety
97    ///
98    /// `tx` must be a transaction from this BackingStorage instance.
99    unsafe fn batch_lookup_data(
100        &self,
101        tx: Option<&Self::ReadTransaction<'_>>,
102        task_ids: &[TaskId],
103        category: SpecificTaskDataCategory,
104    ) -> Result<Vec<TaskStorage>>;
105
106    fn shutdown(&self) -> Result<()> {
107        Ok(())
108    }
109}
110
111impl<L, R> BackingStorage for Either<L, R>
112where
113    L: BackingStorage,
114    R: BackingStorage,
115{
116    fn invalidate(&self, reason_code: &str) -> Result<()> {
117        either::for_both!(self, this => this.invalidate(reason_code))
118    }
119}
120
121impl<L, R> BackingStorageSealed for Either<L, R>
122where
123    L: BackingStorageSealed,
124    R: BackingStorageSealed,
125{
126    type ReadTransaction<'l> = Either<L::ReadTransaction<'l>, R::ReadTransaction<'l>>;
127
128    fn next_free_task_id(&self) -> Result<TaskId> {
129        either::for_both!(self, this => this.next_free_task_id())
130    }
131
132    fn uncompleted_operations(&self) -> Result<Vec<AnyOperation>> {
133        either::for_both!(self, this => this.uncompleted_operations())
134    }
135
136    fn save_snapshot<I>(
137        &self,
138        operations: Vec<Arc<AnyOperation>>,
139        task_cache_updates: Vec<ChunkedVec<(Arc<CachedTaskType>, TaskId)>>,
140        snapshots: Vec<I>,
141    ) -> Result<()>
142    where
143        I: Iterator<Item = SnapshotItem> + Send + Sync,
144    {
145        either::for_both!(self, this => this.save_snapshot(
146            operations,
147            task_cache_updates,
148            snapshots,
149        ))
150    }
151
152    fn start_read_transaction(&self) -> Option<Self::ReadTransaction<'_>> {
153        Some(match self {
154            Either::Left(this) => Either::Left(this.start_read_transaction()?),
155            Either::Right(this) => Either::Right(this.start_read_transaction()?),
156        })
157    }
158
159    unsafe fn lookup_task_candidates(
160        &self,
161        tx: Option<&Self::ReadTransaction<'_>>,
162        key: &CachedTaskType,
163    ) -> Result<SmallVec<[TaskId; 1]>> {
164        match self {
165            Either::Left(this) => {
166                let tx = tx.map(|tx| read_transaction_left_or_panic(tx.as_ref()));
167                // Safety: `tx` is unwrapped from `Either::Left`, so it originated from `this`.
168                unsafe { this.lookup_task_candidates(tx, key) }
169            }
170            Either::Right(this) => {
171                let tx = tx.map(|tx| read_transaction_right_or_panic(tx.as_ref()));
172                // Safety: `tx` is unwrapped from `Either::Right`, so it originated from `this`.
173                unsafe { this.lookup_task_candidates(tx, key) }
174            }
175        }
176    }
177
178    unsafe fn lookup_data(
179        &self,
180        tx: Option<&Self::ReadTransaction<'_>>,
181        task_id: TaskId,
182        category: SpecificTaskDataCategory,
183        storage: &mut TaskStorage,
184    ) -> Result<()> {
185        match self {
186            Either::Left(this) => {
187                let tx = tx.map(|tx| read_transaction_left_or_panic(tx.as_ref()));
188                // Safety: `tx` is unwrapped from `Either::Left`, so it originated from `this`.
189                unsafe { this.lookup_data(tx, task_id, category, storage) }
190            }
191            Either::Right(this) => {
192                let tx = tx.map(|tx| read_transaction_right_or_panic(tx.as_ref()));
193                // Safety: `tx` is unwrapped from `Either::Right`, so it originated from `this`.
194                unsafe { this.lookup_data(tx, task_id, category, storage) }
195            }
196        }
197    }
198
199    unsafe fn batch_lookup_data(
200        &self,
201        tx: Option<&Self::ReadTransaction<'_>>,
202        task_ids: &[TaskId],
203        category: SpecificTaskDataCategory,
204    ) -> Result<Vec<TaskStorage>> {
205        match self {
206            Either::Left(this) => {
207                let tx = tx.map(|tx| read_transaction_left_or_panic(tx.as_ref()));
208                // Safety: `tx` is unwrapped from `Either::Left`, so it originated from `this`.
209                unsafe { this.batch_lookup_data(tx, task_ids, category) }
210            }
211            Either::Right(this) => {
212                let tx = tx.map(|tx| read_transaction_right_or_panic(tx.as_ref()));
213                // Safety: `tx` is unwrapped from `Either::Right`, so it originated from `this`.
214                unsafe { this.batch_lookup_data(tx, task_ids, category) }
215            }
216        }
217    }
218
219    fn shutdown(&self) -> Result<()> {
220        either::for_both!(self, this => this.shutdown())
221    }
222}
223
224// similar to `Either::unwrap_left`, but does not require `R: Debug`.
225fn read_transaction_left_or_panic<L, R>(either: Either<L, R>) -> L {
226    match either {
227        Either::Left(l) => l,
228        Either::Right(_) => panic!(
229            "expected ReadTransaction of Either::Left containing {}, received Either::Right type \
230             of {}",
231            type_name::<L>(),
232            type_name::<R>(),
233        ),
234    }
235}
236
237// similar to `Either::unwrap_right`, but does not require `R: Debug`.
238fn read_transaction_right_or_panic<L, R>(either: Either<L, R>) -> R {
239    match either {
240        Either::Left(_) => panic!(
241            "expected ReadTransaction of Either::Right containing {}, received Either::Left type \
242             of {}",
243            type_name::<R>(),
244            type_name::<L>(),
245        ),
246        Either::Right(r) => r,
247    }
248}