alioth/src/virtio/worker/mio.rs - 4893278

mio.rs81.40%

// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
 
use std::os::fd::AsRawFd;
use std::sync::Arc;
use std::sync::mpsc::Receiver;
use std::thread::JoinHandle;
 
use mio::event::Event;
use mio::unix::SourceFd;
use mio::{Events, Interest, Poll, Registry, Token};
use snafu::ResultExt;
 
use crate::hv::IoeventFd;
use crate::mem::mapped::{Ram, RamBus};
use crate::sync::notifier::Notifier;
use crate::virtio::dev::{
    ActiveBackend, Backend, BackendEvent, Context, StartParam, Virtio, WakeEvent, Worker,
    WorkerState,
};
use crate::virtio::queue::{Queue, QueueReg, VirtQueue};
use crate::virtio::{IrqSender, Result, error};
 
pub trait VirtioMio: Virtio {
    fn activate<'m, Q, S, E>(
        &mut self,
        feature: u128,
        active_mio: &mut ActiveMio<'_, '_, 'm, Q, S, E>,
    ) -> Result<()>
    where
        Q: VirtQueue<'m>,
        S: IrqSender,
        E: IoeventFd;
 
    fn handle_queue<'m, Q, S, E>(
        &mut self,
        index: u16,
        active_mio: &mut ActiveMio<'_, '_, 'm, Q, S, E>,
    ) -> Result<()>
    where
        Q: VirtQueue<'m>,
        S: IrqSender,
        E: IoeventFd;
 
    fn handle_event<'m, Q, S, E>(
        &mut self,
        event: &Event,
        active_mio: &mut ActiveMio<'_, '_, 'm, Q, S, E>,
    ) -> Result<()>
    where
        Q: VirtQueue<'m>,
        S: IrqSender,
        E: IoeventFd;
 
    fn reset(&mut self, registry: &Registry);
}
 
impl BackendEvent for Event {
    fn token(&self) -> u64 {45x
        self.token().0 as u6445x
    }45x
}
 
const TOKEN_QUEUE: u64 = 1 << 62;
 
pub struct Mio {
    poll: Poll,
}
 
impl Mio {
    pub fn spawn_worker<D, S, E>(6x
        dev: D,6x
        event_rx: Receiver<WakeEvent<S, E>>,6x
        memory: Arc<RamBus>,6x
        queue_regs: Arc<[QueueReg]>,6x
    ) -> Result<(JoinHandle<()>, Arc<Notifier>)>6x
    where6x
        D: VirtioMio,6x
        S: IrqSender,6x
        E: IoeventFd,6x
    {
        let poll = Poll::new().context(error::CreatePoll)?;6x
        let m = Mio { poll };6x
        Worker::spawn(dev, m, event_rx, memory, queue_regs)6x
    }6x
}
 
impl<D> Backend<D> for Mio
where
    D: VirtioMio,
{
    fn register_notifier(&mut self, token: u64) -> Result<Arc<Notifier>> {6x
        let mut notifier = Notifier::new()?;6x
        let registry = self.poll.registry();6x
        registry.register(&mut notifier, Token(token as usize), Interest::READABLE)?;6x
        Ok(Arc::new(notifier))6x
    }6x
 
    fn reset(&self, dev: &mut D) -> Result<()> {6x
        dev.reset(self.poll.registry());6x
        Ok(())6x
    }6x
 
    fn event_loop<'m, S, Q, E>(6x
        &mut self,6x
        memory: &'m Ram,6x
        context: &mut Context<D, S, E>,6x
        queues: &mut [Option<Queue<'_, 'm, Q>>],6x
        param: &StartParam<S, E>,6x
    ) -> Result<()>6x
    where6x
        S: IrqSender,6x
        Q: VirtQueue<'m>,6x
        E: IoeventFd,6x
    {
        let mut events = Events::with_capacity(128);6x
        let mut active_mio = ActiveMio {6x
            queues,6x
            irq_sender: &*param.irq_sender,6x
            ioeventfds: param.ioeventfds.as_deref().unwrap_or(&[]),6x
            poll: &mut self.poll,6x
            mem: memory,6x
        };6x
        context.dev.activate(param.feature, &mut active_mio)?;6x
        let registry = active_mio.poll.registry();6x
        for (index, fd) in active_mio.ioeventfds.iter().enumerate() {6x
            if context.dev.ioeventfd_offloaded(index as u16)? {
                continue;
            }
            let token = index as u64 | TOKEN_QUEUE;
            registry
                .register(
                    &mut SourceFd(&fd.as_fd().as_raw_fd()),
                    Token(token as usize),
                    Interest::READABLE,
                )
                .context(error::EventSource)?;
        }
        'out: loop {
            active_mio42x
                .poll42x
                .poll(&mut events, None)42x
                .context(error::PollEvents)?;42x
            for event in events.iter() {45x
                context.handle_event(event, &mut active_mio)?;45x
                if context.state != WorkerState::Running {45x
                    break 'out;6x
                }39x
            }
        }
        let registry = active_mio.poll.registry();6x
        for (index, fd) in active_mio.ioeventfds.iter().enumerate() {6x
            if context.dev.ioeventfd_offloaded(index as u16)? {
                continue;
            }
            registry
                .deregister(&mut SourceFd(&fd.as_fd().as_raw_fd()))
                .context(error::EventSource)?;
        }
        Ok(())6x
    }6x
}
 
pub struct ActiveMio<'a, 'r, 'm, Q, S, E>
where
    Q: VirtQueue<'m>,
{
    pub queues: &'a mut [Option<Queue<'r, 'm, Q>>],
    pub irq_sender: &'a S,
    pub ioeventfds: &'a [E],
    pub poll: &'a mut Poll,
    pub mem: &'m Ram,
}
 
impl<'m, D, Q, S, E> ActiveBackend<D> for ActiveMio<'_, '_, 'm, Q, S, E>
where
    D: VirtioMio,
    Q: VirtQueue<'m>,
    S: IrqSender,
    E: IoeventFd,
{
    type Event = Event;
 
    fn handle_event(&mut self, dev: &mut D, event: &Self::Event) -> Result<()> {9x
        let token = event.token().0 as u64;9x
        if token & TOKEN_QUEUE == TOKEN_QUEUE {9x
            dev.handle_queue(token as u16, self)
        } else {
            dev.handle_event(event, self)9x
        }
    }9x
 
    fn handle_queue(&mut self, dev: &mut D, index: u16) -> Result<()> {30x
        dev.handle_queue(index, self)30x
    }30x
}
 