Skip to content

Commit 4062922

Browse files
committed
feat(virtq): use Rc ptr equality instead generation for pools
Instead of threading generation through pools it's actually simpler to use Rc::ptr_eq to check if pool belongs to current snapshot generation and swap pools on restore. Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>
1 parent ad05110 commit 4062922

8 files changed

Lines changed: 297 additions & 129 deletions

File tree

‎docs/virtio-host-guest-communication.md‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -382,7 +382,7 @@ external `ByteChunks`:
382382
* `VecBytes` values copy into a contiguous `Vec<u8>`.
383383
* Multiple `Bytes` clones or slices backed by one owner keep one slot live.
384384
* The final owner unmaps its alias and returns the virtual range for reuse.
385-
Its lease releases the scratch slot only in the slot's current generation.
385+
Its lease releases the slot in its original pool.
386386

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

396396
Both pools share one guest-global alias allocator. Its state is captured with
397-
the mappings and keeps retained ranges reserved across pool generations.
397+
the mappings and keeps retained ranges reserved across pool swaps.
398398
Each completed `GuestMapping` owns a separate page-aligned range. Mapping occurs
399399
when the guest receives the buffer. Capture, restore, and cloning preserve its
400400
pointer and contents.
401401

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

@@ -497,7 +499,7 @@ leases. Posted chains keep their reservations, preserving the submitted
497499
request. Descriptors and cursors stay unchanged. Retained aliases keep their
498500
captured mappings without payload copying.
499501

500-
Source continuation keeps the pool generation and live scratch leases intact.
502+
Source continuation keeps the pool instances and live scratch leases intact.
501503
Checkpointing alone does not detach retained payloads from scratch.
502504

503505
Result completion prefills free H2G slots. This requires no preparatory guest

‎src/hyperlight_common/src/virtq/buffer.rs‎

Lines changed: 23 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -245,26 +245,18 @@ impl Buf for SegmentsBuf<'_> {
245245
}
246246
}
247247

248-
/// A buffer allocation returned on drop while its pool generation is current.
248+
/// A buffer allocation returned to its original pool on drop.
249249
pub struct BufferLease {
250250
/// The pool that allocated the buffer.
251251
pool: SlotPool,
252-
/// The pool generation at the time of allocation.
253-
generation: u64,
254252
/// The buffer's start address and full allocation capacity.
255253
allocation: Allocation,
256254
}
257255

258256
impl BufferLease {
259257
/// Create a new buffer lease from a pool and allocation.
260258
pub fn new(pool: SlotPool, allocation: Allocation) -> Self {
261-
let generation = pool.generation();
262-
263-
Self {
264-
pool,
265-
allocation,
266-
generation,
267-
}
259+
Self { pool, allocation }
268260
}
269261

270262
/// The buffer's start address and full allocation capacity.
@@ -275,10 +267,6 @@ impl BufferLease {
275267

276268
impl Drop for BufferLease {
277269
fn drop(&mut self) {
278-
if self.generation != self.pool.generation() {
279-
return;
280-
}
281-
282270
if let Err(error) = self.pool.dealloc(self.allocation.addr) {
283271
log::error!("Failed to release a virtqueue buffer: {error}");
284272
debug_assert!(false, "BufferLease deallocation failed: {error}");
@@ -315,38 +303,42 @@ mod tests {
315303
}
316304

317305
#[test]
318-
fn retired_lease_does_not_release_a_reused_slot() {
319-
let pool = SlotPool::new(SlotLayout::new(0x1000, 4096, 1).unwrap()).unwrap();
320-
let allocation = pool.alloc(4096).unwrap();
321-
let retained = BufferLease::new(pool.clone(), allocation);
322-
323-
// SAFETY: This test has no views or queue reservations.
324-
unsafe { pool.retire_generation(core::iter::empty()) }.unwrap();
325-
let reused = pool.alloc(4096).unwrap();
326-
assert_eq!(reused.addr, allocation.addr);
306+
fn old_leases_do_not_release_a_reused_slot() {
307+
let mut pool = SlotPool::new(SlotLayout::new(0x1000, 4096, 1).unwrap()).unwrap();
308+
let mut allocation = pool.alloc(4096).unwrap();
309+
let mut retained = Vec::new();
310+
311+
for _ in 0..3 {
312+
retained.push(BufferLease::new(pool.clone(), allocation));
313+
pool = pool.swap(core::iter::empty()).unwrap();
314+
let reused = pool.alloc(4096).unwrap();
315+
assert_eq!(reused.addr, allocation.addr);
316+
allocation = reused;
317+
}
327318

328319
drop(retained);
329320
assert_eq!(pool.num_free(), 0);
330-
drop(BufferLease::new(pool.clone(), reused));
321+
drop(BufferLease::new(pool.clone(), allocation));
331322
assert_eq!(pool.num_free(), 1);
332323
}
333324

334325
#[test]
335-
fn generation_recovery_preserves_queue_reservations() {
326+
fn swap_preserves_queue_reservations() {
336327
let pool = SlotPool::new(SlotLayout::new(0x1000, 4096, 2).unwrap()).unwrap();
337328
let reserved = pool.alloc(4096).unwrap();
338329
let captured = pool.alloc(4096).unwrap();
339330
let retained = BufferLease::new(pool.clone(), captured);
340331

341-
// SAFETY: No view exists. The iterator names the only queue reservation.
342-
unsafe { pool.retire_generation(core::iter::once(reserved.addr)) }.unwrap();
343-
assert_eq!(pool.num_free(), 1);
344-
assert_eq!(pool.alloc(4096).unwrap().addr, captured.addr);
332+
let swapped = pool.swap(core::iter::once(reserved.addr)).unwrap();
333+
assert_eq!(swapped.num_free(), 1);
334+
assert_eq!(swapped.alloc(4096).unwrap().addr, captured.addr);
345335

346336
drop(retained);
347-
assert_eq!(pool.num_free(), 0);
337+
assert_eq!(swapped.num_free(), 0);
338+
swapped.dealloc(reserved.addr).unwrap();
339+
swapped.dealloc(captured.addr).unwrap();
340+
assert_eq!(swapped.num_free(), 2);
348341
pool.dealloc(reserved.addr).unwrap();
349-
pool.dealloc(captured.addr).unwrap();
350342
assert_eq!(pool.num_free(), 2);
351343
}
352344

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

Lines changed: 37 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ struct Tier {
104104
/// Number of slots in this tier.
105105
count: usize,
106106
/// Free slot addresses, popped/pushed LIFO.
107-
free: SmallVec<[u64; 64]>,
107+
free: SmallVec<[u64; 16]>,
108108
/// One bit per slot index; set means the slot is currently handed out.
109109
allocated: FixedBitSet,
110110
}
@@ -115,10 +115,6 @@ impl Tier {
115115
free.try_reserve_exact(layout.slot_count)
116116
.map_err(|_| AllocError::Bookkeeping)?;
117117

118-
for i in 0..layout.slot_count {
119-
free.push(layout.base_addr + (i * layout.slot_size as usize) as u64);
120-
}
121-
122118
Ok(Self {
123119
base_addr: layout.base_addr,
124120
slot_size: layout.slot_size,
@@ -186,7 +182,7 @@ impl Tier {
186182
Ok(())
187183
}
188184

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

@@ -198,7 +194,7 @@ impl Tier {
198194
Ok(())
199195
}
200196

201-
/// Rebuild the free list from the recovered ownership bitmap.
197+
/// Populate the free list from the allocation bitmap.
202198
fn rebuild_free(&mut self) {
203199
self.free.clear();
204200

@@ -241,7 +237,6 @@ impl Tier {
241237
struct Inner {
242238
lower: Option<Tier>,
243239
upper: Tier,
244-
generation: u64,
245240
}
246241

247242
impl Inner {
@@ -250,7 +245,6 @@ impl Inner {
250245
return Ok(Self {
251246
lower: None,
252247
upper: Tier::from_layout(upper)?,
253-
generation: 0,
254248
});
255249
};
256250

@@ -273,14 +267,12 @@ impl Inner {
273267
return Ok(Self {
274268
lower: None,
275269
upper: Tier::from_layout(layout)?,
276-
generation: 0,
277270
});
278271
}
279272

280273
Ok(Self {
281274
lower: Some(Tier::from_layout(lower)?),
282275
upper: Tier::from_layout(upper)?,
283-
generation: 0,
284276
})
285277
}
286278

@@ -311,47 +303,6 @@ impl Inner {
311303
self.upper.dealloc_addr(addr)
312304
}
313305

314-
/// Start a new generation and rebuild allocation state from `reserved`.
315-
fn retire_generation(
316-
&mut self,
317-
reserved: impl Iterator<Item = u64> + Clone,
318-
) -> Result<(), AllocError> {
319-
let generation = self.generation.checked_add(1).ok_or(AllocError::Overflow)?;
320-
321-
#[cfg(debug_assertions)]
322-
for addr in reserved.clone() {
323-
debug_assert!(
324-
self.allocation_len(addr).is_ok(),
325-
"address {addr:#x} is not a live pool slot"
326-
);
327-
}
328-
329-
self.generation = generation;
330-
self.upper.allocated.clear();
331-
332-
if let Some(lower) = &mut self.lower {
333-
lower.allocated.clear();
334-
}
335-
336-
for addr in reserved {
337-
if let Some(lower) = &mut self.lower
338-
&& lower.contains(addr)
339-
{
340-
lower.reserve_addr(addr)?;
341-
} else {
342-
self.upper.reserve_addr(addr)?;
343-
}
344-
}
345-
346-
self.upper.rebuild_free();
347-
348-
if let Some(lower) = &mut self.lower {
349-
lower.rebuild_free();
350-
}
351-
352-
Ok(())
353-
}
354-
355306
fn allocation_len(&self, addr: u64) -> Result<usize, AllocError> {
356307
if let Some(lower) = &self.lower
357308
&& lower.contains(addr)
@@ -419,23 +370,32 @@ pub struct SlotPool {
419370
}
420371

421372
impl SlotPool {
422-
/// Generation owning the current scratch allocations.
423-
pub(crate) fn generation(&self) -> u64 {
424-
self.inner.borrow().generation
373+
/// Whether both handles share allocation state.
374+
pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
375+
Rc::ptr_eq(&self.inner, &other.inner)
425376
}
426377

427-
/// Retire existing leases and rebuild the pool around the supplied reservations.
428-
///
429-
/// # Safety
378+
/// Create allocation state for a pool swap with the supplied live slots reserved.
430379
///
431-
/// Retired views must already use backing independent of these scratch slots.
432-
/// `reserved` must name every other live allocation exactly once.
433-
/// No allocation or peer access may race this call.
434-
pub(crate) unsafe fn retire_generation(
435-
&self,
436-
reserved: impl Iterator<Item = u64> + Clone,
437-
) -> Result<(), AllocError> {
438-
self.inner.borrow_mut().retire_generation(reserved)
380+
/// Invalid or duplicate reservations leave this pool unchanged.
381+
pub(crate) fn swap(&self, reserved: impl Iterator<Item = u64>) -> Result<Self, AllocError> {
382+
let current = self.inner.borrow();
383+
let (lower, upper) = current.layouts();
384+
let mut inner = Inner::new(lower, upper)?;
385+
386+
for addr in reserved {
387+
// this address is a valid allocation in the current pool
388+
current.allocation_len(addr)?;
389+
if let Some(lower) = &mut inner.lower
390+
&& lower.contains(addr)
391+
{
392+
lower.reserve_addr(addr)?;
393+
} else {
394+
inner.upper.reserve_addr(addr)?;
395+
}
396+
}
397+
398+
Ok(Self::from_inner(inner))
439399
}
440400

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

455415
fn from_layouts(lower: Option<SlotLayout>, upper: SlotLayout) -> Result<Self, AllocError> {
456-
let inner = Inner::new(lower, upper)?;
457-
Ok(Self {
416+
Ok(Self::from_inner(Inner::new(lower, upper)?))
417+
}
418+
419+
fn from_inner(mut inner: Inner) -> Self {
420+
inner.upper.rebuild_free();
421+
if let Some(lower) = &mut inner.lower {
422+
lower.rebuild_free();
423+
}
424+
425+
Self {
458426
inner: Rc::new(RefCell::new(inner)),
459-
})
427+
}
460428
}
461429

462430
/// Return every live slot address in deterministic tier and index order.

0 commit comments

Comments
 (0)