refactor(observe): make Arc internal to proc and cgroup caches (#593)

Author: mccormickt

auraed/src/cells/cell_service/cell_service.rs

@@ -44,20 +44,21 @@ use std::time::Duration;
 use std::{process::ExitStatus, sync::Arc};
 use tokio::sync::Mutex;
 use tonic::{Code, Request, Response, Status};
-use tracing::{info, trace, warn};
+use tracing::{info, instrument, trace, warn};
 
 /**
  * Macro to perform an operation within a cell.
  * It retries the operation with an exponential backoff strategy in case of connection errors.
  */
 macro_rules! do_in_cell {
     ($self:ident, $cell_name:ident, $function:ident, $request:ident) => {{
-        let mut cells = $self.cells.lock().await;
-
         // Retrieve the client socket for the specified cell
-        let client_socket = cells
+        let client_socket = {
+            let mut cells = $self.cells.lock().await;
+            cells
             .get(&$cell_name, |cell| cell.client_socket())
-            .map_err(CellsServiceError::CellsError)?;
+            .map_err(CellsServiceError::CellsError)?
+        };
 
         // Initialize the exponential backoff strategy for retrying the operation
         let mut retry_strategy = backoff::ExponentialBackoffBuilder::new()
@@ -289,22 +290,26 @@ impl CellService {
         assert!(cell_name.is_none());
         info!("CellService: stop() executable_name={:?}", executable_name,);
 
-        let mut executables = self.executables.lock().await;
-
-        // Retrieve the process ID (PID) of the executable to be stopped
-        let pid = executables
-            .get(&executable_name)
-            .map_err(CellsServiceError::ExecutablesError)?
-            .pid()
-            .map_err(CellsServiceError::Io)?
-            .expect("pid")
-            .as_raw();
+        let pid = {
+            let mut executables = self.executables.lock().await;
+
+            // Retrieve the process ID (PID) of the executable to be stopped
+            let pid = executables
+                .get(&executable_name)
+                .map_err(CellsServiceError::ExecutablesError)?
+                .pid()
+                .map_err(CellsServiceError::Io)?
+                .expect("pid")
+                .as_raw();
+
+            // Stop the executable and handle any errors
+            let _: ExitStatus = executables
+                .stop(&executable_name)
+                .await
+                .map_err(CellsServiceError::ExecutablesError)?;
 
-        // Stop the executable and handle any errors
-        let _: ExitStatus = executables
-            .stop(&executable_name)
-            .await
-            .map_err(CellsServiceError::ExecutablesError)?;
+            pid
+        };
 
         // Remove the executable's logs from the observe service.
         if let Err(e) = self
@@ -472,6 +477,7 @@ impl cell_service_server::CellService for CellService {
         Ok(Response::new(self.allocate(request).await?))
     }
 
+    #[instrument(skip(self))]
     async fn free(
         &self,
         request: Request<CellServiceFreeRequest>,
@@ -485,6 +491,7 @@ impl cell_service_server::CellService for CellService {
         Ok(Response::new(self.free(request).await?))
     }
 
+    #[instrument(skip(self))]
     async fn start(
         &self,
         request: Request<CellServiceStartRequest>,
@@ -513,6 +520,7 @@ impl cell_service_server::CellService for CellService {
         }
     }
 
+    #[instrument(skip(self))]
     async fn stop(
         &self,
         request: Request<CellServiceStopRequest>,
@@ -586,7 +594,7 @@ mod tests {
 
         // Create a new instance of CellService for testing
         let service = CellService::new(ObserveService::new(
-            Arc::new(LogChannel::new(String::from("test"))),
+            LogChannel::new(String::from("test")),
             (None, None, None),
         ));
 
@@ -685,7 +693,7 @@ mod tests {
     #[tokio::test]
     async fn start_registers_log_channels_and_returns_uid_gid() {
         let observe_service = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("test"))),
+            LogChannel::new(String::from("test")),
             (None, None, None),
         );
         let service = CellService::new(observe_service.clone());

auraed/src/lib.rs

@@ -86,7 +86,6 @@ use proto::{
     vms::vm_service_server::VmServiceServer,
 };
 use std::path::{Path, PathBuf};
-use std::sync::Arc;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
 use tokio::task::JoinHandle;
@@ -247,7 +246,7 @@ pub async fn run(
             tonic_health::server::health_reporter();
 
         let observe_service = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("auraed"))),
+            LogChannel::new(String::from("auraed")),
             perf_events,
         );
         let observe_service_server =

auraed/src/observe/cgroup_cache.rs

@@ -43,6 +43,7 @@
 \* -------------------------------------------------------------------------- */
 use std::collections::HashMap;
 use std::ffi::OsString;
+use std::sync::{Arc, Mutex};
 use tracing::warn;
 use walkdir::DirEntryExt;
 use walkdir::WalkDir;
@@ -53,30 +54,40 @@ use walkdir::WalkDir;
 /// we should think if inode wraparound is a potential issue. We could look at
 /// how the Linux inode cache is implemented:
 /// https://elixir.bootlin.com/linux/latest/source/fs/inode.c
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct CgroupCache {
     root: OsString,
-    cache: HashMap<u64, OsString>,
+    cache: Arc<Mutex<HashMap<u64, OsString>>>,
 }
 
 impl CgroupCache {
     pub fn new(root: OsString) -> Self {
-        Self { root, cache: HashMap::new() }
+        Self { root, cache: Arc::new(Mutex::new(HashMap::new())) }
     }
 
     pub fn get(&mut self, ino: u64) -> Option<OsString> {
-        if let Some(path) = self.cache.get(&ino) {
-            Some(path.clone())
-        } else {
-            self.refresh_cache();
-            self.cache.get(&ino).cloned()
+        {
+            let cache =
+                self.cache.lock().expect("failed to acquire cgroup lock");
+            if let Some(path) = cache.get(&ino) {
+                return Some(path.clone());
+            }
         }
+
+        self.refresh_cache();
+        self.cache
+            .lock()
+            .expect("failed to acquire cgroup lock")
+            .get(&ino)
+            .cloned()
     }
 
     fn refresh_cache(&mut self) {
+        let mut cache =
+            self.cache.lock().expect("failed to acquire cgroup lock");
         WalkDir::new(&self.root).into_iter().for_each(|res| match res {
             Ok(dir_entry) => {
-                _ = self.cache.insert(dir_entry.ino(), dir_entry.path().into());
+                _ = cache.insert(dir_entry.ino(), dir_entry.path().into());
             }
             Err(e) => {
                 warn!("could not read from {:?}: {}", self.root, e);

auraed/src/observe/observe_service.rs

@@ -31,19 +31,19 @@ use proto::observe::{
     LogItem, Signal as PosixSignal, WorkloadType, observe_service_server,
 };
 use std::collections::HashMap;
+use std::sync::Arc;
 use std::time::Duration;
-use std::{ffi::OsString, sync::Arc};
 use tokio::sync::mpsc;
 use tokio::sync::{Mutex, broadcast::Receiver};
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status};
-use tracing::info;
+use tracing::{info, instrument};
 
 #[derive(Debug, Clone)]
 pub struct ObserveService {
-    aurae_logger: Arc<LogChannel>,
-    cgroup_cache: Arc<Mutex<CgroupCache>>,
-    proc_cache: Option<Arc<Mutex<ProcCache>>>,
+    aurae_logger: LogChannel,
+    cgroup_cache: CgroupCache,
+    proc_cache: Option<ProcCache>,
     posix_signals: Option<PerfEventBroadcast<Signal>>,
     sub_process_consumer_list:
         Arc<Mutex<HashMap<i32, HashMap<LogChannelType, LogChannel>>>>,
@@ -56,24 +56,20 @@ type PerfEvents = (
 );
 
 impl ObserveService {
-    pub fn new(aurae_logger: Arc<LogChannel>, perf_events: PerfEvents) -> Self {
+    pub fn new(aurae_logger: LogChannel, perf_events: PerfEvents) -> Self {
         let proc_cache = match perf_events {
-            (Some(f), Some(e), _) => {
-                Some(Arc::new(Mutex::new(ProcCache::new(
-                    Duration::from_secs(60),
-                    Duration::from_secs(60),
-                    f,
-                    e,
-                    ProcfsProcessInfo {},
-                ))))
-            }
+            (Some(f), Some(e), _) => Some(ProcCache::new(
+                Duration::from_secs(60),
+                Duration::from_secs(60),
+                f,
+                e,
+                ProcfsProcessInfo {},
+            )),
             _ => None,
         };
         Self {
             aurae_logger,
-            cgroup_cache: Arc::new(Mutex::new(CgroupCache::new(
-                OsString::from("/sys/fs/cgroup"),
-            ))),
+            cgroup_cache: CgroupCache::new("/sys/fs/cgroup".into()),
             proc_cache,
             posix_signals: perf_events.2,
             sub_process_consumer_list: Arc::new(Mutex::new(HashMap::new())),
@@ -133,7 +129,8 @@ impl ObserveService {
         self.aurae_logger.subscribe()
     }
 
-    async fn get_posix_signals_stream(
+    #[instrument(skip(self))]
+    fn get_posix_signals_stream(
         &self,
         filter: Option<(WorkloadType, String)>,
     ) -> ReceiverStream<Result<GetPosixSignalsStreamResponse, Status>> {
@@ -212,11 +209,11 @@ impl observe_service_server::ObserveService for ObserveService {
         &self,
         request: Request<GetSubProcessStreamRequest>,
     ) -> Result<Response<Self::GetSubProcessStreamStream>, Status> {
-        let channel = LogChannelType::try_from(request.get_ref().channel_type)
-            .map_err(|_| ObserveServiceError::InvalidLogChannelType {
-                channel_type: request.get_ref().channel_type,
-            })?;
-        let pid: i32 = request.get_ref().process_id;
+        let GetSubProcessStreamRequest { process_id: pid, channel_type } =
+            request.into_inner();
+        let channel = LogChannelType::try_from(channel_type).map_err(|_| {
+            ObserveServiceError::InvalidLogChannelType { channel_type }
+        })?;
 
         println!("Requested Channel {channel:?}");
         println!("Requested Process ID {pid}");
@@ -269,15 +266,9 @@ impl observe_service_server::ObserveService for ObserveService {
             ));
         }
 
-        Ok(Response::new(
-            self.get_posix_signals_stream(
-                request
-                    .into_inner()
-                    .workload
-                    .map(|w| (w.workload_type(), w.id)),
-            )
-            .await,
-        ))
+        Ok(Response::new(self.get_posix_signals_stream(
+            request.into_inner().workload.map(|w| (w.workload_type(), w.id)),
+        )))
     }
 }
 
@@ -286,12 +277,11 @@ mod tests {
     use super::ObserveService;
     use crate::logging::log_channel::LogChannel;
     use proto::observe::LogChannelType;
-    use std::sync::Arc;
 
     #[tokio::test]
     async fn test_register_sub_process_channel_success() {
         let svc = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("auraed"))),
+            LogChannel::new(String::from("auraed")),
             (None, None, None),
         );
         assert!(
@@ -310,7 +300,7 @@ mod tests {
     #[tokio::test]
     async fn test_register_sub_process_channel_duplicate_error() {
         let svc = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("auraed"))),
+            LogChannel::new(String::from("auraed")),
             (None, None, None),
         );
         assert!(
@@ -338,7 +328,7 @@ mod tests {
     #[tokio::test]
     async fn test_unregister_sub_process_channel_success() {
         let svc = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("auraed"))),
+            LogChannel::new(String::from("auraed")),
             (None, None, None),
         );
         assert!(
@@ -362,7 +352,7 @@ mod tests {
     #[tokio::test]
     async fn test_unregister_sub_process_channel_no_pid_error() {
         let svc = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("auraed"))),
+            LogChannel::new(String::from("auraed")),
             (None, None, None),
         );
         assert!(
@@ -377,7 +367,7 @@ mod tests {
     #[tokio::test]
     async fn test_unregister_sub_process_channel_no_channel_type_error() {
         let svc = ObserveService::new(
-            Arc::new(LogChannel::new(String::from("auraed"))),
+            LogChannel::new(String::from("auraed")),
             (None, None, None),
         );
         assert!(