1pub(super) mod leases;
3mod persistence;
4
5use crate::{StableKey, hash};
6use serde::{Deserialize, Serialize};
7use std::collections::{BTreeMap, BTreeSet};
8
9const FRAME: &[u8] = b"ICDIRR01";
10
11#[derive(Clone, Copy, Debug)]
19pub struct StorageConfig {
20 pub allocations: usize,
22 pub extents: usize,
24 pub pages: u64,
26 pub wire: ic_wire::Limits,
28}
29
30impl Default for StorageConfig {
31 fn default() -> Self {
32 Self {
33 allocations: 4_096,
34 extents: 16_384,
35 pages: 1_048_576,
36 wire: ic_wire::Limits {
37 bytes: 4 * 1024 * 1024,
38 allocation: 16 * 1024 * 1024,
39 nodes: 262_144,
40 depth: 16,
41 },
42 }
43 }
44}
45
46#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
53pub struct AllocationId(u64);
54impl AllocationId {
55 #[must_use]
57 pub const fn get(self) -> u64 {
58 self.0
59 }
60}
61
62#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
63struct Extent {
64 start: u64,
65 pub pages: u64,
66}
67
68#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
69struct Entry {
70 id: AllocationId,
71 key: StableKey,
72 extents: Vec<Extent>,
74}
75
76#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
77struct Image {
78 generation: u64,
79 next_id: u64,
80 next_page: u64,
82 metadata: Vec<Extent>,
84 entries: Vec<Entry>,
85}
86
87#[derive(Debug, thiserror::Error)]
88pub enum Error {
89 #[error("allocation page count overflow")]
91 ArithmeticOverflow,
92 #[error("directory budget exceeded")]
93 Budget,
94 #[error("invalid directory: {0}")]
95 Invalid(&'static str),
96 #[error("allocation already has a live lease")]
97 AlreadyOpen,
98 #[error("access exceeds allocation bounds")]
99 Bounds,
100 #[error("backing already has a live owner")]
101 AlreadyClaimed,
102 #[error("allocation does not exist")]
103 Missing,
104 #[error("prepared directory belongs to a stale generation")]
105 Stale,
106 #[error("backing growth refused")]
107 Growth,
108 #[error("directory owner is poisoned after interrupted backing IO")]
109 Poisoned,
110 #[error("Wire directory record: {0}")]
111 Wire(#[from] ic_wire::Error),
112}
113
114#[derive(Clone, Debug)]
115struct Directory {
116 image: Image,
117 by_key: BTreeMap<StableKey, AllocationId>,
118 by_id: BTreeMap<AllocationId, usize>,
119 logical_ends: Vec<Vec<u64>>,
121 budget: StorageConfig,
122 data_start: u64,
123}
124
125#[derive(Debug)]
127struct Prepared {
128 base: Image,
129 candidate: Directory,
130}
131
132impl Directory {
133 fn empty(budget: StorageConfig, data_start: u64) -> Result<Self, Error> {
134 Self::validate(
135 Image {
136 generation: 0,
137 next_id: 1,
138 next_page: data_start,
139 metadata: vec![],
140 entries: vec![],
141 },
142 budget,
143 data_start,
144 )
145 }
146
147 fn validate(image: Image, budget: StorageConfig, data_start: u64) -> Result<Self, Error> {
148 if image.entries.len() > budget.allocations || image.next_page > budget.pages {
149 return Err(Error::Budget);
150 }
151 if data_start == 0 || image.next_page < data_start || image.next_id == 0 {
152 return Err(Error::Invalid("root boundary or identity watermark"));
153 }
154 let mut by_key = BTreeMap::new();
155 let mut by_id = BTreeMap::new();
156 if image.metadata.len() > 128 {
157 return Err(Error::Budget);
158 }
159 let mut ranges = Vec::new();
160 for extent in &image.metadata {
161 let end = extent
162 .start
163 .checked_add(extent.pages)
164 .ok_or(Error::Budget)?;
165 if extent.pages == 0 || extent.start < data_start || end > image.next_page {
166 return Err(Error::Invalid("metadata arena outside backing"));
167 }
168 ranges.push((extent.start, end));
169 }
170 let mut data_extents = 0;
171 let mut previous = None;
172 for (index, entry) in image.entries.iter().enumerate() {
173 if entry.id.0 == 0
174 || entry.id.0 >= image.next_id
175 || previous.is_some_and(|id| id >= entry.id)
176 {
177 return Err(Error::Invalid("noncanonical allocation identities"));
178 }
179 previous = Some(entry.id);
180 if by_key.insert(entry.key.clone(), entry.id).is_some() {
181 return Err(Error::Invalid("duplicate stable key"));
182 }
183 by_id.insert(entry.id, index);
184 if entry.extents.len() > budget.extents.saturating_sub(data_extents) {
185 return Err(Error::Budget);
186 }
187 data_extents += entry.extents.len();
188 for extent in &entry.extents {
189 let end = extent
190 .start
191 .checked_add(extent.pages)
192 .ok_or(Error::Invalid("extent overflow"))?;
193 if extent.pages == 0 || extent.start < data_start || end > image.next_page {
194 return Err(Error::Invalid("extent outside data region"));
195 }
196 ranges.push((extent.start, end));
197 }
198 }
199 ranges.sort_unstable();
200 if ranges.windows(2).any(|pair| pair[0].1 > pair[1].0) {
201 return Err(Error::Invalid("overlapping physical extents"));
202 }
203 let logical_ends = image
204 .entries
205 .iter()
206 .map(|entry| {
207 let mut total = 0;
208 entry
209 .extents
210 .iter()
211 .map(|extent| {
212 total += extent.pages;
214 total
215 })
216 .collect()
217 })
218 .collect();
219 Ok(Self {
220 image,
221 by_key,
222 by_id,
223 logical_ends,
224 budget,
225 data_start,
226 })
227 }
228
229 fn decode(bytes: &[u8], budget: StorageConfig, data_start: u64) -> Result<Self, Error> {
230 if bytes.len() > budget.wire.bytes {
231 return Err(Error::Budget);
232 }
233 let payload = bytes
234 .strip_prefix(FRAME)
235 .ok_or(Error::Invalid("directory frame"))?;
236 let image = ic_wire::native::decode(payload, budget.wire)?;
237 Self::validate(image, budget, data_start)
238 }
239
240 fn encode(&self) -> Result<Vec<u8>, Error> {
241 let mut limits = self.budget.wire;
242 limits.bytes = limits.bytes.checked_sub(FRAME.len()).ok_or(Error::Budget)?;
243 let payload = ic_wire::native::encode(&self.image, limits)?;
244 let mut bytes = Vec::with_capacity(FRAME.len() + payload.len());
245 bytes.extend_from_slice(FRAME);
246 bytes.extend_from_slice(&payload);
247 Ok(bytes)
248 }
249
250 #[cfg(test)]
251 fn prepare(
252 &self,
253 keys: Vec<StableKey>,
254 growth: &[(AllocationId, u64)],
255 ) -> Result<Prepared, Error> {
256 self.prepare_at(keys, growth, self.image.next_page)
257 }
258
259 fn prepare_at(
262 &self,
263 keys: Vec<StableKey>,
264 growth: &[(AllocationId, u64)],
265 physical_end: u64,
266 ) -> Result<Prepared, Error> {
267 let mut image = self.image.clone();
268 image.next_page = image.next_page.max(physical_end);
269 let keys: BTreeSet<_> = keys.into_iter().collect();
270 for key in keys {
271 if self.by_key.contains_key(&key) {
272 continue;
273 }
274 if image.entries.len() >= self.budget.allocations {
275 return Err(Error::Budget);
276 }
277 let id = AllocationId(image.next_id);
278 image.next_id = image.next_id.checked_add(1).ok_or(Error::Budget)?;
279 image.entries.push(Entry {
280 id,
281 key,
282 extents: vec![],
283 });
284 }
285 for &(id, pages) in growth {
286 let index = *self.by_id.get(&id).ok_or(Error::Missing)?;
288 if pages == 0 {
289 continue;
290 }
291 let start = image.next_page;
292 image.next_page = start.checked_add(pages).ok_or(Error::Budget)?;
293 if let Some(last) = image.entries[index].extents.last_mut()
294 && last.start + last.pages == start
295 {
296 last.pages += pages;
297 } else {
298 image.entries[index].extents.push(Extent { start, pages });
299 }
300 }
301 image.generation = image.generation.checked_add(1).ok_or(Error::Budget)?;
302 let candidate = Self::validate(image, self.budget, self.data_start)?;
303 Ok(Prepared {
304 base: self.image.clone(),
305 candidate,
306 })
307 }
308
309 #[cfg(test)]
311 fn publish(&mut self, prepared: Prepared) -> Result<(), Error> {
312 if self.image != prepared.base || self.data_start != prepared.candidate.data_start {
313 return Err(Error::Stale);
314 }
315 let candidate = Self::validate(prepared.candidate.image, self.budget, self.data_start)?;
318 candidate.encode()?;
319 *self = candidate;
320 Ok(())
321 }
322
323 #[cfg(test)]
324 fn page_span(&self, id: AllocationId, logical: u64) -> Option<(u64, u64)> {
325 let index = *self.by_id.get(&id)?;
326 let ends = &self.logical_ends[index];
327 let slot = ends.partition_point(|&end| end <= logical);
328 let extent = self.image.entries[index].extents.get(slot)?;
329 let start = if slot == 0 { 0 } else { ends[slot - 1] };
330 Some((extent.start + logical - start, ends[slot] - logical))
331 }
332
333 #[cfg(test)]
334 fn page(&self, id: AllocationId, logical: u64) -> Option<u64> {
335 self.page_span(id, logical).map(|(page, _)| page)
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342
343 fn key(i: usize) -> StableKey {
344 StableKey::parse(format!("test.store{i}.v1")).unwrap()
345 }
346
347 #[test]
348 fn identities_survive_reorder_omission_and_reopen_beyond_256() {
349 let budget = StorageConfig::default();
350 let mut d = Directory::empty(budget, 1).unwrap();
351 let p = d.prepare((0..1_024).rev().map(key).collect(), &[]).unwrap();
352 d.publish(p).unwrap();
353 let original = d.by_key.clone();
354 let p = d.prepare(vec![key(0), key(2_000)], &[]).unwrap();
355 let recovered = Directory::decode(&p.candidate.encode().unwrap(), budget, 1).unwrap();
356 for (key, id) in original {
357 assert_eq!(recovered.by_key[&key], id);
358 }
359 assert_eq!(recovered.by_key[&key(2_000)], AllocationId(1_025));
360 assert_eq!(recovered.image.entries.len(), 1_025);
361 }
362
363 #[test]
364 fn interleaved_growth_maps_disjoint_ranges_and_rejects_stale_candidates() {
365 let mut d = Directory::empty(StorageConfig::default(), 2).unwrap();
366 d.publish(d.prepare(vec![key(0), key(1)], &[]).unwrap())
367 .unwrap();
368 let a = d.by_key[&key(0)];
369 let b = d.by_key[&key(1)];
370 let first = d.prepare(vec![], &[(a, 2), (b, 3), (a, 1)]).unwrap();
371 let stale = d.prepare(vec![key(2)], &[]).unwrap();
372 d.publish(first).unwrap();
373 assert_eq!(
374 (d.page(a, 0), d.page(a, 1), d.page(a, 2), d.page(a, 3)),
375 (Some(2), Some(3), Some(7), None)
376 );
377 assert_eq!((d.page(b, 0), d.page(b, 2)), (Some(4), Some(6)));
378 assert!(matches!(d.publish(stale), Err(Error::Stale)));
379 }
380
381 #[test]
382 fn rejected_budgets_and_overflows_preserve_the_base() {
383 let budget = StorageConfig {
384 allocations: 1,
385 extents: 1,
386 pages: 3,
387 ..StorageConfig::default()
388 };
389 let mut d = Directory::empty(budget, 1).unwrap();
390 d.publish(d.prepare(vec![key(0)], &[]).unwrap()).unwrap();
391 let original = d.image.clone();
392 let id = d.by_key[&key(0)];
393 assert!(d.prepare(vec![key(1)], &[]).is_err());
394 assert!(d.prepare(vec![], &[(id, 3)]).is_err());
395 assert!(d.prepare(vec![], &[(id, u64::MAX)]).is_err());
396 assert!(d.prepare(vec![], &[(id, 1), (id, 2)]).is_err());
397 assert_eq!(d.image, original);
398 }
399
400 #[test]
401 fn publication_enforces_receiving_budget_and_identity_exhaustion_is_fallible() {
402 let mut strict = Directory::empty(
403 StorageConfig {
404 allocations: 1,
405 ..StorageConfig::default()
406 },
407 1,
408 )
409 .unwrap();
410 let relaxed = Directory::empty(StorageConfig::default(), 1).unwrap();
411 let candidate = relaxed.prepare(vec![key(0), key(1)], &[]).unwrap();
412 let before = strict.image.clone();
413 assert!(matches!(strict.publish(candidate), Err(Error::Budget)));
414 assert_eq!(strict.image, before);
415
416 let mut image = relaxed.image;
417 image.next_id = u64::MAX - 1;
418 let mut d = Directory::validate(image, StorageConfig::default(), 1).unwrap();
419 d.publish(d.prepare(vec![key(0)], &[]).unwrap()).unwrap();
420 assert_eq!(d.by_key[&key(0)], AllocationId(u64::MAX - 1));
421 let before = d.image.clone();
422 assert!(matches!(d.prepare(vec![key(1)], &[]), Err(Error::Budget)));
423 assert_eq!(d.image, before);
424 }
425
426 #[test]
427 fn recovery_rejects_overlap_duplicate_keys_bounds_and_bad_wire() {
428 let budget = StorageConfig::default();
429 let mut d = Directory::empty(budget, 1).unwrap();
430 d.publish(d.prepare(vec![key(0), key(1)], &[]).unwrap())
431 .unwrap();
432 let ids: Vec<_> = d.by_id.keys().copied().collect();
433 d.publish(d.prepare(vec![], &[(ids[0], 2), (ids[1], 2)]).unwrap())
434 .unwrap();
435 for mutation in 0..5 {
436 let mut image = d.image.clone();
437 match mutation {
438 0 => image.entries[1].extents[0].start = 2,
439 1 => image.entries[1].key = image.entries[0].key.clone(),
440 2 => image.entries[0].extents[0].start = 0,
441 3 => image.entries[0].extents[0].pages = u64::MAX,
442 _ => image.entries[1].id = image.entries[0].id,
443 }
444 let payload = ic_wire::native::encode(&image, budget.wire).unwrap();
445 let bytes = [FRAME, &payload].concat();
446 assert!(Directory::decode(&bytes, budget, 1).is_err());
447 }
448 let bytes = d.encode().unwrap();
449 for end in 0..bytes.len() {
450 assert!(Directory::decode(&bytes[..end], budget, 1).is_err());
451 }
452 let mut trailing = bytes.clone();
453 trailing.push(0);
454 assert!(Directory::decode(&trailing, budget, 1).is_err());
455 assert!(
456 Directory::decode(
457 &bytes,
458 StorageConfig {
459 wire: ic_wire::Limits {
460 bytes: 8,
461 ..budget.wire
462 },
463 ..budget
464 },
465 1
466 )
467 .is_err()
468 );
469 }
470}