// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use criterion::{BatchSize, Bencher, Criterion, black_box};
use kvproto::kvrpcpb::Context;
use test_storage::SyncTestStorageBuilderApiV1;
use test_util::KvGenerator;
use tikv::storage::kv::Engine;
use txn_types::{Key, Mutation};

use super::{BenchConfig, DEFAULT_ITERATIONS, EngineFactory};

fn storage_prewrite<E: Engine, F: EngineFactory<E>>(b: &mut Bencher<'_>, config: &BenchConfig<F>) {
    let engine = config.engine_factory.build();
    let store = SyncTestStorageBuilderApiV1::from_engine(engine)
        .build(0)
        .unwrap();
    b.iter_batched(
        || {
            let kvs = KvGenerator::new(config.key_length, config.value_length)
                .generate(DEFAULT_ITERATIONS);

            let data: Vec<(Context, Vec<Mutation>, Vec<u8>)> = kvs
                .iter()
                .map(|(k, v)| {
                    (
                        Context::default(),
                        vec![Mutation::make_put(Key::from_raw(k), v.clone())],
                        k.clone(),
                    )
                })
                .collect();
            (data, &store)
        },
        |(data, store)| {
            for (context, mutations, primary) in data {
                black_box(store.prewrite(context, mutations, primary, 1).unwrap());
            }
        },
        BatchSize::SmallInput,
    );
}

fn storage_commit<E: Engine, F: EngineFactory<E>>(b: &mut Bencher<'_>, config: &BenchConfig<F>) {
    let engine = config.engine_factory.build();
    let store = SyncTestStorageBuilderApiV1::from_engine(engine)
        .build(0)
        .unwrap();
    b.iter_batched(
        || {
            let kvs = KvGenerator::new(config.key_length, config.value_length)
                .generate(DEFAULT_ITERATIONS);

            for (k, v) in &kvs {
                store
                    .prewrite(
                        Context::default(),
                        vec![Mutation::make_put(Key::from_raw(k), v.clone())],
                        k.clone(),
                        1,
                    )
                    .unwrap();
            }

            (kvs, &store)
        },
        |(kvs, store)| {
            for (k, _) in &kvs {
                black_box(store.commit(Context::default(), vec![Key::from_raw(k)], 1, 2)).unwrap();
            }
        },
        BatchSize::SmallInput,
    );
}

pub fn bench_storage<E: Engine, F: EngineFactory<E>>(
    c: &mut Criterion,
    configs: &[BenchConfig<F>],
) {
    let mut group = c.benchmark_group("storage");
    for config in configs {
        group.bench_with_input(
            format!("async_prewrite/{:?}", config),
            config,
            storage_prewrite,
        );
        group.bench_with_input(format!("async_commit/{:?}", config), config, storage_commit);
    }
    group.finish();
}