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
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    /// Logical page count overflowed before any backing access.
102    #[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    // Derived only from validated extents; never persisted as a second authority.
132    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                        // Disjoint ranges fit within the admitted physical watermark.
217                        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    // Extend the validated base directly; no second indexed Directory is needed
244    // merely to account for a native write's uncommitted physical tail.
245    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            // Growth applies only to bindings already present in the base.
271            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    // Model-only application; production validates once after reserving metadata.
290    #[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}