#[macro_use] extern crate serde_json; use clap::{Parser, Subcommand}; use rand::seq::IteratorRandom; use std::io; use std::path::PathBuf; use std::time::Instant; use std::{fs, sync::mpsc, thread}; use ledb::{Options, Storage}; use chgk_ledb_lib::db; use chgk_ledb_lib::questions; use chgk_ledb_lib::source; use crate::questions::{Question, QuestionsConverter}; use crate::source::ReadSourceQuestionsBatches; const ZIP_FILENAME: &str = "json.zip"; const NEW_DB_FILENAME: &str = "db.dat"; const DB_DIR: &str = "db"; #[derive(Subcommand, Debug)] enum Command { Write, Compact, Print { #[clap(value_parser, default_value = "0")] id: u32, }, ZipPrint { #[clap(value_parser, default_value = "0")] file_num: usize, #[clap(value_parser, default_value = "0")] num: usize, }, Write2, Print2 { #[clap(value_parser, default_value = "0")] id: u32, }, } #[derive(Parser, Debug)] #[clap(author, version, about, long_about = None)] #[clap(propagate_version = true)] struct Cli { #[clap(subcommand)] command: Command, #[clap(short, long, action)] measure: bool, } fn zip_reader_task(tx: mpsc::Sender) { let zip_file = fs::File::open(ZIP_FILENAME).unwrap(); let zip_reader = io::BufReader::new(zip_file); let archive = zip::ZipArchive::new(zip_reader).unwrap(); let mut source_questions = archive.source_questions(); let questions = source_questions .convert() .enumerate() .map(|(num, mut question)| { question.num = 1 + num as u32; question }); for question in questions { let res = tx.send(question); if res.is_err() { break; } } println!("read done"); } fn db_writer_task(rx: mpsc::Receiver) { let out_file: PathBuf = [DB_DIR, "data.mdb"].into_iter().collect(); match fs::metadata(&out_file) { Ok(x) if x.is_file() => { fs::remove_file(&out_file).unwrap(); println!(r#""{}" removed"#, out_file.to_str().unwrap()); } _ => {} }; let options: Options = serde_json::from_value(json!({ "map_size": 900 * 1024 * 1024, // 900mb "write_map": true, "map_async": true, "no_lock": true, "no_meta_sync": true, "no_sync": true, })) .unwrap(); let storage = Storage::new(DB_DIR, options).unwrap(); let collection = storage.collection("questions").unwrap(); let count = collection.load(rx).expect("load"); println!("loaded {count}"); println!("syncing to disk..."); storage.sync(true).unwrap(); print!("stats: "); let stats = storage.stat().unwrap(); println!("{:?}", stats); drop(storage); println!("write done"); } fn write_db() { let (tx, rx) = mpsc::channel::(); [ thread::spawn(move || zip_reader_task(tx)), thread::spawn(move || db_writer_task(rx)), ] .into_iter() .for_each(|handle| handle.join().expect("thread panic")); println!("all done"); } fn print_question_from(get_q: F) where F: FnOnce() -> Option, { let q = get_q().unwrap(); println!("{:#?}", q) } fn read_from_zip(file_num: usize, mut num: usize) -> Option { let mut rng = rand::thread_rng(); let zip_file = fs::File::open(ZIP_FILENAME).unwrap(); let zip_reader = io::BufReader::new(zip_file); let archive = zip::ZipArchive::new(zip_reader).unwrap(); let mut source_questions = archive.source_questions(); let (filename, batch) = if file_num == 0 { source_questions.choose(&mut rng).unwrap() } else { source_questions.nth(file_num - 1).unwrap() }; let mut batch = batch.unwrap(); batch.filename = filename; let questions: Vec = batch.into(); if num == 0 { num = (1..=questions.len()).choose(&mut rng).unwrap(); } Some(questions[num - 1].clone()) } fn compact_db() { let options: Options = serde_json::from_value(json!({ "write_map": true, "map_async": true, "no_lock": true, "no_meta_sync": true, "no_sync": true, "compact": true, })) .unwrap(); let storage = Storage::new(DB_DIR, options).unwrap(); storage.sync(true).unwrap(); let stats = storage.stat().unwrap(); println!("{:?}", stats); drop(storage); } fn read_from_db(mut id: u32) -> Option { let options: Options = serde_json::from_value(json!({ "read_only": true, "map_async": true, "no_lock": true, })) .unwrap(); let storage = Storage::new(DB_DIR, options).unwrap(); let collection = storage.collection("questions").unwrap(); let mut rng = rand::thread_rng(); if id == 0 { let last_id = collection.last_id().unwrap(); id = (1..=last_id).choose(&mut rng).unwrap(); } collection.get::(id).unwrap() } // measure and return time elapsed in `func` in seconds pub fn measure(func: F) -> f64 { let start = Instant::now(); func(); let elapsed = start.elapsed(); (elapsed.as_secs() as f64) + (elapsed.subsec_nanos() as f64 / 1_000_000_000.0) } pub fn measure_and_print(func: F) { let m = measure(func); eprintln!("{}", m); } fn main() { let args = Cli::parse(); let mut action: Box = match &args.command { Command::Write => Box::new(write_db), Command::Compact => Box::new(compact_db), Command::Print { id } => { let get_question = Box::new(|| read_from_db(*id)); Box::new(|| print_question_from(get_question)) } Command::ZipPrint { file_num, num } => { let get_question = Box::new(|| read_from_zip(*file_num, *num)); Box::new(|| print_question_from(get_question)) } Command::Write2 => Box::new(write_db2), Command::Print2 { id } => { let get_question = Box::new(|| read_from_db2(*id)); Box::new(|| print_question_from(get_question)) } }; if args.measure { action = Box::new(|| measure_and_print(action)); } action(); } fn read_from_db2(id: u32) -> Option { let reader: db::Reader = db::Reader::new(NEW_DB_FILENAME, 2048).expect("new db reader"); let mut questions = reader.iter(); match id { 0 => { let mut rng = rand::thread_rng(); questions.choose(&mut rng) } _ => questions.nth((id - 1) as usize), } } fn write_db2() { let (tx, rx) = mpsc::channel::(); [ thread::spawn(move || zip_reader_task(tx)), thread::spawn(move || db_writer2_task(rx)), ] .into_iter() .for_each(|handle| handle.join().expect("thread panic")); println!("all done"); } fn db_writer2_task(rx: mpsc::Receiver) { let writer_opts = db::WriterOpts::default(); let mut writer: db::Writer = db::Writer::new(NEW_DB_FILENAME, writer_opts).expect("new db writer"); writer .load(&mut rx.iter()) .unwrap_or_else(|e| panic!("db writer load, {e:#?}")); writer.finish().expect("db writer finish"); println!("write done"); }