profcollectd: Limit storage space used

This prevents profcollect from exhausting all the device storage. We don't
count space used by packed reports or processed data since they are much
smaller (~30KB per trace).

Test: manual
Bug: 207426233
Bug: 207417653
Change-Id: I54c77bf7d5af57b5ab8d3a696b486c951fe79aa8

2 files changed, 36 insertions, 6 deletions

diff --git a/profcollectd/libprofcollectd/config.rs b/profcollectd/libprofcollectd/config.rs
index 4e830c96..0b55d5c9 100644
--- a/profcollectd/libprofcollectd/config.rs
+++ b/profcollectd/libprofcollectd/config.rs
@@ -56,6 +56,8 @@ pub struct Config {
     pub sampling_period: Duration,
     /// An optional filter to limit which binaries to or not to profile.
     pub binary_filter: String,
+    /// Maximum size of the trace directory.
+    pub max_trace_limit: u64,
 }
 
 impl Config {
@@ -70,6 +72,10 @@ impl Config {
             )?),
             sampling_period: Duration::from_millis(get_device_config("sampling_period", 500)?),
             binary_filter: get_device_config("binary_filter", "".to_string())?,
+            max_trace_limit: get_device_config(
+                "max_trace_limit",
+                /* 512MB */ 512 * 1024 * 1024,
+            )?,
         })
     }
 }
diff --git a/profcollectd/libprofcollectd/scheduler.rs b/profcollectd/libprofcollectd/scheduler.rs
index 526345b4..b0b322e1 100644
--- a/profcollectd/libprofcollectd/scheduler.rs
+++ b/profcollectd/libprofcollectd/scheduler.rs
@@ -16,6 +16,8 @@
 
 //! ProfCollect tracing scheduler.
 
+use std::fs;
+use std::path::Path;
 use std::sync::mpsc::{sync_channel, SyncSender};
 use std::sync::Arc;
 use std::sync::Mutex;
@@ -61,11 +63,13 @@ impl Scheduler {
                     Ok(_) => break,
                     Err(_) => {
                         // Did not receive a termination signal, initiate trace event.
-                        trace_provider.lock().unwrap().trace(
-                            &TRACE_OUTPUT_DIR,
-                            "periodic",
-                            &config.sampling_period,
-                        );
+                        if check_space_limit(*TRACE_OUTPUT_DIR, &config).unwrap() {
+                            trace_provider.lock().unwrap().trace(
+                                &TRACE_OUTPUT_DIR,
+                                "periodic",
+                                &config.sampling_period,
+                            );
+                        }
                     }
                 }
             }
@@ -85,7 +89,9 @@ impl Scheduler {
 
     pub fn one_shot(&self, config: &Config, tag: &str) -> Result<()> {
         let trace_provider = self.trace_provider.clone();
-        trace_provider.lock().unwrap().trace(&TRACE_OUTPUT_DIR, tag, &config.sampling_period);
+        if check_space_limit(*TRACE_OUTPUT_DIR, config)? {
+            trace_provider.lock().unwrap().trace(&TRACE_OUTPUT_DIR, tag, &config.sampling_period);
+        }
         Ok(())
     }
 
@@ -107,3 +113,21 @@ impl Scheduler {
         self.trace_provider.lock().unwrap().get_name()
     }
 }
+
+/// Run if space usage is under limit.
+fn check_space_limit(path: &Path, config: &Config) -> Result<bool> {
+    let ret = dir_size(path)? <= config.max_trace_limit;
+    if !ret {
+        log::error!("trace storage exhausted.");
+    }
+    Ok(ret)
+}
+
+/// Returns the size of a directory, non-recursive.
+fn dir_size(path: &Path) -> Result<u64> {
+    fs::read_dir(path)?.try_fold(0, |acc, file| {
+        let metadata = file?.metadata()?;
+        let size = if metadata.is_file() { metadata.len() } else { 0 };
+        Ok(acc + size)
+    })
+}