Skip to content

Commit c6660b6

Browse files
ShadowCurseShadowCurse
committed
refactor(virtio): create VirtioDeviceType enum
Instead of using `u32` values directly from generated `virtio_ids` create an enum with all types we actually use and start using them. Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>
1 parent 3344d0c commit c6660b6

File tree

24 files changed

+152
-123
lines changed

24 files changed

+152
-123
lines changed

‎src/vmm/src/builder.rs‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -753,7 +753,7 @@ pub(crate) mod tests {
753753
use super::*;
754754
use crate::device_manager::tests::default_device_manager;
755755
use crate::devices::virtio::block::CacheType;
756-
use crate::devices::virtio::generated::virtio_ids;
756+
use crate::devices::virtio::device::VirtioDeviceType;
757757
use crate::devices::virtio::rng::device::ENTROPY_DEV_ID;
758758
use crate::devices::virtio::vsock::VSOCK_DEV_ID;
759759
use crate::mmds::data_store::{Mmds, MmdsVersion};
@@ -955,7 +955,7 @@ pub(crate) mod tests {
955955

956956
assert!(
957957
vmm.device_manager
958-
.get_virtio_device(virtio_ids::VIRTIO_ID_VSOCK, &vsock_dev_id)
958+
.get_virtio_device(VirtioDeviceType::Vsock.as_u32(), &vsock_dev_id)
959959
.is_some()
960960
);
961961
}
@@ -980,7 +980,7 @@ pub(crate) mod tests {
980980

981981
assert!(
982982
vmm.device_manager
983-
.get_virtio_device(virtio_ids::VIRTIO_ID_RNG, ENTROPY_DEV_ID)
983+
.get_virtio_device(VirtioDeviceType::Rng.as_u32(), ENTROPY_DEV_ID)
984984
.is_some()
985985
);
986986
}
@@ -1044,7 +1044,7 @@ pub(crate) mod tests {
10441044

10451045
assert!(
10461046
vmm.device_manager
1047-
.get_virtio_device(virtio_ids::VIRTIO_ID_BALLOON, BALLOON_DEV_ID)
1047+
.get_virtio_device(VirtioDeviceType::Balloon.as_u32(), BALLOON_DEV_ID)
10481048
.is_some()
10491049
);
10501050
}
@@ -1095,7 +1095,7 @@ pub(crate) mod tests {
10951095
assert!(cmdline_contains(&cmdline, "root=/dev/vda ro"));
10961096
assert!(
10971097
vmm.device_manager
1098-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
1098+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), drive_id.as_str())
10991099
.is_some()
11001100
);
11011101
}
@@ -1116,7 +1116,7 @@ pub(crate) mod tests {
11161116
assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
11171117
assert!(
11181118
vmm.device_manager
1119-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
1119+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), drive_id.as_str())
11201120
.is_some()
11211121
);
11221122
}
@@ -1138,7 +1138,7 @@ pub(crate) mod tests {
11381138
assert!(!cmdline_contains(&cmdline, "root=/dev/vda"));
11391139
assert!(
11401140
vmm.device_manager
1141-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
1141+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), drive_id.as_str())
11421142
.is_some()
11431143
);
11441144
}
@@ -1175,17 +1175,17 @@ pub(crate) mod tests {
11751175
assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
11761176
assert!(
11771177
vmm.device_manager
1178-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, "root")
1178+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), "root")
11791179
.is_some()
11801180
);
11811181
assert!(
11821182
vmm.device_manager
1183-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, "secondary")
1183+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), "secondary")
11841184
.is_some()
11851185
);
11861186
assert!(
11871187
vmm.device_manager
1188-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, "third")
1188+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), "third")
11891189
.is_some()
11901190
);
11911191

@@ -1214,7 +1214,7 @@ pub(crate) mod tests {
12141214
assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
12151215
assert!(
12161216
vmm.device_manager
1217-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
1217+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), drive_id.as_str())
12181218
.is_some()
12191219
);
12201220
}
@@ -1235,7 +1235,7 @@ pub(crate) mod tests {
12351235
assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 ro"));
12361236
assert!(
12371237
vmm.device_manager
1238-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
1238+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), drive_id.as_str())
12391239
.is_some()
12401240
);
12411241
}
@@ -1256,7 +1256,7 @@ pub(crate) mod tests {
12561256
assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
12571257
assert!(
12581258
vmm.device_manager
1259-
.get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
1259+
.get_virtio_device(VirtioDeviceType::Block.as_u32(), drive_id.as_str())
12601260
.is_some()
12611261
);
12621262
}
@@ -1279,7 +1279,7 @@ pub(crate) mod tests {
12791279
assert!(cmdline_contains(&cmdline, "root=/dev/pmem0 ro"));
12801280
assert!(
12811281
vmm.device_manager
1282-
.get_virtio_device(virtio_ids::VIRTIO_ID_PMEM, id.as_str())
1282+
.get_virtio_device(VirtioDeviceType::Pmem.as_u32(), id.as_str())
12831283
.is_some()
12841284
);
12851285
}

‎src/vmm/src/device_manager/mmio.rs‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ impl MMIODeviceManager {
182182
{
183183
let mmio_device = device.inner.lock().expect("Poisoned lock");
184184
let locked_device = mmio_device.locked_device();
185-
identifier = (locked_device.device_type(), device_id);
185+
identifier = (locked_device.device_type().as_u32(), device_id);
186186
for (i, queue_evt) in locked_device.queue_events().iter().enumerate() {
187187
let io_addr = IoEventAddress::Mmio(
188188
device.resources.addr + u64::from(crate::devices::virtio::NOTIFY_REG_OFFSET),
@@ -430,7 +430,7 @@ pub(crate) mod tests {
430430

431431
use super::*;
432432
use crate::devices::virtio::ActivateError;
433-
use crate::devices::virtio::device::VirtioDevice;
433+
use crate::devices::virtio::device::{VirtioDevice, VirtioDeviceType};
434434
use crate::devices::virtio::queue::Queue;
435435
use crate::devices::virtio::transport::VirtioInterrupt;
436436
use crate::devices::virtio::transport::mmio::IrqTrigger;
@@ -454,7 +454,7 @@ pub(crate) mod tests {
454454
let mmio_device = MmioTransport::new(guest_mem, interrupt, device.clone(), false);
455455
self.register_mmio_virtio_for_boot(vm, dev_id.to_string(), mmio_device, cmdline)?;
456456
Ok(self
457-
.get_virtio_device(device.lock().unwrap().device_type(), dev_id)
457+
.get_virtio_device(device.lock().unwrap().device_type().as_u32(), dev_id)
458458
.unwrap()
459459
.resources
460460
.addr)
@@ -491,7 +491,7 @@ pub(crate) mod tests {
491491
}
492492

493493
impl VirtioDevice for DummyDevice {
494-
impl_device_type!(0);
494+
impl_device_type!(VirtioDeviceType::Net);
495495

496496
fn avail_features(&self) -> u64 {
497497
0
@@ -642,7 +642,7 @@ pub(crate) mod tests {
642642
#[test]
643643
fn test_dummy_device() {
644644
let dummy = DummyDevice::new();
645-
assert_eq!(dummy.device_type(), 0);
645+
assert_eq!(dummy.device_type(), VirtioDeviceType::Net);
646646
assert_eq!(dummy.queues().len(), QUEUE_SIZES.len());
647647
}
648648

@@ -670,23 +670,31 @@ pub(crate) mod tests {
670670
let addr = device_manager
671671
.register_virtio_test_device(&vm, vm.guest_memory().clone(), dummy, &mut cmdline, &id)
672672
.unwrap();
673-
assert!(device_manager.get_virtio_device(type_id, &id).is_some());
673+
assert!(
674+
device_manager
675+
.get_virtio_device(type_id.as_u32(), &id)
676+
.is_some()
677+
);
674678
assert_eq!(
675679
addr,
676-
device_manager.virtio_devices[&(type_id, id.clone())]
680+
device_manager.virtio_devices[&(type_id.as_u32(), id.clone())]
677681
.resources
678682
.addr
679683
);
680684
assert_eq!(
681685
crate::arch::GSI_LEGACY_START,
682-
device_manager.virtio_devices[&(type_id, id)]
686+
device_manager.virtio_devices[&(type_id.as_u32(), id)]
683687
.resources
684688
.gsi
685689
.unwrap()
686690
);
687691

688692
let id = "bar";
689-
assert!(device_manager.get_virtio_device(type_id, id).is_none());
693+
assert!(
694+
device_manager
695+
.get_virtio_device(type_id.as_u32(), id)
696+
.is_none()
697+
);
690698

691699
let dummy2 = Arc::new(Mutex::new(DummyDevice::new()));
692700
let id2 = String::from("foo2");
@@ -697,7 +705,7 @@ pub(crate) mod tests {
697705
let mut count = 0;
698706
let _: Result<(), MmioError> =
699707
device_manager.for_each_virtio_device(|devtype, devid, _| {
700-
assert_eq!(*devtype, type_id);
708+
assert_eq!(*devtype, type_id.as_u32());
701709
match devid.as_str() {
702710
"foo" => count += 1,
703711
"foo2" => count += 2,

‎src/vmm/src/device_manager/mod.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ use crate::devices::legacy::RTCDevice;
3131
use crate::devices::legacy::serial::SerialOut;
3232
use crate::devices::legacy::{IER_RDA_BIT, IER_RDA_OFFSET, SerialDevice};
3333
use crate::devices::pseudo::BootTimer;
34-
use crate::devices::virtio::device::VirtioDevice;
34+
use crate::devices::virtio::device::{VirtioDevice, VirtioDeviceType};
3535
use crate::devices::virtio::transport::mmio::{IrqTrigger, MmioTransport};
3636
use crate::resources::VmResources;
3737
use crate::snapshot::Persist;
@@ -336,7 +336,8 @@ impl DeviceManager {
336336
device_id: &str,
337337
) -> Option<Arc<Mutex<dyn VirtioDevice>>> {
338338
if self.pci_devices.pci_segment.is_some() {
339-
let pci_device = self.pci_devices.get_virtio_device(virtio_type, device_id)?;
339+
let device_type = VirtioDeviceType::from_u32(virtio_type)?;
340+
let pci_device = self.pci_devices.get_virtio_device(device_type, device_id)?;
340341
Some(
341342
pci_device
342343
.lock()
@@ -365,7 +366,7 @@ impl DeviceManager {
365366
T: VirtioDevice + 'static + Debug,
366367
F: FnOnce(&mut T) -> R,
367368
{
368-
if let Some(device) = self.get_virtio_device(T::const_device_type(), id) {
369+
if let Some(device) = self.get_virtio_device(T::const_device_type().as_u32(), id) {
369370
let mut dev = device.lock().expect("Poisoned lock");
370371
Ok(f(dev
371372
.as_mut_any()

‎src/vmm/src/device_manager/pci_mngr.rs‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,7 @@ use crate::devices::virtio::balloon::Balloon;
1616
use crate::devices::virtio::balloon::persist::{BalloonConstructorArgs, BalloonState};
1717
use crate::devices::virtio::block::device::Block;
1818
use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
19-
use crate::devices::virtio::device::VirtioDevice;
20-
use crate::devices::virtio::generated::virtio_ids;
19+
use crate::devices::virtio::device::{VirtioDevice, VirtioDeviceType};
2120
use crate::devices::virtio::mem::VirtioMem;
2221
use crate::devices::virtio::mem::persist::{VirtioMemConstructorArgs, VirtioMemState};
2322
use crate::devices::virtio::net::Net;
@@ -48,7 +47,7 @@ pub struct PciDevices {
4847
/// PCIe segment of the VMM, if PCI is enabled. We currently support a single PCIe segment.
4948
pub pci_segment: Option<PciSegment>,
5049
/// All VirtIO PCI devices of the system
51-
pub virtio_devices: HashMap<(u32, String), Arc<Mutex<VirtioPciDevice>>>,
50+
pub virtio_devices: HashMap<(VirtioDeviceType, String), Arc<Mutex<VirtioPciDevice>>>,
5251
}
5352

5453
#[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -120,7 +119,7 @@ impl PciDevices {
120119
debug!("Allocating BDF: {pci_device_bdf:?} for device");
121120
let mem = vm.guest_memory().clone();
122121

123-
let device_type: u32 = device.lock().expect("Poisoned lock").device_type();
122+
let device_type = device.lock().expect("Poisoned lock").device_type();
124123

125124
// Allocate one MSI vector per queue, plus one for configuration
126125
let msix_num =
@@ -172,7 +171,7 @@ impl PciDevices {
172171
) -> Result<(), PciManagerError> {
173172
// We should only be reaching this point if PCI is enabled
174173
let pci_segment = self.pci_segment.as_ref().unwrap();
175-
let device_type: u32 = device.lock().expect("Poisoned lock").device_type();
174+
let device_type = device.lock().expect("Poisoned lock").device_type();
176175

177176
let virtio_device = Arc::new(Mutex::new(VirtioPciDevice::new_from_state(
178177
device_id.to_string(),
@@ -207,7 +206,7 @@ impl PciDevices {
207206
/// Gets the specified device.
208207
pub fn get_virtio_device(
209208
&self,
210-
device_type: u32,
209+
device_type: VirtioDeviceType,
211210
device_id: &str,
212211
) -> Option<&Arc<Mutex<VirtioPciDevice>>> {
213212
self.virtio_devices
@@ -290,7 +289,7 @@ impl<'a> Persist<'a> for PciDevices {
290289
let pci_device_bdf = transport_state.pci_device_bdf.into();
291290

292291
match locked_virtio_dev.device_type() {
293-
virtio_ids::VIRTIO_ID_BALLOON => {
292+
VirtioDeviceType::Balloon => {
294293
let balloon_device = locked_virtio_dev
295294
.as_any()
296295
.downcast_ref::<Balloon>()
@@ -305,7 +304,7 @@ impl<'a> Persist<'a> for PciDevices {
305304
transport_state,
306305
});
307306
}
308-
virtio_ids::VIRTIO_ID_BLOCK => {
307+
VirtioDeviceType::Block => {
309308
let block_dev = locked_virtio_dev
310309
.as_mut_any()
311310
.downcast_mut::<Block>()
@@ -326,7 +325,7 @@ impl<'a> Persist<'a> for PciDevices {
326325
});
327326
}
328327
}
329-
virtio_ids::VIRTIO_ID_NET => {
328+
VirtioDeviceType::Net => {
330329
let net_dev = locked_virtio_dev
331330
.as_mut_any()
332331
.downcast_mut::<Net>()
@@ -348,7 +347,7 @@ impl<'a> Persist<'a> for PciDevices {
348347
transport_state,
349348
})
350349
}
351-
virtio_ids::VIRTIO_ID_VSOCK => {
350+
VirtioDeviceType::Vsock => {
352351
let vsock_dev = locked_virtio_dev
353352
.as_mut_any()
354353
// Currently, VsockUnixBackend is the only implementation of VsockBackend.
@@ -379,7 +378,7 @@ impl<'a> Persist<'a> for PciDevices {
379378
transport_state,
380379
});
381380
}
382-
virtio_ids::VIRTIO_ID_RNG => {
381+
VirtioDeviceType::Rng => {
383382
let rng_dev = locked_virtio_dev
384383
.as_mut_any()
385384
.downcast_mut::<Entropy>()
@@ -393,7 +392,7 @@ impl<'a> Persist<'a> for PciDevices {
393392
transport_state,
394393
})
395394
}
396-
virtio_ids::VIRTIO_ID_PMEM => {
395+
VirtioDeviceType::Pmem => {
397396
let pmem_dev = locked_virtio_dev
398397
.as_mut_any()
399398
.downcast_mut::<Pmem>()
@@ -406,7 +405,7 @@ impl<'a> Persist<'a> for PciDevices {
406405
transport_state,
407406
});
408407
}
409-
virtio_ids::VIRTIO_ID_MEM => {
408+
VirtioDeviceType::Mem => {
410409
let mem_dev = locked_virtio_dev
411410
.as_mut_any()
412411
.downcast_mut::<VirtioMem>()
@@ -420,7 +419,6 @@ impl<'a> Persist<'a> for PciDevices {
420419
transport_state,
421420
})
422421
}
423-
_ => unreachable!(),
424422
}
425423
}
426424

0 commit comments

Comments
 (0)