Skip to main content

ic_storage/runtime/
directory_model.rs

1//! Owned named-allocation directory and durable extent validation.
2pub(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///
12/// StorageConfig
13///
14/// Resource ceilings checked when creating, opening and extending the directory.
15/// Physical page accounting includes retained metadata and abandoned tails.
16///
17
18#[derive(Clone, Copy, Debug)]
19pub struct StorageConfig {
20    /// Maximum retained named allocations, including currently omitted names.
21    pub allocations: usize,
22    /// Maximum total data extents across every allocation.
23    pub extents: usize,
24    /// Maximum physical backing extent, in 64 KiB Wasm pages.
25    pub pages: u64,
26    /// Encoding and decoding budgets for each directory snapshot.
27    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///
47/// AllocationId
48///
49/// Opaque durable identity assigned by the directory. It grants no access.
50///
51
52#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
53pub struct AllocationId(u64);
54impl AllocationId {
55    /// Opaque diagnostic identity; callers cannot choose an allocation ID.
56    #[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    // Logical page order is explicit; physical ranges need not be contiguous.
73    extents: Vec<Extent>,
74}
75
76#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
77struct Image {
78    generation: u64,
79    next_id: u64,
80    // High-water mark, including space reserved for directory persistence.
81    next_page: u64,
82    // Reserved snapshot arenas, including the active arena. Never application data.
83    metadata: Vec<Extent>,
84    entries: Vec<Entry>,
85}
86
87#[derive(Debug, thiserror::Error)]
88pub enum Error {
89    /// Logical page count overflowed before any backing access.
90    #[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    // Derived only from validated extents; never persisted as a second authority.
120    logical_ends: Vec<Vec<u64>>,
121    budget: StorageConfig,
122    data_start: u64,
123}
124
125// A candidate owns all tentative changes. Preparation never mutates its base.
126#[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                        // Disjoint ranges fit within the admitted physical watermark.
213                        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    // Extend the validated base directly; no second indexed Directory is needed
260    // merely to account for a native write's uncommitted physical tail.
261    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            // Growth applies only to bindings already present in the base.
287            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    // Model-only publication; durable publication belongs to Store.
310    #[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        // Re-admit the candidate under this owner's budgets, even when two
316        // owners began from the same image with different local limits.
317        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}