Skip to content

Commit

Permalink
Add moving average
Browse files Browse the repository at this point in the history
  • Loading branch information
@lbeder
lbeder committed Nov 25, 2024
1 parent 40dbbd5 commit 5ee2aa3
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 27 deletions.
17 changes: 16 additions & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,8 @@ Options:
Specifies the number of most recent checkpoints to keep, while automatically deleting older ones [default: 1]
--restore-from-checkpoint <RESTORE_FROM_CHECKPOINT>
Path to an existing checkpoint from which to resume the derivation process
--iteration-moving-window <ITERATION_MOVING_WINDOW>
Iteration time sampling moving window size [default: 10]
-h, --help
Print help
```
Expand Down Expand Up @@ -213,6 +215,8 @@ Salt is: s...t
Password is: p...d
████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ 1/10 10% (54s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
```
Final result:
Expand All @@ -228,6 +232,8 @@ Password is: p...d
████████████████████████████████████████████████████████████████████████████████ 10/10 100% (0s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
Key is (please highlight to see): 0xad9aa031287b42f45c40a5caf3b3ed47f795d9315d22ab50a25652b3f2a6b716
Start time: 2024-09-18 18:57:34
Expand Down Expand Up @@ -313,6 +319,7 @@ Password is: p...d
████████████████████████████████████████████████████████████████░░░░░░░░░░░░░░░░ 5/10 80% (10s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
Created checkpoint #5 with data hash 0x3c0c7ab8bb2001c1efd67ce049a437c760cf95d4cc2967160b708fb7216d74d1
```
Expand Down Expand Up @@ -408,14 +415,16 @@ The password, salt and internal data are correct
████████████████████████████████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ 5/10 0% (4s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
```
Final result:
```sh
████████████████████████████████████████████████████████████████████████████████ 10/10 100% (0s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
Key is (please highlight to see): 0xad9aa031287b42f45c40a5caf3b3ed47f795d9315d22ab50a25652b3f2a6b716
Start time: 2024-09-18 19:00:59
Expand Down Expand Up @@ -464,6 +473,8 @@ Final result:
```sh
████████████████████████████████████████████████████████████████████████████████ 20/20 100% (0s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
Key is (please highlight to see): 0x07eee820a3f92c5577dedd07e7d325dc58bb1064f9ae05af30be9863ec6e7354
Start time: 2024-09-18 19:00:59
Expand Down Expand Up @@ -496,6 +507,8 @@ Password is: p...d
████████████████████████████████████████████████████████████████████████████████ 10/10 100% (0s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
Key is (please highlight to see): 0xad9aa031287b42f45c40a5caf3b3ed47f795d9315d22ab50a25652b3f2a6b716
Key (base64) is (please highlight to see): rZqgMSh7QvRcQKXK87PtR/eV2TFdIqtQolZSs/KmtxY+W4
Key (base58) is (please highlight to see): CggHSjC3rpDdCGcbL2uB28qpFeBsWVsUMph1iGpbnDGy
Expand Down Expand Up @@ -530,6 +543,8 @@ SlowKey Parameters:
████████████████████████████████████████████████████████████████████████████████ 10/10 100% (0s)
Iteration time moving average (10): 4s 425ms, last iteration time: 4s 322ms
Key is (please highlight to see): 0xad9aa031287b42f45c40a5caf3b3ed47f795d9315d22ab50a25652b3f2a6b716
Saved encrypted output to "~/output.enc"
Expand Down
72 changes: 46 additions & 26 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use mimalloc::MiMalloc;
use sha2::{Digest, Sha512};
use std::{
cmp::Ordering,
collections::VecDeque,
env,
path::PathBuf,
str::from_utf8,
Expand Down Expand Up @@ -152,6 +153,9 @@ enum Commands {
help = "Path to an existing checkpoint from which to resume the derivation process"
)]
restore_from_checkpoint: Option<PathBuf>,

#[arg(long, default_value = "10", help = "Iteration time sampling moving window size")]
iteration_moving_window: u32,
},

#[command(about = "Decrypt and print a checkpoint")]
Expand Down Expand Up @@ -457,6 +461,7 @@ fn main() {
checkpoint_dir,
restore_from_checkpoint,
max_checkpoints_to_keep,
iteration_moving_window,
}) => {
println!(
"Please input all data either in raw or hex format starting with the {} prefix\n",
Expand Down Expand Up @@ -574,8 +579,8 @@ fn main() {

let mb = MultiProgress::new();

// Please note that we are using a custom message, instead of percents, since we want a higher resolution
// that the default one
// Create the main progress bar. Please note that we are using a custom message, instead of percents, since
// we want a higher resolution that the default one
let pb = mb.add(ProgressBar::new(iterations as u64)).with_style(
ProgressStyle::with_template("{bar:80.cyan/blue} {pos:>8}/{len:8} {msg}% ({eta})").unwrap(),
);
Expand All @@ -587,27 +592,20 @@ fn main() {
// Set the percent using a custom message
pb.set_message(format!("{}", (offset * 100) as f64 / iterations as f64));

let mut cpb: Option<ProgressBar> = None;

if checkpoint.is_some() && checkpointing_interval != 0 {
cpb = Some(
mb.add(ProgressBar::new(
((iterations - offset) / checkpointing_interval) as u64,
))
.with_style(ProgressStyle::with_template("{msg}").unwrap()),
);

if let Some(ref mut cpb) = &mut cpb {
cpb.set_position((offset / checkpointing_interval) as u64);
cpb.reset_eta();
cpb.enable_steady_tick(Duration::from_secs(1));
}
}
// Create a progress bar to track iteration times and checkpoints
let ipb = mb
.add(ProgressBar::new(iterations as u64))
.with_style(ProgressStyle::with_template("{msg}").unwrap());

let start_time = SystemTime::now();
let running_time = Instant::now();
let mut iteration_time = Instant::now();
let mut samples: VecDeque<u128> = VecDeque::new();

let slowkey = SlowKey::new(&slowkey_opts);

let mut checkpoint_info = String::new();

let prev_data_mutex = Arc::new(Mutex::new(prev_data));
let prev_data_thread = Arc::clone(&prev_data_mutex);

Expand All @@ -618,6 +616,31 @@ fn main() {
&offset_data,
offset,
|current_iteration, current_data| {
// Track iteration times
let last_iteration_time = iteration_time.elapsed().as_millis();
iteration_time = Instant::now();

samples.push_back(last_iteration_time);

// If we have more than the required samples, remove the oldest one
if samples.len() > iteration_moving_window as usize {
samples.pop_front();
}

// Calculate the moving average
let moving_average = samples.iter().sum::<u128>() as f64 / samples.len() as f64;

let iteration_info = format!(
"\nIteration time moving average ({}): {}, last iteration time: {}",
iteration_moving_window.to_string().cyan(),
format_duration(Duration::from_millis(moving_average as u64))
.to_string()
.cyan(),
format_duration(Duration::from_millis(last_iteration_time as u64))
.to_string()
.cyan(),
);

let mut prev_data = prev_data_thread.lock().unwrap();

// Create a checkpoint if we've reached the checkpoint interval
Expand All @@ -626,16 +649,14 @@ fn main() {

if let Some(checkpoint) = &mut checkpoint {
checkpoint.create_checkpoint(current_iteration, current_data, prev_data);
}

if let Some(ref mut cpb) = &mut cpb {
let hash = Checkpoint::hash_checkpoint(current_iteration, current_data, prev_data);

cpb.set_message(format!(
checkpoint_info = format!(
"\nCreated checkpoint #{} with data hash {}",
(current_iteration + 1).to_string().cyan(),
format!("0x{}", hex::encode(hash)).cyan()
));
);
}
}

Expand All @@ -652,14 +673,13 @@ fn main() {
"{}",
((current_iteration + 1) * 100) as f64 / iterations as f64
));

ipb.set_message(format!("{}{}", iteration_info, checkpoint_info));
},
);

pb.finish();

if let Some(ref mut cpb) = &mut cpb {
cpb.finish();
}
ipb.finish();

key
});
Expand Down

0 comments on commit 5ee2aa3

Please sign in to comment.