thingvellir 0.0.14

use std::{
    collections::HashMap,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    thread,
    time::Duration,
};

use anyhow::Error;
use thingvellir::{
    service_builder, Commit, CommitToUpstream, DataCommitRequest, DataLoadRequest,
    LoadFromUpstream, MutableServiceHandle, ServiceData, ShardShutdownStats, ShutdownResult,
};

#[derive(Clone, Default, Debug)]
struct CountersUpstream {
    num_commits: Arc<AtomicU64>,
    upstream_data: HashMap<u64, CounterData>,
}

#[derive(Clone, Default, Copy, Debug)]
struct CounterData {
    count: u64,
}

impl ServiceData for CounterData {}

impl CommitToUpstream<u64, CounterData> for CountersUpstream {
    fn commit(&mut self, request: DataCommitRequest<u64, CounterData>) {
        println!(
            "[COMMIT TO UPSTREAM] key {} committed with count {}",
            request.key(),
            request.data().count
        );
        self.upstream_data.insert(*request.key(), *request.data());
        self.num_commits.fetch_add(1, Ordering::SeqCst);
        request.resolve()
    }
}

impl LoadFromUpstream<u64, CounterData> for CountersUpstream {
    fn load(&mut self, request: DataLoadRequest<u64, CounterData>) {
        if let Some(value) = self.upstream_data.get(request.key()) {
            println!(
                "[LOAD FROM UPSTREAM] key {} loaded with count {}",
                request.key(),
                value.count
            );
            request.resolve(*value);
        } else {
            println!(
                "[LOAD FROM UPSTREAM] key {} not found upstream. returning count of 0",
                request.key()
            );
            request.resolve(CounterData { count: 0 })
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    async fn increment_counter(
        key: u64,
        service: &mut MutableServiceHandle<u64, CounterData>,
    ) -> Option<u64> {
        let exec = service.execute_mut(key, move |counter_data| {
            println!(
                "Counter {} was {} and is being changed to {}",
                key,
                counter_data.count,
                counter_data.count + 1
            );
            counter_data.count += 1;
            Commit::default(counter_data.count)
        });
        exec.await.ok()
    }

    async fn load_counter(
        key: u64,
        service: &mut MutableServiceHandle<u64, CounterData>,
    ) -> Option<u64> {
        let exec = service.execute(key, |x| x.count);
        exec.await.ok()
    }

    let num_commits = Arc::new(AtomicU64::new(0));

    let (mut service, shutdown) = service_builder::<u64, CounterData>(1000).build_mutable(
        CountersUpstream {
            num_commits: Arc::clone(&num_commits),
            upstream_data: HashMap::new(),
        },
        thingvellir::DefaultCommitPolicy::Within(Duration::from_secs(3)),
    );

    // Increment keys 1, 2, and 3 until they each have counts of 500
    for key in 1..4 {
        for _ in 1..501 {
            increment_counter(key, &mut service).await;
        }
    }

    // Check all the keys have counts of 500
    for key in 1..4 {
        assert_eq!(
            load_counter(key, &mut service)
                .await
                .unwrap_or_else(|| panic!("No count for key {key}")),
            500
        );
    }

    // Wait for the accumulated commit policy to occur
    println!("Sleeping for 5 seconds");
    thread::sleep(Duration::from_secs(5));

    // Check to ensure the upstream commits were coalesced into 1 per key, totalling 3
    let num_commits_unwrapped = num_commits.load(Ordering::SeqCst);
    assert_eq!(num_commits_unwrapped, 3);

    // Increment keys 1, 2, and 3 500 more times, until they each have counts of 1000
    for key in 1..4 {
        for _ in 1..501 {
            increment_counter(key, &mut service).await;
        }
    }

    // Check all the keys have counts of 1000
    for key in 1..4 {
        assert_eq!(
            load_counter(key, &mut service)
                .await
                .unwrap_or_else(|| panic!("No count for key {key}")),
            1000
        );
    }

    // Shutdown and verify that the commit queue is coalesced & flushed properly
    let shutdown_result = shutdown.gracefully_shutdown().join().await;
    match shutdown_result {
        ShutdownResult::GracefullyShutdown(vec) => {
            // merge all of the ShardShutdownStats into a single ShardShutdownStats
            let mut final_stats = ShardShutdownStats::default();
            for stats in vec {
                final_stats.initial_commits_enqueued += stats.initial_commits_enqueued;
                final_stats.commits_completed += stats.commits_completed;
                final_stats.loads_completed += stats.loads_completed;
            }
            println!("Graceful shutdown stats: {:?}", final_stats);
            assert_eq!(final_stats.initial_commits_enqueued, 3);
            assert_eq!(final_stats.loads_completed, 0);
            assert_eq!(final_stats.commits_completed, 3);
        }
        ShutdownResult::HardShutdown => unreachable!(), // This shouldn't happen!
        ShutdownResult::AlreadyShutdown => unreachable!(), // This shouldn't happen!
    }

    // Total commits should now be 6
    let num_commits_unwrapped = num_commits.load(Ordering::SeqCst);
    assert_eq!(num_commits_unwrapped, 6);

    Ok(())
}