add DeviceUsage and DeviceUsageKind for Instance.device_usage

agent/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "agent"
-version = "0.11.4"
+version = "0.11.5"
 authors = ["Kate Goldenring <[email protected]>", "<[email protected]>"]
 edition = "2018"
 rust-version = "1.68.1"

agent/src/util/crictl_containers.rs

@@ -1,5 +1,6 @@
-use akri_shared::akri::AKRI_SLOT_ANNOTATION_NAME_PREFIX;
-use std::collections::{HashMap, HashSet};
+use akri_shared::akri::{instance::device_usage::NodeUsage, AKRI_SLOT_ANNOTATION_NAME_PREFIX};
+use std::collections::HashMap;
+use std::str::FromStr;
 
 /// Output from crictl query
 #[derive(Serialize, Deserialize, Clone, Debug)]
@@ -16,19 +17,21 @@ struct CriCtlContainer {
 }
 
 /// This gets the usage slots for an instance by getting the annotations that were stored at id `AKRI_SLOT_ANNOTATION_NAME_PREFIX` during allocate.
-pub fn get_container_slot_usage(crictl_output: &str) -> HashSet<String> {
+pub fn get_container_slot_usage(crictl_output: &str) -> HashMap<String, NodeUsage> {
     match serde_json::from_str::<CriCtlOutput>(crictl_output) {
         Ok(crictl_output_parsed) => crictl_output_parsed
             .containers
             .iter()
             .flat_map(|container| &container.annotations)
             .filter_map(|(key, value)| {
-                if key.starts_with(AKRI_SLOT_ANNOTATION_NAME_PREFIX)
-                    && value.eq(key
+                if key.starts_with(AKRI_SLOT_ANNOTATION_NAME_PREFIX) {
+                    let slot_id = key
                         .strip_prefix(AKRI_SLOT_ANNOTATION_NAME_PREFIX)
-                        .unwrap_or_default())
-                {
-                    Some(value.clone())
+                        .unwrap_or_default();
+                    match NodeUsage::from_str(value) {
+                        Ok(node_usage) => Some((slot_id.to_string(), node_usage)),
+                        Err(_) => None,
+                    }
                 } else {
                     None
                 }
@@ -40,14 +43,15 @@ pub fn get_container_slot_usage(crictl_output: &str) -> HashSet<String> {
                 e,
                 &crictl_output
             );
-            HashSet::default()
+            HashMap::default()
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use akri_shared::akri::instance::device_usage::DeviceUsageKind;
 
     fn get_container_str(annotation: &str) -> String {
         format!("{{ \
@@ -86,55 +90,70 @@ mod tests {
         let _ = env_logger::builder().is_test(true).try_init();
 
         // Empty output
-        assert_eq!(HashSet::<String>::new(), get_container_slot_usage(r#""#));
+        assert_eq!(
+            HashMap::<String, NodeUsage>::new(),
+            get_container_slot_usage(r#""#)
+        );
         // Empty json output
-        assert_eq!(HashSet::<String>::new(), get_container_slot_usage(r#"{}"#));
+        assert_eq!(
+            HashMap::<String, NodeUsage>::new(),
+            get_container_slot_usage(r#"{}"#)
+        );
         // Expected output with no containers
         assert_eq!(
-            HashSet::<String>::new(),
+            HashMap::<String, NodeUsage>::new(),
             get_container_slot_usage(r#"{\"containers\": []}"#)
         );
         // Output with syntax error
         assert_eq!(
-            HashSet::<String>::new(),
+            HashMap::<String, NodeUsage>::new(),
             get_container_slot_usage(r#"{ddd}"#)
         ); // syntax error
            // Expected output with no slot
         assert_eq!(
-            HashSet::<String>::new(),
+            HashMap::<String, NodeUsage>::new(),
             get_container_slot_usage(&format!(
                 "{{ \"containers\": [ {} ] }}",
                 &get_container_str("")
             ))
         );
         // Expected output with slot (including unexpected property)
-        let mut expected = HashSet::new();
-        expected.insert("foo".to_string());
+        let mut expected = HashMap::new();
+        expected.insert(
+            "foo".to_string(),
+            NodeUsage::create(&DeviceUsageKind::Instance, "node-a").unwrap(),
+        );
         assert_eq!(
             expected,
             get_container_slot_usage(&format!(
                 "{{ \"ddd\": \"\", \"containers\": [ {} ] }}",
-                &get_container_str("\"akri.agent.slot-foo\": \"foo\",")
+                &get_container_str("\"akri.agent.slot-foo\": \"node-a\",")
             ))
         );
         // Expected output with slot
         assert_eq!(
             expected,
             get_container_slot_usage(&format!(
                 "{{ \"containers\": [ {} ] }}",
-                &get_container_str("\"akri.agent.slot-foo\": \"foo\",")
+                &get_container_str("\"akri.agent.slot-foo\": \"node-a\",")
             ))
         );
         // Expected output with multiple containers
-        let mut expected_2 = HashSet::new();
-        expected_2.insert("foo1".to_string());
-        expected_2.insert("foo2".to_string());
+        let mut expected_2 = HashMap::new();
+        expected_2.insert(
+            "foo1".to_string(),
+            NodeUsage::create(&DeviceUsageKind::Instance, "node-a").unwrap(),
+        );
+        expected_2.insert(
+            "foo2".to_string(),
+            NodeUsage::create(&DeviceUsageKind::Instance, "node-b").unwrap(),
+        );
         assert_eq!(
             expected_2,
             get_container_slot_usage(&format!(
                 "{{ \"containers\": [ {}, {} ] }}",
-                &get_container_str("\"akri.agent.slot-foo1\": \"foo1\","),
-                &get_container_str("\"akri.agent.slot-foo2\": \"foo2\","),
+                &get_container_str("\"akri.agent.slot-foo1\": \"node-a\","),
+                &get_container_str("\"akri.agent.slot-foo2\": \"node-b\","),
             ))
         );
     }