From fe2a9bd7ba95e58a4e04aa6bc2de54f0baa93718 Mon Sep 17 00:00:00 2001
From: Quentin <qmatillat@gmail.com>
Date: Wed, 11 Dec 2024 00:50:48 +0100
Subject: [PATCH] feat(warmup)!: Add warmup wait command (#379)

This allows to programmatically (with script) wait for elements to be
ready, instead of hardcoding a duration.

<details><summary>Scripts examples</summary>

Warm-up:
```bash
KEY="data/${1:0:2}/$1";
if [ "$(aws s3api head-object --bucket cold --key "$KEY" --query StorageClass)" == '"GLACIER"' ]
then
    aws s3api restore-object --restore-request Days=7 --bucket cold --key "$KEY"
fi
```

Warm-up wait:
```bash
KEY="data/${1:0:2}/$1";
until [ "$(aws s3api head-object --bucket cold --key "$KEY" --query StorageClass)" != '"GLACIER"' ]
do
    sleep 60s
done
```
</details>

This is a simple version that will wait for each object sequentially
(but that shouldn't be a problem as all warm up request are done in a
first pass, and thus runs in parallel).

I'm up for discussion if another model looks good, like
- moving the loop in the lib, and the `wait` script just check if the
file is ready
- having only one script, that warm-up and then return the status

---------

Signed-off-by: MATILLAT Quentin <qmatillat@gmail.com>
---
 crates/core/src/repository.rs         | 13 ++++++++++++-
 crates/core/src/repository/warm_up.rs | 27 +++++++++++++++++++++------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/crates/core/src/repository.rs b/crates/core/src/repository.rs
index 703490ed..309810c7 100644
--- a/crates/core/src/repository.rs
+++ b/crates/core/src/repository.rs
@@ -157,6 +157,14 @@ pub struct RepositoryOptions {
     #[cfg_attr(feature = "merge", merge(strategy = conflate::option::overwrite_none))]
     pub warm_up_command: Option<CommandInput>,
 
+    /// Wait for end of warm up by running the command with %id replaced by pack id
+    #[cfg_attr(
+        feature = "clap",
+        clap(long, global = true, conflicts_with = "warm_up_wait",)
+    )]
+    #[cfg_attr(feature = "merge", merge(strategy = conflate::option::overwrite_none))]
+    pub warm_up_wait_command: Option<CommandInput>,
+
     /// Duration (e.g. 10m) to wait after warm up
     #[cfg_attr(feature = "clap", clap(long, global = true, value_name = "DURATION"))]
     #[serde_as(as = "Option<DisplayFromStr>")]
@@ -748,7 +756,10 @@ impl<P: ProgressBars, S> Repository<P, S> {
     ///
     /// * If the command could not be parsed.
     /// * If the thread pool could not be created.
-    pub fn warm_up_wait(&self, packs: impl ExactSizeIterator<Item = PackId>) -> RusticResult<()> {
+    pub(crate) fn warm_up_wait(
+        &self,
+        packs: impl ExactSizeIterator<Item = PackId> + Clone,
+    ) -> RusticResult<()> {
         warm_up_wait(self, packs)
     }
 }
diff --git a/crates/core/src/repository/warm_up.rs b/crates/core/src/repository/warm_up.rs
index a1a74a6b..a592706c 100644
--- a/crates/core/src/repository/warm_up.rs
+++ b/crates/core/src/repository/warm_up.rs
@@ -31,10 +31,13 @@ pub(super) mod constants {
 /// * If the thread pool could not be created.
 pub(crate) fn warm_up_wait<P: ProgressBars, S>(
     repo: &Repository<P, S>,
-    packs: impl ExactSizeIterator<Item = PackId>,
+    packs: impl ExactSizeIterator<Item = PackId> + Clone,
 ) -> RusticResult<()> {
-    warm_up(repo, packs)?;
-    if let Some(wait) = repo.opts.warm_up_wait {
+    warm_up(repo, packs.clone())?;
+
+    if let Some(warm_up_wait_cmd) = &repo.opts.warm_up_wait_command {
+        warm_up_command(packs, warm_up_wait_cmd, &repo.pb, &WarmUpType::WaitPack)?;
+    } else if let Some(wait) = repo.opts.warm_up_wait {
         let p = repo.pb.progress_spinner(format!("waiting {wait}..."));
         sleep(*wait);
         p.finish();
@@ -58,13 +61,19 @@ pub(crate) fn warm_up<P: ProgressBars, S>(
     packs: impl ExactSizeIterator<Item = PackId>,
 ) -> RusticResult<()> {
     if let Some(warm_up_cmd) = &repo.opts.warm_up_command {
-        warm_up_command(packs, warm_up_cmd, &repo.pb)?;
+        warm_up_command(packs, warm_up_cmd, &repo.pb, &WarmUpType::WarmUp)?;
     } else if repo.be.needs_warm_up() {
         warm_up_repo(repo, packs)?;
     }
     Ok(())
 }
 
+#[derive(Debug)]
+enum WarmUpType {
+    WarmUp,
+    WaitPack,
+}
+
 /// Warm up the repository using a command.
 ///
 /// # Arguments
@@ -80,8 +89,12 @@ fn warm_up_command<P: ProgressBars>(
     packs: impl ExactSizeIterator<Item = PackId>,
     command: &CommandInput,
     pb: &P,
+    ty: &WarmUpType,
 ) -> RusticResult<()> {
-    let p = pb.progress_counter("warming up packs...");
+    let p = pb.progress_counter(match ty {
+        WarmUpType::WarmUp => "warming up packs...",
+        WarmUpType::WaitPack => "waiting for packs to be ready...",
+    });
     p.set_length(packs.len() as u64);
     for pack in packs {
         let args: Vec<_> = command
@@ -102,11 +115,13 @@ fn warm_up_command<P: ProgressBars>(
                     err,
                 )
                 .attach_context("command", command.to_string())
+                .attach_context("type", format!("{ty:?}"))
             })?;
 
         if !status.success() {
-            warn!("warm-up command was not successful for pack {pack:?}. {status}");
+            warn!("{ty:?} command was not successful for pack {pack:?}. {status}");
         }
+        p.inc(1);
     }
     p.finish();
     Ok(())