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
26pub trait BackingStorage: BackingStorageSealed {
36 fn invalidate(&self, reason_code: &str) -> Result<()>;
48}
49
50pub 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 unsafe fn lookup_task_candidates(
79 &self,
80 tx: Option<&Self::ReadTransaction<'_>>,
81 key: &CachedTaskType,
82 ) -> Result<SmallVec<[TaskId; 1]>>;
83 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 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 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 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 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 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 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 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
224fn 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
237fn 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}