feat: initial support for idm send in daemon

crates/daemon/src/idm.rs

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