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 {
    ($name:ident) => {
        #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
        #[repr(transparent)]
        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!(ShortRoomId);
short_id_type!(ShortEventId);
short_id_type!(ShortStateHash);
short_id_type!(ShortStateKey);

mod data;

pub(crate) use data::Data;

pub(crate) struct Service {
    db: &'static dyn Data,
    shorteventid_cache: Mutex<LruCache<ShortEventId, Arc<EventId>>>,
    eventidshort_cache: Mutex<LruCache<OwnedEventId, ShortEventId>>,
    statekeyshort_cache:
        Mutex<LruCache<(StateEventType, String), ShortStateKey>>,
    shortstatekey_cache:
        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: Mutex::new(LruCache::new(
                shorteventid_cache_size,
            )),
            eventidshort_cache: Mutex::new(LruCache::new(
                eventidshort_cache_size,
            )),
            statekeyshort_cache: Mutex::new(LruCache::new(
                statekeyshort_cache_size,
            )),
            shortstatekey_cache: 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(short) =
            self.eventidshort_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);
        }

        self.eventidshort_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(short) = self
            .statekeyshort_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);

            self.statekeyshort_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(short) = self
            .statekeyshort_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);
        }

        self.statekeyshort_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(id) =
            self.shorteventid_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);

        self.shorteventid_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(id) =
            self.shortstatekey_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);

        self.shortstatekey_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)
    }
}