diff --git a/src/agent/src/rpc.rs b/src/agent/src/rpc.rs index 8210c0722e..3475a7e26f 100644 --- a/src/agent/src/rpc.rs +++ b/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()); } } diff --git a/src/agent/src/sandbox.rs b/src/agent/src/sandbox.rs index 7ffd81b56b..759d932065 100644 --- a/src/agent/src/sandbox.rs +++ b/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, oneshot::Sender); -#[derive(Clone, Debug)] +pub type StorageDeviceObject = Arc>; + +#[derive(Clone)] pub struct StorageState { - inner: Arc>, + 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 { + 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() { diff --git a/src/libs/kata-types/src/mount.rs b/src/libs/kata-types/src/mount.rs index 91a061797a..44ccd26045 100644 --- a/src/libs/kata-types/src/mount.rs +++ b/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 } + } +} + +impl StorageDevice for StorageDeviceGeneric { + fn path(&self) -> &str { + &self.path } - /// Get reference count. - pub fn ref_count(&self) -> u32 { + 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 { + handlers: HashMap, +} + +impl Default for StorageHandlerManager { + fn default() -> Self { + Self::new() + } +} + +impl StorageHandlerManager { + /// 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());