feat: idm v2 (#102)

* feat: rebuild idm to separate transport from content

* feat: fast guest lookup table and host identification

crates/krata/build.rs

@@ -6,12 +6,20 @@ fn main() -> Result<()> {
         .descriptor_pool("crate::DESCRIPTOR_POOL")
         .configure(
             &mut config,
-            &["proto/krata/v1/control.proto", "proto/krata/bus/idm.proto"],
+            &[
+                "proto/krata/v1/control.proto",
+                "proto/krata/idm/transport.proto",
+                "proto/krata/idm/internal.proto",
+            ],
             &["proto/"],
         )?;
     tonic_build::configure().compile_with_config(
         config,
-        &["proto/krata/v1/control.proto", "proto/krata/bus/idm.proto"],
+        &[
+            "proto/krata/v1/control.proto",
+            "proto/krata/idm/transport.proto",
+            "proto/krata/idm/internal.proto",
+        ],
         &["proto/"],
     )?;
     Ok(())

crates/krata/proto/krata/bus/idm.proto

@@ -1,67 +0,0 @@
-syntax = "proto3";
-
-package krata.bus.idm;
-
-option java_multiple_files = true;
-option java_package = "dev.krata.proto.bus.idm";
-option java_outer_classname = "IdmProto";
-
-import "google/protobuf/struct.proto";
-
-message IdmPacket {
-    oneof content {
-        IdmEvent event = 1;
-        IdmRequest request = 2;
-        IdmResponse response = 3;
-    }
-}
-
-message IdmEvent {
-    oneof event {
-        IdmExitEvent exit = 1;
-    }
-}
-
-message IdmExitEvent {
-    int32 code = 1;
-}
-
-message IdmRequest {
-    uint64 id = 1;
-    oneof request {
-        IdmPingRequest ping = 2;
-        IdmMetricsRequest metrics = 3;
-    }
-}
-
-message IdmPingRequest {}
-
-message IdmMetricsRequest {}
-
-message IdmResponse {
-    uint64 id = 1;
-    oneof response {
-        IdmPingResponse ping = 2;
-        IdmMetricsResponse metrics = 3;
-    }
-}
-
-message IdmPingResponse {}
-
-message IdmMetricsResponse {
-    IdmMetricNode root = 1;
-}
-
-message IdmMetricNode {
-    string name = 1;
-    google.protobuf.Value value = 2;
-    IdmMetricFormat format = 3;
-    repeated IdmMetricNode children = 4;
-}
-
-enum IdmMetricFormat {
-    IDM_METRIC_FORMAT_UNKNOWN = 0;
-    IDM_METRIC_FORMAT_BYTES = 1;
-    IDM_METRIC_FORMAT_INTEGER = 2;
-    IDM_METRIC_FORMAT_DURATION_SECONDS = 3;
-}

crates/krata/proto/krata/idm/internal.proto

@@ -0,0 +1,57 @@
+syntax = "proto3";
+
+package krata.idm.internal;
+
+option java_multiple_files = true;
+option java_package = "dev.krata.proto.idm.internal";
+option java_outer_classname = "IdmInternalProto";
+
+import "google/protobuf/struct.proto";
+
+message ExitEvent {
+    int32 code = 1;
+}
+
+message PingRequest {}
+
+message PingResponse {}
+
+message MetricsRequest {}
+
+message MetricsResponse {
+    MetricNode root = 1;
+}
+
+message MetricNode {
+    string name = 1;
+    google.protobuf.Value value = 2;
+    MetricFormat format = 3;
+    repeated MetricNode children = 4;
+}
+
+enum MetricFormat {
+    METRIC_FORMAT_UNKNOWN = 0;
+    METRIC_FORMAT_BYTES = 1;
+    METRIC_FORMAT_INTEGER = 2;
+    METRIC_FORMAT_DURATION_SECONDS = 3;
+}
+
+message Event {
+    oneof event {
+        ExitEvent exit = 1;
+    }
+}
+
+message Request {
+    oneof request {
+        PingRequest ping = 1;
+        MetricsRequest metrics = 2;
+    }
+}
+
+message Response {
+    oneof response {
+        PingResponse ping = 1;
+        MetricsResponse metrics = 2;
+    }
+}

crates/krata/proto/krata/idm/transport.proto

@@ -0,0 +1,22 @@
+syntax = "proto3";
+
+package krata.idm.transport;
+
+option java_multiple_files = true;
+option java_package = "dev.krata.proto.idm.transport";
+option java_outer_classname = "IdmTransportProto";
+
+message IdmTransportPacket {
+    uint64 id = 1;
+    uint64 channel = 2;
+    IdmTransportPacketForm form = 3;
+    bytes data = 4;
+}
+
+enum IdmTransportPacketForm {
+    IDM_TRANSPORT_PACKET_FORM_UNKNOWN = 0;
+    IDM_TRANSPORT_PACKET_FORM_RAW = 1;
+    IDM_TRANSPORT_PACKET_FORM_EVENT = 2;
+    IDM_TRANSPORT_PACKET_FORM_REQUEST = 3;
+    IDM_TRANSPORT_PACKET_FORM_RESPONSE = 4;
+}

crates/krata/proto/krata/v1/common.proto

@@ -62,7 +62,8 @@ message GuestState {
     GuestNetworkState network = 2;
     GuestExitInfo exit_info = 3;
     GuestErrorInfo error_info = 4;
-    uint32 domid = 5;
+    string host = 5;
+    uint32 domid = 6;
 }
 
 enum GuestStatus {

crates/krata/proto/krata/v1/control.proto

@@ -6,10 +6,12 @@ option java_multiple_files = true;
 option java_package = "dev.krata.proto.v1.control";
 option java_outer_classname = "ControlProto";
 
-import "krata/bus/idm.proto";
+import "krata/idm/transport.proto";
 import "krata/v1/common.proto";
 
 service ControlService {
+    rpc IdentifyHost(IdentifyHostRequest) returns (IdentifyHostReply);
+
     rpc CreateGuest(CreateGuestRequest) returns (CreateGuestReply);
     rpc DestroyGuest(DestroyGuestRequest) returns (DestroyGuestReply);
     rpc ResolveGuest(ResolveGuestRequest) returns (ResolveGuestReply);
@@ -24,6 +26,14 @@ service ControlService {
     rpc PullImage(PullImageRequest) returns (stream PullImageReply);
 }
 
+message IdentifyHostRequest {}
+
+message IdentifyHostReply {
+    string host_uuid = 1;
+    uint32 host_domid = 2;
+    string krata_version = 3;
+}
+
 message CreateGuestRequest {
     krata.v1.common.GuestSpec spec = 1;
 }
@@ -84,9 +94,9 @@ message ReadGuestMetricsReply {
 message SnoopIdmRequest {}
 
 message SnoopIdmReply {
-    uint32 from = 1;
-    uint32 to = 2;
-    krata.bus.idm.IdmPacket packet = 3;
+    string from = 1;
+    string to = 2;
+    krata.idm.transport.IdmTransportPacket packet = 3;
 }
 
 message ImageProgress {

crates/krata/src/bus/mod.rs

@@ -1 +0,0 @@
-pub mod idm;

crates/krata/src/idm/client.rs

@@ -8,10 +8,6 @@ use std::{
     time::Duration,
 };
 
-use super::protocol::{
-    idm_packet::Content, idm_request::Request, idm_response::Response, IdmEvent, IdmPacket,
-    IdmRequest, IdmResponse,
-};
 use anyhow::{anyhow, Result};
 use log::{debug, error};
 use nix::sys::termios::{cfmakeraw, tcgetattr, tcsetattr, SetArg};
@@ -22,14 +18,21 @@ use tokio::{
     select,
     sync::{
         broadcast,
-        mpsc::{channel, Receiver, Sender},
+        mpsc::{self, Receiver, Sender},
         oneshot, Mutex,
     },
     task::JoinHandle,
     time::timeout,
 };
 
-type RequestMap = Arc<Mutex<HashMap<u64, oneshot::Sender<IdmResponse>>>>;
+use super::{
+    internal,
+    serialize::{IdmRequest, IdmSerializable},
+    transport::{IdmTransportPacket, IdmTransportPacketForm},
+};
+
+type RequestMap<R> = Arc<Mutex<HashMap<u64, oneshot::Sender<<R as IdmRequest>::Response>>>>;
+pub type IdmInternalClient = IdmClient<internal::Request, internal::Event>;
 
 const IDM_PACKET_QUEUE_LEN: usize = 100;
 const IDM_REQUEST_TIMEOUT_SECS: u64 = 10;
@@ -37,8 +40,8 @@ const IDM_PACKET_MAX_SIZE: usize = 20 * 1024 * 1024;
 
 #[async_trait::async_trait]
 pub trait IdmBackend: Send {
-    async fn recv(&mut self) -> Result<IdmPacket>;
-    async fn send(&mut self, packet: IdmPacket) -> Result<()>;
+    async fn recv(&mut self) -> Result<IdmTransportPacket>;
+    async fn send(&mut self, packet: IdmTransportPacket) -> Result<()>;
 }
 
 pub struct IdmFileBackend {
@@ -66,30 +69,30 @@ impl IdmFileBackend {
 
 #[async_trait::async_trait]
 impl IdmBackend for IdmFileBackend {
-    async fn recv(&mut self) -> Result<IdmPacket> {
+    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(IdmPacket::default());
+            return Ok(IdmTransportPacket::default());
         }
         let b2 = guard.get_inner_mut().read_u8().await?;
         if b2 != 0xff {
-            return Ok(IdmPacket::default());
+            return Ok(IdmTransportPacket::default());
         }
         let size = guard.get_inner_mut().read_u32_le().await?;
         if size == 0 {
-            return Ok(IdmPacket::default());
+            return Ok(IdmTransportPacket::default());
         }
         let mut buffer = vec![0u8; size as usize];
         guard.get_inner_mut().read_exact(&mut buffer).await?;
-        match IdmPacket::decode(buffer.as_slice()) {
+        match IdmTransportPacket::decode(buffer.as_slice()) {
             Ok(packet) => Ok(packet),
             Err(error) => Err(anyhow!("received invalid idm packet: {}", error)),
         }
     }
 
-    async fn send(&mut self, packet: IdmPacket) -> Result<()> {
+    async fn send(&mut self, packet: IdmTransportPacket) -> Result<()> {
         let mut file = self.write.lock().await;
         let data = packet.encode_to_vec();
         file.write_all(&[0xff, 0xff]).await?;
@@ -100,16 +103,17 @@ impl IdmBackend for IdmFileBackend {
 }
 
 #[derive(Clone)]
-pub struct IdmClient {
-    request_backend_sender: broadcast::Sender<IdmRequest>,
+pub struct IdmClient<R: IdmRequest, E: IdmSerializable> {
+    channel: u64,
+    request_backend_sender: broadcast::Sender<(u64, R)>,
     next_request_id: Arc<Mutex<u64>>,
-    event_receiver_sender: broadcast::Sender<IdmEvent>,
-    tx_sender: Sender<IdmPacket>,
-    requests: RequestMap,
+    event_receiver_sender: broadcast::Sender<E>,
+    tx_sender: Sender<IdmTransportPacket>,
+    requests: RequestMap<R>,
     task: Arc<JoinHandle<()>>,
 }
 
-impl Drop for IdmClient {
+impl<R: IdmRequest, E: IdmSerializable> Drop for IdmClient<R, E> {
     fn drop(&mut self) {
         if Arc::strong_count(&self.task) <= 1 {
             self.task.abort();
@@ -117,12 +121,12 @@ impl Drop for IdmClient {
     }
 }
 
-impl IdmClient {
-    pub async fn new(backend: Box<dyn IdmBackend>) -> Result<IdmClient> {
+impl<R: IdmRequest, E: IdmSerializable> IdmClient<R, E> {
+    pub async fn new(channel: u64, backend: Box<dyn IdmBackend>) -> Result<Self> {
         let requests = Arc::new(Mutex::new(HashMap::new()));
         let (event_sender, event_receiver) = broadcast::channel(IDM_PACKET_QUEUE_LEN);
         let (internal_request_backend_sender, _) = broadcast::channel(IDM_PACKET_QUEUE_LEN);
-        let (tx_sender, tx_receiver) = channel(IDM_PACKET_QUEUE_LEN);
+        let (tx_sender, tx_receiver) = mpsc::channel(IDM_PACKET_QUEUE_LEN);
         let backend_event_sender = event_sender.clone();
         let request_backend_sender = internal_request_backend_sender.clone();
         let requests_for_client = requests.clone();
@@ -141,6 +145,7 @@ impl IdmClient {
             }
         });
         Ok(IdmClient {
+            channel,
             next_request_id: Arc::new(Mutex::new(0)),
             event_receiver_sender: event_sender.clone(),
             request_backend_sender,
@@ -150,7 +155,7 @@ impl IdmClient {
         })
     }
 
-    pub async fn open<P: AsRef<Path>>(path: P) -> Result<IdmClient> {
+    pub async fn open<P: AsRef<Path>>(channel: u64, path: P) -> Result<Self> {
         let read_file = File::options()
             .read(true)
             .write(false)
@@ -164,39 +169,48 @@ impl IdmClient {
             .open(path)
             .await?;
         let backend = IdmFileBackend::new(read_file, write_file).await?;
-        IdmClient::new(Box::new(backend) as Box<dyn IdmBackend>).await
+        IdmClient::new(channel, Box::new(backend) as Box<dyn IdmBackend>).await
     }
 
-    pub async fn emit(&self, event: IdmEvent) -> Result<()> {
+    pub async fn emit<T: IdmSerializable>(&self, event: T) -> Result<()> {
+        let id = {
+            let mut guard = self.next_request_id.lock().await;
+            let req = *guard;
+            *guard = req.wrapping_add(1);
+            req
+        };
         self.tx_sender
-            .send(IdmPacket {
-                content: Some(Content::Event(event)),
+            .send(IdmTransportPacket {
+                id,
+                form: IdmTransportPacketForm::Event.into(),
+                channel: self.channel,
+                data: event.encode()?,
             })
             .await?;
         Ok(())
     }
 
-    pub async fn requests(&self) -> Result<broadcast::Receiver<IdmRequest>> {
+    pub async fn requests(&self) -> Result<broadcast::Receiver<(u64, R)>> {
         Ok(self.request_backend_sender.subscribe())
     }
 
-    pub async fn respond(&self, id: u64, response: Response) -> Result<()> {
-        let packet = IdmPacket {
-            content: Some(Content::Response(IdmResponse {
-                id,
-                response: Some(response),
-            })),
+    pub async fn respond<T: IdmSerializable>(&self, id: u64, response: T) -> Result<()> {
+        let packet = IdmTransportPacket {
+            id,
+            form: IdmTransportPacketForm::Response.into(),
+            channel: self.channel,
+            data: response.encode()?,
         };
         self.tx_sender.send(packet).await?;
         Ok(())
     }
 
-    pub async fn subscribe(&self) -> Result<broadcast::Receiver<IdmEvent>> {
+    pub async fn subscribe(&self) -> Result<broadcast::Receiver<E>> {
         Ok(self.event_receiver_sender.subscribe())
     }
 
-    pub async fn send(&self, request: Request) -> Result<Response> {
-        let (sender, receiver) = oneshot::channel::<IdmResponse>();
+    pub async fn send(&self, request: R) -> Result<R::Response> {
+        let (sender, receiver) = oneshot::channel::<R::Response>();
         let req = {
             let mut guard = self.next_request_id.lock().await;
             let req = *guard;
@@ -217,49 +231,52 @@ impl IdmClient {
             });
         });
         self.tx_sender
-            .send(IdmPacket {
-                content: Some(Content::Request(IdmRequest {
-                    id: req,
-                    request: Some(request),
-                })),
+            .send(IdmTransportPacket {
+                id: req,
+                channel: self.channel,
+                form: IdmTransportPacketForm::Request.into(),
+                data: request.encode()?,
             })
             .await?;
 
         let response = timeout(Duration::from_secs(IDM_REQUEST_TIMEOUT_SECS), receiver).await??;
         success.store(true, Ordering::Release);
-        if let Some(response) = response.response {
-            Ok(response)
-        } else {
-            Err(anyhow!("response did not contain any content"))
-        }
+        Ok(response)
     }
 
     async fn process(
         mut backend: Box<dyn IdmBackend>,
-        event_sender: broadcast::Sender<IdmEvent>,
-        requests: RequestMap,
-        request_backend_sender: broadcast::Sender<IdmRequest>,
-        _event_receiver: broadcast::Receiver<IdmEvent>,
-        mut receiver: Receiver<IdmPacket>,
+        event_sender: broadcast::Sender<E>,
+        requests: RequestMap<R>,
+        request_backend_sender: broadcast::Sender<(u64, R)>,
+        _event_receiver: broadcast::Receiver<E>,
+        mut receiver: Receiver<IdmTransportPacket>,
     ) -> Result<()> {
         loop {
             select! {
                 x = backend.recv() => match x {
                     Ok(packet) => {
-                        match packet.content {
-                            Some(Content::Event(event)) => {
-                                let _ = event_sender.send(event);
+                        match packet.form() {
+                            IdmTransportPacketForm::Event => {
+                                if let Ok(event) = E::decode(&packet.data) {
+                                    let _ = event_sender.send(event);
+                                }
                             },
 
-                            Some(Content::Request(request)) => {
-                                let _ = request_backend_sender.send(request);
+                            IdmTransportPacketForm::Request => {
+                                if let Ok(request) = R::decode(&packet.data) {
+                                    let _ = request_backend_sender.send((packet.id, request));
+                                }
                             },
 
-                            Some(Content::Response(response)) => {
+                            IdmTransportPacketForm::Response => {
                                 let mut requests = requests.lock().await;
-                                if let Some(sender) = requests.remove(&response.id) {
+                                if let Some(sender) = requests.remove(&packet.id) {
                                     drop(requests);
-                                    let _ = sender.send(response);
+
+                                    if let Ok(response) = R::Response::decode(&packet.data) {
+                                        let _ = sender.send(response);
+                                    }
                                 }
                             },
 

crates/krata/src/idm/internal.rs

@@ -1,26 +1,66 @@
+use anyhow::Result;
+use prost::Message;
 use prost_types::{ListValue, Value};
 
-include!(concat!(env!("OUT_DIR"), "/krata.bus.idm.rs"));
+use super::serialize::{IdmRequest, IdmSerializable};
+
+include!(concat!(env!("OUT_DIR"), "/krata.idm.internal.rs"));
+
+pub const INTERNAL_IDM_CHANNEL: u64 = 0;
+
+impl IdmSerializable for Event {
+    fn encode(&self) -> Result<Vec<u8>> {
+        Ok(self.encode_to_vec())
+    }
+
+    fn decode(bytes: &[u8]) -> Result<Self> {
+        Ok(<Self as prost::Message>::decode(bytes)?)
+    }
+}
+
+impl IdmSerializable for Request {
+    fn encode(&self) -> Result<Vec<u8>> {
+        Ok(self.encode_to_vec())
+    }
+
+    fn decode(bytes: &[u8]) -> Result<Self> {
+        Ok(<Self as prost::Message>::decode(bytes)?)
+    }
+}
+
+impl IdmRequest for Request {
+    type Response = Response;
+}
+
+impl IdmSerializable for Response {
+    fn encode(&self) -> Result<Vec<u8>> {
+        Ok(self.encode_to_vec())
+    }
+
+    fn decode(bytes: &[u8]) -> Result<Self> {
+        Ok(<Self as prost::Message>::decode(bytes)?)
+    }
+}
 
 pub trait AsIdmMetricValue {
     fn as_metric_value(&self) -> Value;
 }
 
-impl IdmMetricNode {
-    pub fn structural<N: AsRef<str>>(name: N, children: Vec<IdmMetricNode>) -> IdmMetricNode {
-        IdmMetricNode {
+impl MetricNode {
+    pub fn structural<N: AsRef<str>>(name: N, children: Vec<MetricNode>) -> MetricNode {
+        MetricNode {
             name: name.as_ref().to_string(),
             value: None,
-            format: IdmMetricFormat::Unknown.into(),
+            format: MetricFormat::Unknown.into(),
             children,
         }
     }
 
-    pub fn raw_value<N: AsRef<str>, V: AsIdmMetricValue>(name: N, value: V) -> IdmMetricNode {
-        IdmMetricNode {
+    pub fn raw_value<N: AsRef<str>, V: AsIdmMetricValue>(name: N, value: V) -> MetricNode {
+        MetricNode {
             name: name.as_ref().to_string(),
             value: Some(value.as_metric_value()),
-            format: IdmMetricFormat::Unknown.into(),
+            format: MetricFormat::Unknown.into(),
             children: vec![],
         }
     }
@@ -28,9 +68,9 @@ impl IdmMetricNode {
     pub fn value<N: AsRef<str>, V: AsIdmMetricValue>(
         name: N,
         value: V,
-        format: IdmMetricFormat,
-    ) -> IdmMetricNode {
-        IdmMetricNode {
+        format: MetricFormat,
+    ) -> MetricNode {
+        MetricNode {
             name: name.as_ref().to_string(),
             value: Some(value.as_metric_value()),
             format: format.into(),

crates/krata/src/idm/mod.rs

@@ -1,3 +1,5 @@
 #[cfg(unix)]
 pub mod client;
-pub use crate::bus::idm as protocol;
+pub mod internal;
+pub mod serialize;
+pub mod transport;

crates/krata/src/idm/serialize.rs

@@ -0,0 +1,10 @@
+use anyhow::Result;
+
+pub trait IdmSerializable: Sized + Clone + Send + Sync + 'static {
+    fn decode(bytes: &[u8]) -> Result<Self>;
+    fn encode(&self) -> Result<Vec<u8>>;
+}
+
+pub trait IdmRequest: IdmSerializable {
+    type Response: IdmSerializable;
+}

crates/krata/src/idm/transport.rs

@@ -0,0 +1 @@
+include!(concat!(env!("OUT_DIR"), "/krata.idm.transport.rs"));

crates/krata/src/lib.rs

@@ -1,7 +1,6 @@
 use once_cell::sync::Lazy;
 use prost_reflect::DescriptorPool;
 
-pub mod bus;
 pub mod v1;
 
 pub mod client;