1use super::{DecodeLimits, GenerationRoot, decode_generation_root, encode_generation_root};
4use crate::cancellation::CancellationToken;
5use crate::foundation::GenerationId;
6use crate::performance::{OperationFailure, WorkBudget, WorkCounters, WorkError};
7use crate::storage::{
8 OBJECT_DIGEST_ENVELOPE_BYTES, ObjectId, ObjectKind, ObjectStore, ObjectStoreError,
9 object_digest,
10};
11use bytes::Bytes;
12use std::mem::size_of;
13use thiserror::Error;
14
15#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub struct CheckpointRequest {
18 pub base: ObjectId,
20 pub file_table: ObjectId,
22 pub merge_parent: Option<ObjectId>,
24}
25
26#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct CheckpointReceipt {
29 pub root: ObjectId,
31 pub generation_id: GenerationId,
33 pub reused_base: bool,
35 pub work: WorkCounters,
37}
38
39pub type CheckpointFailure = OperationFailure<CheckpointError>;
41
42pub fn build_checkpoint<S: ObjectStore>(
56 store: &S,
57 request: CheckpointRequest,
58 limits: DecodeLimits,
59 budget: WorkBudget,
60) -> Result<CheckpointReceipt, CheckpointFailure> {
61 validate_request(request)?;
62 let (base, mut work) = read_root(store, request.base, limits, budget, WorkCounters::default())?;
63 let merge = match request.merge_parent {
64 Some(parent_id) => {
65 let (parent, next) = read_root(store, parent_id, limits, budget, work)?;
66 work = next;
67 validate_merge_parents(&base, &parent, work)?;
68 Some(parent_id)
69 }
70 None => None,
71 };
72 if merge.is_none() && request.file_table == base.file_table {
73 return Ok(reused_checkpoint(request.base, work));
74 }
75 let (object, encoded, prospective) = encode_checkpoint(request, &base, merge, work, budget)?;
76 let remaining = prospective
77 .remaining(budget)
78 .map_err(|error| OperationFailure::new(error.into(), work))?;
79 let receipt = store
80 .put(object, Bytes::from(encoded), remaining)
81 .map_err(|failure| failure.map_with_prior_work(prospective, Into::into))?;
82 work = prospective
83 .checked_add(receipt.work)
84 .map_err(|error| OperationFailure::new(error.into(), prospective))?;
85 work.verify(budget)
86 .map_err(|error| OperationFailure::new(error.into(), work))?;
87 Ok(CheckpointReceipt {
88 root: object,
89 generation_id: GenerationId::new(object.digest),
90 reused_base: false,
91 work,
92 })
93}
94
95pub async fn build_checkpoint_async<S: crate::async_storage::AsyncObjectStore>(
106 store: &S,
107 request: CheckpointRequest,
108 limits: DecodeLimits,
109 budget: WorkBudget,
110 cancellation: &CancellationToken,
111) -> Result<CheckpointReceipt, CheckpointFailure> {
112 validate_request(request)?;
113 let (base, mut work) = read_root_async(
114 store,
115 request.base,
116 limits,
117 budget,
118 WorkCounters::default(),
119 cancellation,
120 )
121 .await?;
122 let merge = match request.merge_parent {
123 Some(parent_id) => {
124 let (parent, next) =
125 read_root_async(store, parent_id, limits, budget, work, cancellation).await?;
126 work = next;
127 validate_merge_parents(&base, &parent, work)?;
128 Some(parent_id)
129 }
130 None => None,
131 };
132 if merge.is_none() && request.file_table == base.file_table {
133 return Ok(reused_checkpoint(request.base, work));
134 }
135 let (object, encoded, prospective) = encode_checkpoint(request, &base, merge, work, budget)?;
136 let remaining = prospective
137 .remaining(budget)
138 .map_err(|error| OperationFailure::new(error.into(), work))?;
139 let receipt = crate::async_storage::AsyncObjectStore::put(
140 store,
141 object,
142 Bytes::from(encoded),
143 remaining,
144 cancellation,
145 )
146 .await
147 .map_err(|failure| failure.map_with_prior_work(prospective, Into::into))?;
148 work = prospective
149 .checked_add(receipt.work)
150 .map_err(|error| OperationFailure::new(error.into(), prospective))?;
151 work.verify(budget)
152 .map_err(|error| OperationFailure::new(error.into(), work))?;
153 Ok(CheckpointReceipt {
154 root: object,
155 generation_id: GenerationId::new(object.digest),
156 reused_base: false,
157 work,
158 })
159}
160
161fn validate_request(request: CheckpointRequest) -> Result<(), CheckpointFailure> {
162 if request.base.kind != ObjectKind::GenerationRoot
163 || request
164 .merge_parent
165 .is_some_and(|parent| parent.kind != ObjectKind::GenerationRoot)
166 {
167 return Err(OperationFailure::before_work(
168 CheckpointError::WrongGenerationKind,
169 ));
170 }
171 if request.file_table.kind != ObjectKind::FileTablePage {
172 return Err(OperationFailure::before_work(
173 CheckpointError::WrongFileTableKind,
174 ));
175 }
176 if request.merge_parent == Some(request.base) {
177 return Err(OperationFailure::before_work(
178 CheckpointError::DuplicateParent,
179 ));
180 }
181 Ok(())
182}
183
184fn validate_merge_parents(
185 base: &GenerationRoot,
186 parent: &GenerationRoot,
187 work: WorkCounters,
188) -> Result<(), CheckpointFailure> {
189 if parent.volume_id != base.volume_id {
190 return Err(OperationFailure::new(
191 CheckpointError::CrossVolumeMerge,
192 work,
193 ));
194 }
195 if parent.root_file_id != base.root_file_id {
196 return Err(OperationFailure::new(
197 CheckpointError::RootIdentityMismatch,
198 work,
199 ));
200 }
201 Ok(())
202}
203
204fn reused_checkpoint(root: ObjectId, work: WorkCounters) -> CheckpointReceipt {
205 CheckpointReceipt {
206 root,
207 generation_id: GenerationId::new(root.digest),
208 reused_base: true,
209 work,
210 }
211}
212
213fn encode_checkpoint(
214 request: CheckpointRequest,
215 base: &GenerationRoot,
216 merge: Option<ObjectId>,
217 work: WorkCounters,
218 budget: WorkBudget,
219) -> Result<(ObjectId, Vec<u8>, WorkCounters), CheckpointFailure> {
220 let parent_count = 1_usize + usize::from(merge.is_some());
221 let mut parents = Vec::with_capacity(parent_count);
222 parents.push(GenerationId::new(request.base.digest));
223 if let Some(parent) = merge {
224 parents.push(GenerationId::new(parent.digest));
225 }
226 let root = GenerationRoot {
227 volume_id: base.volume_id,
228 root_file_id: base.root_file_id,
229 file_table: request.file_table,
230 parents,
231 required_features: base.required_features,
232 };
233 let encoded =
234 encode_generation_root(&root).map_err(|error| OperationFailure::new(error.into(), work))?;
235 let object = ObjectId {
236 kind: ObjectKind::GenerationRoot,
237 digest: object_digest(ObjectKind::GenerationRoot, &encoded),
238 };
239 let encoded_capacity = u64::try_from(encoded.capacity()).unwrap_or(u64::MAX);
240 let parent_capacity = u64::try_from(root.parents.capacity())
241 .unwrap_or(u64::MAX)
242 .checked_mul(u64::try_from(size_of::<GenerationId>()).unwrap_or(u64::MAX))
243 .ok_or_else(|| OperationFailure::new(WorkError::Overflow.into(), work))?;
244 let peak_allocation_bytes = encoded_capacity
245 .checked_add(parent_capacity)
246 .ok_or_else(|| OperationFailure::new(WorkError::Overflow.into(), work))?;
247 let encoded_length = u64::try_from(encoded.len()).unwrap_or(u64::MAX);
248 let prospective = work
249 .checked_add(WorkCounters {
250 bytes_encoded: encoded_length,
251 bytes_hashed: encoded_length.saturating_add(OBJECT_DIGEST_ENVELOPE_BYTES),
252 allocation_operations: 2,
253 peak_allocation_bytes,
254 ..WorkCounters::default()
255 })
256 .map_err(|error| OperationFailure::new(error.into(), work))?;
257 prospective
258 .verify(budget)
259 .map_err(|error| OperationFailure::new(error.into(), work))?;
260 Ok((object, encoded, prospective))
261}
262
263fn read_root<S: ObjectStore>(
264 store: &S,
265 object: ObjectId,
266 limits: DecodeLimits,
267 budget: WorkBudget,
268 work: WorkCounters,
269) -> Result<(GenerationRoot, WorkCounters), CheckpointFailure> {
270 let remaining = work
271 .remaining(budget)
272 .map_err(|error| OperationFailure::new(error.into(), work))?;
273 let receipt = store
274 .read(object, limits.maximum_object_bytes, remaining)
275 .map_err(|failure| failure.map_with_prior_work(work, Into::into))?;
276 let combined = work
277 .checked_add(receipt.work)
278 .map_err(|error| OperationFailure::new(error.into(), work))?;
279 let root = decode_generation_root(&receipt.value, limits)
280 .map_err(|error| OperationFailure::new(error.into(), combined))?;
281 Ok((root, combined))
282}
283
284async fn read_root_async<S: crate::async_storage::AsyncObjectStore>(
285 store: &S,
286 object: ObjectId,
287 limits: DecodeLimits,
288 budget: WorkBudget,
289 work: WorkCounters,
290 cancellation: &CancellationToken,
291) -> Result<(GenerationRoot, WorkCounters), CheckpointFailure> {
292 let remaining = work
293 .remaining(budget)
294 .map_err(|error| OperationFailure::new(error.into(), work))?;
295 let receipt = crate::async_storage::AsyncObjectStore::read(
296 store,
297 object,
298 limits.maximum_object_bytes,
299 remaining,
300 cancellation,
301 )
302 .await
303 .map_err(|failure| failure.map_with_prior_work(work, Into::into))?;
304 let combined = work
305 .checked_add(receipt.work)
306 .map_err(|error| OperationFailure::new(error.into(), work))?;
307 let root = decode_generation_root(&receipt.value.bytes, limits)
308 .map_err(|error| OperationFailure::new(error.into(), combined))?;
309 Ok((root, combined))
310}
311
312#[derive(Debug, Error)]
314pub enum CheckpointError {
315 #[error("checkpoint parent has the wrong object kind")]
317 WrongGenerationKind,
318 #[error("checkpoint file table has the wrong object kind")]
320 WrongFileTableKind,
321 #[error("checkpoint merge parents are duplicated")]
323 DuplicateParent,
324 #[error("checkpoint cannot merge generations from different volumes")]
326 CrossVolumeMerge,
327 #[error("checkpoint merge parents have different root identities")]
329 RootIdentityMismatch,
330 #[error(transparent)]
332 Storage(#[from] ObjectStoreError),
333 #[error(transparent)]
335 Decode(#[from] super::CanonicalDecodeError),
336 #[error(transparent)]
338 Work(#[from] WorkError),
339}
340
341#[cfg(all(test, feature = "memory"))]
342#[path = "tests/checkpoint.rs"]
343mod tests;