fix(idm): reimplement packet processing algorithm (#330)

* chore(xen): rewrite event channel code * fix(idm): repair idm bugs on the file backend
edera-dev · Aug 13, 2024 · 1cf03a4 · 1cf03a4
1 parent ffc9dcc
commit 1cf03a4
Show file tree

Hide file tree

Showing 16 changed files with 270 additions and 168 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/ctl/src/cli/zone/attach.rs b/crates/ctl/src/cli/zone/attach.rs
@@ -26,7 +26,7 @@ impl ZoneAttachCommand {
         let input = StdioConsoleStream::stdin_stream(zone_id.clone()).await;
         let output = client.attach_zone_console(input).await?.into_inner();
         let stdout_handle =
-            tokio::task::spawn(async move { StdioConsoleStream::stdout(output).await });
+            tokio::task::spawn(async move { StdioConsoleStream::stdout(output, true).await });
         let exit_hook_task = StdioConsoleStream::zone_exit_hook(zone_id.clone(), events).await?;
         let code = select! {
             x = stdout_handle => {

diff --git a/crates/ctl/src/cli/zone/launch.rs b/crates/ctl/src/cli/zone/launch.rs
@@ -155,7 +155,7 @@ impl ZoneLaunchCommand {
             let input = StdioConsoleStream::stdin_stream(id.clone()).await;
             let output = client.attach_zone_console(input).await?.into_inner();
             let stdout_handle =
-                tokio::task::spawn(async move { StdioConsoleStream::stdout(output).await });
+                tokio::task::spawn(async move { StdioConsoleStream::stdout(output, true).await });
             let exit_hook_task = StdioConsoleStream::zone_exit_hook(id.clone(), events).await?;
             select! {
                 x = stdout_handle => {

diff --git a/crates/ctl/src/cli/zone/logs.rs b/crates/ctl/src/cli/zone/logs.rs
@@ -43,7 +43,7 @@ impl ZoneLogsCommand {
         };
         let output = client.attach_zone_console(input).await?.into_inner();
         let stdout_handle =
-            tokio::task::spawn(async move { StdioConsoleStream::stdout(output).await });
+            tokio::task::spawn(async move { StdioConsoleStream::stdout(output, false).await });
         let exit_hook_task = StdioConsoleStream::zone_exit_hook(zone_id.clone(), events).await?;
         let code = select! {
             x = stdout_handle => {

diff --git a/crates/ctl/src/console.rs b/crates/ctl/src/console.rs
@@ -73,8 +73,8 @@ impl StdioConsoleStream {
         }
     }
 
-    pub async fn stdout(mut stream: Streaming<ZoneConsoleReply>) -> Result<()> {
-        if stdin().is_tty() {
+    pub async fn stdout(mut stream: Streaming<ZoneConsoleReply>, raw: bool) -> Result<()> {
+        if raw && stdin().is_tty() {
             enable_raw_mode()?;
             StdioConsoleStream::register_terminal_restore_hook()?;
         }

diff --git a/crates/daemon/src/idm.rs b/crates/daemon/src/idm.rs
@@ -254,9 +254,9 @@ pub struct IdmDaemonBackend {
 
 #[async_trait::async_trait]
 impl IdmBackend for IdmDaemonBackend {
-    async fn recv(&mut self) -> Result<IdmTransportPacket> {
+    async fn recv(&mut self) -> Result<Vec<IdmTransportPacket>> {
         if let Some(packet) = self.rx_receiver.recv().await {
-            Ok(packet)
+            Ok(vec![packet])
         } else {
             Err(anyhow!("idm receive channel closed"))
         }

diff --git a/crates/krata/src/idm/client.rs b/crates/krata/src/idm/client.rs
@@ -9,13 +9,13 @@ use std::{
 };
 
 use anyhow::{anyhow, Result};
-use bytes::{BufMut, BytesMut};
+use bytes::{Buf, BufMut, BytesMut};
 use log::{debug, error};
 use nix::sys::termios::{cfmakeraw, tcgetattr, tcsetattr, SetArg};
 use prost::Message;
 use tokio::{
     fs::File,
-    io::{unix::AsyncFd, AsyncReadExt, AsyncWriteExt},
+    io::{AsyncReadExt, AsyncWriteExt},
     select,
     sync::{
         broadcast,
@@ -43,12 +43,13 @@ const IDM_PACKET_MAX_SIZE: usize = 20 * 1024 * 1024;
 
 #[async_trait::async_trait]
 pub trait IdmBackend: Send {
-    async fn recv(&mut self) -> Result<IdmTransportPacket>;
+    async fn recv(&mut self) -> Result<Vec<IdmTransportPacket>>;
     async fn send(&mut self, packet: IdmTransportPacket) -> Result<()>;
 }
 
 pub struct IdmFileBackend {
-    read_fd: Arc<Mutex<AsyncFd<File>>>,
+    read: Arc<Mutex<File>>,
+    read_buffer: BytesMut,
     write: Arc<Mutex<File>>,
 }
 
@@ -57,7 +58,8 @@ impl IdmFileBackend {
         IdmFileBackend::set_raw_port(&read_file)?;
         IdmFileBackend::set_raw_port(&write_file)?;
         Ok(IdmFileBackend {
-            read_fd: Arc::new(Mutex::new(AsyncFd::new(read_file)?)),
+            read: Arc::new(Mutex::new(read_file)),
+            read_buffer: BytesMut::new(),
             write: Arc::new(Mutex::new(write_file)),
         })
     }
@@ -72,26 +74,58 @@ impl IdmFileBackend {
 
 #[async_trait::async_trait]
 impl IdmBackend for IdmFileBackend {
-    async fn recv(&mut self) -> Result<IdmTransportPacket> {
-        let mut fd = self.read_fd.lock().await;
-        let mut guard = fd.readable_mut().await?;
-        let b1 = guard.get_inner_mut().read_u8().await?;
-        if b1 != 0xff {
-            return Ok(IdmTransportPacket::default());
-        }
-        let b2 = guard.get_inner_mut().read_u8().await?;
-        if b2 != 0xff {
-            return Ok(IdmTransportPacket::default());
-        }
-        let size = guard.get_inner_mut().read_u32_le().await?;
-        if size == 0 {
-            return Ok(IdmTransportPacket::default());
-        }
-        let mut buffer = vec![0u8; size as usize];
-        guard.get_inner_mut().read_exact(&mut buffer).await?;
-        match IdmTransportPacket::decode(buffer.as_slice()) {
-            Ok(packet) => Ok(packet),
-            Err(error) => Err(anyhow!("received invalid idm packet: {}", error)),
+    async fn recv(&mut self) -> Result<Vec<IdmTransportPacket>> {
+        let mut data = vec![0; 8192];
+        let mut first = true;
+        'read_more: loop {
+            let mut packets = Vec::new();
+            if !first {
+                if !packets.is_empty() {
+                    return Ok(packets);
+                }
+                let size = self.read.lock().await.read(&mut data).await?;
+                self.read_buffer.extend_from_slice(&data[0..size]);
+            }
+            first = false;
+            loop {
+                if self.read_buffer.len() < 6 {
+                    continue 'read_more;
+                }
+
+                let b1 = self.read_buffer[0];
+                let b2 = self.read_buffer[1];
+
+                if b1 != 0xff || b2 != 0xff {
+                    self.read_buffer.clear();
+                    continue 'read_more;
+                }
+
+                let size = (self.read_buffer[2] as u32
+                    | (self.read_buffer[3] as u32) << 8
+                    | (self.read_buffer[4] as u32) << 16
+                    | (self.read_buffer[5] as u32) << 24) as usize;
+                let needed = size + 6;
+                if self.read_buffer.len() < needed {
+                    continue 'read_more;
+                }
+
+                let mut packet = self.read_buffer.split_to(needed);
+                packet.advance(6);
+
+                match IdmTransportPacket::decode(packet) {
+                    Ok(packet) => {
+                        packets.push(packet);
+                    }
+                    Err(error) => {
+                        return Err(anyhow!("received invalid idm packet: {}", error));
+                    }
+                }
+
+                if self.read_buffer.is_empty() {
+                    break;
+                }
+            }
+            return Ok(packets);
         }
     }
 
@@ -403,8 +437,9 @@ impl<R: IdmRequest, E: IdmSerializable> IdmClient<R, E> {
         loop {
             select! {
                 x = backend.recv() => match x {
-                    Ok(packet) => {
-                        if packet.channel != channel {
+                    Ok(packets) => {
+                        for packet in packets {
+                         if packet.channel != channel {
                             continue;
                         }
 
@@ -478,6 +513,7 @@ impl<R: IdmRequest, E: IdmSerializable> IdmClient<R, E> {
 
                             _ => {},
                         }
+                        }
                     },
 
                     Err(error) => {

diff --git a/crates/runtime/src/channel.rs b/crates/runtime/src/channel.rs
@@ -8,14 +8,11 @@ use anyhow::{anyhow, Result};
 use log::{debug, error};
 use tokio::{
     select,
-    sync::{
-        broadcast,
-        mpsc::{channel, Receiver, Sender},
-    },
+    sync::mpsc::{channel, Receiver, Sender},
     task::JoinHandle,
     time::sleep,
 };
-use xenevtchn::EventChannel;
+use xenevtchn::EventChannelService;
 use xengnt::{sys::GrantRef, GrantTab, MappedMemory};
 use xenstore::{XsdClient, XsdInterface};
 
@@ -43,7 +40,7 @@ pub struct ChannelService {
     typ: String,
     use_reserved_ref: Option<u64>,
     backends: HashMap<u32, ChannelBackend>,
-    evtchn: EventChannel,
+    evtchn: EventChannelService,
     store: XsdClient,
     gnttab: GrantTab,
     input_receiver: Receiver<(u32, Vec<u8>)>,
@@ -64,7 +61,7 @@ impl ChannelService {
         let (output_sender, output_receiver) = channel(GROUPED_CHANNEL_QUEUE_LEN);
 
         debug!("opening Xen event channel");
-        let evtchn = EventChannel::open().await?;
+        let evtchn = EventChannelService::open().await?;
         debug!("opening XenStore");
         let store = XsdClient::open().await?;
         debug!("opening GrantTab");
@@ -234,7 +231,7 @@ impl ChannelBackend {
         domid: u32,
         id: u32,
         store: XsdClient,
-        evtchn: EventChannel,
+        evtchn: EventChannelService,
         gnttab: GrantTab,
         output_sender: Sender<(u32, Option<Vec<u8>>)>,
         use_reserved_ref: Option<u64>,
@@ -273,7 +270,7 @@ pub struct KrataChannelBackendProcessor {
     id: u32,
     domid: u32,
     store: XsdClient,
-    evtchn: EventChannel,
+    evtchn: EventChannelService,
     gnttab: GrantTab,
 }
 
@@ -492,25 +489,18 @@ impl KrataChannelBackendProcessor {
                 },
 
                 x = channel.receiver.recv() => match x {
-                    Ok(_) => {
+                    Some(_) => {
                         unsafe {
                             let buffer = self.read_output_buffer(channel.local_port, &memory).await?;
                             if !buffer.is_empty() {
                                 sender.send((self.domid, Some(buffer))).await?;
                             }
                         };
-                        channel.unmask_sender.send(channel.local_port).await?;
+                        channel.unmask().await?;
                     },
 
-                    Err(error) => {
-                        match error {
-                            broadcast::error::RecvError::Closed => {
-                                break;
-                            },
-                            error => {
-                                return Err(anyhow!("failed to receive event notification: {}", error));
-                            }
-                        }
+                    None => {
+                        break;
                     }
                 }
             };

diff --git a/crates/runtime/src/ip.rs b/crates/runtime/src/ip.rs
@@ -7,7 +7,7 @@ use std::{
 
 use anyhow::{anyhow, Result};
 use ipnetwork::{Ipv4Network, Ipv6Network};
-use log::{debug, error, trace};
+use log::{debug, error};
 use tokio::sync::RwLock;
 use uuid::Uuid;
 use xenstore::{XsdClient, XsdInterface};

diff --git a/crates/runtime/src/power.rs b/crates/runtime/src/power.rs
@@ -167,7 +167,10 @@ impl PowerManagementContext {
             .set_cpufreq_gov(CpuId::All, policy)
             .await
             .unwrap_or_else(|error| {
-                info!("non-fatal error while setting scheduler policy: {:?}", error);
+                info!(
+                    "non-fatal error while setting scheduler policy: {:?}",
+                    error
+                );
             });
         Ok(())
     }

diff --git a/crates/xen/xenevtchn/Cargo.toml b/crates/xen/xenevtchn/Cargo.toml
@@ -9,6 +9,7 @@ edition = "2021"
 resolver = "2"
 
 [dependencies]
+byteorder = { workspace = true }
 libc = { workspace = true }
 log = { workspace = true }
 thiserror = { workspace = true }

diff --git a/crates/xen/xenevtchn/examples/simple.rs b/crates/xen/xenevtchn/examples/simple.rs
@@ -1,9 +1,9 @@
 use xenevtchn::error::Result;
-use xenevtchn::EventChannel;
+use xenevtchn::EventChannelService;
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let channel = EventChannel::open().await?;
+    let channel = EventChannelService::open().await?;
     println!("channel opened");
     let port = channel.bind_unbound_port(0).await?;
     println!("port: {}", port);

diff --git a/crates/xen/xenevtchn/src/error.rs b/crates/xen/xenevtchn/src/error.rs
@@ -8,6 +8,10 @@ pub enum Error {
     Io(#[from] io::Error),
     #[error("failed to send event channel wake: {0}")]
     WakeSend(tokio::sync::broadcast::error::SendError<u32>),
+    #[error("failed to acquire lock")]
+    LockAcquireFailed,
+    #[error("event port already in use")]
+    PortInUse,
 }
 
 pub type Result<T> = std::result::Result<T, Error>;