Skip to content

Commit

Permalink
feat(warmup)!: Add warmup wait command (#379)
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
tinou98 authored Dec 10, 2024
1 parent 70bc278 commit fe2a9bd
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 7 deletions.
13 changes: 12 additions & 1 deletion crates/core/src/repository.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>")]
Expand Down Expand Up @@ -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)
}
}
Expand Down
27 changes: 21 additions & 6 deletions crates/core/src/repository/warm_up.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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(())
Expand Down

0 comments on commit fe2a9bd

Please sign in to comment.