matrix-rust-sdk

use std::{sync::Arc, time::Duration};

use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};

use matrix_sdk::{

    SqliteEventCacheStore,

    linked_chunk::{LinkedChunk, LinkedChunkId, Update, lazy_loader},

};

use matrix_sdk_base::event_cache::{

    Event, Gap,

    store::{DEFAULT_CHUNK_CAPACITY, DynEventCacheStore, IntoEventCacheStore, MemoryStore},

};

use matrix_sdk_test::{ALICE, event_factory::EventFactory};

use ruma::room_id;

use tempfile::tempdir;

use tokio::runtime::Builder;

#[derive(Clone, Debug)]

enum Operation {

    PushItemsBack(Vec<Event>),

    PushGapBack(Gap),

}

#[cfg(not(feature = "codspeed"))]

const NUMBER_OF_EVENTS: &[u64] = &[10, 100, 1000, 10_000, 100_000];

#[cfg(feature = "codspeed")]

const NUMBER_OF_EVENTS: &[u64] = &[10, 100, 1000];

fn writing(c: &mut Criterion) {

    // Create a new asynchronous runtime.

    let runtime = Builder::new_multi_thread()

        .enable_time()

        .enable_io()

        .build()

        .expect("Failed to create an asynchronous runtime");

    let room_id = room_id!("!fabricandofitfaber:bar.baz");

    let linked_chunk_id = LinkedChunkId::Room(room_id);

    let event_factory = EventFactory::new().room(room_id).sender(&ALICE);

    let mut group = c.benchmark_group("Linked chunk writing");

    group.sample_size(10).measurement_time(Duration::from_secs(30));

    for &number_of_events in NUMBER_OF_EVENTS {

        let sqlite_temp_dir = tempdir().unwrap();

        // Declare new stores for this set of events.

        let stores: [(&str, Option<Arc<DynEventCacheStore>>); 3] = [

            ("none", None),

            ("memory store", Some(MemoryStore::default().into_event_cache_store())),

            (

                "sqlite store",

                runtime.block_on(async {

                    Some(

                        SqliteEventCacheStore::open(sqlite_temp_dir.path().join("bench"), None)

                            .await

                            .unwrap()

                            .into_event_cache_store(),

                    )

                }),

            ),

        ];

        for (store_name, store) in stores {

            // Create the operations we want to bench.

            let mut operations = Vec::new();

            {

                let mut events = (0..number_of_events)

                    .map(|nth| event_factory.text_msg(format!("foo {nth}")).into_event())

                    .peekable();

                let mut gap_nth = 0;

                while events.peek().is_some() {

                    {

                        let events_to_push_back = events.by_ref().take(80).collect::<Vec<_>>();

                        if events_to_push_back.is_empty() {

                            break;

                        }

                        operations.push(Operation::PushItemsBack(events_to_push_back));

                    }

                    {

                        operations

                            .push(Operation::PushGapBack(Gap { token: format!("gap{gap_nth}") }));

                        gap_nth += 1;

                    }

                }

            }

            // Define the throughput.

            group.throughput(Throughput::Elements(number_of_events));

            // Get a bencher.

            group.bench_with_input(

                BenchmarkId::new(format!("Linked chunk writing [{store_name}]"), number_of_events),

                &operations,

                |bencher, operations| {

                    // Bench the routine.

                    bencher.to_async(&runtime).iter_batched(

                        || operations.clone(),

                        |operations| async {

                            // The routine to bench!

                            let mut linked_chunk = LinkedChunk::<DEFAULT_CHUNK_CAPACITY, Event, Gap>::new_with_update_history();

                            for operation in operations {

                                match operation {

                                    Operation::PushItemsBack(events) => linked_chunk.push_items_back(events),

                                    Operation::PushGapBack(gap) => linked_chunk.push_gap_back(gap),

                                }

                            }

                            if let Some(store) = &store {

                                let updates = linked_chunk.updates().unwrap().take();

                                store.handle_linked_chunk_updates(linked_chunk_id, updates).await.unwrap();

                                // Empty the store.

                                store.handle_linked_chunk_updates(linked_chunk_id, vec![Update::Clear]).await.unwrap();

                            }

                        },

                        BatchSize::SmallInput

                    )

                },

            );

            {

                let _guard = runtime.enter();

                drop(store);

            }

        }

    }

    group.finish()

}

fn reading(c: &mut Criterion) {

    // Create a new asynchronous runtime.

    let runtime = Builder::new_multi_thread()

        .enable_time()

        .enable_io()

        .build()

        .expect("Failed to create an asynchronous runtime");

    let room_id = room_id!("!fabricandofitfaber:bar.baz");

    let linked_chunk_id = LinkedChunkId::Room(room_id);

    let event_factory = EventFactory::new().room(room_id).sender(&ALICE);

    let mut group = c.benchmark_group("Linked chunk reading");

    group.sample_size(10);

    for &num_events in NUMBER_OF_EVENTS {

        let sqlite_temp_dir = tempdir().unwrap();

        // Declare new stores for this set of events.

        let stores: [(&str, Arc<DynEventCacheStore>); 2] = [

            ("memory store", MemoryStore::default().into_event_cache_store()),

            (

                "sqlite store",

                runtime.block_on(async {

                    SqliteEventCacheStore::open(sqlite_temp_dir.path().join("bench"), None)

                        .await

                        .unwrap()

                        .into_event_cache_store()

                }),

            ),

        ];

        for (store_name, store) in stores {

            // Store some events and gap chunks in the store.

            {

                let mut events = (0..num_events)

                    .map(|nth| event_factory.text_msg(format!("foo {nth}")).into_event())

                    .peekable();

                let mut lc =

                    LinkedChunk::<DEFAULT_CHUNK_CAPACITY, Event, Gap>::new_with_update_history();

                let mut num_gaps = 0;

                while events.peek().is_some() {

                    let events_chunk = events.by_ref().take(80).collect::<Vec<_>>();

                    if events_chunk.is_empty() {

                        break;

                    }

                    lc.push_items_back(events_chunk);

                    lc.push_gap_back(Gap { token: format!("gap{num_gaps}") });

                    num_gaps += 1;

                }

                // Now persist the updates to recreate this full linked chunk.

                let updates = lc.updates().unwrap().take();

                runtime

                    .block_on(store.handle_linked_chunk_updates(linked_chunk_id, updates))

                    .unwrap();

            }

            // Define the throughput.

            group.throughput(Throughput::Elements(num_events));

            // Bench the lazy loader.

            group.bench_function(

                BenchmarkId::new(format!("Linked chunk lazy loader[{store_name}]"), num_events),

                |bencher| {

                    // Bench the routine.

                    bencher.to_async(&runtime).iter(|| async {

                        // Load the last chunk first,

                        let (last_chunk, chunk_id_gen) =

                            store.load_last_chunk(linked_chunk_id).await.unwrap();

                        let mut lc =

                            lazy_loader::from_last_chunk::<128, _, _>(last_chunk, chunk_id_gen)

                                .expect("no error when reconstructing the linked chunk")

                                .expect("there is a linked chunk in the store");

                        // Then load until the start of the linked chunk.

                        let mut cur_chunk_id = lc.chunks().next().unwrap().identifier();

                        while let Some(prev) =

                            store.load_previous_chunk(linked_chunk_id, cur_chunk_id).await.unwrap()

                        {

                            cur_chunk_id = prev.identifier;

                            lazy_loader::insert_new_first_chunk(&mut lc, prev)

                                .expect("no error when linking the previous lazy-loaded chunk");

                        }

                    })

                },

            );

            // Bench the metadata loader.

            group.bench_function(

                BenchmarkId::new(format!("Linked chunk metadata loader[{store_name}]"), num_events),

                |bencher| {

                    // Bench the routine.

                    bencher.to_async(&runtime).iter(|| async {

                        let _metadata = store

                            .load_all_chunks_metadata(linked_chunk_id)

                            .await

                            .expect("metadata must load");

                    })

                },

            );

            {

                let _guard = runtime.enter();

                drop(store);

            }

        }

    }

    group.finish()

}

criterion_group! {

    name = event_cache;

    config = Criterion::default();

    targets = writing, reading,

}

criterion_main!(event_cache);