|
| 1 | +/* |
| 2 | +Copyright 2026 The Hyperlight Authors. |
| 3 | +
|
| 4 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +you may not use this file except in compliance with the License. |
| 6 | +You may obtain a copy of the License at |
| 7 | +
|
| 8 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | +Unless required by applicable law or agreed to in writing, software |
| 11 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +See the License for the specific language governing permissions and |
| 14 | +limitations under the License. |
| 15 | +*/ |
| 16 | + |
| 17 | +//! Buffer allocation traits and shared types for virtqueue buffer management. |
| 18 | +
|
| 19 | +use alloc::rc::Rc; |
| 20 | +use alloc::sync::Arc; |
| 21 | + |
| 22 | +use thiserror::Error; |
| 23 | + |
| 24 | +use super::access::MemOps; |
| 25 | + |
| 26 | +#[derive(Debug, Error, Copy, Clone)] |
| 27 | +pub enum AllocError { |
| 28 | + #[error("Invalid region addr {0}")] |
| 29 | + InvalidAlign(u64), |
| 30 | + #[error("Invalid free addr {0} and size {1}")] |
| 31 | + InvalidFree(u64, usize), |
| 32 | + #[error("Invalid argument")] |
| 33 | + InvalidArg, |
| 34 | + #[error("Empty region")] |
| 35 | + EmptyRegion, |
| 36 | + #[error("Out of memory")] |
| 37 | + OutOfMemory, |
| 38 | + #[error("Overflow")] |
| 39 | + Overflow, |
| 40 | +} |
| 41 | + |
| 42 | +/// Allocation result |
| 43 | +#[derive(Debug, Clone, Copy)] |
| 44 | +pub struct Allocation { |
| 45 | + /// Starting address of the allocation |
| 46 | + pub addr: u64, |
| 47 | + /// Length of the allocation in bytes rounded up to slab size |
| 48 | + pub len: usize, |
| 49 | +} |
| 50 | + |
| 51 | +/// Trait for buffer providers. |
| 52 | +pub trait BufferProvider { |
| 53 | + /// Allocate at least `len` bytes. |
| 54 | + fn alloc(&self, len: usize) -> Result<Allocation, AllocError>; |
| 55 | + |
| 56 | + /// Free a previously allocated block. |
| 57 | + fn dealloc(&self, alloc: Allocation) -> Result<(), AllocError>; |
| 58 | + |
| 59 | + /// Resize by trying in-place grow; otherwise reserve a new block and free old. |
| 60 | + fn resize(&self, old_alloc: Allocation, new_len: usize) -> Result<Allocation, AllocError>; |
| 61 | + |
| 62 | + /// Reset the pool to initial state. |
| 63 | + fn reset(&self) {} |
| 64 | +} |
| 65 | + |
| 66 | +impl<T: BufferProvider> BufferProvider for Rc<T> { |
| 67 | + fn alloc(&self, len: usize) -> Result<Allocation, AllocError> { |
| 68 | + (**self).alloc(len) |
| 69 | + } |
| 70 | + fn dealloc(&self, alloc: Allocation) -> Result<(), AllocError> { |
| 71 | + (**self).dealloc(alloc) |
| 72 | + } |
| 73 | + fn resize(&self, old_alloc: Allocation, new_len: usize) -> Result<Allocation, AllocError> { |
| 74 | + (**self).resize(old_alloc, new_len) |
| 75 | + } |
| 76 | + fn reset(&self) { |
| 77 | + (**self).reset() |
| 78 | + } |
| 79 | +} |
| 80 | + |
| 81 | +impl<T: BufferProvider> BufferProvider for Arc<T> { |
| 82 | + fn alloc(&self, len: usize) -> Result<Allocation, AllocError> { |
| 83 | + (**self).alloc(len) |
| 84 | + } |
| 85 | + fn dealloc(&self, alloc: Allocation) -> Result<(), AllocError> { |
| 86 | + (**self).dealloc(alloc) |
| 87 | + } |
| 88 | + fn resize(&self, old_alloc: Allocation, new_len: usize) -> Result<Allocation, AllocError> { |
| 89 | + (**self).resize(old_alloc, new_len) |
| 90 | + } |
| 91 | + fn reset(&self) { |
| 92 | + (**self).reset() |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +/// The owner of a mapped buffer, ensuring its lifetime. |
| 97 | +/// |
| 98 | +/// Holds a pool allocation and provides direct access to the underlying |
| 99 | +/// shared memory via [`MemOps::as_slice`]. Implements `AsRef<[u8]>` so it |
| 100 | +/// can be used with [`Bytes::from_owner`](bytes::Bytes::from_owner) for |
| 101 | +/// zero-copy `Bytes` backed by shared memory. |
| 102 | +/// |
| 103 | +/// When dropped, the allocation is returned to the pool. |
| 104 | +#[derive(Debug, Clone)] |
| 105 | +pub struct BufferOwner<P: BufferProvider, M: MemOps> { |
| 106 | + pub(crate) pool: P, |
| 107 | + pub(crate) mem: M, |
| 108 | + pub(crate) alloc: Allocation, |
| 109 | + pub(crate) written: usize, |
| 110 | +} |
| 111 | + |
| 112 | +impl<P: BufferProvider, M: MemOps> Drop for BufferOwner<P, M> { |
| 113 | + fn drop(&mut self) { |
| 114 | + let _ = self.pool.dealloc(self.alloc); |
| 115 | + } |
| 116 | +} |
| 117 | + |
| 118 | +impl<P: BufferProvider, M: MemOps> AsRef<[u8]> for BufferOwner<P, M> { |
| 119 | + fn as_ref(&self) -> &[u8] { |
| 120 | + let len = self.written.min(self.alloc.len); |
| 121 | + // Safety: BufferOwner keeps both the pool allocation and the M |
| 122 | + // alive, so the memory region is valid. Protocol-level descriptor |
| 123 | + // ownership transfer guarantees no concurrent writes. |
| 124 | + match unsafe { self.mem.as_slice(self.alloc.addr, len) } { |
| 125 | + Ok(slice) => slice, |
| 126 | + Err(_) => &[], |
| 127 | + } |
| 128 | + } |
| 129 | +} |
| 130 | + |
| 131 | +/// A guard that runs a cleanup function when dropped, unless dismissed. |
| 132 | +pub struct AllocGuard<F: FnOnce(Allocation)>(Option<(Allocation, F)>); |
| 133 | + |
| 134 | +impl<F: FnOnce(Allocation)> AllocGuard<F> { |
| 135 | + pub fn new(alloc: Allocation, cleanup: F) -> Self { |
| 136 | + Self(Some((alloc, cleanup))) |
| 137 | + } |
| 138 | + |
| 139 | + pub fn release(mut self) -> Allocation { |
| 140 | + self.0.take().unwrap().0 |
| 141 | + } |
| 142 | +} |
| 143 | + |
| 144 | +impl<F: FnOnce(Allocation)> core::ops::Deref for AllocGuard<F> { |
| 145 | + type Target = Allocation; |
| 146 | + |
| 147 | + fn deref(&self) -> &Allocation { |
| 148 | + &self.0.as_ref().unwrap().0 |
| 149 | + } |
| 150 | +} |
| 151 | + |
| 152 | +impl<F: FnOnce(Allocation)> Drop for AllocGuard<F> { |
| 153 | + fn drop(&mut self) { |
| 154 | + if let Some((alloc, cleanup)) = self.0.take() { |
| 155 | + cleanup(alloc) |
| 156 | + } |
| 157 | + } |
| 158 | +} |
0 commit comments