grapevine/src/service/rooms/short.rs

use std::sync::{Arc, Mutex};

use lru_cache::LruCache;
use ruma::{events::StateEventType, EventId, OwnedEventId, RoomId};

use crate::{
    observability::{FoundIn, Lookup, METRICS},
    utils::error::Result,
};

macro_rules! short_id_type {
    ($($(#[$doc:meta])* struct $name:ident(u64);)*) => {
        $(
            #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
            #[repr(transparent)]
            $(#[$doc])*
            pub(crate) struct $name(u64);

            impl $name {
                pub(crate) fn new(id: u64) -> Self {
                    Self(id)
                }

                pub(crate) fn get(&self) -> u64 {
                    self.0
                }
            }
        )*
    };
}

short_id_type!(
    /// Interned [RoomId].
    ///
    /// Created using [`get_shortroomid()`](Service::get_shortroomid) or
    /// [`get_or_create_shortroomid()`](Service::get_or_create_shortroomid).
    struct ShortRoomId(u64);
    /// Interned [EventId].
    ///
    /// Created using
    /// [`get_or_create_shorteventid()`](Service::get_or_create_shorteventid),
    /// resolved using
    /// [`get_eventid_from_short()`](Service::get_eventid_from_short).
    struct ShortEventId(u64);
    /// Interned hash of concatenated state events.
    ///
    /// Equal state sets do not necessarily correspond to equal short state
    /// hashes, because the calculated hash is dependent on `HashSet`
    /// iteration order.
    ///
    /// Created using
    /// [`get_or_create_shortstatehash()`](Service::get_or_create_shortstatehash).
    struct ShortStateHash(u64);
    /// Interned `(event type, state key)` tuple.
    ///
    /// Created using [`get_shortstatekey()`](Service::get_shortstatekey) or
    /// [`get_or_create_shortstatekey()`](Service::get_or_create_shortstatekey),
    /// resolved using
    /// [`get_statekey_from_short()`](Service::get_statekey_from_short).
    struct ShortStateKey(u64);
);

mod data;

pub(crate) use data::Data;

pub(crate) struct Service {
    db: &'static dyn Data,
    shorteventid_cache: Option<Mutex<LruCache<ShortEventId, Arc<EventId>>>>,
    eventidshort_cache: Option<Mutex<LruCache<OwnedEventId, ShortEventId>>>,
    statekeyshort_cache:
        Option<Mutex<LruCache<(StateEventType, String), ShortStateKey>>>,
    shortstatekey_cache:
        Option<Mutex<LruCache<ShortStateKey, (StateEventType, String)>>>,
}

impl Service {
    pub(crate) fn new(
        db: &'static dyn Data,
        shorteventid_cache_size: usize,
        eventidshort_cache_size: usize,
        statekeyshort_cache_size: usize,
        shortstatekey_cache_size: usize,
    ) -> Self {
        Self {
            db,
            shorteventid_cache: (shorteventid_cache_size > 0)
                .then(|| Mutex::new(LruCache::new(shorteventid_cache_size))),
            eventidshort_cache: (eventidshort_cache_size > 0)
                .then(|| Mutex::new(LruCache::new(eventidshort_cache_size))),
            statekeyshort_cache: (statekeyshort_cache_size > 0)
                .then(|| Mutex::new(LruCache::new(statekeyshort_cache_size))),
            shortstatekey_cache: (shortstatekey_cache_size > 0)
                .then(|| Mutex::new(LruCache::new(shortstatekey_cache_size))),
        }
    }

    pub(crate) fn get_or_create_shorteventid(
        &self,
        event_id: &EventId,
    ) -> Result<ShortEventId> {
        let lookup = Lookup::CreateEventIdToShort;

        if let Some(cache) = &self.eventidshort_cache {
            if let Some(short) = cache.lock().unwrap().get_mut(event_id) {
                METRICS.record_lookup(lookup, FoundIn::Cache);
                return Ok(*short);
            }
        }

        let (short, created) = self.db.get_or_create_shorteventid(event_id)?;

        if created {
            METRICS.record_lookup(lookup, FoundIn::Nothing);
        } else {
            METRICS.record_lookup(lookup, FoundIn::Database);
        }

        if let Some(cache) = &self.eventidshort_cache {
            cache.lock().unwrap().insert(event_id.to_owned(), short);
        }

        Ok(short)
    }

    pub(crate) fn get_shortstatekey(
        &self,
        event_type: &StateEventType,
        state_key: &str,
    ) -> Result<Option<ShortStateKey>> {
        let lookup = Lookup::StateKeyToShort;

        if let Some(cache) = &self.statekeyshort_cache {
            if let Some(short) = cache
                .lock()
                .unwrap()
                .get_mut(&(event_type.clone(), state_key.to_owned()))
            {
                METRICS.record_lookup(lookup, FoundIn::Cache);
                return Ok(Some(*short));
            }
        }

        let short = self.db.get_shortstatekey(event_type, state_key)?;

        if let Some(short) = short {
            METRICS.record_lookup(lookup, FoundIn::Database);

            if let Some(cache) = &self.statekeyshort_cache {
                cache
                    .lock()
                    .unwrap()
                    .insert((event_type.clone(), state_key.to_owned()), short);
            }
        } else {
            METRICS.record_lookup(lookup, FoundIn::Nothing);
        }

        Ok(short)
    }

    pub(crate) fn get_or_create_shortstatekey(
        &self,
        event_type: &StateEventType,
        state_key: &str,
    ) -> Result<ShortStateKey> {
        let lookup = Lookup::CreateStateKeyToShort;

        if let Some(cache) = &self.statekeyshort_cache {
            if let Some(short) = cache
                .lock()
                .unwrap()
                .get_mut(&(event_type.clone(), state_key.to_owned()))
            {
                METRICS.record_lookup(lookup, FoundIn::Cache);
                return Ok(*short);
            }
        }

        let (short, created) =
            self.db.get_or_create_shortstatekey(event_type, state_key)?;

        if created {
            METRICS.record_lookup(lookup, FoundIn::Nothing);
        } else {
            METRICS.record_lookup(lookup, FoundIn::Database);
        }

        if let Some(cache) = &self.statekeyshort_cache {
            cache
                .lock()
                .unwrap()
                .insert((event_type.clone(), state_key.to_owned()), short);
        }

        Ok(short)
    }

    pub(crate) fn get_eventid_from_short(
        &self,
        shorteventid: ShortEventId,
    ) -> Result<Arc<EventId>> {
        let lookup = Lookup::ShortToEventId;

        if let Some(cache) = &self.shorteventid_cache {
            if let Some(id) = cache.lock().unwrap().get_mut(&shorteventid) {
                METRICS.record_lookup(lookup, FoundIn::Cache);
                return Ok(Arc::clone(id));
            }
        }
        let event_id = self.db.get_eventid_from_short(shorteventid)?;

        METRICS.record_lookup(lookup, FoundIn::Database);

        if let Some(cache) = &self.shorteventid_cache {
            cache.lock().unwrap().insert(shorteventid, Arc::clone(&event_id));
        }

        Ok(event_id)
    }

    pub(crate) fn get_statekey_from_short(
        &self,
        shortstatekey: ShortStateKey,
    ) -> Result<(StateEventType, String)> {
        let lookup = Lookup::ShortToStateKey;

        if let Some(cache) = &self.shortstatekey_cache {
            if let Some(id) = cache.lock().unwrap().get_mut(&shortstatekey) {
                METRICS.record_lookup(lookup, FoundIn::Cache);
                return Ok(id.clone());
            }
        }
        let x = self.db.get_statekey_from_short(shortstatekey)?;

        METRICS.record_lookup(lookup, FoundIn::Database);

        if let Some(cache) = &self.shortstatekey_cache {
            cache.lock().unwrap().insert(shortstatekey, x.clone());
        }

        Ok(x)
    }

    /// Returns `(shortstatehash, already_existed)`
    pub(crate) fn get_or_create_shortstatehash(
        &self,
        state_hash: &[u8],
    ) -> Result<(ShortStateHash, bool)> {
        self.db.get_or_create_shortstatehash(state_hash)
    }

    pub(crate) fn get_shortroomid(
        &self,
        room_id: &RoomId,
    ) -> Result<Option<ShortRoomId>> {
        self.db.get_shortroomid(room_id)
    }

    pub(crate) fn get_or_create_shortroomid(
        &self,
        room_id: &RoomId,
    ) -> Result<ShortRoomId> {
        self.db.get_or_create_shortroomid(room_id)
    }
}