iob/
blob_id_allocator.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use core::mem::{MaybeUninit, size_of};
6use core::slice;
7use core::sync::atomic::{AtomicU64, Ordering};
8use zx_status::Status;
9
10/// Possible error conditions when attempting to allocate a blob ID.
11#[derive(Copy, Clone, Debug, PartialEq, Eq)]
12pub enum AllocateError {
13    /// The header was invalid (purportedly either with indices extending past the blob head
14    /// offset or the blob head offset extending past the length of the buffer). Suggests invalid
15    /// initialization or corruption.
16    InvalidHeader,
17    /// The would-be bookkeeping index slot for the new allocated ID is non-empty (i.e., not in
18    /// the initialized state). Suggests corruption.
19    NonEmptyIndex,
20    /// Out Of Memory: there is insufficient memory available for the requested allocation.
21    /// (What is available can be queried with `remaining_bytes()`).
22    OutOfMemory,
23}
24
25impl From<AllocateError> for Status {
26    fn from(err: AllocateError) -> Self {
27        match err {
28            AllocateError::OutOfMemory => Status::NO_MEMORY,
29            AllocateError::InvalidHeader | AllocateError::NonEmptyIndex => {
30                Status::IO_DATA_INTEGRITY
31            }
32        }
33    }
34}
35
36/// Error type for `allocate_with` combining allocator errors and copy errors.
37#[derive(Copy, Clone, Debug, PartialEq, Eq)]
38pub enum AllocateErrorWith<E> {
39    /// Allocation error within the ID allocator structure.
40    Error(AllocateError),
41    /// Error returned by the caller-provided copy closure.
42    Copy(E),
43}
44
45impl<E> From<AllocateErrorWith<E>> for Status
46where
47    Status: From<E>,
48{
49    fn from(err: AllocateErrorWith<E>) -> Self {
50        match err {
51            AllocateErrorWith::Error(e) => e.into(),
52            AllocateErrorWith::Copy(e) => Status::from(e),
53        }
54    }
55}
56
57/// Possible error conditions when retrieving a blob by ID, either via iteration or `get_blob()`.
58#[derive(Copy, Clone, Debug, PartialEq, Eq)]
59pub enum BlobError {
60    /// The header was invalid (purportedly either with indices extending past the blob head
61    /// offset or the blob head offset extending past the length of the buffer). Suggests invalid
62    /// initialization or corruption.
63    InvalidHeader,
64    /// The requested ID has not yet been allocated.
65    UnallocatedId,
66    /// Suggests a lost race in which the requested ID is valid and the header has been updated
67    /// to reflect that, but the bookkeeping index has not yet been committed. Such a case
68    /// unfortunately is indistinguishable from the other possibility that the index was already
69    /// committed but then subsequently corrupted with zeroes. Where there is confidence in the
70    /// first case, this call should be retried.
71    UncommittedIndex,
72    /// The corresponding bookkeeping index is invalid, with the blob purportedly not being
73    /// contained within `[blob head, end of region)`. Suggests corruption.
74    InvalidIndex,
75}
76
77impl From<BlobError> for Status {
78    fn from(err: BlobError) -> Self {
79        match err {
80            BlobError::UnallocatedId => Status::NOT_FOUND,
81            BlobError::UncommittedIndex => Status::SHOULD_WAIT,
82            BlobError::InvalidHeader | BlobError::InvalidIndex => Status::IO_DATA_INTEGRITY,
83        }
84    }
85}
86
87/// Specifies whether memory should be zeroed when initializing a `BlobIdAllocator`.
88#[derive(Copy, Clone, Debug, PartialEq, Eq)]
89pub enum ZeroFill {
90    No,
91    Yes,
92}
93
94/// Header for the ID allocator region.
95#[repr(C)]
96#[derive(Copy, Clone, Debug, PartialEq, Eq)]
97pub struct Header {
98    pub next_id: u32,
99    pub blob_head: u32,
100}
101
102const _: () = assert!(size_of::<Header>() == size_of::<u64>());
103
104impl Header {
105    pub const SIZE: usize = size_of::<Self>();
106
107    pub fn from_u64(val: u64) -> Self {
108        Self { next_id: val as u32, blob_head: (val >> 32) as u32 }
109    }
110
111    pub fn to_u64(self) -> u64 {
112        (self.next_id as u64) | ((self.blob_head as u64) << 32)
113    }
114
115    /// Returns the byte offset just past the last bookkeeping index, or `None` if the computation
116    /// would overflow a `u32`.
117    pub fn index_end(self) -> Option<u32> {
118        self.next_id.checked_mul(Index::SIZE as u32)?.checked_add(Header::SIZE as u32)
119    }
120
121    /// Returns `true` if the header reflects a valid state for a buffer of `length` bytes.
122    pub fn is_valid(self, length: usize) -> bool {
123        self.remaining_bytes(length).is_some()
124    }
125
126    /// Returns the remaining number of available bytes in the region, given the length of the
127    /// region, or `None` in the event of an invalid header.
128    pub fn remaining_bytes(self, length: usize) -> Option<usize> {
129        let end = self.index_end()?;
130        if end <= self.blob_head && (self.blob_head as usize) <= length {
131            Some((self.blob_head - end) as usize)
132        } else {
133            None
134        }
135    }
136}
137
138/// Represents a blob bookkeeping index slot.
139#[repr(C)]
140#[derive(Copy, Clone, Debug, PartialEq, Eq)]
141pub struct Index {
142    pub size: u32,
143    pub offset: u32,
144}
145
146const _: () = assert!(size_of::<Index>() == size_of::<u64>());
147
148impl Index {
149    pub const SIZE: usize = size_of::<Self>();
150
151    pub fn from_u64(val: u64) -> Self {
152        Self { size: val as u32, offset: (val >> 32) as u32 }
153    }
154
155    pub fn to_u64(self) -> u64 {
156        (self.size as u64) | ((self.offset as u64) << 32)
157    }
158
159    /// See `BlobError::InvalidIndex`. The current blob head offset and length of the region must
160    /// be provided, and are assumed to have already been validated.
161    pub fn is_valid(self, blob_head: u32, length: usize) -> bool {
162        debug_assert!((blob_head as usize) <= length);
163        (blob_head <= self.offset)
164            && (self.offset as usize <= length)
165            && (self.size as usize <= length - (self.offset as usize))
166    }
167}
168
169/// Helper struct implementing the lock-free blob-id allocator over memory.
170///
171/// Represents a thread-safe view into an IOBuffer region of "ID allocator" discipline
172/// (`ZX_IOB_DISCIPLINE_TYPE_ID_ALLOCATOR`), used to map sized data blobs to sequentially-allocated
173/// numeric IDs.
174///
175/// Suppose there are N mapped blobs. The memory is laid out as follows, with copies of the blobs
176/// growing down and their corresponding bookkeeping indices growing up:
177/// ```text
178/// --------------------------------
179///   next available blob ID (4 bytes)
180///   blob head offset (4 bytes)
181///   ----------------------------
182///   blob 0 size (4 bytes)       } <-- bookkeeping index
183///   blob 0 offset (4 bytes)     }
184///   ...
185///   blob N-1 size (4 bytes)
186///   blob N-1 offset (4 bytes)
187///   ----------------------------
188///   zero-initialized memory      <-- remaining bytes available
189///   ---------------------------- <-- blob head offset
190///   blob N-1
191///   ...
192///   blob 0
193/// --------------------------------
194/// ```
195///
196/// This struct takes care of the atomic nuance required of accessing and updating such a
197/// structure.
198#[derive(Clone, Copy, Debug)]
199pub struct BlobIdAllocator<'a> {
200    bytes: &'a [u8],
201}
202
203impl<'a> BlobIdAllocator<'a> {
204    /// Constructs a view from a byte slice.
205    ///
206    /// The provided slice must be at least 8-byte-aligned and at least 8 bytes in size.
207    pub fn from_slice(slice: &'a [u8]) -> Self {
208        debug_assert!(slice.as_ptr().cast::<Header>().is_aligned());
209        debug_assert!((Header::SIZE..=u32::MAX as usize).contains(&slice.len()));
210        Self { bytes: slice }
211    }
212
213    /// Initializes the backing memory as an ID allocator region with no blobs yet mapped,
214    /// returning an initialized `BlobIdAllocator`.
215    ///
216    /// If the region is already known to be zero-filled, `zero_fill` may be `ZeroFill::No`.
217    ///
218    /// The provided slice must be at least 8-byte-aligned and at least 8 bytes in size.
219    pub fn init_from_slice(slice: &'a mut [u8], zero_fill: ZeroFill) -> Self {
220        debug_assert!((Header::SIZE..=u32::MAX as usize).contains(&slice.len()));
221        if zero_fill == ZeroFill::Yes {
222            slice[Header::SIZE..].fill(0);
223        }
224        let allocator = Self::from_slice(slice);
225        allocator.store_header(Header { next_id: 0, blob_head: allocator.bytes.len() as u32 });
226        allocator
227    }
228
229    /// The next ID to be allocated.
230    pub fn next_id(&self) -> u32 {
231        self.load_header().next_id
232    }
233
234    /// The remaining number of available bytes in the allocator (including those that might be
235    /// used for bookkeeping). `None` is returned in the case of an invalid header (see
236    /// `AllocateError::InvalidHeader` for more detail).
237    pub fn remaining_bytes(&self) -> Option<usize> {
238        self.load_header().remaining_bytes(self.bytes.len())
239    }
240
241    /// Attempts to store the provided blob and allocate its ID.
242    pub fn allocate(&self, blob: &[u8]) -> Result<u32, AllocateError> {
243        self.allocate_with(blob.len(), |dest| {
244            for (d, s) in dest.iter_mut().zip(blob) {
245                d.write(*s);
246            }
247            Ok::<(), core::convert::Infallible>(())
248        })
249        .map_err(|e| match e {
250            AllocateErrorWith::Error(err) => err,
251            AllocateErrorWith::Copy(infallible) => match infallible {},
252        })
253    }
254
255    /// A variation of the allocation routine that abstracts the representation of the supplied
256    /// blob and the manner in which it is copied. This is of particular value to the use of this
257    /// library in kernel, which requires care in dealing with user-supplied memory.
258    ///
259    /// `copy`, which performs the copy of blob to a specified destination, is a callable of input
260    /// signature `(dest: &mut [MaybeUninit<u8>]) -> Result<(), E>`.
261    pub fn allocate_with<E>(
262        &self,
263        blob_size: usize,
264        copy: impl FnOnce(&mut [MaybeUninit<u8>]) -> Result<(), E>,
265    ) -> Result<u32, AllocateErrorWith<E>> {
266        let header_atomic = self.header_atomic();
267        let mut raw_hdr = header_atomic.load(Ordering::Acquire);
268        let (id, offset) = loop {
269            let hdr = Header::from_u64(raw_hdr);
270            let remaining = hdr
271                .remaining_bytes(self.bytes.len())
272                .ok_or(AllocateErrorWith::Error(AllocateError::InvalidHeader))?;
273            if remaining < Index::SIZE || remaining - Index::SIZE < blob_size {
274                return Err(AllocateErrorWith::Error(AllocateError::OutOfMemory));
275            }
276            let offset = hdr.blob_head - (blob_size as u32);
277            let updated = Header { next_id: hdr.next_id + 1, blob_head: offset };
278            match header_atomic.compare_exchange_weak(
279                raw_hdr,
280                updated.to_u64(),
281                Ordering::AcqRel,
282                Ordering::Acquire,
283            ) {
284                Ok(_) => break (hdr.next_id, offset),
285                Err(actual) => raw_hdr = actual,
286            }
287        };
288
289        // We store the blob and then store the index with release semantics to ensure the
290        // following:
291        // (1) The write of the index stays ordered after the previous store of the blob.
292        // (2) The write of the index stays ordered before subsequent reads with acquire semantics.
293        //
294        // SAFETY: The space check and atomic CAS ensure `[offset, offset + blob_size)` is
295        // in-bounds, disjoint from the index table, and claimed exclusively.
296        let dest_slice = unsafe {
297            let dest_ptr = self.bytes.as_ptr().add(offset as usize).cast_mut().cast();
298            slice::from_raw_parts_mut(dest_ptr, blob_size)
299        };
300        copy(dest_slice).map_err(AllocateErrorWith::Copy)?;
301
302        // Before overwriting, to be safe, check that the index is in the initial state (empty).
303        let new_index = Index { size: blob_size as u32, offset };
304        // SAFETY: `remaining_bytes` ensured `offset >= index_end`, so `id`'s index slot is
305        // in-bounds and disjoint from `dest_slice`.
306        unsafe { self.index_atomic(id) }
307            .compare_exchange(0, new_index.to_u64(), Ordering::Release, Ordering::Relaxed)
308            .map_err(|_| AllocateErrorWith::Error(AllocateError::NonEmptyIndex))?;
309        Ok(id)
310    }
311
312    /// Returns the blob corresponding to a given ID.
313    pub fn get_blob(&self, id: u32) -> Result<&'a [u8], BlobError> {
314        let hdr = self.load_header();
315        if !hdr.is_valid(self.bytes.len()) {
316            return Err(BlobError::InvalidHeader);
317        }
318        if id >= hdr.next_id {
319            return Err(BlobError::UnallocatedId);
320        }
321        // We load the index with acquire semantics as this ensures the following:
322        // (1) The read of the index stays ordered before the subsequent load of the blob.
323        // (2) The read of the index stays ordered after previous updates, which were written with
324        // release semantics.
325        //
326        // SAFETY: `hdr.is_valid` and `id < hdr.next_id` ensure `id`'s index slot is in-bounds and
327        // disjoint from payload data.
328        let index_raw = unsafe { self.index_atomic(id) }.load(Ordering::Acquire);
329        if index_raw == 0 {
330            return Err(BlobError::UncommittedIndex);
331        }
332        let index = Index::from_u64(index_raw);
333        if !index.is_valid(hdr.blob_head, self.bytes.len()) {
334            return Err(BlobError::InvalidIndex);
335        }
336        let offset = index.offset as usize;
337        let size = index.size as usize;
338        Ok(&self.bytes[offset..offset + size])
339    }
340
341    /// Provides an iterator through all allocated blobs and IDs.
342    pub fn iter(&self) -> Iter<'a> {
343        Iter { allocator: *self, next_id: 0 }
344    }
345
346    fn header_atomic(&self) -> &AtomicU64 {
347        // SAFETY:
348        // (1) `self.bytes` is at least `Header::SIZE` (8 bytes) and 8-byte aligned per the
349        //     preconditions of `from_slice` and `init_from_slice`.
350        // (2) The header is exclusively accessed through atomic operations, so casting to a
351        //     shared `&AtomicU64` reference is sound.
352        unsafe { &*self.bytes.as_ptr().cast::<AtomicU64>() }
353    }
354
355    /// Returns a reference to the `AtomicU64` index slot for `id`.
356    ///
357    /// # Safety
358    ///
359    /// The caller must ensure `id` corresponds to an in-bounds index slot that does not overlap
360    /// with any active mutable borrows.
361    unsafe fn index_atomic(&self, id: u32) -> &AtomicU64 {
362        let index_offset = (id as usize + 1) * Index::SIZE;
363        // SAFETY:
364        // (1) The caller guarantees `index_offset` is in-bounds and disjoint from mutable borrows.
365        // (2) `self.bytes` is 8-byte aligned and `index_offset` is a multiple of `Index::SIZE`
366        //     (8 bytes), so the pointer is properly aligned for `AtomicU64`.
367        // (3) The index slot is exclusively accessed through atomic operations.
368        unsafe { &*self.bytes.as_ptr().add(index_offset).cast::<AtomicU64>() }
369    }
370
371    fn load_header(&self) -> Header {
372        Header::from_u64(self.header_atomic().load(Ordering::Relaxed))
373    }
374
375    fn store_header(&self, header: Header) {
376        self.header_atomic().store(header.to_u64(), Ordering::Release);
377    }
378}
379
380/// Iterator over blobs in a `BlobIdAllocator`.
381pub struct Iter<'a> {
382    allocator: BlobIdAllocator<'a>,
383    next_id: u32,
384}
385
386impl<'a> Iterator for Iter<'a> {
387    type Item = Result<(u32, &'a [u8]), BlobError>;
388
389    fn next(&mut self) -> Option<Self::Item> {
390        let max_id = self.allocator.next_id();
391        if self.next_id >= max_id {
392            return None;
393        }
394        let id = self.next_id;
395        self.next_id += 1;
396        Some(self.allocator.get_blob(id).map(|blob| (id, blob)))
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403    use core::ptr;
404    use std::vec::Vec;
405    use std::{thread, vec};
406
407    #[test]
408    fn test_header_and_index() {
409        let hdr = Header { next_id: 42, blob_head: 4096 };
410        let raw = hdr.to_u64();
411        let decoded = Header::from_u64(raw);
412        assert_eq!(hdr, decoded);
413
414        assert_eq!(Header { next_id: 0, blob_head: 4096 }.index_end(), Some(8));
415        assert_eq!(Header { next_id: 1, blob_head: 4096 }.index_end(), Some(16));
416        assert_eq!(Header { next_id: 5, blob_head: 4096 }.index_end(), Some(48));
417
418        const MAX_NEXT_ID: u32 = (u32::MAX - 8) / 8;
419        assert_eq!(
420            Header { next_id: MAX_NEXT_ID, blob_head: 4096 }.index_end(),
421            Some(8 + MAX_NEXT_ID * 8)
422        );
423        assert_eq!(Header { next_id: MAX_NEXT_ID + 1, blob_head: 4096 }.index_end(), None);
424
425        assert_eq!(Header { next_id: 0, blob_head: 1024 }.remaining_bytes(1024), Some(1024 - 8));
426        assert_eq!(Header { next_id: 1, blob_head: 1000 }.remaining_bytes(1024), Some(1000 - 16));
427        assert_eq!(Header { next_id: 0, blob_head: 8 }.remaining_bytes(8), Some(0));
428        assert_eq!(Header { next_id: 1, blob_head: 15 }.remaining_bytes(1024), None);
429        assert_eq!(Header { next_id: 0, blob_head: 2000 }.remaining_bytes(1024), None);
430
431        let idx = Index { size: 100, offset: 500 };
432        let raw_idx = idx.to_u64();
433        assert_eq!(raw_idx, 100 | (500u64 << 32));
434        assert_eq!(Index::from_u64(raw_idx), idx);
435        assert!(idx.is_valid(400, 1024));
436        assert!(!idx.is_valid(600, 1024));
437        assert!(!idx.is_valid(400, 550));
438    }
439
440    #[test]
441    fn test_single_threaded() {
442        let blob_a = [b'a'; 51];
443        let blob_b = [b'b'; 17];
444        let blob_c = [b'c'; 1];
445
446        let mut buffer = [0u64; 100 / 8 + 1];
447        let slice = unsafe { slice::from_raw_parts_mut(buffer.as_mut_ptr().cast::<u8>(), 100) };
448        let allocator = BlobIdAllocator::init_from_slice(slice, ZeroFill::Yes);
449
450        assert_eq!(allocator.remaining_bytes(), Some(92)); // 100 - 8 (header)
451
452        let id_a = allocator.allocate(&blob_a).expect("allocate A");
453        assert_eq!(id_a, 0);
454        assert_eq!(allocator.remaining_bytes(), Some(33)); // 92 - 8 (index) - 51 (blob)
455
456        let res_a = allocator.get_blob(0).expect("get A");
457        assert_eq!(res_a, &blob_a[..]);
458
459        let id_b = allocator.allocate(&blob_b).expect("allocate B");
460        assert_eq!(id_b, 1);
461        assert_eq!(allocator.remaining_bytes(), Some(8)); // 33 - 8 (index) - 17 (blob)
462
463        let res_b = allocator.get_blob(1).expect("get B");
464        assert_eq!(res_b, &blob_b[..]);
465
466        // Allocating blob C requires 8 (index) + 1 (data) = 9 bytes, but only 8 remain.
467        assert_eq!(allocator.allocate(&blob_c), Err(AllocateError::OutOfMemory));
468
469        let items: Vec<(u32, &[u8])> = allocator.iter().map(Result::unwrap).collect();
470        assert_eq!(items.len(), 2);
471        assert_eq!(items[0], (0, &blob_a[..]));
472        assert_eq!(items[1], (1, &blob_b[..]));
473    }
474
475    #[test]
476    fn test_multi_threaded() {
477        const NUM_THREADS: usize = 100;
478        const BUF_SIZE: usize = 8 + NUM_THREADS * 8 + NUM_THREADS * 1;
479        let mut buffer = vec![0u64; BUF_SIZE / 8 + 1];
480        let slice =
481            unsafe { slice::from_raw_parts_mut(buffer.as_mut_ptr().cast::<u8>(), BUF_SIZE) };
482        let allocator = BlobIdAllocator::init_from_slice(slice, ZeroFill::Yes);
483
484        let mut ids: Vec<u32> = thread::scope(|s| {
485            let mut handles = Vec::with_capacity(NUM_THREADS);
486            for i in 0..NUM_THREADS {
487                handles.push(s.spawn(move || {
488                    let byte = [i as u8];
489                    allocator.allocate(&byte).expect("allocate in thread")
490                }));
491            }
492            handles.into_iter().map(|h| h.join().unwrap()).collect()
493        });
494        ids.sort();
495        for (expected, actual) in ids.iter().enumerate() {
496            assert_eq!(expected as u32, *actual);
497        }
498
499        let mut blob_values: Vec<u8> = allocator
500            .iter()
501            .map(|res| {
502                let (_id, blob) = res.unwrap();
503                assert_eq!(blob.len(), 1);
504                blob[0]
505            })
506            .collect();
507        assert_eq!(blob_values.len(), NUM_THREADS);
508        blob_values.sort();
509        for (i, val) in blob_values.iter().enumerate() {
510            assert_eq!(i as u8, *val);
511        }
512    }
513
514    #[test]
515    fn test_corruption_detection() {
516        let mut buffer = [0u64; 16];
517        let slice = unsafe { slice::from_raw_parts_mut(buffer.as_mut_ptr().cast::<u8>(), 128) };
518        let allocator = BlobIdAllocator::init_from_slice(slice, ZeroFill::No);
519
520        // Corrupt header by setting blob_head < index_end
521        let corrupted_hdr = Header { next_id: 1, blob_head: 8 };
522        unsafe {
523            ptr::write_volatile(buffer.as_mut_ptr(), corrupted_hdr.to_u64());
524        }
525
526        assert_eq!(allocator.allocate(&[0u8; 4]), Err(AllocateError::InvalidHeader));
527        assert_eq!(allocator.get_blob(0), Err(BlobError::InvalidHeader));
528
529        // Re-initialize and corrupt index slot to cause CAS failure
530        let slice = unsafe { slice::from_raw_parts_mut(buffer.as_mut_ptr().cast::<u8>(), 128) };
531        let allocator = BlobIdAllocator::init_from_slice(slice, ZeroFill::No);
532        unsafe {
533            // Slot for index 0 is at offset 8 (index 1 in u64 array). Write non-zero value to it.
534            ptr::write_volatile(buffer.as_mut_ptr().add(1), 0x1234);
535        }
536        assert_eq!(allocator.allocate(&[0u8; 4]), Err(AllocateError::NonEmptyIndex));
537    }
538
539    #[test]
540    fn test_overflowing_next_id_is_invalid() {
541        let mut buffer = [0u64; 512];
542        let slice = unsafe { slice::from_raw_parts_mut(buffer.as_mut_ptr().cast::<u8>(), 4096) };
543        let allocator = BlobIdAllocator::init_from_slice(slice, ZeroFill::Yes);
544
545        // Corrupt next_id to the overflow-triggering value.
546        // next_id = 0x20000000 causes: (uint32_t)(8 + 0x20000000*8) = 8 (wraps!)
547        unsafe {
548            let raw = buffer.as_mut_ptr() as *mut u32;
549            ptr::write_volatile(raw, 0x2000_0000);
550        }
551
552        assert_eq!(allocator.remaining_bytes(), None);
553        assert_eq!(allocator.allocate(&[0u8; 8]), Err(AllocateError::InvalidHeader));
554        assert_eq!(allocator.get_blob(0), Err(BlobError::InvalidHeader));
555    }
556}
iob/blob_id_allocator.rs

iob/
blob_id_allocator.rs
