//! A simple end-to-end replication benchmark. It sets up N Ditto instances
//! connected in a chain, and measures how long it takes for one peer to
//! insert a number of documents, and for the last peer in the chain to receive them.
//!
//! This benchmark is parameterized by the size of the documents being sent,
//! and the number of documents being sent. This is purely an insertion benchmark,
//! and does not measure mutations to existing documents.

use std::sync::{
    atomic::{AtomicUsize, Ordering},
    Arc,
};

use criterion::*;
use rand::distributions::{Alphanumeric, DistString};
use serde_json::json;

mod common;
use common::setup;

/// Generate a random payload of `len` bytes. It's important to use random data
/// so that the compression middleware in the transports doesn't skew the results.
fn generate_payload(len: usize) -> String {
    let mut rng = rand::thread_rng();
    Alphanumeric.sample_string(&mut rng, len)
}

const PEERS: usize = 3;
const DOCS: usize = 100;

fn simple(c: &mut Criterion) {
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();

    let mut group = c.benchmark_group("replication_e2e");

    let ditto = runtime.block_on(setup(PEERS)).unwrap();

    // wildcard subscription to make sure data floodfills
    for d in &ditto {
        d.sync()
            .register_subscription_v2("SELECT * FROM data")
            .unwrap();

        d.sync().start().unwrap();
    }

    let received_docs = Arc::new(AtomicUsize::new(0));

    ditto[PEERS - 1]
        .store()
        .register_observer_v2("SELECT * FROM data", {
            let received_docs = received_docs.clone();
            move |items| {
                received_docs.fetch_add(items.item_count(), Ordering::SeqCst);
            }
        })
        .unwrap();

    let chunk_sizes = [
        1usize, // Mostly sending overhead
        64,     // Small user payloads
        512,
        1024,
        4096,
        8192,
        16 * 1024,
    ];
    for payload_size in chunk_sizes {
        for d in &ditto {
            runtime.block_on(async {
                d.store().execute_v2("DELETE FROM data").await.unwrap();
            })
        }

        // this isn't quite accurate but it's close enough (we don't account for id or document
        // wrapper overhead)
        group.throughput(Throughput::Bytes((DOCS * payload_size) as _));

        group.bench_function(format!("replication_{}_bytes", payload_size), |b| {
            b.to_async(&runtime).iter_batched(
                || {
                    // reset counter
                    received_docs.store(0, Ordering::SeqCst);
                    (&ditto[0], received_docs.clone())
                },
                |(src, received_docs)| async move {
                    let docs = (0..DOCS)
                        .map(|_| json! ({"data": generate_payload(payload_size)}))
                        .collect::<Vec<_>>();

                    src.store()
                        .execute_v2((
                            "INSERT INTO data DOCUMENTS (:docs)",
                            json! ({ "docs": docs }),
                        ))
                        .await
                        .unwrap();

                    while received_docs.load(Ordering::SeqCst) < DOCS {
                        tokio::task::yield_now().await;
                    }
                },
                criterion::BatchSize::LargeInput,
            );
        });
    }

    group.finish();
}

criterion_group!(replication, simple);
criterion_main!(replication);
