chgk_ledb/app_async/src/main.rs

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;
use chgk_ledb_lib::util::ErrorToString;

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,
}

#[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;
}

// 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);
}

async fn print_question_from<F>(get_q: F)
where
    F: Future<Output = Result<Question, String>>,
{
    let q = get_q.await.expect("question not found");
    println!("{:#?}", q)
}

async fn read_from_zip(file_num: usize, mut num: usize) -> Result<Question, String> {
    let mut rng = thread_rng();
    let zip_file = fs::File::open(ZIP_FILENAME).await.str_err()?;
    let mut zip_reader = io::BufReader::new(zip_file);
    let archive = ZipFileReader::with_tokio(&mut zip_reader).await.str_err()?;

    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
        .map_err(|e| format!("get file {file_num} => {e}"))?;
    let src = stream::once(async { src });
    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)
        .ok_or(format!("get question #{num} => None"))?
        .clone();
    question.num = num as u32;
    Ok(question)
}

async fn read_from_db(id: u32) -> Result<Question, String> {
    let reader: async_db::Reader<Question> = async_db::Reader::new(NEW_DB_FILENAME).await?;

    let len = reader.len();

    let index = if id == 0 {
        let mut rng = thread_rng();
        let questions = Uniform::new(0, len);
        rng.sample(questions)
    } else {
        id as usize - 1
    };

    reader
        .get(index)
        .await
        .map_err(|e| format!("get #{index} => {e}"))
}

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 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);

    source_questions
        .converter()
        .convert()
        .enumerate()
        .map(|(num, mut question)| {
            question.num = 1 + (num as u32);
            question
        })
        .for_each_concurrent(None, |question| async {
            tx.send(question).expect("send");
        })
        .await;

    println!("read 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");
}