feat: initial support for idm send in daemon

This commit is contained in:
Alex Zenla
2024-04-10 09:26:00 +00:00
parent ec9060d872
commit 092a243a83
5 changed files with 75 additions and 64 deletions

View File

@ -7,8 +7,9 @@ use kratart::channel::ChannelService;
use log::{error, warn}; use log::{error, warn};
use prost::Message; use prost::Message;
use tokio::{ use tokio::{
select,
sync::{ sync::{
mpsc::{Receiver, Sender}, mpsc::{channel, Receiver, Sender},
Mutex, Mutex,
}, },
task::JoinHandle, task::JoinHandle,
@ -19,16 +20,23 @@ type ListenerMap = Arc<Mutex<HashMap<u32, Sender<(u32, IdmPacket)>>>>;
#[derive(Clone)] #[derive(Clone)]
pub struct DaemonIdmHandle { pub struct DaemonIdmHandle {
listeners: ListenerMap, listeners: ListenerMap,
tx_sender: Sender<(u32, IdmPacket)>,
task: Arc<JoinHandle<()>>, task: Arc<JoinHandle<()>>,
} }
#[derive(Clone)] #[derive(Clone)]
pub struct DaemonIdmSubscribeHandle { pub struct DaemonIdmSubscribeHandle {
domid: u32, domid: u32,
tx_sender: Sender<(u32, IdmPacket)>,
listeners: ListenerMap, listeners: ListenerMap,
} }
impl DaemonIdmSubscribeHandle { impl DaemonIdmSubscribeHandle {
pub async fn send(&self, packet: IdmPacket) -> Result<()> {
self.tx_sender.send((self.domid, packet)).await?;
Ok(())
}
pub async fn unsubscribe(&self) -> Result<()> { pub async fn unsubscribe(&self) -> Result<()> {
let mut guard = self.listeners.lock().await; let mut guard = self.listeners.lock().await;
let _ = guard.remove(&self.domid); let _ = guard.remove(&self.domid);
@ -37,6 +45,11 @@ impl DaemonIdmSubscribeHandle {
} }
impl DaemonIdmHandle { impl DaemonIdmHandle {
pub async fn send(&self, domid: u32, packet: IdmPacket) -> Result<()> {
self.tx_sender.send((domid, packet)).await?;
Ok(())
}
pub async fn subscribe( pub async fn subscribe(
&self, &self,
domid: u32, domid: u32,
@ -46,6 +59,7 @@ impl DaemonIdmHandle {
guard.insert(domid, sender); guard.insert(domid, sender);
Ok(DaemonIdmSubscribeHandle { Ok(DaemonIdmSubscribeHandle {
domid, domid,
tx_sender: self.tx_sender.clone(),
listeners: self.listeners.clone(), listeners: self.listeners.clone(),
}) })
} }
@ -61,17 +75,25 @@ impl Drop for DaemonIdmHandle {
pub struct DaemonIdm { pub struct DaemonIdm {
listeners: ListenerMap, listeners: ListenerMap,
receiver: Receiver<(u32, Vec<u8>)>, tx_sender: Sender<(u32, IdmPacket)>,
tx_raw_sender: Sender<(u32, Vec<u8>)>,
tx_receiver: Receiver<(u32, IdmPacket)>,
rx_receiver: Receiver<(u32, Vec<u8>)>,
task: JoinHandle<()>, task: JoinHandle<()>,
} }
impl DaemonIdm { impl DaemonIdm {
pub async fn new() -> Result<DaemonIdm> { pub async fn new() -> Result<DaemonIdm> {
let (service, _, receiver) = ChannelService::new("krata-channel".to_string(), None).await?; let (service, tx_raw_sender, rx_receiver) =
ChannelService::new("krata-channel".to_string(), None).await?;
let (tx_sender, tx_receiver) = channel(100);
let task = service.launch().await?; let task = service.launch().await?;
let listeners = Arc::new(Mutex::new(HashMap::new())); let listeners = Arc::new(Mutex::new(HashMap::new()));
Ok(DaemonIdm { Ok(DaemonIdm {
receiver, rx_receiver,
tx_receiver,
tx_sender,
tx_raw_sender,
task, task,
listeners, listeners,
}) })
@ -79,6 +101,7 @@ impl DaemonIdm {
pub async fn launch(mut self) -> Result<DaemonIdmHandle> { pub async fn launch(mut self) -> Result<DaemonIdmHandle> {
let listeners = self.listeners.clone(); let listeners = self.listeners.clone();
let tx_sender = self.tx_sender.clone();
let task = tokio::task::spawn(async move { let task = tokio::task::spawn(async move {
let mut buffers: HashMap<u32, BytesMut> = HashMap::new(); let mut buffers: HashMap<u32, BytesMut> = HashMap::new();
if let Err(error) = self.process(&mut buffers).await { if let Err(error) = self.process(&mut buffers).await {
@ -87,16 +110,16 @@ impl DaemonIdm {
}); });
Ok(DaemonIdmHandle { Ok(DaemonIdmHandle {
listeners, listeners,
tx_sender,
task: Arc::new(task), task: Arc::new(task),
}) })
} }
async fn process(&mut self, buffers: &mut HashMap<u32, BytesMut>) -> Result<()> { async fn process(&mut self, buffers: &mut HashMap<u32, BytesMut>) -> Result<()> {
loop { loop {
let Some((domid, data)) = self.receiver.recv().await else { select! {
break; x = self.rx_receiver.recv() => match x {
}; Some((domid, data)) => {
let buffer = buffers.entry(domid).or_insert_with_key(|_| BytesMut::new()); let buffer = buffers.entry(domid).or_insert_with_key(|_| BytesMut::new());
buffer.extend_from_slice(&data); buffer.extend_from_slice(&data);
if buffer.len() < 2 { if buffer.len() < 2 {
@ -123,6 +146,23 @@ impl DaemonIdm {
warn!("received invalid packet from domain {}: {}", domid, packet); warn!("received invalid packet from domain {}: {}", domid, packet);
} }
} }
},
None => {
break;
}
},
x = self.tx_receiver.recv() => match x {
Some((domid, packet)) => {
let data = packet.encode_to_vec();
self.tx_raw_sender.send((domid, data)).await?;
},
None => {
break;
}
}
};
} }
Ok(()) Ok(())
} }

View File

@ -23,6 +23,8 @@ async fn main() -> Result<()> {
if let Err(error) = guest.init().await { if let Err(error) = guest.init().await {
error!("failed to initialize guest: {}", error); error!("failed to initialize guest: {}", error);
death(127).await?; death(127).await?;
return Ok(());
} }
death(1).await?;
Ok(()) Ok(())
} }

View File

@ -8,7 +8,7 @@ use krata::idm::{
client::IdmClient, client::IdmClient,
protocol::{idm_event::Event, IdmEvent, IdmExitEvent, IdmPacket}, protocol::{idm_event::Event, IdmEvent, IdmExitEvent, IdmPacket},
}; };
use log::error; use log::debug;
use nix::unistd::Pid; use nix::unistd::Pid;
use tokio::select; use tokio::select;
@ -38,7 +38,7 @@ impl GuestBackground {
}, },
None => { None => {
error!("idm packet channel closed"); debug!("idm packet channel closed");
break; break;
} }
}, },

View File

@ -1,18 +0,0 @@
use anyhow::Result;
use tokio::fs::File;
pub struct XenConsole {
pub read_handle: File,
pub write_handle: File,
}
impl XenConsole {
pub async fn new(tty: &str) -> Result<XenConsole> {
let read_handle = File::options().read(true).write(false).open(tty).await?;
let write_handle = File::options().read(false).write(true).open(tty).await?;
Ok(XenConsole {
read_handle,
write_handle,
})
}
}

View File

@ -15,7 +15,6 @@ use xenstore::{XsdClient, XsdInterface};
use self::{ use self::{
autoloop::AutoLoop, autoloop::AutoLoop,
console::XenConsole,
launch::{GuestLaunchRequest, GuestLauncher}, launch::{GuestLaunchRequest, GuestLauncher},
}; };
use krataoci::cache::ImageCache; use krataoci::cache::ImageCache;
@ -23,7 +22,6 @@ use krataoci::cache::ImageCache;
pub mod autoloop; pub mod autoloop;
pub mod cfgblk; pub mod cfgblk;
pub mod channel; pub mod channel;
pub mod console;
pub mod launch; pub mod launch;
pub struct GuestLoopInfo { pub struct GuestLoopInfo {
@ -321,17 +319,6 @@ impl Runtime {
Ok(uuid) Ok(uuid)
} }
pub async fn console(&self, uuid: Uuid) -> Result<XenConsole> {
let info = self
.context
.resolve(uuid)
.await?
.ok_or_else(|| anyhow!("unable to resolve guest: {}", uuid))?;
let domid = info.domid;
let tty = self.context.xen.get_console_path(domid).await?;
XenConsole::new(&tty).await
}
pub async fn list(&self) -> Result<Vec<GuestInfo>> { pub async fn list(&self) -> Result<Vec<GuestInfo>> {
self.context.list().await self.context.list().await
} }