use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use vm_memory::{Bytes, GuestAddress, GuestMemoryMmap};
use vmm_sys_util::eventfd::EventFd;
use virtio_bindings::virtio_blk::{
VIRTIO_BLK_S_IOERR, VIRTIO_BLK_S_UNSUPP, VIRTIO_BLK_T_FLUSH, VIRTIO_BLK_T_GET_ID,
VIRTIO_BLK_T_IN, VIRTIO_BLK_T_OUT,
};
use virtio_bindings::virtio_config::VIRTIO_CONFIG_S_NEEDS_RESET;
use virtio_bindings::virtio_mmio::{VIRTIO_MMIO_INT_CONFIG, VIRTIO_MMIO_INT_VRING};
use virtio_queue::Error as VirtioQueueError;
use virtio_queue::{QueueOwnedT, QueueT};
use super::{
Backing, BlkQueue, BlkWorkerState, ChainDescriptor, NUM_QUEUES, REQ_QUEUE,
VIRTIO_BLK_OUTHDR_SIZE, VIRTIO_BLK_SECTOR_SIZE, VIRTIO_BLK_SEG_MAX, VIRTIO_BLK_SIZE_MAX,
VirtioBlk, VirtioBlkOutHdr, publish_completion,
};
use crate::vmm::PiMutex;
use crate::vmm::virtio_msix::{IrqSource, MsixState};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum DrainOutcome {
Done,
ThrottleStalled { wait_nanos: u64 },
}
pub(crate) fn drain_bracket_impl(
state: &mut BlkWorkerState,
queues: &mut [BlkQueue; NUM_QUEUES],
mem: &GuestMemoryMmap,
irq_evt: &EventFd,
interrupt_status: &AtomicU32,
device_status: &AtomicU32,
msix: Option<&Arc<PiMutex<MsixState>>>,
) -> DrainOutcome {
if !queues[REQ_QUEUE].ready() {
return DrainOutcome::Done;
}
if state.queue_poisoned {
return DrainOutcome::Done;
}
let mut signal_needed = false;
let mut stall_outcome: Option<u64> = None;
'outer: loop {
if let Err(e) = queues[REQ_QUEUE].disable_notification(mem) {
tracing::warn!(%e, "virtio-blk disable_notification failed");
}
loop {
let iter_outcome = {
let q = &mut queues[REQ_QUEUE];
#[allow(unused_mut)]
let mut guard = q.lock();
match guard.iter(mem) {
Ok(mut iter) => Ok(iter.next()),
Err(e) => Err(e),
}
};
let chain = match iter_outcome {
Ok(Some(c)) => c,
Ok(None) => break,
Err(VirtioQueueError::InvalidAvailRingIndex) => {
state.queue_poisoned = true;
state.counters.record_invalid_avail_idx();
tracing::warn!(
"virtio-blk avail.idx exceeds next_avail by more \
than queue.size (virtio-v1.2 §2.7.13.3 \
avail.idx-distance violation); poisoning queue \
until guest reset"
);
debug_assert_eq!(
VIRTIO_CONFIG_S_NEEDS_RESET.count_ones(),
1,
"VIRTIO_CONFIG_S_NEEDS_RESET must be a single bit; \
multi-bit fetch_or would expand set_status's CAS retry universe"
);
device_status.fetch_or(VIRTIO_CONFIG_S_NEEDS_RESET, Ordering::SeqCst);
interrupt_status.fetch_or(VIRTIO_MMIO_INT_CONFIG, Ordering::Release);
signal_or_intx(IrqSource::Config, msix, irq_evt);
break 'outer;
}
Err(e) => {
state.queue_poisoned = true;
state.counters.record_io_error();
tracing::warn!(
%e,
"virtio-blk iter() failed (non-InvalidAvailRingIndex); \
poisoning queue until guest reset"
);
debug_assert_eq!(
VIRTIO_CONFIG_S_NEEDS_RESET.count_ones(),
1,
"VIRTIO_CONFIG_S_NEEDS_RESET must be a single bit; \
multi-bit fetch_or would expand set_status's CAS retry universe"
);
device_status.fetch_or(VIRTIO_CONFIG_S_NEEDS_RESET, Ordering::SeqCst);
interrupt_status.fetch_or(VIRTIO_MMIO_INT_CONFIG, Ordering::Release);
signal_or_intx(IrqSource::Config, msix, irq_evt);
break 'outer;
}
};
let q = &mut queues[REQ_QUEUE];
let head = chain.head_index();
state.all_descs_scratch.clear();
const SCRATCH_CAP: usize = VIRTIO_BLK_SEG_MAX as usize + 2;
let mut chain_len: usize = 0;
for desc in chain {
chain_len += 1;
if state.all_descs_scratch.len() < SCRATCH_CAP {
state.all_descs_scratch.push(ChainDescriptor {
addr: desc.addr(),
len: desc.len(),
is_write_only: desc.is_write_only(),
});
}
}
if chain_len > SCRATCH_CAP {
tracing::warn!(
head,
desc_count = chain_len,
"virtio-blk chain exceeds seg_max + 2; dropping (oversized chain \
prevents reliable status_addr identification, leaving the guest \
to surface the request via blk-mq timer / hung-task watchdog)"
);
state.counters.record_io_error();
continue;
}
let mut header_addr: Option<GuestAddress> = None;
let mut status_addr: Option<GuestAddress> = None;
if let Some((first, rest)) = state.all_descs_scratch.split_first() {
if !first.is_write_only && (first.len as usize) >= VIRTIO_BLK_OUTHDR_SIZE {
header_addr = Some(first.addr);
}
if let Some((last, _middle)) = rest.split_last() {
if last.is_write_only && last.len >= 1 {
status_addr = last
.addr
.0
.checked_add(last.len as u64 - 1)
.map(GuestAddress);
}
}
}
let Some(status_addr) = status_addr else {
tracing::warn!(head, "virtio-blk request without status descriptor");
state.counters.record_io_error();
continue;
};
let Some(header_addr) = header_addr else {
tracing::warn!(head, "virtio-blk request without valid header descriptor");
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"bad header",
) {
signal_needed = true;
}
continue;
};
let hdr: VirtioBlkOutHdr = match mem.read_obj(header_addr) {
Ok(h) => h,
Err(_) => {
tracing::warn!(head, "virtio-blk header read failed");
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"bad hdr read",
) {
signal_needed = true;
}
continue;
}
};
let req_type = hdr.type_;
let sector = hdr.sector;
let data_segments: &[ChainDescriptor] = &state.all_descs_scratch[1..chain_len - 1];
if data_segments.iter().any(|d| d.len > VIRTIO_BLK_SIZE_MAX) {
tracing::warn!(head, "virtio-blk descriptor exceeds size_max");
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"size_max reject",
) {
signal_needed = true;
}
continue;
}
let data_len: u64 = data_segments.iter().map(|d| d.len as u64).sum();
if matches!(
req_type,
VIRTIO_BLK_T_IN | VIRTIO_BLK_T_OUT | VIRTIO_BLK_T_GET_ID
) && data_segments.is_empty()
{
tracing::warn!(
head,
req_type,
"virtio-blk T_IN/T_OUT/T_GET_ID with no data segments"
);
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"zero-data",
) {
signal_needed = true;
}
continue;
}
if req_type == VIRTIO_BLK_T_FLUSH && !data_segments.is_empty() {
tracing::warn!(
head,
seg_count = data_segments.len(),
"virtio-blk T_FLUSH with non-empty data segments"
);
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"flush-with-data",
) {
signal_needed = true;
}
continue;
}
if matches!(req_type, VIRTIO_BLK_T_IN | VIRTIO_BLK_T_OUT)
&& !data_len.is_multiple_of(VIRTIO_BLK_SECTOR_SIZE as u64)
{
tracing::warn!(
head,
req_type,
data_len,
"virtio-blk T_IN/T_OUT data_len not a multiple of 512"
);
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"sub-sector",
) {
signal_needed = true;
}
continue;
}
let backing: &dyn Backing = &*state.backing;
let counters = state.counters.as_ref();
let cap_bytes = state.capacity_bytes;
let read_only = state.read_only;
let pre_throttle = VirtioBlk::classify_pre_throttle(req_type, read_only, counters);
let direction_violation = pre_throttle.is_none()
&& match req_type {
VIRTIO_BLK_T_IN | VIRTIO_BLK_T_GET_ID => {
data_segments.iter().any(|d| !d.is_write_only)
}
VIRTIO_BLK_T_OUT => data_segments.iter().any(|d| d.is_write_only),
_ => false,
};
if direction_violation {
tracing::warn!(
head,
req_type,
"virtio-blk T_IN/T_OUT/T_GET_ID data segment direction mismatch"
);
state.counters.record_io_error();
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
VIRTIO_BLK_S_IOERR as u8,
1,
"direction",
) {
signal_needed = true;
}
continue;
}
if pre_throttle.is_none() {
let ops_ok = state.ops_bucket.can_consume(1);
let bytes_ok = state.bytes_bucket.can_consume(data_len);
if !ops_ok || !bytes_ok {
state.counters.record_throttled();
if !state.currently_stalled {
state.currently_stalled = true;
state.counters.record_throttle_pending_inc();
}
let prev = queues[REQ_QUEUE].next_avail();
queues[REQ_QUEUE].set_next_avail(prev.wrapping_sub(1));
let ops_wait = if !ops_ok {
state.ops_bucket.nanos_until_n_tokens(1)
} else {
0
};
let bytes_wait = if !bytes_ok {
state.bytes_bucket.nanos_until_n_tokens(data_len)
} else {
0
};
let wait_nanos = ops_wait.max(bytes_wait);
tracing::trace!(
head,
ops_ok,
bytes_ok,
wait_nanos,
"virtio-blk throttle stall; rolling back chain"
);
stall_outcome = Some(wait_nanos);
break;
}
let ops_consumed = state.ops_bucket.consume(1);
let bytes_consumed = state.bytes_bucket.consume(data_len);
debug_assert!(
ops_consumed && bytes_consumed,
"throttle invariant: can_consume must imply consume",
);
if state.currently_stalled {
state.currently_stalled = false;
state.counters.record_throttle_pending_dec();
}
}
let (status_byte, used_len) = if let Some(out) = pre_throttle {
out
} else {
match req_type {
VIRTIO_BLK_T_IN => VirtioBlk::handle_read_vectored_impl(
backing,
cap_bytes,
counters,
mem,
sector,
data_segments,
data_len,
),
VIRTIO_BLK_T_OUT => VirtioBlk::handle_write_vectored_impl(
backing,
cap_bytes,
counters,
mem,
sector,
data_segments,
data_len,
),
VIRTIO_BLK_T_FLUSH => VirtioBlk::handle_flush_impl(backing, counters),
VIRTIO_BLK_T_GET_ID => {
VirtioBlk::handle_get_id_impl(counters, mem, data_segments)
}
_ => {
counters.record_io_error();
(VIRTIO_BLK_S_UNSUPP as u8, 1)
}
}
};
let req_type_name = match req_type {
VIRTIO_BLK_T_IN => "in",
VIRTIO_BLK_T_OUT => "out",
VIRTIO_BLK_T_FLUSH => "flush",
VIRTIO_BLK_T_GET_ID => "get_id",
_ => "unsupp",
};
tracing::trace!(
req_type = req_type_name,
req_type_raw = req_type,
sector,
head,
status = status_byte,
used_len,
"virtio-blk request done"
);
if publish_completion(
mem,
q,
&state.counters,
head,
status_addr,
status_byte,
used_len,
"publish completion",
) {
signal_needed = true;
}
}
if stall_outcome.is_some() {
if let Err(e) = queues[REQ_QUEUE].enable_notification(mem) {
tracing::warn!(
%e,
"virtio-blk enable_notification failed on throttle stall"
);
}
break 'outer;
}
match queues[REQ_QUEUE].enable_notification(mem) {
Ok(true) => continue 'outer,
Ok(false) => break 'outer,
Err(e) => {
tracing::warn!(%e, "virtio-blk enable_notification failed");
break 'outer;
}
}
}
if signal_needed {
interrupt_status.fetch_or(VIRTIO_MMIO_INT_VRING, Ordering::Release);
let q = &mut queues[REQ_QUEUE];
if q.needs_notification(mem)
.inspect_err(
|e| tracing::warn!(%e, "needs_notification failed; firing IRQ as fail-safe"),
)
.unwrap_or(true)
{
signal_or_intx(IrqSource::Vring { queue: REQ_QUEUE }, msix, irq_evt);
}
}
match stall_outcome {
Some(wait_nanos) => DrainOutcome::ThrottleStalled { wait_nanos },
None => DrainOutcome::Done,
}
}
fn signal_or_intx(source: IrqSource, msix: Option<&Arc<PiMutex<MsixState>>>, irq_evt: &EventFd) {
if let Some(m) = msix {
let mut g = m.lock();
if g.enabled() {
g.signal(source);
return;
}
}
if let Err(e) = irq_evt.write(1) {
tracing::warn!(%e, "virtio-blk irq_evt.write failed");
}
}