thingvellir 0.0.14

use std::borrow::BorrowMut;
use std::fmt::Display;
use std::hash::Hash;
use std::ops::Add;
use std::sync::Arc;

use futures::Future;
use log::{error, info, warn};

use tokio::sync::oneshot;
use tokio::sync::{mpsc, Mutex};
use tokio::task::{JoinError, JoinSet};
use tokio::time::{interval, Instant, Interval};

use super::commit_queue::{EnqueuedCommit, TakeOwnershipResult};
use super::messages::{InternalJoinSetResult, InternalMpscMessage, ShutdownSender};
use super::types::ExecuteIfCachedFn;
use super::{
    utils, AccumulatedCommitPolicy, DataMap, DataState, PendingCallback, ServiceHandleMessage,
    ShardConfig, ShardShutdownStats, ShardStats, TakeDataSender, TakenData, TryTakeResult,
    UpstreamManager,
};
use crate::{LoadFromUpstream, ServiceData, ShardError, UpstreamError};

pub(super) struct Shard<Key, Data, Upstream> {
    config: ShardConfig,

    shard_receiver: mpsc::Receiver<ServiceHandleMessage<Key, Data>>,

    internal_receiver: mpsc::UnboundedReceiver<InternalMpscMessage>,
    internal_sender: mpsc::UnboundedSender<InternalMpscMessage>,
    internal_join_set: JoinSet<InternalJoinSetResult<Key, Data>>,

    data_map: DataMap<Key, Data>,
    upstream_manager: UpstreamManager<Key, Data, Upstream>,
    stats: ShardStats,
    expedited_expiration_probe_enqueued: bool,
    cache_expiration_probe_interval: Interval,
}

impl<Key, Data, Upstream, MaybeMutableUpstream> Shard<Key, Data, MaybeMutableUpstream>
where
    Key: Clone + Hash + Eq + Display + Send + 'static,
    Data: ServiceData,
    Upstream: LoadFromUpstream<Key, Data>,
    MaybeMutableUpstream: super::marker::MaybeMutableUpstream<Key, Data, Immutable = Upstream>,
{
    pub(super) fn new(
        shard_receiver: mpsc::Receiver<ServiceHandleMessage<Key, Data>>,
        upstream: MaybeMutableUpstream,
        config: ShardConfig,
    ) -> Self {
        // we use an unbounded channel here, as we want the upstream implementation to be "fire and forget" which is to say,
        // the data loaded callback should not need to be awaited (which is the case with the regular Sender),
        // and we don't necessarily need back-pressure here, as the demand for loaded data originates from the external
        // shard message.
        let (internal_sender, internal_receiver) = mpsc::unbounded_channel::<InternalMpscMessage>();
        let internal_join_set = JoinSet::new();

        let cache_expiration_probe_interval = interval(config.cache_expiration_probe_interval);
        let data_map = DataMap::new(&config);
        let upstream_manager = UpstreamManager::new(upstream, &config);

        Self {
            config,
            shard_receiver,
            internal_receiver,
            internal_sender,
            internal_join_set,
            cache_expiration_probe_interval,
            upstream_manager,
            data_map,
            stats: Default::default(),
            expedited_expiration_probe_enqueued: false,
        }
    }

    pub(super) async fn run(mut self) {
        // for logs („• ᴗ •„)
        let hostname = hostname::get()
            .map(|h| {
                h.into_string()
                    .expect("hostname was not a valid utf-8 string")
            })
            .unwrap_or("localhost".to_owned());

        let shard_id = self.config.shard_id;

        let _abort_on_panic = super::utils::AbortOnPanic::new(format!(
            "(hostname={}) shard={} panicked!",
            hostname, shard_id
        ));

        info!("(hostname={}) shard={} is starting", hostname, shard_id);

        let shutdown_sender = match self.shard_loop().await {
            Err(error) => {
                error!(
                    "(hostname={}) shard={} has failed with error: {:?}",
                    hostname, shard_id, error
                );
                std::process::abort();
            }
            Ok(shutdown_sender) => shutdown_sender,
        };

        info!(
            "(hostname={}) shard={} is closing the internal receiver",
            hostname, shard_id
        );

        // Right now we do not need to drain this, because it only handles TTLing stuff out,
        // which does not matter because we are shutting down.
        self.internal_receiver.close();

        info!(
            "(hostname={}) shard={} is persisting writes to storage",
            hostname, shard_id
        );

        let persist_all_start = Instant::now();
        let shutdown_stats = self.drain(&hostname).await;

        info!(
            "(hostname={}) shard={} has finished persisting, took {:?}. stats: {:?}",
            hostname,
            shard_id,
            persist_all_start.elapsed(),
            shutdown_stats,
        );

        // Indicate to the shutdown API that we have finished
        shutdown_sender.send(shutdown_stats).ok();
    }

    async fn shard_loop(&mut self) -> Result<ShutdownSender, anyhow::Error> {
        loop {
            // This struct exists because we want autocomplete + auto-formatting, and that does not work in `tokio::select!`.
            enum Action<Key, Data> {
                ShardMessage(ServiceHandleMessage<Key, Data>),
                InternalMpscMessage(InternalMpscMessage),
                JoinSetResult(Result<InternalJoinSetResult<Key, Data>, JoinError>),
                ProbeExpiredEntries,
                HandleCommitData(EnqueuedCommit<Key, Data>),
            }

            let action = tokio::select! {
                Some(shard_message) = self.shard_receiver.recv() => Action::ShardMessage(shard_message),
                Some(internal_message) = self.internal_receiver.recv() => Action::InternalMpscMessage(internal_message),
                Some(join_set_result) = self.internal_join_set.join_next(), if !self.internal_join_set.is_empty() => Action::JoinSetResult(join_set_result),
                _interval = self.cache_expiration_probe_interval.tick(), if !self.expedited_expiration_probe_enqueued => Action::ProbeExpiredEntries,
                Ok(enqueued_commit) = self.upstream_manager.poll_commit_queue_ready() => Action::HandleCommitData(enqueued_commit),
            };

            match action {
                Action::ShardMessage(shard_message) => {
                    if let Some(shutdown_sender) = self.handle_shard_message(shard_message) {
                        return Ok(shutdown_sender);
                    }
                }
                Action::InternalMpscMessage(internal_message) => {
                    self.handle_internal_mpsc_message(internal_message)
                }
                Action::JoinSetResult(Ok(join_set_result)) => {
                    self.handle_internal_join_set_result(join_set_result)
                }
                Action::JoinSetResult(Err(join_set_err)) => {
                    std::panic::resume_unwind(join_set_err.into_panic())
                }
                Action::ProbeExpiredEntries => self.probe_expired_entries(),
                Action::HandleCommitData(enqueued_commit) => {
                    self.handle_commit_data(enqueued_commit)
                }
            }
        }
    }

    #[inline]
    fn handle_commit_data(&mut self, enqueued_commit: EnqueuedCommit<Key, Data>) {
        let (key, mut data) = enqueued_commit.into_inner();
        match &mut data {
            Some(data) => {
                self.upstream_manager
                    .do_commit_data(key, data, &mut self.internal_join_set)
            }
            None => {
                let mut guard = self
                    .data_map
                    .get_loaded_data(&key)
                    .expect("invariant: referenced key from persist queue does not exist.");

                self.upstream_manager.do_commit_data(
                    key.clone(),
                    guard.as_mut(),
                    &mut self.internal_join_set,
                )
            }
        };
    }

    #[inline]
    fn handle_internal_mpsc_message(&mut self, message: InternalMpscMessage) {
        match message {
            InternalMpscMessage::DoExpeditedExpirationProbe => {
                self.probe_expired_entries();
            }
        }
    }

    #[inline]
    fn handle_internal_join_set_result(&mut self, message: InternalJoinSetResult<Key, Data>) {
        match message {
            InternalJoinSetResult::DataLoadResult(key, result) => match result {
                Ok(data) => self.handle_data_load_result_success(key, data),
                Err(error) => self.handle_data_load_result_error(key, error.into()),
            },
            InternalJoinSetResult::DataCommitResult(key, result) => match result {
                Ok(()) => self.handle_data_commit_result_success(key),
                Err(error) => self.handle_data_commit_result_error(key, error),
            },
        }
    }

    #[inline]
    fn handle_shard_message(
        &mut self,
        message: ServiceHandleMessage<Key, Data>,
    ) -> Option<ShutdownSender> {
        match message {
            ServiceHandleMessage::Execute(key, func) => {
                self.execute_or_enqueue_load(key, PendingCallback::Execute(func), false);
            }
            ServiceHandleMessage::ExecuteMut(key, func) => {
                self.execute_or_enqueue_load(key, PendingCallback::ExecuteMut(func), false);
            }
            ServiceHandleMessage::ExecuteIfCached(key, func) => {
                self.execute_if_cached(key, func);
            }
            ServiceHandleMessage::GetStats(result_tx) => {
                self.stats.data_size = self.data_map.len();
                self.stats.expiring_keys = self.data_map.expiring_keys_len();
                self.stats.internal_tasks = self.internal_join_set.len();
                result_tx.send(self.stats.clone()).ok();
            }
            ServiceHandleMessage::TakeData(key, result_tx) => {
                self.take_data(key, result_tx);
            }
            ServiceHandleMessage::Shutdown(result_tx) => return Some(result_tx),
        }

        None
    }

    #[inline]
    fn execute_if_cached(&mut self, key: Key, func: ExecuteIfCachedFn<Data>) {
        let expired = {
            let guard = self.data_map.get_loaded_data(&key);
            match guard {
                None => {
                    (func)(None);
                    None
                }
                Some(guard) => {
                    if utils::is_expired(guard.as_ref()) {
                        Some((guard.into_cloned_key(), func))
                    } else {
                        (func)(Some(guard.as_ref()));
                        None
                    }
                }
            }
        };

        if let Some((key, func)) = expired {
            self.evict_key(&key, EvictionReason::Ttl);
            self.upstream_manager.cancel_commit(&key);
            self.execute_if_cached(key, func);
        }
    }

    fn handle_data_load_result_success(&mut self, key: Key, data: Data) {
        if utils::is_expired(&data) {
            self.handle_data_load_result_error(key, ShardError::DataImmediatelyExpired);
            return;
        }

        let (mut guard, pending_actions) = match self.data_map.swap_loaded_data(key, data) {
            // The demand for the data is gone? Wat.
            None => return,
            Some((guard, pending_actions)) => (guard, pending_actions),
        };

        let pending_executions_completed = pending_actions.callbacks.len();
        let default_commit_policy = self.config.default_commit_policy;

        let accumulated_commit_policy = pending_actions.callbacks.into_iter().fold(
            AccumulatedCommitPolicy::new(),
            |acc, pending_callback| {
                acc.accumulate(pending_callback.resolve(&mut guard, default_commit_policy))
            },
        );

        self.stats.record_load_complete(Ok(()));
        self.stats
            .record_pending_executions_completed(pending_executions_completed);

        // We have finished processing the data load, and performed all pending callbacks,
        // do we have a request to now take ownership of the data elsewhere? if so, let's
        // go ahead and respect that request.
        if let Some(take_data_sender) = pending_actions.take_data_sender {
            let key = guard.into_cloned_key();
            self.take_data(key, take_data_sender);
        } else if accumulated_commit_policy.did_mutate() && guard.as_ref().should_persist() {
            self.upstream_manager
                .enqueue_persist_with_accumulated_commit_policy(
                    &guard.into_cloned_key(),
                    accumulated_commit_policy,
                );
        }
    }

    fn handle_data_load_result_error(&mut self, key: Key, error: ShardError) {
        let pending_actions = match self.data_map.take_pending_actions(key) {
            None => return,
            Some(pending_actions) => pending_actions,
        };

        let pending_executions_completed = pending_actions.callbacks.len();
        for pending_callback in pending_actions.callbacks {
            pending_callback.reject(error.clone(), self.config.default_commit_policy);
        }

        self.stats.record_load_complete(Err(&error));
        self.stats
            .record_pending_executions_completed(pending_executions_completed);

        // If the load failed, there's nothing to take.
        if let Some(take_data_sender) = pending_actions.take_data_sender {
            take_data_sender.send(None).ok();
        }
    }

    fn handle_data_commit_result_success(&mut self, key: Key) {
        self.upstream_manager.temp_mark_commit_complete(key);
        // todo!()
    }

    fn handle_data_commit_result_error(&mut self, key: Key, error: UpstreamError) {
        self.upstream_manager.temp_mark_commit_complete(key);
        // todo!()
    }

    #[inline]
    fn execute_or_enqueue_load(
        &mut self,
        key: Key,
        pending_callback: PendingCallback<Data>,
        is_expire_retry: bool,
    ) {
        // If we are at capacity, and we are unable to evict, and this isn't a known key,
        // we will reject the request. Future executions (once loads are complete),
        // should allow us to evict items.
        if self.data_map.would_exceed_capacity_if_inserted(&key) && !self.evict_lru() {
            pending_callback.reject(
                ShardError::ShardAtCapacity,
                self.config.default_commit_policy,
            );
            return;
        }

        let key_is_expired = {
            let upstream_manager = &mut self.upstream_manager;
            let data_state = self.data_map
                .get_or_insert_loaded_or_loading_data(key, |key_ref| {
                    let take_ownership_result =
                        upstream_manager.try_take_ownership_from_commit_queue(key_ref);

                    match take_ownership_result {
                        // The persist queue did not own the, we need to schedule it to be fetched from upstream.
                        TakeOwnershipResult::NotEnqueued => None,
                        // We have successfully taken the ownership of the data, so insert it into the
                        // data map.
                        TakeOwnershipResult::Transferred(data) => Some(data),
                        TakeOwnershipResult::NotOwned => panic!(
                            "invariant: tried to relinquish ownership of key that was already in the hash-map"
                        ),
                    }
                });

            match data_state {
                DataState::MustLoad(key, pending_callbacks) => {
                    pending_callbacks.push(pending_callback);
                    self.upstream_manager
                        .do_load_data(key, &mut self.internal_join_set);

                    self.stats.loads_in_progress += 1;
                    self.stats.executions_pending += 1;
                    None
                }

                DataState::Loading(pending_callbacks) => {
                    pending_callbacks.push(pending_callback);
                    self.stats.executions_pending += 1;
                    None
                }

                DataState::Loaded(mut guard) => {
                    if utils::is_expired(guard.as_ref()) {
                        if is_expire_retry {
                            pending_callback.reject(
                                ShardError::DataImmediatelyExpired,
                                self.config.default_commit_policy,
                            );
                            None
                        } else {
                            Some((guard.into_cloned_key(), pending_callback))
                        }
                    } else {
                        let accumulated_commit_policy: AccumulatedCommitPolicy = pending_callback
                            .resolve(&mut guard, self.config.default_commit_policy)
                            .into();

                        if accumulated_commit_policy.did_mutate() && guard.as_ref().should_persist()
                        {
                            self.upstream_manager
                                .enqueue_persist_with_accumulated_commit_policy(
                                    &guard.into_cloned_key(),
                                    accumulated_commit_policy,
                                );
                        }

                        self.stats.executions_complete += 1;
                        None
                    }
                }
            }
        };

        // If the key is expired, we'll retry this operation, first evicting the key,
        // and then re-trying the execute enqueue or load.
        if let Some((key, pending_callback)) = key_is_expired {
            self.evict_key(&key, EvictionReason::Ttl);
            // If we're expired, let's go ahead and cancel the commit as well. This will prevent us from getting into a race.
            self.upstream_manager.cancel_commit(&key);
            self.execute_or_enqueue_load(key, pending_callback, true);
        }
    }

    fn take_data(&mut self, key: Key, sender: TakeDataSender<Key, Data>) {
        match self.data_map.try_take(key, sender) {
            TryTakeResult::NotFound(key, sender) => {
                let taken_data = self
                    .upstream_manager
                    .cancel_and_take_data_from_commit_queue(&key);

                if taken_data.is_some() {
                    self.stats.keys_taken += 1;
                }

                sender.send(taken_data).ok();
            }
            TryTakeResult::Taken(key, data, sender) => {
                // big todo: if we are persisting data due to an immediate write commit policy,
                // we need to ensure that take doesn't move the data before it's persisted.
                let was_enqueued_at = self
                    .upstream_manager
                    .cancel_commit(&key)
                    .map(|data| data.at());

                let taken_data = TakenData {
                    key,
                    data,
                    was_enqueued_at,
                };
                self.stats.keys_taken += 1;
                sender.send(Some(taken_data)).ok();
            }
            TryTakeResult::TakeAlreadyEnqueued(sender) => {
                // We already have an existing sender, so return `None` here.
                // todo: maybe an AlreadyTakenError?? none is fine for now tho.
                sender.send(None).ok();
            }
            TryTakeResult::Enqueued => {}
        }
    }

    /// Evicts what is "probably" the least recently used data from storage.
    ///
    /// Returns true if data was able to be successfully evicted, otherwise, false.
    fn evict_lru(&mut self) -> bool {
        match self.data_map.probe_and_take_best_lru_candidate() {
            Some(key_to_evict) => {
                self.evict_key(&key_to_evict, EvictionReason::Lru);
                true
            }
            None => false,
        }
    }

    fn evict_key(&mut self, key_to_evict: &Key, eviction_reason: EvictionReason) {
        let (key, data) = self
            .data_map
            .remove_loaded_data(key_to_evict)
            .expect("evict_key tried to evict a key that does not exist.");

        self.upstream_manager
            .give_ownership_to_commit_queue_if_enqueued(key, data);

        match eviction_reason {
            EvictionReason::Ttl => {
                self.stats.keys_ttl_evicted += 1;
            }
            EvictionReason::Lru => {
                self.stats.keys_lru_evicted += 1;
            }
        }
    }

    /// Probes the expiring keys to see if any have expired. We do this rather than maintain a rather (costly) delay queue of all the
    /// expiring keys, applying a probabilistic approach to key expiry.
    ///
    /// Every time this probe is called, we look at a random sample of N many keys within the expiring keys set and expire them. If the
    /// number of keys expired during the probe (X) exceed 25% of N, we immediately schedule another probe, until the % of expired keys
    /// is under 25% per probe, then we go back to probing at the regular tick interval.
    fn probe_expired_entries(&mut self) {
        let mut num_expired_keys = 0;
        let now = Instant::now();

        for _ in 0..self.config.cache_expiration_probe_keys_per_tick {
            match self.data_map.probe_and_take_expiring_key(&now) {
                // There are no keys to evict, so we have nothing left to do.
                Err(_probe_empty) => break,
                // The probe returned nothing.
                Ok(None) => continue,
                Ok(Some(expiring_key)) => {
                    self.evict_key(&expiring_key, EvictionReason::Ttl);
                    num_expired_keys += 1;
                }
            }
        }

        // Did we expire enough keys to make us enter the expedited loop?
        // The thought here is that if we've expired over 25% of the keys when we probed, we should expedite the next tick, hopefully,
        // to the next tick of the event loop for this shard, going until the number of expired keys is probably under 25%. Then we'll
        // go back to probing at the configured probe interval. We do this instead of looping, so we can avoid "blocking" the event
        // loop processing expiration probes for too long.
        let expedited_expiration_threshold = self.config.cache_expiration_probe_keys_per_tick / 4;
        self.expedited_expiration_probe_enqueued =
            num_expired_keys >= expedited_expiration_threshold;

        self.stats.expiration_probes_ran += 1;
        if self.expedited_expiration_probe_enqueued {
            self.stats.expiration_probes_expedited += 1;
            self.internal_sender
                .send(InternalMpscMessage::DoExpeditedExpirationProbe)
                .ok();
        }
    }

    /// Drains the shard's commit queue, persisting all data upstream as soon as possible and
    /// returning stats to the caller once the operation is complete.
    async fn drain(&mut self, hostname: &String) -> ShardShutdownStats {
        self.upstream_manager.start_draining();

        let initial_commits_enqueued = self.upstream_manager.commit_queue_len();
        let mut excess_executes = 0;
        let mut excess_execute_if_cacheds = 0;
        let mut excess_execute_muts = 0;
        let mut loads_completed = 0;
        let mut commits_completed = 0;
        let mut loads_failed = 0;
        let mut commits_failed = 0;
        let mut take_data_requests = 0;
        let mut has_closed_receiver = false;

        while self.upstream_manager.commit_queue_len() > 0 || !self.internal_join_set.is_empty() {
            if !has_closed_receiver && self.upstream_manager.commit_queue_len() == 0 {
                info!("(hostname={}) shard={} has no more commits in the commit queue and is now closing the shard receiver", hostname, self.config.shard_id);
                self.shard_receiver.close();
                has_closed_receiver = true;
            }

            enum Action<Key, Data> {
                ShardMessage(ServiceHandleMessage<Key, Data>),
                JoinSetResult(Result<InternalJoinSetResult<Key, Data>, JoinError>),
                HandleCommitData(EnqueuedCommit<Key, Data>),
            }

            let action = tokio::select! {
                Some(shard_message) = self.shard_receiver.recv() => Action::ShardMessage(shard_message),
                Some(join_set_result) = self.internal_join_set.join_next(), if !self.internal_join_set.is_empty() => Action::JoinSetResult(join_set_result),
                Ok(enqueued_commit) = self.upstream_manager.poll_commit_queue_ready() => Action::HandleCommitData(enqueued_commit),
            };

            match action {
                Action::ShardMessage(shard_message) => {
                    match shard_message {
                        ServiceHandleMessage::TakeData(..) => take_data_requests += 1,
                        ServiceHandleMessage::Execute(..) => excess_executes += 1,
                        ServiceHandleMessage::ExecuteIfCached(..) => excess_execute_if_cacheds += 1,
                        ServiceHandleMessage::ExecuteMut(..) => excess_execute_muts += 1,
                        ServiceHandleMessage::GetStats(..) => {}
                        ServiceHandleMessage::Shutdown(..) => unreachable!(),
                    }
                    self.handle_shard_message(shard_message);
                }
                Action::JoinSetResult(Ok(join_set_result)) => {
                    match &join_set_result {
                        InternalJoinSetResult::DataLoadResult(_, result) => match result {
                            Ok(_) => loads_completed += 1,
                            _ => loads_failed += 1,
                        },
                        InternalJoinSetResult::DataCommitResult(_, result) => match result {
                            Ok(()) => commits_completed += 1,
                            _ => commits_failed += 1,
                        },
                    }
                    self.handle_internal_join_set_result(join_set_result);
                }
                Action::JoinSetResult(Err(join_set_err)) => {
                    std::panic::resume_unwind(join_set_err.into_panic());
                }
                Action::HandleCommitData(enqueued_commit) => {
                    self.handle_commit_data(enqueued_commit);
                }
            };
        }

        ShardShutdownStats {
            initial_commits_enqueued,
            excess_execute_muts,
            excess_executes,
            excess_execute_if_cacheds,
            take_data_requests,
            loads_completed,
            loads_failed,
            commits_completed,
            commits_failed,
        }
    }
}

enum EvictionReason {
    Lru,
    Ttl,
}