chgk_ledb/app/benches/questions_bench.rs

#[macro_use]
extern crate criterion;
extern crate serde;
extern crate serde_derive;
extern crate serde_json;
extern crate tempfile;

use chgk_ledb_lib::db::{Reader, Writer, WriterOpts};
use chgk_ledb_lib::questions::{binary::Question, QuestionsConverter};
use chgk_ledb_lib::source::ReadSourceQuestionsBatches;

use std::path::Path;
use std::time::Duration;
use std::{fs, io};

use criterion::{BatchSize, Criterion};
use tempfile::{tempdir, NamedTempFile};

const ZIP_FILENAME: &str = "../json.zip";
const NEW_DB_FILENAME: &str = "../db.dat";

const N: usize = 4096;

fn read_sample() -> Vec<Question> {
    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();

    source_questions
        .convert()
        .take(N)
        .enumerate()
        .map(|(num, mut question)| {
            question.num = 1 + num as u32;
            question
        })
        .collect()
}

fn prepare_db_writer<P: AsRef<Path>>(path: P) -> Writer<Question> {
    let opts = WriterOpts {
        compress_lvl: 1,
        data_buf_size: 100 * 1024 * 1024,
        out_buf_size: 100 * 1024 * 1024,
        current_buf_size: 10240,
    };

    Writer::new(path, opts).expect("new writer")
}

fn questions_read(c: &mut Criterion) {
    c.bench_function("questions_read", |b| {
        b.iter_batched(
            || {
                let reader: Reader<Question> =
                    Reader::new(NEW_DB_FILENAME, 4096).expect("new reader");
                reader.into_iter().take(N)
            },
            |reader| {
                for item in reader {
                    drop(item);
                }
            },
            BatchSize::SmallInput,
        )
    });
}

fn questions_write(c: &mut Criterion) {
    let dir = tempdir().expect("tempdir");

    c.bench_function("questions_write", |b| {
        b.iter_batched(
            || {
                let tmpfile = NamedTempFile::new_in(dir.path())
                    .expect("new tempfile")
                    .into_temp_path();
                let src = read_sample().into_iter();
                let writer = prepare_db_writer(&tmpfile);
                (src, writer)
            },
            |(mut src, mut writer)| {
                writer.load(&mut src).unwrap();
                writer.finish().unwrap();
            },
            BatchSize::SmallInput,
        )
    });
}

fn config() -> Criterion {
    Criterion::default()
        .sample_size(40)
        .warm_up_time(Duration::from_secs(7))
        .measurement_time(Duration::from_secs(20))
}

criterion_group! {name=benches; config = config(); targets = questions_read, questions_write}
criterion_main!(benches);