From de0e0360a55c28e56932e8e9f3514466b28b04e7 Mon Sep 17 00:00:00 2001 From: Changyuan Lyu Date: Sat, 10 Oct 2026 23:45:28 -0700 Subject: [PATCH] refactor(virtio): use atomic types in queue structs Assisted-by: Antigravity:Gemini-4-Argon Signed-off-by: Changyuan Lyu --- alioth/src/virtio/queue/packed.rs | 42 +++++++++++--------------- alioth/src/virtio/queue/packed_test.rs | 33 ++++++++++---------- alioth/src/virtio/queue/split.rs | 37 ++++++++++------------- alioth/src/virtio/queue/split_test.rs | 34 +++++++++++---------- 4 files changed, 69 insertions(+), 77 deletions(-) diff --git a/alioth/src/virtio/queue/packed.rs b/alioth/src/virtio/queue/packed.rs index 61dabeed..a4e2126a 100644 --- a/alioth/src/virtio/queue/packed.rs +++ b/alioth/src/virtio/queue/packed.rs @@ -16,7 +16,7 @@ use std::marker::PhantomData; use std::sync::atomic::{AtomicU16, Ordering}; use bitfield::bitfield; -use zerocopy::{FromBytes, Immutable, IntoBytes}; +use zerocopy::{FromBytes, IntoBytes}; use crate::consts; use crate::mem::mapped::Ram; @@ -24,12 +24,12 @@ use crate::virtio::queue::{DescChain, DescFlag, QueueReg, VirtQueue}; use crate::virtio::{Result, error}; #[repr(C, align(16))] -#[derive(Debug, Copy, Clone, Default, FromBytes, Immutable, IntoBytes)] +#[derive(Debug, Default, FromBytes, IntoBytes)] struct Desc { pub addr: u64, pub len: u32, pub id: u16, - pub flag: u16, + pub flag: AtomicU16, } bitfield! { @@ -80,8 +80,8 @@ consts! { #[repr(C)] struct DescEvent { - index: WrappedIndex, - flag: EventFlag, + index: AtomicU16, + flag: AtomicU16, } #[derive(Debug)] @@ -142,10 +142,8 @@ impl<'m> VirtQueue<'m> for PackedQueue<'m> { const INIT_INDEX: WrappedIndex = WrappedIndex::INIT; fn desc_avail(&self, index: WrappedIndex) -> bool { - let flag = unsafe { - AtomicU16::from_ptr(&raw mut (*self.desc.offset(index.offset() as isize)).flag) - .load(Ordering::Acquire) - }; + let flag = + unsafe { &(*self.desc.offset(index.offset() as isize)).flag }.load(Ordering::Acquire); self.flag_is_avail(DescFlag::from_bits_retain(flag), index.wrap_counter()) } @@ -159,12 +157,12 @@ impl<'m> VirtQueue<'m> for PackedQueue<'m> { let mut offset = index.offset(); let id = loop { let desc = unsafe { self.desc.offset(offset as isize).read_volatile() }; - let flag = DescFlag::from_bits_retain(desc.flag); + let flag = DescFlag::from_bits_retain(desc.flag.into_inner()); if flag.contains(DescFlag::INDIRECT) { for i in 0..(desc.len as usize / size_of::()) { let addr = desc.addr + (i * size_of::()) as u64; let desc: Desc = ram.read_t(addr)?; - let flag = DescFlag::from_bits_retain(desc.flag); + let flag = DescFlag::from_bits_retain(desc.flag.into_inner()); if flag.contains(DescFlag::WRITE) { writeable.push((desc.addr, desc.len as u64)); } else { @@ -195,10 +193,10 @@ impl<'m> VirtQueue<'m> for PackedQueue<'m> { let first = self.desc.offset(index.offset() as isize); (&raw mut (*first).id).write_volatile(id); (&raw mut (*first).len).write_volatile(len); - let flag_ptr = AtomicU16::from_ptr(&raw mut (*first).flag); - let mut flag = DescFlag::from_bits_retain(flag_ptr.load(Ordering::Relaxed)); + let flag_ref = &(*first).flag; + let mut flag = DescFlag::from_bits_retain(flag_ref.load(Ordering::Relaxed)); self.set_flag_used(&mut flag, index.wrap_counter()); - flag_ptr.store(flag.bits(), Ordering::Release); + flag_ref.store(flag.bits(), Ordering::Release); } } @@ -208,20 +206,16 @@ impl<'m> VirtQueue<'m> for PackedQueue<'m> { } else { EventFlag::DISABLE }; - unsafe { - AtomicU16::from_ptr(&raw mut (*self.notification).flag.0) - .store(flag.raw(), Ordering::Relaxed); - } + unsafe { &*self.notification } + .flag + .store(flag.raw(), Ordering::Relaxed) } fn interrupt_enabled(&self, index: Self::Index, delta: u16) -> bool { - let flag = EventFlag(unsafe { - AtomicU16::from_ptr(&raw mut (*self.interrupt).flag.0).load(Ordering::Acquire) - }); + let flag = EventFlag(unsafe { &*self.interrupt }.flag.load(Ordering::Acquire)); if self.enable_event_idx && flag == EventFlag::DESC { - let event_index = WrappedIndex(unsafe { - AtomicU16::from_ptr(&raw mut (*self.interrupt).index.0).load(Ordering::Relaxed) - }); + let event_index = + WrappedIndex(unsafe { &*self.interrupt }.index.load(Ordering::Relaxed)); let prev_used_index = index.wrapping_sub(delta, self.size); let base = prev_used_index.offset(); let end = base + delta; diff --git a/alioth/src/virtio/queue/packed_test.rs b/alioth/src/virtio/queue/packed_test.rs index 95d15183..4d88a391 100644 --- a/alioth/src/virtio/queue/packed_test.rs +++ b/alioth/src/virtio/queue/packed_test.rs @@ -95,7 +95,7 @@ impl<'m> VirtQueueGuest<'m> for PackedQueue<'m> { desc.addr = *addr; desc.len = *len; desc.id = id; - desc.flag = flag.bits(); + desc.flag.store(flag.bits(), Ordering::Release); } ids.len() as u16 } @@ -105,8 +105,8 @@ impl<'m> VirtQueueGuest<'m> for PackedQueue<'m> { index: Self::Index, chains: &HashMap>, ) -> Option { - let desc = unsafe { &mut *self.desc.offset(index.offset() as isize) }; - let flag = DescFlag::from_bits_retain(desc.flag); + let desc = unsafe { &*self.desc.offset(index.offset() as isize) }; + let flag = DescFlag::from_bits_retain(desc.flag.load(Ordering::Acquire)); if !self.flag_is_used(flag, index.wrap_counter()) { return None; } @@ -118,9 +118,7 @@ impl<'m> VirtQueueGuest<'m> for PackedQueue<'m> { } fn notification_enabled(&self) -> bool { - let flag = EventFlag(unsafe { - AtomicU16::from_ptr(&raw mut (*self.notification).flag.0).load(Ordering::Relaxed) - }); + let flag = EventFlag(unsafe { &*self.notification }.flag.load(Ordering::Relaxed)); flag != EventFlag::DISABLE } @@ -132,12 +130,9 @@ impl<'m> VirtQueueGuest<'m> for PackedQueue<'m> { } else { EventFlag::ENABLE }; - unsafe { - AtomicU16::from_ptr(&raw mut (*self.interrupt).index.0) - .store(index.0, Ordering::Relaxed); - AtomicU16::from_ptr(&raw mut (*self.interrupt).flag.0) - .store(flag.raw(), Ordering::Release); - } + let interrupt = unsafe { &*self.interrupt }; + interrupt.index.store(index.0, Ordering::Relaxed); + interrupt.flag.store(flag.raw(), Ordering::Release); } } @@ -234,9 +229,15 @@ fn enable_notification() { let q = PackedQueue::new(reg, &ram, false).unwrap().unwrap(); q.enable_notification(false); - assert_eq!(unsafe { &*q.notification }.flag, EventFlag::DISABLE); + assert_eq!( + unsafe { &*q.notification }.flag.load(Ordering::Acquire), + EventFlag::DISABLE.raw() + ); q.enable_notification(true); - assert_eq!(unsafe { &*q.notification }.flag, EventFlag::ENABLE); + assert_eq!( + unsafe { &*q.notification }.flag.load(Ordering::Acquire), + EventFlag::ENABLE.raw() + ); } #[rstest] @@ -270,8 +271,8 @@ fn is_interrupt_enabled( .unwrap(); *unsafe { &mut *q.interrupt } = DescEvent { - index: WrappedIndex(event_index), - flag: event_flag, + index: AtomicU16::new(event_index), + flag: AtomicU16::new(event_flag.raw()), }; assert_eq!( diff --git a/alioth/src/virtio/queue/split.rs b/alioth/src/virtio/queue/split.rs index efc347d4..87d2131d 100644 --- a/alioth/src/virtio/queue/split.rs +++ b/alioth/src/virtio/queue/split.rs @@ -40,10 +40,10 @@ bitflags! { } #[repr(C, align(2))] -#[derive(Debug, Clone, Layout, Immutable, FromBytes, IntoBytes)] +#[derive(Debug, Layout)] pub struct AvailHeader { - flags: u16, - idx: u16, + flags: AtomicU16, + idx: AtomicU16, } bitflags! { @@ -53,10 +53,10 @@ bitflags! { } #[repr(C, align(4))] -#[derive(Debug, Clone, Layout)] +#[derive(Debug, Layout)] pub struct UsedHeader { - flags: u16, - idx: u16, + flags: AtomicU16, + idx: AtomicU16, } #[repr(C)] @@ -71,34 +71,32 @@ pub struct SplitQueue<'m> { size: u16, avail_hdr: *mut AvailHeader, avail_ring: *mut u16, - used_event: Option<*mut u16>, + used_event: Option<*mut AtomicU16>, used_hdr: *mut UsedHeader, used_ring: *mut UsedElem, - avail_event: Option<*mut u16>, + avail_event: Option<*mut AtomicU16>, desc: *mut Desc, _phantom: PhantomData<&'m ()>, } impl SplitQueue<'_> { pub fn avail_index(&self) -> u16 { - unsafe { AtomicU16::from_ptr(&raw mut (*self.avail_hdr).idx).load(Ordering::Acquire) } + unsafe { &*self.avail_hdr }.idx.load(Ordering::Acquire) } pub fn set_used_index(&self, val: u16) { - unsafe { - AtomicU16::from_ptr(&raw mut (*self.used_hdr).idx).store(val, Ordering::Release); - } + unsafe { &*self.used_hdr }.idx.store(val, Ordering::Release) } pub fn used_event(&self) -> Option { self.used_event - .map(|event| unsafe { AtomicU16::from_ptr(event).load(Ordering::Relaxed) }) + .map(|event| unsafe { &*event }.load(Ordering::Relaxed)) } pub fn set_avail_event(&self, op: impl FnOnce(&AtomicU16)) -> bool { match self.avail_event { Some(avail_event) => { - op(unsafe { AtomicU16::from_ptr(avail_event) }); + op(unsafe { &*avail_event }); true } None => false, @@ -106,16 +104,13 @@ impl SplitQueue<'_> { } pub fn set_flag_notification(&self, enabled: bool) { - unsafe { - AtomicU16::from_ptr(&raw mut (*self.used_hdr).flags) - .store((!enabled) as u16, Ordering::Relaxed); - } + unsafe { &*self.used_hdr } + .flags + .store(!enabled as u16, Ordering::Relaxed); } pub fn flag_interrupt_enabled(&self) -> bool { - unsafe { - AtomicU16::from_ptr(&raw mut (*self.avail_hdr).flags).load(Ordering::Relaxed) == 0 - } + unsafe { &*self.avail_hdr }.flags.load(Ordering::Relaxed) == 0 } fn get_desc(&self, id: u16) -> Result { diff --git a/alioth/src/virtio/queue/split_test.rs b/alioth/src/virtio/queue/split_test.rs index 239befa5..2d9a71cf 100644 --- a/alioth/src/virtio/queue/split_test.rs +++ b/alioth/src/virtio/queue/split_test.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::collections::HashMap; -use std::sync::atomic::{AtomicU16, Ordering}; +use std::sync::atomic::Ordering; use assert_matches::assert_matches; @@ -54,12 +54,14 @@ impl<'m> VirtQueueGuest<'m> for SplitQueue<'m> { let avail_idx = self.avail_index(); assert_eq!(index, avail_idx); *unsafe { &mut *self.avail_ring.offset((avail_idx % self.size) as isize) } = last_id; - unsafe { &mut *self.avail_hdr }.idx = avail_idx.wrapping_add(1); + unsafe { &*self.avail_hdr } + .idx + .store(avail_idx.wrapping_add(1), Ordering::Release); 1 } fn get_used(&mut self, index: u16, _: &HashMap>) -> Option { - let used_idx = unsafe { &*self.used_hdr }.idx; + let used_idx = unsafe { &*self.used_hdr }.idx.load(Ordering::Acquire); if !(index < used_idx || index - used_idx >= !(self.size - 1)) { return None; } @@ -71,12 +73,10 @@ impl<'m> VirtQueueGuest<'m> for SplitQueue<'m> { fn notification_enabled(&self) -> bool { match self.avail_event { - Some(avail_event) => unsafe { - AtomicU16::from_ptr(avail_event).load(Ordering::Relaxed) == self.avail_index() - }, - None => unsafe { - AtomicU16::from_ptr(&raw mut (*self.used_hdr).flags).load(Ordering::Relaxed) == 0 - }, + Some(avail_event) => { + unsafe { &*avail_event }.load(Ordering::Relaxed) == self.avail_index() + } + None => unsafe { &*self.used_hdr }.flags.load(Ordering::Relaxed) == 0, } } @@ -87,12 +87,11 @@ impl<'m> VirtQueueGuest<'m> for SplitQueue<'m> { } else { index.wrapping_sub(1) }; - unsafe { AtomicU16::from_ptr(used_event).store(idx, Ordering::Relaxed) }; + unsafe { &*used_event }.store(idx, Ordering::Relaxed); } else { - unsafe { - AtomicU16::from_ptr(&raw mut (*self.avail_hdr).flags) - .store((!enabled) as u16, Ordering::Relaxed) - }; + unsafe { &*self.avail_hdr } + .flags + .store(!enabled as u16, Ordering::Relaxed) } } } @@ -186,9 +185,12 @@ fn event_idx_enabled() { let ram = ram_bus.lock_layout(); let reg = &queues[0]; let q = SplitQueue::new(reg, &ram, true).unwrap().unwrap(); - unsafe { *q.used_event.unwrap() = 1 }; + unsafe { &*q.used_event.unwrap() }.store(1, Ordering::Release); assert_eq!(q.used_event(), Some(1)); assert!(q.set_avail_event(|event| event.store(12, Ordering::Relaxed))); - assert_eq!(unsafe { *q.avail_event.unwrap() }, 12); + assert_eq!( + unsafe { &*q.avail_event.unwrap() }.load(Ordering::Acquire), + 12 + ); }