feat(virtq): introduce pool alloc for managing allocation lifetime

Author: andreiltd

src/hyperlight_common/src/virtq/buffer.rs

@@ -97,52 +97,26 @@ impl<T: BufferProvider> BufferProvider for Arc<T> {
 
 /// The owner of a mapped buffer, ensuring its lifetime.
 ///
-/// Holds a pool allocation and provides direct access to the underlying
+/// Holds a [`PoolAlloc`] and provides direct access to the underlying
 /// shared memory via [`MemOps::as_slice`]. Implements `AsRef<[u8]>` so it
 /// can be used with [`Bytes::from_owner`](bytes::Bytes::from_owner) for
 /// zero-copy `Bytes` backed by shared memory.
 ///
 /// When dropped, the allocation is returned to the pool.
 #[derive(Debug)]
 pub struct BufferOwner<P: BufferProvider, M: MemOps> {
-    pub(crate) pool: P,
+    pub(crate) alloc: PoolAlloc<P>,
     pub(crate) mem: M,
-    pub(crate) alloc: Allocation,
     pub(crate) written: usize,
 }
 
-impl<P: BufferProvider, M: MemOps> Drop for BufferOwner<P, M> {
-    fn drop(&mut self) {
-        let _ = self.pool.dealloc(self.alloc);
-    }
-}
-
-impl<P: BufferProvider, M: MemOps> BufferOwner<P, M> {
-    pub(crate) fn try_new(
-        pool: P,
-        mem: M,
-        alloc: Allocation,
-        written: usize,
-    ) -> Result<Self, M::Error> {
-        // Pre check direct access before handing the owner to Bytes::from_owner
-        let len = written.min(alloc.len);
-        let _ = unsafe { mem.as_slice(alloc.addr, len) }?;
-
-        Ok(Self {
-            pool,
-            mem,
-            alloc,
-            written,
-        })
-    }
-}
-
 impl<P: BufferProvider, M: MemOps> AsRef<[u8]> for BufferOwner<P, M> {
     fn as_ref(&self) -> &[u8] {
-        let len = self.written.min(self.alloc.len);
+        let alloc = self.alloc.allocation();
+        let len = self.written.min(alloc.len);
         // Safety: BufferOwner keeps both the pool allocation and the M alive,
         // so the memory region is valid.
-        match unsafe { self.mem.as_slice(self.alloc.addr, len) } {
+        match unsafe { self.mem.as_slice(alloc.addr, len) } {
             Ok(slice) => slice,
             Err(_) => {
                 debug_assert!(false, "BufferOwner direct slice failed");
@@ -152,41 +126,74 @@ impl<P: BufferProvider, M: MemOps> AsRef<[u8]> for BufferOwner<P, M> {
     }
 }
 
-/// A guard that runs a cleanup function when dropped, unless dismissed.
-pub struct AllocGuard<F: FnOnce(Allocation)>(Option<(Allocation, F)>);
+/// Pool-owned allocation that is returned to the pool on drop.
+///
+/// Use [`into_raw`](Self::into_raw) to transfer ownership to a descriptor
+/// state that will deallocate the raw [`Allocation`] through another path.
+#[derive(Debug)]
+pub struct PoolAlloc<P: BufferProvider> {
+    inner: Option<PoolAllocInner<P>>,
+}
+
+#[derive(Debug)]
+struct PoolAllocInner<P: BufferProvider> {
+    pool: P,
+    alloc: Allocation,
+}
 
-impl<F: FnOnce(Allocation)> AllocGuard<F> {
-    pub fn new(alloc: Allocation, cleanup: F) -> Self {
-        Self(Some((alloc, cleanup)))
+impl<P: BufferProvider> PoolAlloc<P> {
+    /// Wrap an existing allocation with its owning pool.
+    pub fn new(pool: P, alloc: Allocation) -> Self {
+        Self {
+            inner: Some(PoolAllocInner { pool, alloc }),
+        }
+    }
+
+    /// Allocate from `pool` and return an owning guard.
+    pub fn allocate(pool: P, len: usize) -> Result<Self, AllocError> {
+        let alloc = pool.alloc(len)?;
+        Ok(Self::new(pool, alloc))
+    }
+
+    /// The raw allocation currently owned by this guard.
+    pub fn allocation(&self) -> Allocation {
+        self.inner
+            .as_ref()
+            .map(|inner| inner.alloc)
+            .unwrap_or_else(|| {
+                unreachable!("PoolAlloc::allocation called after ownership transfer")
+            })
     }
 
-    pub fn release(mut self) -> Allocation {
-        // Safety: AllocGuard is always constructed with Some, and release is only called once
-        self.0
+    /// Release ownership and return the raw allocation.
+    pub fn into_raw(mut self) -> Allocation {
+        self.inner
             .take()
-            .map(|(alloc, _)| alloc)
-            .unwrap_or_else(|| unreachable!("AllocGuard::release called on dismissed guard"))
+            .map(|inner| inner.alloc)
+            .unwrap_or_else(|| unreachable!("PoolAlloc::into_raw called after ownership transfer"))
     }
-}
 
-impl<F: FnOnce(Allocation)> core::ops::Deref for AllocGuard<F> {
-    type Target = Allocation;
+    pub(crate) fn into_buffer_owner<M: MemOps>(
+        self,
+        mem: M,
+        written: usize,
+    ) -> Result<BufferOwner<P, M>, M::Error> {
+        let alloc = self.allocation();
+        let len = written.min(alloc.len);
+        let _ = unsafe { mem.as_slice(alloc.addr, len) }?;
 
-    fn deref(&self) -> &Allocation {
-        // Safety: AllocGuard is always constructed with Some, and the inner value is only
-        // taken by release() or Drop.
-        &self
-            .0
-            .as_ref()
-            .unwrap_or_else(|| unreachable!("AllocGuard::deref called on dismissed guard"))
-            .0
+        Ok(BufferOwner {
+            alloc: self,
+            mem,
+            written,
+        })
     }
 }
 
-impl<F: FnOnce(Allocation)> Drop for AllocGuard<F> {
+impl<P: BufferProvider> Drop for PoolAlloc<P> {
     fn drop(&mut self) {
-        if let Some((alloc, cleanup)) = self.0.take() {
-            cleanup(alloc)
+        if let Some(PoolAllocInner { pool, alloc }) = self.inner.take() {
+            let _ = pool.dealloc(alloc);
         }
     }
 }

src/hyperlight_common/src/virtq/producer.rs

@@ -253,30 +253,21 @@ where
             self.pool.dealloc(entry)?;
         }
 
-        let completion_guard = inf.completion().map(|buf| {
-            let pool = self.pool.clone();
-            AllocGuard::new(buf, move |a| {
-                let _ = pool.dealloc(a);
-            })
-        });
+        let completion_guard = inf
+            .completion()
+            .map(|buf| PoolAlloc::new(self.pool.clone(), buf));
 
         // Read completion data
         let has_completion = completion_guard.is_some();
         let data = match completion_guard {
-            Some(buf) => {
-                if written > buf.len {
-                    return Err(VirtqError::InvalidState);
-                }
-                let owner = BufferOwner::try_new(
-                    self.pool.clone(),
-                    self.inner.mem().clone(),
-                    *buf,
-                    written,
-                )
-                .map_err(|_| VirtqError::MemoryReadError)?;
-                let _ = buf.release();
-                Bytes::from_owner(owner)
+            Some(buf) if written > buf.allocation().len => {
+                // This is a protocol violation
+                return Err(VirtqError::InvalidState);
             }
+            Some(buf) => Bytes::from_owner(
+                buf.into_buffer_owner(self.inner.mem().clone(), written)
+                    .map_err(|_| VirtqError::MemoryReadError)?,
+            ),
             None => Bytes::new(),
         };
 
@@ -642,16 +633,8 @@ impl<M: MemOps, P: BufferProvider + Clone> ChainBuilder<M, P> {
         }
     }
 
-    fn alloc(
-        &self,
-        size: usize,
-    ) -> Result<AllocGuard<impl FnOnce(Allocation) + use<M, P>>, VirtqError> {
-        let alloc = self.pool.alloc(size)?;
-        let pool = self.pool.clone();
-
-        Ok(AllocGuard::new(alloc, move |a| {
-            let _ = pool.dealloc(a);
-        }))
+    fn alloc(&self, size: usize) -> Result<PoolAlloc<P>, VirtqError> {
+        Ok(PoolAlloc::allocate(self.pool.clone(), size)?)
     }
 
     /// Request an entry buffer of `cap` bytes.
@@ -688,14 +671,14 @@ impl<M: MemOps, P: BufferProvider + Clone> ChainBuilder<M, P> {
 
         let inflight = match (entry_alloc, completion_alloc) {
             (Some(entry), Some(cqe)) => Inflight::ReadWrite {
-                entry: entry.release(),
-                completion: cqe.release(),
+                entry: entry.into_raw(),
+                completion: cqe.into_raw(),
             },
             (Some(entry), None) => Inflight::ReadOnly {
-                entry: entry.release(),
+                entry: entry.into_raw(),
             },
             (None, Some(cqe)) => Inflight::WriteOnly {
-                completion: cqe.release(),
+                completion: cqe.into_raw(),
             },
             (None, None) => unreachable!(),
         };