Added background scan option (#7)

MindFlavor · web-flow · commit abcf22d3f112 · 2020-12-24T20:15:34.000+01:00
* Added background thread option

* added startup block to avoid returning empty data

* Updated README
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name           = "prometheus_folder_size_exporter"
-version        = "0.3.0"
+version        = "0.4.0"
 authors        = ["Francesco Cogno <francesco.cogno@outlook.com>", "Guido Scatena <guido.scatena@unipi.it>"]
 description    = "Prometheus Folder Size Exporter"
 edition        = "2018"
diff --git a/README.md b/README.md
@@ -41,6 +41,7 @@ Start the binary with `-h` to get the complete syntax. The parameters are:
 | `-v` | no | <switch> | | Enable verbose mode.
 | `-p` | no | any valid port number | 9974 | Specify the serivce port. This is the port your Prometheus instance should point to.
 | `-i` | yes | a valid config json file | - | The configuration file. This json is detailed below (you can find an example here: [example.json](example.json)).
+| `-b` | no | Any number > 0 | Off | Enables the async storage calculation. The value specifies how often the calculation will be done. If not specified, the values will be calculated synchronously at each HTTP Get.
 
 Once started, the tool will listen on the specified port (or the default one, 9974, if not specified) and return a Prometheus valid response at the url `/metrics`. So to check if the tool is working properly simply browse the `http://localhost:9974` (or whichever port you choose).
 
diff --git a/src/main.rs b/src/main.rs
@@ -7,17 +7,33 @@ use hyper::Body;
 use std::env;
 use std::sync::Arc;
 mod options;
+mod state;
 use options::Options;
+pub use state::State;
 mod exporter_error;
 mod folder_scanner;
 mod render_to_prometheus;
 use prometheus_exporter_base::prelude::*;
+use std::time::Duration;
 
 async fn perform_request(
     _req: http::request::Request<Body>,
-    options: Arc<Options>,
+    state: Arc<Arc<State>>,
 ) -> Result<String, failure::Error> {
-    let results = options.folders_to_scan.scan()?;
+    let results = if state.options.background_poll_seconds.is_some() {
+        // loop until we have some data.
+        // This is needed because the first scan can take a lot of time
+        // and we must block until we have something. It will happen only
+        // at startup though.
+        let mut results = state.shared_vec.read().unwrap().clone();
+        while results.is_empty() {
+            std::thread::sleep(Duration::from_millis(100));
+            results = state.shared_vec.read().unwrap().clone();
+        }
+        results
+    } else {
+        state.options.folders_to_scan.scan()?
+    };
 
     let mut metric = PrometheusMetric::build()
         .with_name("folder_size")
@@ -28,7 +44,7 @@ async fn perform_request(
         .build();
 
     for result in results {
-        if let Some(user) = result.folder.user {
+        if let Some(user) = &result.folder.user {
             metric.render_and_append_instance(
                 &PrometheusInstance::new()
                     .with_label("path", result.folder.path.as_str())
@@ -88,27 +104,48 @@ async fn main() {
                 .help("verbose logging")
                 .takes_value(false),
         )
+        .arg(
+            Arg::with_name("background_poll_seconds")
+                .short("b")
+                .help("enables background scanning every <sec> seconds")
+                .required(false)
+                .takes_value(true),
+        )
         .get_matches();
 
     let options = Options::from_claps(&matches);
-
     if options.verbose {
         env::set_var("RUST_LOG", "prometheus_folder_size_exporter=trace");
     } else {
         env::set_var("RUST_LOG", "prometheus_folder_size_exporter=info");
     }
     env_logger::init();
-
     log::info!("using options: {:?}", options);
 
+    let state = Arc::new(State::new(options));
+
+    // start the background thread only
+    // if requested
+    if let Some(background_poll_seconds) = state.options.background_poll_seconds {
+        let state = state.clone();
+        std::thread::spawn(move || loop {
+            log::info!("starting background folder structure update");
+            let results = state.options.folders_to_scan.scan().unwrap();
+            *state.shared_vec.write().unwrap() = results;
+            log::info!("background folder structure update completed");
+
+            std::thread::sleep(background_poll_seconds);
+        });
+    }
+
     let bind = matches.value_of("port").unwrap();
     let bind = u16::from_str_radix(&bind, 10).expect("port must be a valid number");
     let addr = ([0, 0, 0, 0], bind).into();
 
     log::info!("starting exporter on {}", addr);
 
-    prometheus_exporter_base::render_prometheus(addr, options, |request, options| {
-        Box::pin(perform_request(request, options))
+    prometheus_exporter_base::render_prometheus(addr, state, |request, state| {
+        Box::pin(perform_request(request, state))
     })
     .await
 }
diff --git a/src/options.rs b/src/options.rs
@@ -7,11 +7,15 @@ pub(crate) struct Options {
     pub folders_file: String,
     pub verbose: bool,
     pub folders_to_scan: FolderScanner,
+    pub background_poll_seconds: Option<std::time::Duration>,
 }
 
 impl Options {
     pub fn from_claps(matches: &clap::ArgMatches<'_>) -> Options {
         let folders_file = matches.value_of("folders_file").unwrap().to_owned();
+        let background_poll_seconds = matches
+            .value_of("background_poll_seconds")
+            .map(|s| std::time::Duration::from_secs(s.parse().unwrap()));
 
         let mut file = File::open(&folders_file).unwrap();
         let mut file_contents = String::new();
@@ -22,6 +26,7 @@ impl Options {
             folders_file,
             verbose: matches.is_present("verbose"),
             folders_to_scan: FolderScanner::from_json(&file_contents).unwrap(),
+            background_poll_seconds,
         }
     }
 }
diff --git a/src/state.rs b/src/state.rs
@@ -0,0 +1,19 @@
+use crate::folder_scanner::FolderWithSize;
+use crate::Options;
+use std::sync::Arc;
+use std::sync::RwLock;
+
+#[derive(Debug, Clone)]
+pub struct State {
+    pub(crate) options: Options,
+    pub(crate) shared_vec: Arc<RwLock<Vec<FolderWithSize>>>,
+}
+
+impl State {
+    pub(crate) fn new(options: Options) -> Self {
+        Self {
+            options,
+            shared_vec: Arc::new(RwLock::new(Vec::new())),
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -7,11 +7,15 @@ pub(crate) struct Options {`
`7`	`7`	`pub folders_file: String,`
`8`	`8`	`pub verbose: bool,`
`9`	`9`	`pub folders_to_scan: FolderScanner,`
	`10`	`+ pub background_poll_seconds: Option<std::time::Duration>,`
`10`	`11`	`}`
`11`	`12`
`12`	`13`	`impl Options {`
`13`	`14`	`pub fn from_claps(matches: &clap::ArgMatches<'_>) -> Options {`
`14`	`15`	`let folders_file = matches.value_of("folders_file").unwrap().to_owned();`
	`16`	`+ let background_poll_seconds = matches`
	`17`	`+ .value_of("background_poll_seconds")`
	`18`	`+ .map(\|s\| std::time::Duration::from_secs(s.parse().unwrap()));`
`15`	`19`
`16`	`20`	`let mut file = File::open(&folders_file).unwrap();`
`17`	`21`	`let mut file_contents = String::new();`
`@@ -22,6 +26,7 @@ impl Options {`
`22`	`26`	`folders_file,`
`23`	`27`	`verbose: matches.is_present("verbose"),`
`24`	`28`	`folders_to_scan: FolderScanner::from_json(&file_contents).unwrap(),`
	`29`	`+ background_poll_seconds,`
`25`	`30`	`}`
`26`	`31`	`}`
`27`	`32`	`}`