[pdm-devel] [PATCH proxmox-datacenter-manager v7 23/24] metric collection: allow to wait until completion when triggering collection manually

[Lukas Wagner]

This allows us to request the latest metrics for a single remote in the
rrddata API calls, closing the gap in data the results from the longer
10min poll interval.

Signed-off-by: Lukas Wagner <l.wagner at proxmox.com>
---

Notes:
    New in v7.

 server/src/api/metric_collection.rs            |  2 +-
 server/src/api/remotes.rs                      |  2 +-
 .../src/metric_collection/collection_task.rs   | 16 +++++++++++-----
 server/src/metric_collection/mod.rs            | 18 +++++++++++++++---
 4 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/server/src/api/metric_collection.rs b/server/src/api/metric_collection.rs
index 845cc0e6..0658fb1f 100644
--- a/server/src/api/metric_collection.rs
+++ b/server/src/api/metric_collection.rs
@@ -34,7 +34,7 @@ const SUBDIRS: SubdirMap = &sorted!([
 )]
 /// Trigger metric collection for a provided remote or for all remotes if no remote is passed.
 pub async fn trigger_metric_collection(remote: Option<String>) -> Result<(), Error> {
-    crate::metric_collection::trigger_metric_collection(remote).await?;
+    crate::metric_collection::trigger_metric_collection(remote, false).await?;
 
     Ok(())
 }
diff --git a/server/src/api/remotes.rs b/server/src/api/remotes.rs
index c2489e60..033aa7c9 100644
--- a/server/src/api/remotes.rs
+++ b/server/src/api/remotes.rs
@@ -186,7 +186,7 @@ pub async fn add_remote(mut entry: Remote, create_token: Option<String>) -> Resu
 
     pdm_config::remotes::save_config(&remotes)?;
 
-    if let Err(e) = metric_collection::trigger_metric_collection(Some(name)).await {
+    if let Err(e) = metric_collection::trigger_metric_collection(Some(name), false).await {
         log::error!("could not trigger metric collection after adding remote: {e}");
     }
 
diff --git a/server/src/metric_collection/collection_task.rs b/server/src/metric_collection/collection_task.rs
index 1e23fa88..5f67d65d 100644
--- a/server/src/metric_collection/collection_task.rs
+++ b/server/src/metric_collection/collection_task.rs
@@ -41,7 +41,7 @@ pub const MIN_COLLECTION_INTERVAL: u64 = 10;
 
 /// Control messages for the metric collection task.
 pub(super) enum ControlMsg {
-    TriggerMetricCollection(Option<String>),
+    TriggerMetricCollection(Option<String>, oneshot::Sender<()>),
 }
 
 /// Task which periodically collects metrics from all remotes and stores
@@ -136,20 +136,26 @@ impl MetricCollectionTask {
     /// Handle a control message for force-triggered collection.
     async fn handle_control_message(&mut self, message: ControlMsg) {
         if let Some(remotes) = Self::load_remote_config() {
-            match message {
-                ControlMsg::TriggerMetricCollection(Some(remote)) => {
+            let done_tx = match message {
+                ControlMsg::TriggerMetricCollection(Some(remote), done_tx) => {
                     log::debug!("starting metric collection for remote '{remote}'- triggered by control message");
                     self.fetch_remotes(&remotes, &[remote]).await;
+                    done_tx
                 }
-                ControlMsg::TriggerMetricCollection(None) => {
+                ControlMsg::TriggerMetricCollection(None, done_tx) => {
                     log::debug!("starting metric collection from all remotes - triggered by control message");
                     let to_fetch = remotes
                         .iter()
                         .map(|(name, _)| name.into())
                         .collect::<Vec<String>>();
                     self.fetch_remotes(&remotes, &to_fetch).await;
+                    done_tx
                 }
-            }
+            };
+
+            // We don't care about the result, if the caller does not wait for the result, it
+            // might have dropped the receiver already.
+            let _ = done_tx.send(());
         }
     }
 
diff --git a/server/src/metric_collection/mod.rs b/server/src/metric_collection/mod.rs
index 2ddffda5..0e6860fc 100644
--- a/server/src/metric_collection/mod.rs
+++ b/server/src/metric_collection/mod.rs
@@ -5,6 +5,7 @@ use std::sync::OnceLock;
 use anyhow::{bail, Error};
 use nix::sys::stat::Mode;
 use tokio::sync::mpsc::{self, Sender};
+use tokio::sync::oneshot;
 
 use pdm_api_types::MetricCollectionStatus;
 use pdm_buildcfg::PDM_STATE_DIR_M;
@@ -66,15 +67,26 @@ pub fn start_task() -> Result<(), Error> {
     Ok(())
 }
 
-/// Schedule metric collection for a given remote as soon as possible.
+/// Schedule metric collection as soon as possible.
+///
+/// If `remote` is `Some(String)`, then the remote with the given ID is
+/// collected. If remote is `None`, all remotes are scheduled for collection.
+/// If `wait` is `true`, this function waits for the completion of the requested
+/// metric collection run.
 ///
 /// Has no effect if the tx end of the channel has not been initialized yet.
 /// Returns an error if the mpsc channel has been closed already.
-pub async fn trigger_metric_collection(remote: Option<String>) -> Result<(), Error> {
+pub async fn trigger_metric_collection(remote: Option<String>, wait: bool) -> Result<(), Error> {
+    let (done_sender, done_receiver) = oneshot::channel();
+
     if let Some(sender) = CONTROL_MESSAGE_TX.get() {
         sender
-            .send(ControlMsg::TriggerMetricCollection(remote))
+            .send(ControlMsg::TriggerMetricCollection(remote, done_sender))
             .await?;
+
+        if wait {
+            done_receiver.await?;
+        }
     }
 
     Ok(())
-- 
2.47.2