extern crate serde; extern crate serde_derive; #[macro_use] extern crate serde_json; extern crate ledb; extern crate ledb_types; extern crate zip; use clap::{Parser, Subcommand}; use rand::seq::IteratorRandom; use std::io::{Read, Seek, Write}; use std::path::PathBuf; use std::time::Instant; use std::{fs, io, sync::mpsc, thread}; use ledb::{Options, Storage}; mod questions; mod source; use crate::questions::{Question, QuestionsConverter}; use crate::source::ReadSourceQuestionsBatches; const ZIP_FILENAME: &str = "json.zip"; 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 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(); archive .source_questions() .convert() .for_each(|question| tx.send(question).expect("send question")); 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 mut count: usize = 0; let count = &mut count; rx.into_iter().for_each(|question| { let result = collection.insert(&question); if result.is_err() { println!("-- {:#?}", question); panic!("{:#?}", result); } else { *count += 1; } }); println!("inserted {}\nwriting...", count); 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 || 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 { use std::io::Cursor; const LEN_SIZE: usize = std::mem::size_of::(); let cfg = bincode::config::standard().with_fixed_int_encoding(); let input = fs::File::open("test.bin").expect("open input"); let mut input = std::io::BufReader::with_capacity(100 * 1024 * 1024, input); let mut len_data: [u8; LEN_SIZE] = [0; LEN_SIZE]; input.read_exact(&mut len_data).expect("read len"); let len = u64::from_le_bytes(len_data) as usize; //println!("read len done"); let mut zdata = vec![0u8; len]; input.read_exact(&mut zdata[..len]).expect("read ztab"); let tab_data = zstd::decode_all(Cursor::new(zdata)).expect("zstd decode table"); let tab: (Vec, usize) = bincode::decode_from_slice(&tab_data, cfg).expect("bincode decode tab"); let tab = tab.0; //println!("read tab done, len {}", tab.len()); let index = match id { 0 => { let mut rng = rand::thread_rng(); (1..tab.len()).into_iter().choose(&mut rng).unwrap() } _ => (id - 1) as usize, }; //println!("index {}", index); let pos = *tab.get(index).expect("get pos"); let pos_next = *tab.get((index + 1) as usize).expect("get pos next"); let len = pos_next - pos; //println!("pos {} | next {} | len {}", pos, pos_next, len); input.seek_relative(pos as i64).expect("q seek"); let reader = input.take(len as u64); let data = zstd::decode_all(reader).expect("zstd decode data"); //println!("zstd decoded len {}", data.len()); let question: (Question, usize) = bincode::decode_from_slice(&data, cfg).expect("bincode decode q"); let question = question.0; //println!("read done"); Some(question) } fn write_db2() { let (tx, rx) = mpsc::channel::(); [ thread::spawn(move || 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) { const LEN_SIZE: usize = std::mem::size_of::(); const COMP_DATA_LEVEL: i32 = 2; const COMP_HDR_LEVEL: i32 = 2; let cfg = bincode::config::standard().with_fixed_int_encoding(); let mut table: Vec = vec![]; let buf_data: Vec = Vec::with_capacity(500 * 1024 * 1024); use std::io::Cursor; let mut buf = Cursor::new(buf_data); let mut pos: u32 = 0; let mut num = 1; rx.into_iter().for_each(|mut q| { q.num = num; let data = bincode::encode_to_vec(q, cfg).expect("bincode q encode"); let data = zstd::encode_all(Cursor::new(data), COMP_DATA_LEVEL).expect("zstd q encode"); let len = buf.write(&data).expect("write question"); table.push(pos); pos += len as u32; num += 1; }); table.push(pos); println!( "zbuf done, tab len {}, buf size {}", table.len(), buf.position() ); buf.set_position(0); let tab_data = bincode::encode_to_vec(&table, cfg).expect("encode table"); let zdata = zstd::encode_all(Cursor::new(tab_data), COMP_HDR_LEVEL).expect("zstd enc table"); let zlen = zdata.len() as u64; println!("z tab done, tab data_len = {}", zlen); let out = fs::File::create("test.bin").expect("out create"); let mut out = std::io::BufWriter::with_capacity(500 * 1024 * 1024, out); let len_writed = out.write(&zlen.to_le_bytes()).expect("write zlen"); assert_eq!(len_writed, LEN_SIZE); let ztab_writed = out.write(&zdata).expect("write tab zdata"); assert_eq!(ztab_writed, zdata.len()); drop(zdata); println!("header write done, pos: {}", out.stream_position().unwrap()); std::io::copy(&mut buf, &mut out).expect("copy z buf"); drop(buf); out.flush().expect("out flush"); println!("write done"); }