108 lines
2.8 KiB
Rust
108 lines
2.8 KiB
Rust
use std::{
|
|
error::Error,
|
|
sync::{Arc, Mutex},
|
|
};
|
|
|
|
use clap::{arg, Parser};
|
|
use kdam::BarExt;
|
|
use kdam::{tqdm, Bar};
|
|
use rand::prelude::SliceRandom;
|
|
use reqwest::blocking::{multipart, Client};
|
|
use walkdir::WalkDir;
|
|
|
|
#[derive(Parser, Debug, Clone)]
|
|
struct Args {
|
|
#[arg(long)]
|
|
hint_name: String,
|
|
#[arg(long)]
|
|
fixture_dir: String,
|
|
#[arg(long)]
|
|
limit: Option<usize>,
|
|
#[arg(long)]
|
|
num_threads: Option<usize>,
|
|
}
|
|
|
|
fn main() -> Result<(), Box<dyn Error>> {
|
|
let args = Args::parse();
|
|
let hint_name = args.hint_name;
|
|
let num_threads = args.num_threads.unwrap_or(1);
|
|
let pb = Arc::new(Mutex::new(tqdm!()));
|
|
|
|
let mut entries = WalkDir::new(&args.fixture_dir)
|
|
.into_iter()
|
|
.filter_map(Result::ok)
|
|
.filter(|entry| entry.file_type().is_file())
|
|
.filter_map(|entry| {
|
|
let path = entry.path();
|
|
let name = path.to_str()?.to_string();
|
|
Some(name)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
// shuffle the entries
|
|
entries.shuffle(&mut rand::thread_rng());
|
|
|
|
// if there's a limit, drop the rest
|
|
if let Some(limit) = args.limit {
|
|
entries.truncate(limit);
|
|
}
|
|
|
|
let entry_slices = entries.chunks(entries.len() / num_threads);
|
|
|
|
let join_handles = entry_slices.map(|entry_slice| {
|
|
let hint_name = hint_name.clone();
|
|
let pb = pb.clone();
|
|
let entry_slice = entry_slice.to_vec();
|
|
std::thread::spawn(move || {
|
|
let client = Client::new();
|
|
for entry in entry_slice.into_iter() {
|
|
store_file(&client, &hint_name, &entry, pb.clone()).unwrap();
|
|
}
|
|
})
|
|
});
|
|
|
|
for join_handle in join_handles {
|
|
join_handle.join().unwrap();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn store_file(
|
|
client: &Client,
|
|
hint_name: &str,
|
|
file_name: &str,
|
|
pb: Arc<Mutex<Bar>>,
|
|
) -> Result<(), Box<dyn Error>> {
|
|
let file_bytes = std::fs::read(file_name)?;
|
|
|
|
let content_type = if file_name.ends_with(".html") {
|
|
"text/html"
|
|
} else if file_name.ends_with(".json") {
|
|
"application/json"
|
|
} else if file_name.ends_with(".jpg") || file_name.ends_with(".jpeg") {
|
|
"image/jpeg"
|
|
} else if file_name.ends_with(".png") {
|
|
"image/png"
|
|
} else if file_name.ends_with(".txt") {
|
|
"text/plain"
|
|
} else if file_name.ends_with(".kindle.images") {
|
|
"application/octet-stream"
|
|
} else {
|
|
"text/html"
|
|
};
|
|
|
|
let form = multipart::Form::new()
|
|
.text("content_type", content_type)
|
|
.text("compression_hint", hint_name.to_string())
|
|
.part("data", multipart::Part::bytes(file_bytes));
|
|
|
|
let _ = client
|
|
.post("http://localhost:7692/store")
|
|
.multipart(form)
|
|
.send();
|
|
|
|
let mut pb = pb.lock().unwrap();
|
|
pb.update(1)?;
|
|
Ok(())
|
|
}
|