190 lines
5.4 KiB
Rust
190 lines
5.4 KiB
Rust
extern crate serde_json;
|
|
use clap::{Parser, Subcommand};
|
|
use futures::{pin_mut, Future};
|
|
use rand::distributions::Uniform;
|
|
use rand::seq::IteratorRandom;
|
|
use rand::{thread_rng, Rng};
|
|
|
|
use async_zip::tokio::read::seek::ZipFileReader;
|
|
use futures::stream::{self, StreamExt};
|
|
use std::time::Instant;
|
|
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
|
|
|
|
use async_db::WriterOpts;
|
|
|
|
use tokio::{fs, io};
|
|
use tokio_stream::wrappers::UnboundedReceiverStream;
|
|
|
|
use chgk_ledb_lib::async_db;
|
|
use chgk_ledb_lib::questions::Question;
|
|
use chgk_ledb_lib::questions::QuestionsConverterAsyncForStream;
|
|
use chgk_ledb_lib::source::ReadSourceQuestionsBatchesAsync;
|
|
|
|
const ZIP_FILENAME: &str = "json.zip";
|
|
const NEW_DB_FILENAME: &str = "db.dat";
|
|
|
|
#[derive(Subcommand, Debug)]
|
|
enum Command {
|
|
Write,
|
|
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,
|
|
},
|
|
}
|
|
|
|
#[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,
|
|
}
|
|
|
|
async fn zip_reader_task(tx: UnboundedSender<Question>) {
|
|
let mut file = fs::File::open(ZIP_FILENAME).await.expect("open zip");
|
|
let archive = ZipFileReader::with_tokio(&mut file)
|
|
.await
|
|
.expect("open zip file reader");
|
|
let mut source_questions = archive.source_questions();
|
|
let source_questions = source_questions.stream();
|
|
pin_mut!(source_questions);
|
|
|
|
let converter = source_questions.converter().convert();
|
|
pin_mut!(converter);
|
|
|
|
let mut num: u32 = 0;
|
|
while let Some(mut question) = converter.next().await {
|
|
question.num = 1 + num;
|
|
tx.send(question).expect("send");
|
|
num += 1;
|
|
}
|
|
|
|
println!("read done");
|
|
}
|
|
|
|
async fn print_question_from<F>(get_q: F)
|
|
where
|
|
F: Future<Output = Option<Question>>,
|
|
{
|
|
let q = get_q.await.expect("question not found");
|
|
println!("{:#?}", q)
|
|
}
|
|
|
|
async fn read_from_zip(file_num: usize, mut num: usize) -> Option<Question> {
|
|
let mut rng = thread_rng();
|
|
let zip_file = fs::File::open(ZIP_FILENAME).await.expect("open zip file");
|
|
let mut zip_reader = io::BufReader::new(zip_file);
|
|
let archive = ZipFileReader::with_tokio(&mut zip_reader)
|
|
.await
|
|
.expect("open zip file reader");
|
|
|
|
let mut source = archive.source_questions();
|
|
let files_count = source.len();
|
|
let file_index = if file_num == 0 {
|
|
let files = Uniform::new(0, files_count);
|
|
rng.sample(files)
|
|
} else {
|
|
file_num - 1
|
|
};
|
|
|
|
let src = source.get(file_index).await;
|
|
let src = stream::once(async { src.expect("get source file") });
|
|
pin_mut!(src);
|
|
let converter = src.converter();
|
|
let questions: Vec<_> = converter.convert().collect().await;
|
|
if num == 0 {
|
|
num = (1..=questions.len()).choose(&mut rng).unwrap();
|
|
}
|
|
let mut question = questions.get(num - 1).expect("get question").clone();
|
|
question.num = num as u32;
|
|
Some(question)
|
|
}
|
|
|
|
// measure and return time elapsed in `fut` in seconds
|
|
pub async fn measure<F: Future>(fut: F) -> f64 {
|
|
let start = Instant::now();
|
|
fut.await;
|
|
let elapsed = start.elapsed();
|
|
(elapsed.as_secs() as f64) + (elapsed.subsec_nanos() as f64 / 1_000_000_000.0)
|
|
}
|
|
|
|
pub async fn measure_and_print<F: Future>(fut: F) {
|
|
let m = measure(fut).await;
|
|
eprintln!("{}", m);
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let args = Cli::parse();
|
|
|
|
let mut action: Box<dyn Future<Output = _>> = match &args.command {
|
|
Command::Write => Box::new(write_db()),
|
|
Command::Print { id } => {
|
|
let get_question = read_from_db(*id);
|
|
Box::new(print_question_from(get_question))
|
|
}
|
|
Command::ZipPrint { file_num, num } => {
|
|
let get_question = read_from_zip(*file_num, *num);
|
|
Box::new(print_question_from(get_question))
|
|
}
|
|
};
|
|
|
|
if args.measure {
|
|
action = Box::new(measure_and_print(Box::into_pin(action)));
|
|
}
|
|
|
|
Box::into_pin(action).await;
|
|
}
|
|
|
|
async fn read_from_db(id: u32) -> Option<Question> {
|
|
let reader: async_db::Reader<Question> = async_db::Reader::new(NEW_DB_FILENAME)
|
|
.await
|
|
.expect("new db reader");
|
|
|
|
let len = reader.len();
|
|
|
|
let index = if len == 0 {
|
|
let mut rng = thread_rng();
|
|
let questions = Uniform::new(0, len);
|
|
rng.sample(questions)
|
|
} else {
|
|
id as usize - 1
|
|
};
|
|
|
|
match reader.get(index).await {
|
|
Ok(question) => Some(question),
|
|
Err(_) => None,
|
|
}
|
|
}
|
|
async fn write_db() {
|
|
let (tx, rx) = mpsc::unbounded_channel::<Question>();
|
|
tokio::try_join!(
|
|
tokio::spawn(zip_reader_task(tx)),
|
|
tokio::spawn(db_writer_task(rx))
|
|
)
|
|
.expect("tokio join");
|
|
println!("all done");
|
|
}
|
|
async fn db_writer_task(rx: UnboundedReceiver<Question>) {
|
|
let writer_opts = WriterOpts::default();
|
|
let mut writer: async_db::Writer<Question> =
|
|
async_db::Writer::new(NEW_DB_FILENAME, writer_opts)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("db writer load, {e:#?}"));
|
|
|
|
let stream: UnboundedReceiverStream<_> = rx.into();
|
|
let stream = stream;
|
|
writer.load(stream).await.expect("load");
|
|
writer.finish().await.expect("db writer finish");
|
|
|
|
println!("write done");
|
|
}
|