Skip to content

feat(virtq): use Rc pointer equality instead if generation for pools #1945

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andreiltd wants to merge 1 commit into
base: main
Choose a base branch
from
+297 −129
Open

feat(virtq): use Rc pointer equality instead if generation for pools #1945

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions docs/virtio-host-guest-communication.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ external `ByteChunks`:
* `VecBytes` values copy into a contiguous `Vec<u8>`.
* Multiple `Bytes` clones or slices backed by one owner keep one slot live.
* 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.
Its lease releases the slot in its original pool.

Producer reset releases allocations still owned by queue bookkeeping. After
both producers reset and before H2G prefill, every live pool slot belongs to
Expand All @@ -394,13 +394,15 @@ handles. The host resets both consumers before processing more queue traffic.
### Retained virtual addresses

Both pools share one guest-global alias allocator. Its state is captured with
the mappings and keeps retained ranges reserved across pool generations.
the mappings and keeps retained ranges reserved across pool swaps.
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.

Retained aliases reach captured memory after restore. The first transport
entry advances the pool generations and recycles slots outside posted chains.
entry swaps allocators that hold retained slots, preserving posted reservations.
Retained leases keep their original pools, so their release cannot free slots
in the new pool.
Retained values keep their aliases until their final owner drops. Reusing
freed virtual ranges bounds page-table growth by the alias high-water mark.

Expand Down Expand Up @@ -497,7 +499,7 @@ 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.
Source continuation keeps the pool instances 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
Expand Down
54 changes: 23 additions & 31 deletions src/hyperlight_common/src/virtq/buffer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,26 +245,18 @@ impl Buf for SegmentsBuf<'_> {
}
}

/// A buffer allocation returned on drop while its pool generation is current.
/// A buffer allocation returned to its original pool on drop.
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 {
let generation = pool.generation();

Self {
pool,
allocation,
generation,
}
Self { pool, allocation }
}

/// The buffer's start address and full allocation capacity.
Expand All @@ -275,10 +267,6 @@ 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}");
Expand Down Expand Up @@ -315,38 +303,42 @@ mod tests {
}

#[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);
fn old_leases_do_not_release_a_reused_slot() {
let mut pool = SlotPool::new(SlotLayout::new(0x1000, 4096, 1).unwrap()).unwrap();
let mut allocation = pool.alloc(4096).unwrap();
let mut retained = Vec::new();

for _ in 0..3 {
retained.push(BufferLease::new(pool.clone(), allocation));
pool = pool.swap(core::iter::empty()).unwrap();
let reused = pool.alloc(4096).unwrap();
assert_eq!(reused.addr, allocation.addr);
allocation = reused;
}

drop(retained);
assert_eq!(pool.num_free(), 0);
drop(BufferLease::new(pool.clone(), reused));
drop(BufferLease::new(pool.clone(), allocation));
assert_eq!(pool.num_free(), 1);
}

#[test]
fn generation_recovery_preserves_queue_reservations() {
fn swap_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);
let swapped = pool.swap(core::iter::once(reserved.addr)).unwrap();
assert_eq!(swapped.num_free(), 1);
assert_eq!(swapped.alloc(4096).unwrap().addr, captured.addr);

drop(retained);
assert_eq!(pool.num_free(), 0);
assert_eq!(swapped.num_free(), 0);
swapped.dealloc(reserved.addr).unwrap();
swapped.dealloc(captured.addr).unwrap();
assert_eq!(swapped.num_free(), 2);
pool.dealloc(reserved.addr).unwrap();
pool.dealloc(captured.addr).unwrap();
assert_eq!(pool.num_free(), 2);
}

Expand Down
106 changes: 37 additions & 69 deletions src/hyperlight_common/src/virtq/pool/slot.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ struct Tier {
/// Number of slots in this tier.
count: usize,
/// Free slot addresses, popped/pushed LIFO.
free: SmallVec<[u64; 64]>,
free: SmallVec<[u64; 16]>,
/// One bit per slot index; set means the slot is currently handed out.
allocated: FixedBitSet,
}
Expand All @@ -115,10 +115,6 @@ impl Tier {
free.try_reserve_exact(layout.slot_count)
.map_err(|_| AllocError::Bookkeeping)?;

for i in 0..layout.slot_count {
free.push(layout.base_addr + (i * layout.slot_size as usize) as u64);
}

Ok(Self {
base_addr: layout.base_addr,
slot_size: layout.slot_size,
Expand Down Expand Up @@ -186,7 +182,7 @@ impl Tier {
Ok(())
}

/// Keep one queue-owned slot unavailable during generation recovery.
/// Keep one queue-owned slot unavailable for a pool swap.
fn reserve_addr(&mut self, addr: u64) -> Result<(), AllocError> {
let slot = self.slot_of(addr)?;

Expand All @@ -198,7 +194,7 @@ impl Tier {
Ok(())
}

/// Rebuild the free list from the recovered ownership bitmap.
/// Populate the free list from the allocation bitmap.
fn rebuild_free(&mut self) {
self.free.clear();

Expand Down Expand Up @@ -241,7 +237,6 @@ impl Tier {
struct Inner {
lower: Option<Tier>,
upper: Tier,
generation: u64,
}

impl Inner {
Expand All @@ -250,7 +245,6 @@ impl Inner {
return Ok(Self {
lower: None,
upper: Tier::from_layout(upper)?,
generation: 0,
});
};

Expand All @@ -273,14 +267,12 @@ 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,
})
}

Expand Down Expand Up @@ -311,47 +303,6 @@ 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)?;

#[cfg(debug_assertions)]
for addr in reserved.clone() {
debug_assert!(
self.allocation_len(addr).is_ok(),
"address {addr:#x} is not a live pool slot"
);
}

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)
Expand Down Expand Up @@ -419,23 +370,32 @@ pub struct SlotPool {
}

impl SlotPool {
/// Generation owning the current scratch allocations.
pub(crate) fn generation(&self) -> u64 {
self.inner.borrow().generation
/// Whether both handles share allocation state.
pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
Rc::ptr_eq(&self.inner, &other.inner)
}

/// Retire existing leases and rebuild the pool around the supplied reservations.
///
/// # Safety
/// Create allocation state for a pool swap with the supplied live slots reserved.
///
/// 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)
/// Invalid or duplicate reservations leave this pool unchanged.
pub(crate) fn swap(&self, reserved: impl Iterator<Item = u64>) -> Result<Self, AllocError> {
let current = self.inner.borrow();
let (lower, upper) = current.layouts();
let mut inner = Inner::new(lower, upper)?;

for addr in reserved {
if let Some(lower) = &mut inner.lower
&& lower.contains(addr)
{
lower.allocation_len(addr)?;
lower.reserve_addr(addr)?;
} else {
inner.upper.allocation_len(addr)?;
inner.upper.reserve_addr(addr)?;
}
}
Comment on lines +386 to +396

Ok(Self::from_inner(inner))
}

/// Create a single-tier recycling pool from exact slot placement.
Expand All @@ -453,10 +413,18 @@ impl SlotPool {
}

fn from_layouts(lower: Option<SlotLayout>, upper: SlotLayout) -> Result<Self, AllocError> {
let inner = Inner::new(lower, upper)?;
Ok(Self {
Ok(Self::from_inner(Inner::new(lower, upper)?))
}

fn from_inner(mut inner: Inner) -> Self {
inner.upper.rebuild_free();
if let Some(lower) = &mut inner.lower {
lower.rebuild_free();
}

Self {
inner: Rc::new(RefCell::new(inner)),
})
}
}

/// Return every live slot address in deterministic tier and index order.
Expand Down
Loading
Loading