feat(virtq): add support for multi-descriptor payloads

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

Author: andreiltd

src/hyperlight_common/src/virtq/consumer.rs

@@ -255,6 +255,7 @@ pub struct VirtqConsumer<M, N> {
     inner: RingConsumer<M>,
     notifier: N,
     inflight: FixedBitSet,
+    next_token: u32,
 }
 
 impl<M: MemOps + Clone, N: Notifier> VirtqConsumer<M, N> {
@@ -273,6 +274,7 @@ impl<M: MemOps + Clone, N: Notifier> VirtqConsumer<M, N> {
             inner,
             notifier,
             inflight,
+            next_token: 0,
         }
     }
 
@@ -323,7 +325,8 @@ impl<M: MemOps + Clone, N: Notifier> VirtqConsumer<M, N> {
         }
 
         self.inflight.insert(id_idx);
-        let token = Token(id);
+        let token = Token(self.next_token, id);
+        self.next_token = self.next_token.wrapping_add(1);
 
         // Copy entry data from shared memory
         let data = entry_elem

src/hyperlight_common/src/virtq/mod.rs

@@ -335,11 +335,11 @@ pub enum SuppressionKind {
 
 /// A token representing a sent entry in the virtqueue.
 ///
-/// Tokens uniquely identify in-flight requests and are used to correlate
-/// requests with their responses. The token value corresponds to the
-/// descriptor ID in the underlying ring.
+/// Tokens uniquely identify in-flight requests and are used to correlate requests with their responses.
+/// The first element is a monotonically increasing generation counter. The second element is the
+/// underlying descriptor ID
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
-pub struct Token(pub u16);
+pub struct Token(pub u32, pub u16);
 
 impl From<BufferElement> for Allocation {
     fn from(value: BufferElement) -> Self {
@@ -972,6 +972,54 @@ mod tests {
         assert_eq!(cqe2.token, tok_rw);
         assert_eq!(&cqe2.data[..], b"reply");
     }
+
+    /// Regression test: reclaim + submit must not cause token collisions.
+    ///
+    /// Before the monotonic generation counter, Token wrapped the descriptor
+    /// ID which gets recycled. This caused stale pending completions to
+    /// match newly submitted entries with the same recycled descriptor ID.
+    #[test]
+    fn test_reclaim_submit_no_token_collision() {
+        let ring = make_ring(8);
+        let (mut producer, mut consumer, _) = make_test_producer(&ring);
+
+        // Submit and complete a ReadOnly entry
+        let tok_old = send_readonly(&mut producer, b"log");
+
+        let (_, c) = consumer.poll(1024).unwrap().unwrap();
+        consumer.complete(c).unwrap();
+
+        // Reclaim pushes the completion to pending (token = tok_old)
+        let count = producer.reclaim().unwrap();
+        assert_eq!(count, 1);
+
+        // Submit a new ReadWrite entry - may reuse the same descriptor ID
+        let tok_new = send_readwrite(&mut producer, b"call", 64);
+
+        // Tokens must differ even if the descriptor ID was recycled
+        assert_ne!(
+            tok_old, tok_new,
+            "tokens must be unique across reclaim/submit cycles"
+        );
+
+        // Complete the ReadWrite entry
+        let (_, c) = consumer.poll(1024).unwrap().unwrap();
+        let SendCompletion::Writable(mut wc) = c else {
+            panic!("expected writable");
+        };
+        wc.write_all(b"result").unwrap();
+        consumer.complete(wc.into()).unwrap();
+
+        // Poll should return the stale ReadOnly completion first (wrong token)
+        let cqe1 = producer.poll().unwrap().unwrap();
+        assert_eq!(cqe1.token, tok_old);
+        assert!(cqe1.data.is_empty());
+
+        // Then the new ReadWrite completion (matching token)
+        let cqe2 = producer.poll().unwrap().unwrap();
+        assert_eq!(cqe2.token, tok_new);
+        assert_eq!(&cqe2.data[..], b"result");
+    }
 }
 #[cfg(all(test, loom))]
 mod fuzz {

src/hyperlight_common/src/virtq/msg.rs

@@ -20,6 +20,8 @@ limitations under the License.
 //! fixed 8-byte header, enabling message type discrimination and
 //! request/response correlation.
 
+use bitflags::bitflags;
+
 /// Message types for the virtqueue wire protocol.
 #[repr(u8)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -54,24 +56,33 @@ impl TryFrom<u8> for MsgKind {
     }
 }
 
+bitflags! {
+    #[repr(transparent)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+    pub struct MsgFlags: u8 {
+        /// More descriptors follow for this message.
+        const MORE = 1 << 0;
+    }
+}
+
 /// Wire header for all virtqueue messages
 #[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
 #[repr(C)]
 pub struct VirtqMsgHeader {
     /// Discriminates the message type.
     pub kind: u8,
-    /// Per-type flags TODO(ring): add flags type.
+    /// Per-message flags (see [`MsgFlags`]).
     pub flags: u8,
     /// Caller-assigned correlation ID. Responses echo the request's ID.
     pub req_id: u16,
-    /// Byte length of the payload following this header.
+    /// Byte length of the payload following this header in this descriptor.
     pub payload_len: u32,
 }
 
 impl VirtqMsgHeader {
     pub const SIZE: usize = core::mem::size_of::<Self>();
 
-    /// Create a new message header.
+    /// Create a new message header with no flags set.
     pub const fn new(kind: MsgKind, req_id: u16, payload_len: u32) -> Self {
         Self {
             kind: kind as u8,
@@ -82,10 +93,10 @@ impl VirtqMsgHeader {
     }
 
     /// Create a new header with flags.
-    pub const fn with_flags(kind: MsgKind, flags: u8, req_id: u16, payload_len: u32) -> Self {
+    pub const fn with_flags(kind: MsgKind, flags: MsgFlags, req_id: u16, payload_len: u32) -> Self {
         Self {
             kind: kind as u8,
-            flags,
+            flags: flags.bits(),
             req_id,
             payload_len,
         }
@@ -95,4 +106,15 @@ impl VirtqMsgHeader {
     pub fn msg_kind(&self) -> Result<MsgKind, u8> {
         MsgKind::try_from(self.kind)
     }
+
+    /// Interpret the raw flags field as [`MsgFlags`].
+    pub fn msg_flags(&self) -> MsgFlags {
+        MsgFlags::from_bits_truncate(self.flags)
+    }
+
+    /// Returns true if [`MsgFlags::MORE`] is set, indicating more
+    /// descriptors follow for this message.
+    pub const fn has_more(&self) -> bool {
+        self.flags & MsgFlags::MORE.bits() != 0
+    }
 }

src/hyperlight_common/src/virtq/pool.rs

@@ -150,10 +150,10 @@ impl<const N: usize> Slab<N> {
         }
 
         // Fallback to full search
+        let total = self.used_slots.len();
         self.used_slots.zeroes().find(|&next_free| {
-            self.used_slots
-                .count_zeroes(next_free..next_free + slots_num)
-                == slots_num
+            let end = next_free + slots_num;
+            end <= total && self.used_slots.count_zeroes(next_free..end) == slots_num
         })
     }
 
@@ -416,6 +416,11 @@ impl<const L: usize, const U: usize> BufferPool<L, U> {
             inner: SyncWrap(Rc::new(RefCell::new(inner))),
         })
     }
+
+    /// Upper slab slot size in bytes.
+    pub const fn upper_slot_size() -> usize {
+        U
+    }
 }
 
 #[cfg(all(test, loom))]
@@ -821,6 +826,40 @@ mod tests {
         assert!(matches!(result, Err(AllocError::InvalidFree(_, _))));
     }
 
+    #[test]
+    fn test_slab_multi_slot_alloc_near_end() {
+        let mut slab = make_slab::<256>(1792); // 7 slots
+        let a0 = slab.alloc(256).unwrap();
+        let a1 = slab.alloc(256).unwrap();
+        let _a2 = slab.alloc(256).unwrap();
+        let _a3 = slab.alloc(256).unwrap();
+        let _a4 = slab.alloc(256).unwrap();
+        let _a5 = slab.alloc(256).unwrap();
+        let _a6 = slab.alloc(256).unwrap();
+
+        slab.dealloc(a0).unwrap();
+        slab.dealloc(a1).unwrap();
+
+        // 2-slot run fits at indices 0..2 but the search visits index 6
+        // (a free zero) first if slots 0-1 are not found before it.
+        // Actually slots 0-1 are free, so it should find them.
+        let run = slab.alloc(300).unwrap(); // needs 2 slots
+        assert_eq!(run.len, 512);
+    }
+
+    #[test]
+    fn test_slab_multi_slot_alloc_no_room_at_end() {
+        // Only the last slot is free but a 2-slot run is requested.
+        // find_slots must not panic when checking beyond the bitset.
+        let mut slab = make_slab::<256>(1792); // 7 slots
+        let allocs: Vec<_> = (0..7).map(|_| slab.alloc(256).unwrap()).collect();
+        // Free only the last slot (index 6)
+        slab.dealloc(allocs[6]).unwrap();
+
+        let result = slab.alloc(300); // needs 2 slots, only 1 free
+        assert!(matches!(result, Err(AllocError::NoSpace)));
+    }
+
     #[test]
     fn test_slab_free_invalid_address() {
         let mut slab = make_slab::<256>(1024);

src/hyperlight_common/src/virtq/producer.rs

@@ -125,7 +125,8 @@ pub struct VirtqProducer<M, N, P> {
     inner: RingProducer<M>,
     notifier: N,
     pool: P,
-    inflight: Vec<Option<Inflight>>,
+    next_token: u32,
+    inflight: Vec<Option<(Token, Inflight)>>,
     pending: VecDeque<RecvCompletion>,
 }
 
@@ -152,6 +153,7 @@ where
             pool,
             notifier,
             inflight,
+            next_token: 0,
             pending: VecDeque::new(),
         }
     }
@@ -218,13 +220,20 @@ where
         };
 
         let id = used.id as usize;
-        let inf = self
+        let (token, inf) = self
             .inflight
             .get_mut(id)
             .ok_or(VirtqError::InvalidState)?
             .take()
             .ok_or(VirtqError::InvalidState)?;
 
+        // the token's descriptor ID must match the ring's
+        debug_assert_eq!(
+            token.1, used.id,
+            "ring returned desc_id={} but inflight slot {} has token with desc_id={}",
+            used.id, id, token.1,
+        );
+
         let written = used.len as usize;
 
         // Free entry buffers (request data no longer needed)
@@ -250,10 +259,7 @@ where
             None => Bytes::new(),
         };
 
-        Ok(Some(RecvCompletion {
-            token: Token(used.id),
-            data,
-        }))
+        Ok(Some(RecvCompletion { token, data }))
     }
 
     /// Drain all available completions, calling the provided closure for each.
@@ -310,6 +316,9 @@ where
         let chain = inflight.try_into_chain(written)?;
         let id = self.inner.submit_available(&chain)?;
 
+        let token = Token(self.next_token, id);
+        self.next_token = self.next_token.wrapping_add(1);
+
         let slot = self
             .inflight
             .get_mut(id as usize)
@@ -319,7 +328,7 @@ where
             return Err(VirtqError::InvalidState);
         }
 
-        *slot = Some(inflight);
+        *slot = Some((token, inflight));
 
         let should_notify = self.inner.should_notify_since(cursor_before)?;
 
@@ -336,7 +345,7 @@ where
             });
         }
 
-        Ok(Token(id))
+        Ok(token)
     }
 
     /// Signal backpressure to the consumer.
@@ -474,12 +483,18 @@ where
                 .slot_addr(pos as usize)
                 .ok_or(VirtqError::InvalidState)?;
 
-            self.inflight[id as usize] = Some(Inflight::WriteOnly {
-                completion: Allocation {
-                    addr,
-                    len: slot_size,
+            let token = Token(self.next_token, id);
+            self.next_token = self.next_token.wrapping_add(1);
+
+            self.inflight[id as usize] = Some((
+                token,
+                Inflight::WriteOnly {
+                    completion: Allocation {
+                        addr,
+                        len: slot_size,
+                    },
                 },
-            });
+            ));
 
             ids.push(id);
         }
@@ -869,7 +884,7 @@ mod tests {
         // Ring should still be fully usable
         let se = producer.chain().entry(64).completion(128).build().unwrap();
         let tok = producer.submit(se).unwrap();
-        assert!(tok.0 < 16);
+        assert!(tok.1 < 16);
     }
 
     #[test]
@@ -885,7 +900,7 @@ mod tests {
         // Ring should still be fully usable
         let se = producer.chain().entry(64).completion(128).build().unwrap();
         let tok = producer.submit(se).unwrap();
-        assert!(tok.0 < 16);
+        assert!(tok.1 < 16);
     }
 
     #[test]