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