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    fn prepare(
251        &self,
252        keys: Vec<StableKey>,
253        growth: &[(AllocationId, u64)],
254    ) -> Result<Prepared, Error> {
255        let mut image = self.image.clone();
256        let keys: BTreeSet<_> = keys.into_iter().collect();
257        for key in keys {
258            if self.by_key.contains_key(&key) {
259                continue;
260            }
261            if image.entries.len() >= self.budget.allocations {
262                return Err(Error::Budget);
263            }
264            let id = AllocationId(image.next_id);
265            image.next_id = image.next_id.checked_add(1).ok_or(Error::Budget)?;
266            image.entries.push(Entry {
267                id,
268                key,
269                extents: vec![],
270            });
271        }
272        for &(id, pages) in growth {
273            // Growth applies only to bindings already present in the base.
274            let index = *self.by_id.get(&id).ok_or(Error::Missing)?;
275            if pages == 0 {
276                continue;
277            }
278            let start = image.next_page;
279            image.next_page = start.checked_add(pages).ok_or(Error::Budget)?;
280            if let Some(last) = image.entries[index].extents.last_mut()
281                && last.start + last.pages == start
282            {
283                last.pages += pages;
284            } else {
285                image.entries[index].extents.push(Extent { start, pages });
286            }
287        }
288        image.generation = image.generation.checked_add(1).ok_or(Error::Budget)?;
289        let candidate = Self::validate(image, self.budget, self.data_start)?;
290        Ok(Prepared {
291            base: self.image.clone(),
292            candidate,
293        })
294    }
295
296    // Model-only publication; durable publication belongs to Store.
297    #[cfg(test)]
298    fn publish(&mut self, prepared: Prepared) -> Result<(), Error> {
299        if self.image != prepared.base || self.data_start != prepared.candidate.data_start {
300            return Err(Error::Stale);
301        }
302        // Re-admit the candidate under this owner's budgets, even when two
303        // owners began from the same image with different local limits.
304        let candidate = Self::validate(prepared.candidate.image, self.budget, self.data_start)?;
305        candidate.encode()?;
306        *self = candidate;
307        Ok(())
308    }
309
310    fn page_span(&self, id: AllocationId, logical: u64) -> Option<(u64, u64)> {
311        let index = *self.by_id.get(&id)?;
312        let ends = &self.logical_ends[index];
313        let slot = ends.partition_point(|&end| end <= logical);
314        let extent = self.image.entries[index].extents.get(slot)?;
315        let start = if slot == 0 { 0 } else { ends[slot - 1] };
316        Some((extent.start + logical - start, ends[slot] - logical))
317    }
318
319    #[cfg(test)]
320    fn page(&self, id: AllocationId, logical: u64) -> Option<u64> {
321        self.page_span(id, logical).map(|(page, _)| page)
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328
329    fn key(i: usize) -> StableKey {
330        StableKey::parse(format!("test.store{i}.v1")).unwrap()
331    }
332
333    #[test]
334    fn identities_survive_reorder_omission_and_reopen_beyond_256() {
335        let budget = StorageConfig::default();
336        let mut d = Directory::empty(budget, 1).unwrap();
337        let p = d.prepare((0..1_024).rev().map(key).collect(), &[]).unwrap();
338        d.publish(p).unwrap();
339        let original = d.by_key.clone();
340        let p = d.prepare(vec![key(0), key(2_000)], &[]).unwrap();
341        let recovered = Directory::decode(&p.candidate.encode().unwrap(), budget, 1).unwrap();
342        for (key, id) in original {
343            assert_eq!(recovered.by_key[&key], id);
344        }
345        assert_eq!(recovered.by_key[&key(2_000)], AllocationId(1_025));
346        assert_eq!(recovered.image.entries.len(), 1_025);
347    }
348
349    #[test]
350    fn interleaved_growth_maps_disjoint_ranges_and_rejects_stale_candidates() {
351        let mut d = Directory::empty(StorageConfig::default(), 2).unwrap();
352        d.publish(d.prepare(vec![key(0), key(1)], &[]).unwrap())
353            .unwrap();
354        let a = d.by_key[&key(0)];
355        let b = d.by_key[&key(1)];
356        let first = d.prepare(vec![], &[(a, 2), (b, 3), (a, 1)]).unwrap();
357        let stale = d.prepare(vec![key(2)], &[]).unwrap();
358        d.publish(first).unwrap();
359        assert_eq!(
360            (d.page(a, 0), d.page(a, 1), d.page(a, 2), d.page(a, 3)),
361            (Some(2), Some(3), Some(7), None)
362        );
363        assert_eq!((d.page(b, 0), d.page(b, 2)), (Some(4), Some(6)));
364        assert!(matches!(d.publish(stale), Err(Error::Stale)));
365    }
366
367    #[test]
368    fn rejected_budgets_and_overflows_preserve_the_base() {
369        let budget = StorageConfig {
370            allocations: 1,
371            extents: 1,
372            pages: 3,
373            ..StorageConfig::default()
374        };
375        let mut d = Directory::empty(budget, 1).unwrap();
376        d.publish(d.prepare(vec![key(0)], &[]).unwrap()).unwrap();
377        let original = d.image.clone();
378        let id = d.by_key[&key(0)];
379        assert!(d.prepare(vec![key(1)], &[]).is_err());
380        assert!(d.prepare(vec![], &[(id, 3)]).is_err());
381        assert!(d.prepare(vec![], &[(id, u64::MAX)]).is_err());
382        assert!(d.prepare(vec![], &[(id, 1), (id, 2)]).is_err());
383        assert_eq!(d.image, original);
384    }
385
386    #[test]
387    fn publication_enforces_receiving_budget_and_identity_exhaustion_is_fallible() {
388        let mut strict = Directory::empty(
389            StorageConfig {
390                allocations: 1,
391                ..StorageConfig::default()
392            },
393            1,
394        )
395        .unwrap();
396        let relaxed = Directory::empty(StorageConfig::default(), 1).unwrap();
397        let candidate = relaxed.prepare(vec![key(0), key(1)], &[]).unwrap();
398        let before = strict.image.clone();
399        assert!(matches!(strict.publish(candidate), Err(Error::Budget)));
400        assert_eq!(strict.image, before);
401
402        let mut image = relaxed.image;
403        image.next_id = u64::MAX - 1;
404        let mut d = Directory::validate(image, StorageConfig::default(), 1).unwrap();
405        d.publish(d.prepare(vec![key(0)], &[]).unwrap()).unwrap();
406        assert_eq!(d.by_key[&key(0)], AllocationId(u64::MAX - 1));
407        let before = d.image.clone();
408        assert!(matches!(d.prepare(vec![key(1)], &[]), Err(Error::Budget)));
409        assert_eq!(d.image, before);
410    }
411
412    #[test]
413    fn recovery_rejects_overlap_duplicate_keys_bounds_and_bad_wire() {
414        let budget = StorageConfig::default();
415        let mut d = Directory::empty(budget, 1).unwrap();
416        d.publish(d.prepare(vec![key(0), key(1)], &[]).unwrap())
417            .unwrap();
418        let ids: Vec<_> = d.by_id.keys().copied().collect();
419        d.publish(d.prepare(vec![], &[(ids[0], 2), (ids[1], 2)]).unwrap())
420            .unwrap();
421        for mutation in 0..5 {
422            let mut image = d.image.clone();
423            match mutation {
424                0 => image.entries[1].extents[0].start = 2,
425                1 => image.entries[1].key = image.entries[0].key.clone(),
426                2 => image.entries[0].extents[0].start = 0,
427                3 => image.entries[0].extents[0].pages = u64::MAX,
428                _ => image.entries[1].id = image.entries[0].id,
429            }
430            let payload = ic_wire::native::encode(&image, budget.wire).unwrap();
431            let bytes = [FRAME, &payload].concat();
432            assert!(Directory::decode(&bytes, budget, 1).is_err());
433        }
434        let bytes = d.encode().unwrap();
435        for end in 0..bytes.len() {
436            assert!(Directory::decode(&bytes[..end], budget, 1).is_err());
437        }
438        let mut trailing = bytes.clone();
439        trailing.push(0);
440        assert!(Directory::decode(&trailing, budget, 1).is_err());
441        assert!(
442            Directory::decode(
443                &bytes,
444                StorageConfig {
445                    wire: ic_wire::Limits {
446                        bytes: 8,
447                        ..budget.wire
448                    },
449                    ..budget
450                },
451                1
452            )
453            .is_err()
454        );
455    }
456}