kata-types: introduce StorageDevice and StorageHandlerManager

Introduce StorageDevice and StorageHandlerManager, which will be used
to refine storage device management for kata-agent.

Signed-off-by: Jiang Liu <gerry@linux.alibaba.com>

src/agent/src/rpc.rs

@@ -1686,7 +1686,7 @@ async fn remove_container_resources(sandbox: &mut Sandbox, cid: &str) -> Result<
     let mounts = sandbox.container_mounts.get(cid);
     if let Some(mounts) = mounts {
         for m in mounts.iter() {
-            if sandbox.storages.get(m).is_some() {
+            if sandbox.storages.contains_key(m) {
                 cmounts.push(m.to_string());
             }
         }

src/agent/src/sandbox.rs

@@ -5,6 +5,7 @@
 
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
+use std::fmt::{Debug, Formatter};
 use std::fs;
 use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
@@ -14,7 +15,7 @@ use std::{thread, time};
 
 use anyhow::{anyhow, Context, Result};
 use kata_types::cpu::CpuSet;
-use kata_types::mount::StorageDeviceGeneric;
+use kata_types::mount::{StorageDevice, StorageDeviceGeneric};
 use libc::pid_t;
 use oci::{Hook, Hooks};
 use protocols::agent::OnlineCPUMemRequest;
@@ -42,18 +43,31 @@ pub const ERR_INVALID_CONTAINER_ID: &str = "Invalid container id";
 
 type UeventWatcher = (Box<dyn UeventMatcher>, oneshot::Sender<Uevent>);
 
-#[derive(Clone, Debug)]
+pub type StorageDeviceObject = Arc<Mutex<dyn StorageDevice>>;
+
+#[derive(Clone)]
 pub struct StorageState {
-    inner: Arc<Mutex<StorageDeviceGeneric>>,
+    inner: StorageDeviceObject,
+}
+
+impl Debug for StorageState {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("StorageState").finish()
+    }
 }
 
 impl StorageState {
     fn new() -> Self {
         StorageState {
-            inner: Arc::new(Mutex::new(StorageDeviceGeneric::new())),
+            inner: Arc::new(Mutex::new(StorageDeviceGeneric::new("".to_string()))),
         }
     }
 
+    #[allow(dead_code)]
+    pub fn from_device(device: StorageDeviceObject) -> Self {
+        Self { inner: device }
+    }
+
     pub async fn ref_count(&self) -> u32 {
         self.inner.lock().await.ref_count()
     }
@@ -145,7 +159,25 @@ impl Sandbox {
         }
     }
 
+    /// Update the storage device associated with a path.
+    #[allow(dead_code)]
+    pub fn update_sandbox_storage(
+        &mut self,
+        path: &str,
+        device: StorageDeviceObject,
+    ) -> std::result::Result<StorageDeviceObject, StorageDeviceObject> {
+        if !self.storages.contains_key(path) {
+            return Err(device);
+        }
+
+        let state = StorageState::from_device(device);
+        // Safe to unwrap() because we have just ensured existence of entry.
+        let state = self.storages.insert(path.to_string(), state).unwrap();
+        Ok(state.inner)
+    }
+
     // Clean mount and directory of a mountpoint.
+    // This is actually StorageDeviceGeneric::cleanup(), kept here due to dependency chain.
     #[instrument]
     fn cleanup_sandbox_storage(&mut self, path: &str) -> Result<()> {
         if path.is_empty() {

src/libs/kata-types/src/mount.rs

@@ -5,6 +5,7 @@
 //
 
 use anyhow::{anyhow, Context, Error, Result};
+use std::collections::hash_map::Entry;
 use std::convert::TryFrom;
 use std::fmt::Formatter;
 use std::{collections::HashMap, fs, path::PathBuf};
@@ -431,6 +432,7 @@ impl TryFrom<&NydusExtraOptions> for KataVirtualVolume {
 /// An implementation of generic storage device.
 pub struct StorageDeviceGeneric {
     refcount: u32,
+    path: String,
 }
 
 impl std::fmt::Debug for StorageDeviceGeneric {
@@ -443,26 +445,85 @@ impl std::fmt::Debug for StorageDeviceGeneric {
 
 impl StorageDeviceGeneric {
     /// Create a new instance of `StorageStateCommon`.
-    pub fn new() -> Self {
-        StorageDeviceGeneric { refcount: 0 }
+    pub fn new(path: String) -> Self {
+        StorageDeviceGeneric { refcount: 1, path }
+    }
 }
 
-    /// Get reference count.
-    pub fn ref_count(&self) -> u32 {
+impl StorageDevice for StorageDeviceGeneric {
+    fn path(&self) -> &str {
+        &self.path
+    }
+
+    fn ref_count(&self) -> u32 {
         self.refcount
     }
 
-    /// Decrease reference count and return true if it reaches zero.
-    pub fn inc_ref_count(&mut self) {
+    fn inc_ref_count(&mut self) {
         self.refcount += 1;
     }
 
-    /// Decrease reference count and return true if it reaches zero.
-    pub fn dec_and_test_ref_count(&mut self) -> bool {
+    fn dec_and_test_ref_count(&mut self) -> bool {
         assert!(self.refcount > 0);
         self.refcount -= 1;
         self.refcount == 0
     }
+
+    fn cleanup(&self) {}
+}
+
+/// Trait object for storage device.
+pub trait StorageDevice: Send + Sync {
+    /// Path
+    fn path(&self) -> &str;
+
+    /// Get reference count.
+    fn ref_count(&self) -> u32;
+
+    /// Increase reference count.
+    fn inc_ref_count(&mut self);
+
+    /// Decrease reference count and return true if it reaches zero.
+    fn dec_and_test_ref_count(&mut self) -> bool;
+
+    /// Clean up resources related to the storage device.
+    fn cleanup(&self);
+}
+
+/// Manager to manage registered storage device handlers.
+pub struct StorageHandlerManager<H> {
+    handlers: HashMap<String, H>,
+}
+
+impl<H> Default for StorageHandlerManager<H> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<H> StorageHandlerManager<H> {
+    /// Create a new instance of `StorageHandlerManager`.
+    pub fn new() -> Self {
+        Self {
+            handlers: HashMap::new(),
+        }
+    }
+
+    /// Register a storage device handler.
+    pub fn add_handler(&mut self, id: &str, handler: H) -> Result<()> {
+        match self.handlers.entry(id.to_string()) {
+            Entry::Occupied(_) => Err(anyhow!("storage handler for {} already exists", id)),
+            Entry::Vacant(entry) => {
+                entry.insert(handler);
+                Ok(())
+            }
+        }
+    }
+
+    /// Get storage handler with specified `id`.
+    pub fn handler(&self, id: &str) -> Option<&H> {
+        self.handlers.get(id)
+    }
 }
 
 /// Join user provided volume path with kata direct-volume root path.
@@ -700,12 +761,14 @@ mod tests {
 
     #[test]
     fn test_storage_state_common() {
-        let mut state = StorageDeviceGeneric::new();
-        assert_eq!(state.ref_count(), 0);
-        state.inc_ref_count();
+        let mut state = StorageDeviceGeneric::new("".to_string());
         assert_eq!(state.ref_count(), 1);
         state.inc_ref_count();
         assert_eq!(state.ref_count(), 2);
+        state.inc_ref_count();
+        assert_eq!(state.ref_count(), 3);
+        assert!(!state.dec_and_test_ref_count());
+        assert_eq!(state.ref_count(), 2);
         assert!(!state.dec_and_test_ref_count());
         assert_eq!(state.ref_count(), 1);
         assert!(state.dec_and_test_ref_count());