feat: map virtq buffers lazily

Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>

Author: andreiltd

CHANGELOG.md

@@ -20,6 +20,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   snapshot capture, restore, and cloning.
 
 ### Changed
+* Guest transport aliases follow buffer ownership. Restored buffers keep
+  captured backing while their scratch slots become reusable.
 * `Sandbox` is the primary initialized sandbox type. `MultiUseSandbox` remains
   as a deprecated alias.
 * `Sandbox` lives in the private `sandbox::initialized` module and is reached

Cargo.lock

@@ -364,8 +364,8 @@ drains and acknowledges them during the same VM exit.
 Guest `SlotPool` instances own all transport buffers. Pool clones share one
 allocation bitmap with each producer.
 Each producer pairs its backend with that pool. Completion leases carry the
-original slot addresses and full capacities. Backing owns fixed scratch-to-alias
-translation and paging. The backend holds only scratch bounds.
+original slot addresses and full capacities. Backing owns alias allocation and
+paging. The backend holds only scratch bounds.
 Guest completion mapping checks scratch bounds. Allocation ownership and
 initialized lengths follow the `BufferMap` safety contract.
 
@@ -381,8 +381,8 @@ external `ByteChunks`:
 * G2H host responses can become owner-backed guest `Bytes`.
 * `VecBytes` values copy into a contiguous `Vec<u8>`.
 * Multiple `Bytes` clones or slices backed by one owner keep one slot live.
-* The final owner returns its original `BufferLease`. Active pool mappings
-  remain available across ordinary calls.
+* The final owner unmaps its alias and returns the virtual range for reuse.
+  Its lease releases the scratch slot only in the slot's current generation.
 
 Producer reset releases allocations still owned by queue bookkeeping. After
 both producers reset and before H2G prefill, every live pool slot belongs to
@@ -393,14 +393,16 @@ handles. The host resets both consumers before processing more queue traffic.
 
 ### Retained virtual addresses
 
-Each pool has a stable virtual alias range, mapped eagerly at initialization.
-Completed `GuestMapping` owners keep their original leases.
-Capture, restore, and cloning preserve their pointers and contents.
+Both pools share one guest-global alias allocator. Its state is captured with
+the mappings and keeps retained ranges reserved across pool generations.
+Each completed `GuestMapping` owns a separate page-aligned range. Mapping occurs
+when the guest receives the buffer. Capture, restore, and cloning preserve its
+pointer and contents.
 
-Checkpoint preparation records the live slots. After restore, aliases of their
-pages reach captured data. The first transport entry copies the recorded slots
-into scratch and maps every pool alias to scratch. Only recorded slots are
-copied, so neighboring slots keep host writes such as the first request.
+Retained aliases reach captured memory after restore. The first transport
+entry advances the pool generations and recycles slots outside posted chains.
+Retained values keep their aliases until their final owner drops. Reusing
+freed virtual ranges bounds page-table growth by the alias high-water mark.
 
 ### Trust boundary
 
@@ -437,8 +439,6 @@ snapshot uses this flow:
   |                                     | reclaim completed G2H work
   |                                     | reset G2H producer
   |                                     | reset H2G producer
-  |                                     | record live slots
-  |                                     | unmap pages outside live slots
   |                                     | prefill free H2G slots
   |                                     | mailbox = CheckpointComplete
   |<------------------------------------| halt
@@ -459,7 +459,7 @@ The canonical state is:
   bounded by queue size. Each descriptor names a distinct, configured-size
   slot aligned relative to the pool start. Available descriptors form a prefix
   followed by zeroed descriptors.
-* Guest producer and pool bookkeeping matches the rings and retained leases.
+* Guest producer and pool bookkeeping matches the rings and current leases.
 * Driver and device event suppression is normalized.
 * Host consumers start at cursor zero.
 
@@ -490,11 +490,13 @@ The first request fits the H2G capacity posted at checkpoint. Retained slots
 reduce pool capacity but hold no ring descriptors. Framing, slot rounding, ring
 size, and the external-byte control reserve still apply.
 
-The first transport entry after checkpoint maps missing pool pages to scratch.
-After restore, it also copies recorded slots from captured aliases into scratch
-and remaps their pages. This applies to logs and host callbacks before request
-dispatch as well as ordinary H2G calls. Pool ownership, descriptors, and
-cursors stay unchanged, preserving the submitted request.
+The first transport entry after restore recycles slots held by captured
+leases. Posted chains keep their reservations, preserving the submitted
+request. Descriptors and cursors stay unchanged. Retained aliases keep their
+captured mappings without payload copying.
+
+Source continuation keeps the pool generation and live scratch leases intact.
+Checkpointing alone does not detach retained payloads from scratch.
 
 Result completion prefills free H2G slots. This requires no preparatory guest
 entry or application warmup for the checkpoint-posted capacity.

docs/virtio-host-guest-communication.md

@@ -18,6 +18,8 @@ pub const WHP_GITS_TRANSLATOR_BASE_GPA: u64 = 0xeff6_8000;
 pub const WHP_GICR_BASE_GPA: u64 = 0xeffe_e000;
 pub const WHP_GICD_BASE_GPA: u64 = 0xffff_0000;
 
+const _: () = assert!(VIRTQ_BUFFER_GVA_END <= 1 << 48);
+
 pub const fn io_page() -> Option<(crate::vmem::PhysAddr, crate::vmem::VirtAddr)> {
     Some((IO_PAGE_GPA, IO_PAGE_GVA))
 }

src/hyperlight_common/src/arch/aarch64/layout.rs

@@ -20,6 +20,11 @@ pub const VIRTQ_BUFFER_GVA_END: u64 = 0xffff_fd00_0000_0000;
 /// ever memory-constrained.
 pub const SCRATCH_TOP_GPA: usize = 0x0000_000f_ffff_ffff;
 
+const _: () = {
+    assert!(VIRTQ_BUFFER_GVA_START >> 47 == 0x1ffff);
+    assert!((VIRTQ_BUFFER_GVA_END - 1) >> 47 == 0x1ffff);
+};
+
 pub fn io_page() -> Option<(u64, u64)> {
     None
 }

src/hyperlight_common/src/arch/amd64/layout.rs

@@ -109,13 +109,6 @@ const _: () = {
     assert!((VIRTQ_BUFFER_GVA_END as usize).is_multiple_of(crate::vmem::PAGE_SIZE));
     assert!(VIRTQ_BUFFER_GVA_START > SNAPSHOT_PT_GVA_MAX as u64);
 
-    #[cfg(target_arch = "x86_64")]
-    assert!(VIRTQ_BUFFER_GVA_START >> 47 == 0x1ffff);
-    #[cfg(target_arch = "x86_64")]
-    assert!((VIRTQ_BUFFER_GVA_END - 1) >> 47 == 0x1ffff);
-    #[cfg(target_arch = "aarch64")]
-    assert!(VIRTQ_BUFFER_GVA_END <= 1 << 48);
-
     assert!(VIRTQ_BUFFER_GVA_END < scratch_base_gva(16 * 1024 * 1024 * 1024));
     assert!((SCRATCH_TOP_GPA as u64) < VIRTQ_BUFFER_GVA_END - VIRTQ_BUFFER_GVA_START);
 };

src/hyperlight_common/src/layout.rs

@@ -245,18 +245,26 @@ impl Buf for SegmentsBuf<'_> {
     }
 }
 
-/// An exclusively owned buffer allocation returned to its pool on drop.
+/// A buffer allocation returned on drop while its pool generation is current.
 pub struct BufferLease {
     /// The pool that allocated the buffer.
     pool: SlotPool,
+    /// The pool generation at the time of allocation.
+    generation: u64,
     /// The buffer's start address and full allocation capacity.
     allocation: Allocation,
 }
 
 impl BufferLease {
     /// Create a new buffer lease from a pool and allocation.
     pub fn new(pool: SlotPool, allocation: Allocation) -> Self {
-        Self { pool, allocation }
+        let generation = pool.generation();
+
+        Self {
+            pool,
+            allocation,
+            generation,
+        }
     }
 
     /// The buffer's start address and full allocation capacity.
@@ -267,6 +275,10 @@ impl BufferLease {
 
 impl Drop for BufferLease {
     fn drop(&mut self) {
+        if self.generation != self.pool.generation() {
+            return;
+        }
+
         if let Err(error) = self.pool.dealloc(self.allocation.addr) {
             log::error!("Failed to release a virtqueue buffer: {error}");
             debug_assert!(false, "BufferLease deallocation failed: {error}");
@@ -302,6 +314,42 @@ mod tests {
         pool.dealloc(reused.addr).unwrap();
     }
 
+    #[test]
+    fn retired_lease_does_not_release_a_reused_slot() {
+        let pool = SlotPool::new(SlotLayout::new(0x1000, 4096, 1).unwrap()).unwrap();
+        let allocation = pool.alloc(4096).unwrap();
+        let retained = BufferLease::new(pool.clone(), allocation);
+
+        // SAFETY: This test has no views or queue reservations.
+        unsafe { pool.retire_generation(core::iter::empty()) }.unwrap();
+        let reused = pool.alloc(4096).unwrap();
+        assert_eq!(reused.addr, allocation.addr);
+
+        drop(retained);
+        assert_eq!(pool.num_free(), 0);
+        drop(BufferLease::new(pool.clone(), reused));
+        assert_eq!(pool.num_free(), 1);
+    }
+
+    #[test]
+    fn generation_recovery_preserves_queue_reservations() {
+        let pool = SlotPool::new(SlotLayout::new(0x1000, 4096, 2).unwrap()).unwrap();
+        let reserved = pool.alloc(4096).unwrap();
+        let captured = pool.alloc(4096).unwrap();
+        let retained = BufferLease::new(pool.clone(), captured);
+
+        // SAFETY: No view exists. The iterator names the only queue reservation.
+        unsafe { pool.retire_generation(core::iter::once(reserved.addr)) }.unwrap();
+        assert_eq!(pool.num_free(), 1);
+        assert_eq!(pool.alloc(4096).unwrap().addr, captured.addr);
+
+        drop(retained);
+        assert_eq!(pool.num_free(), 0);
+        pool.dealloc(reserved.addr).unwrap();
+        pool.dealloc(captured.addr).unwrap();
+        assert_eq!(pool.num_free(), 2);
+    }
+
     #[test]
     fn segments_cursor_advances_across_segments() {
         let segments = Segments::new([

src/hyperlight_common/src/virtq/buffer.rs

@@ -186,6 +186,28 @@ impl Tier {
         Ok(())
     }
 
+    /// Keep one queue-owned slot unavailable during generation recovery.
+    fn reserve_addr(&mut self, addr: u64) -> Result<(), AllocError> {
+        let slot = self.slot_of(addr)?;
+
+        if self.allocated.contains(slot) {
+            return Err(AllocError::InvalidArg);
+        }
+
+        self.allocated.insert(slot);
+        Ok(())
+    }
+
+    /// Rebuild the free list from the recovered ownership bitmap.
+    fn rebuild_free(&mut self) {
+        self.free.clear();
+
+        for slot in self.allocated.zeroes() {
+            self.free
+                .push(self.base_addr + slot as u64 * u64::from(self.slot_size));
+        }
+    }
+
     fn allocation_len(&self, addr: u64) -> Result<usize, AllocError> {
         self.live_slot_of(addr)?;
         Ok(self.slot_size as usize)
@@ -219,6 +241,7 @@ impl Tier {
 struct Inner {
     lower: Option<Tier>,
     upper: Tier,
+    generation: u64,
 }
 
 impl Inner {
@@ -227,6 +250,7 @@ impl Inner {
             return Ok(Self {
                 lower: None,
                 upper: Tier::from_layout(upper)?,
+                generation: 0,
             });
         };
 
@@ -249,12 +273,14 @@ impl Inner {
             return Ok(Self {
                 lower: None,
                 upper: Tier::from_layout(layout)?,
+                generation: 0,
             });
         }
 
         Ok(Self {
             lower: Some(Tier::from_layout(lower)?),
             upper: Tier::from_layout(upper)?,
+            generation: 0,
         })
     }
 
@@ -285,6 +311,43 @@ impl Inner {
         self.upper.dealloc_addr(addr)
     }
 
+    /// Start a new generation and rebuild allocation state from `reserved`.
+    fn retire_generation(
+        &mut self,
+        reserved: impl Iterator<Item = u64> + Clone,
+    ) -> Result<(), AllocError> {
+        let generation = self.generation.checked_add(1).ok_or(AllocError::Overflow)?;
+
+        for addr in reserved.clone() {
+            self.allocation_len(addr)?;
+        }
+
+        self.generation = generation;
+        self.upper.allocated.clear();
+
+        if let Some(lower) = &mut self.lower {
+            lower.allocated.clear();
+        }
+
+        for addr in reserved {
+            if let Some(lower) = &mut self.lower
+                && lower.contains(addr)
+            {
+                lower.reserve_addr(addr)?;
+            } else {
+                self.upper.reserve_addr(addr)?;
+            }
+        }
+
+        self.upper.rebuild_free();
+
+        if let Some(lower) = &mut self.lower {
+            lower.rebuild_free();
+        }
+
+        Ok(())
+    }
+
     fn allocation_len(&self, addr: u64) -> Result<usize, AllocError> {
         if let Some(lower) = &self.lower
             && lower.contains(addr)
@@ -352,6 +415,25 @@ pub struct SlotPool {
 }
 
 impl SlotPool {
+    /// Generation owning the current scratch allocations.
+    pub(crate) fn generation(&self) -> u64 {
+        self.inner.borrow().generation
+    }
+
+    /// Retire existing leases and rebuild the pool around the supplied reservations.
+    ///
+    /// # Safety
+    ///
+    /// Retired views must already use backing independent of these scratch slots.
+    /// `reserved` must name every other live allocation exactly once.
+    /// No allocation or peer access may race this call.
+    pub(crate) unsafe fn retire_generation(
+        &self,
+        reserved: impl Iterator<Item = u64> + Clone,
+    ) -> Result<(), AllocError> {
+        self.inner.borrow_mut().retire_generation(reserved)
+    }
+
     /// Create a single-tier recycling pool from exact slot placement.
     pub fn new(layout: SlotLayout) -> Result<Self, AllocError> {
         Self::from_layouts(None, layout)

src/hyperlight_common/src/virtq/pool/slot.rs

@@ -240,6 +240,30 @@ where
         self.inner.mem()
     }
 
+    /// Reclaim retained scratch slots while preserving in-flight chains.
+    ///
+    /// # Safety
+    ///
+    /// Every live allocation outside this producer's in-flight chains must
+    /// belong to a returned view with backing independent of reused scratch.
+    /// No unsent chains or builders may remain. Peer accesses must be stopped.
+    pub unsafe fn reclaim_slots(&mut self) -> Result<(), VirtqError> {
+        if !self.pending.is_empty() {
+            return Err(VirtqError::InvalidState);
+        }
+
+        let reserved = self
+            .inflight
+            .live
+            .iter()
+            .flat_map(|inflight| inflight.chain.buffers.iter().map(|buffer| buffer.addr));
+
+        // SAFETY: In-flight records preserve all non-retired allocations.
+        // The caller guarantees independent backing and excludes untracked chains.
+        unsafe { self.pool.retire_generation(reserved)? };
+        Ok(())
+    }
+
     /// Begin building a descriptor chain for submission.
     ///
     /// The builder captures the current free-descriptor budget.

src/hyperlight_common/src/virtq/producer.rs

@@ -15,9 +15,9 @@ Provides only the essential building blocks for interacting with the host enviro
 anyhow = { version = "1.0.102", default-features = false }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 hyperlight-common = { workspace = true, default-features = false }
-fixedbitset = { version = "0.5.7", default-features = false }
 flatbuffers = { version= "25.12.19", default-features = false }
 tracing = { version = "0.1.44", default-features = false, features = ["attributes"] }
+itertools = { version = "0.15.0", default-features = false }
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 hyperlight-guest-tracing = { workspace = true, default-features = false, optional = true }