143 lines
4 KiB
Rust
143 lines
4 KiB
Rust
use std::{collections::BTreeSet, ffi::OsString, time::Duration};
|
|
|
|
use clap::Parser;
|
|
use rand::{thread_rng, Rng};
|
|
use rand_distr::Normal;
|
|
use sqlx::{
|
|
sqlite::{SqliteConnectOptions, SqlitePoolOptions},
|
|
SqlitePool,
|
|
};
|
|
use tokio::task::JoinSet;
|
|
use tokio_retry::Retry;
|
|
use witch_watch::{
|
|
get_db_pool,
|
|
import_utils::{add_omega_watches, add_user, add_watch_quest},
|
|
DbId, WatchQuest,
|
|
};
|
|
|
|
#[derive(Debug, Parser)]
|
|
struct Cli {
|
|
/// path to the movie database
|
|
#[clap(long = "database", short)]
|
|
pub db_path: OsString,
|
|
|
|
/// number of users to create
|
|
#[clap(long, short, default_value_t = 1000)]
|
|
pub users: usize,
|
|
|
|
/// expected gaussian value for number of movies per use
|
|
#[clap(long = "movies", short, default_value_t = 100)]
|
|
pub movies_per_user: u32,
|
|
|
|
/// path to the dictionary to be used for usernames [default:
|
|
/// /usr/share/dict/words]
|
|
#[clap(long, short)]
|
|
pub words: Option<OsString>,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let cli = Cli::parse();
|
|
let path = cli.db_path;
|
|
let num_users = cli.users;
|
|
let mpu = cli.movies_per_user as f32;
|
|
let dict = if let Some(dict) = cli.words {
|
|
dict
|
|
} else {
|
|
"/usr/share/dict/words".into()
|
|
};
|
|
|
|
let words = std::fs::read_to_string(dict).expect("tried to open {dict:?}");
|
|
let words: Vec<&str> = words.split('\n').collect();
|
|
|
|
let opts = SqliteConnectOptions::new().filename(&path).read_only(true);
|
|
let movie_db = SqlitePoolOptions::new()
|
|
.idle_timeout(Duration::from_secs(90))
|
|
.connect_with(opts)
|
|
.await
|
|
.expect("could not open movies db");
|
|
let ww_db = get_db_pool().await;
|
|
|
|
let users = &gen_users(num_users, &words, &ww_db).await;
|
|
let movies = &add_omega_watches(&ww_db, &movie_db).await;
|
|
|
|
let normal = Normal::new(mpu, mpu / 10.0).unwrap();
|
|
let rng = &mut thread_rng();
|
|
for user in users {
|
|
let mut joinset = JoinSet::new();
|
|
|
|
let mut mset = BTreeSet::new();
|
|
let num_movies = rng.sample(normal) as usize;
|
|
|
|
while mset.len() < num_movies {
|
|
let idx = rng.gen_range(0..10_000usize);
|
|
mset.insert(idx);
|
|
}
|
|
dbg!("done with mset pop");
|
|
|
|
let retry_strategy = tokio_retry::strategy::ExponentialBackoff::from_millis(100)
|
|
.map(tokio_retry::strategy::jitter)
|
|
.take(4);
|
|
|
|
for movie in mset.iter() {
|
|
let movie = movies[*movie];
|
|
let quest = WatchQuest {
|
|
id: DbId::new(),
|
|
user: *user,
|
|
watch: movie,
|
|
is_public: true,
|
|
already_watched: false,
|
|
};
|
|
|
|
let retry_strategy = retry_strategy.clone();
|
|
let db = ww_db.clone();
|
|
let key = quest.id.as_string();
|
|
joinset.spawn(async move {
|
|
(
|
|
key,
|
|
Retry::spawn(retry_strategy, || async {
|
|
add_watch_quest(&db, quest).await
|
|
})
|
|
.await,
|
|
)
|
|
});
|
|
}
|
|
|
|
// stragglers
|
|
while (joinset.join_next().await).is_some() {}
|
|
}
|
|
}
|
|
|
|
async fn gen_users(num: usize, words: &[&str], pool: &SqlitePool) -> Vec<DbId> {
|
|
let mut rng = thread_rng();
|
|
let rng = &mut rng;
|
|
let range = 0usize..(words.len());
|
|
let mut users = Vec::with_capacity(num);
|
|
for _ in 0..num {
|
|
let n1 = rng.gen_range(range.clone());
|
|
let n2 = rng.gen_range(range.clone());
|
|
let n3 = rng.gen_range(range.clone());
|
|
let nn = rng.gen_range(0..200);
|
|
|
|
let n1 = &words[n1];
|
|
let n2 = &words[n2];
|
|
let email_domain = &words[n3];
|
|
|
|
let username = format!("{n1}{n2}{nn}");
|
|
let displayname = Some(format!("{n1} {n2}"));
|
|
let email = Some(format!("{username}@{email_domain}"));
|
|
let id = DbId::new();
|
|
|
|
add_user(
|
|
pool,
|
|
&username,
|
|
displayname.as_deref(),
|
|
email.as_deref(),
|
|
Some(id),
|
|
)
|
|
.await;
|
|
(&mut users).push(id);
|
|
}
|
|
|
|
users
|
|
}
|