grapevine/src/database/key_value/rooms/short.rs

use std::sync::Arc;

use ruma::{events::StateEventType, EventId, RoomId};

use crate::{
    database::KeyValueDatabase,
    observability::prelude::*,
    service::{
        self,
        rooms::short::{
            ShortEventId, ShortRoomId, ShortStateHash, ShortStateKey,
        },
    },
    services, utils, Error, Result,
};

impl service::rooms::short::Data for KeyValueDatabase {
    #[t::instrument(skip(self))]
    fn get_or_create_shorteventid(
        &self,
        event_id: &EventId,
    ) -> Result<(ShortEventId, bool)> {
        let (short, created) = if let Some(shorteventid) =
            self.eventid_shorteventid.get(event_id.as_bytes())?
        {
            let shorteventid =
                utils::u64_from_bytes(&shorteventid).map_err(|_| {
                    Error::bad_database("Invalid shorteventid in db.")
                })?;

            (shorteventid, false)
        } else {
            let shorteventid = services().globals.next_count()?;
            self.eventid_shorteventid
                .insert(event_id.as_bytes(), &shorteventid.to_be_bytes())?;
            self.shorteventid_eventid
                .insert(&shorteventid.to_be_bytes(), event_id.as_bytes())?;
            (shorteventid, true)
        };

        Ok((ShortEventId::new(short), created))
    }

    #[t::instrument(skip(self), fields(cache_result))]
    fn get_shortstatekey(
        &self,
        event_type: &StateEventType,
        state_key: &str,
    ) -> Result<Option<ShortStateKey>> {
        let mut db_key = event_type.to_string().as_bytes().to_vec();
        db_key.push(0xFF);
        db_key.extend_from_slice(state_key.as_bytes());

        let short = self
            .statekey_shortstatekey
            .get(&db_key)?
            .map(|shortstatekey| {
                utils::u64_from_bytes(&shortstatekey)
                    .map_err(|_| {
                        Error::bad_database("Invalid shortstatekey in db.")
                    })
                    .map(ShortStateKey::new)
            })
            .transpose()?;

        Ok(short)
    }

    #[t::instrument(skip(self))]
    fn get_or_create_shortstatekey(
        &self,
        event_type: &StateEventType,
        state_key: &str,
    ) -> Result<(ShortStateKey, bool)> {
        let mut db_key = event_type.to_string().as_bytes().to_vec();
        db_key.push(0xFF);
        db_key.extend_from_slice(state_key.as_bytes());

        let (short, created) = if let Some(shortstatekey) =
            self.statekey_shortstatekey.get(&db_key)?
        {
            (
                utils::u64_from_bytes(&shortstatekey).map_err(|_| {
                    Error::bad_database("Invalid shortstatekey in db.")
                })?,
                false,
            )
        } else {
            let shortstatekey = services().globals.next_count()?;
            self.statekey_shortstatekey
                .insert(&db_key, &shortstatekey.to_be_bytes())?;
            self.shortstatekey_statekey
                .insert(&shortstatekey.to_be_bytes(), &db_key)?;
            (shortstatekey, true)
        };

        let short = ShortStateKey::new(short);

        Ok((short, created))
    }

    #[t::instrument(skip(self))]
    fn get_eventid_from_short(
        &self,
        shorteventid: ShortEventId,
    ) -> Result<Arc<EventId>> {
        let bytes = self
            .shorteventid_eventid
            .get(&shorteventid.get().to_be_bytes())?
            .ok_or_else(|| {
                Error::bad_database("Shorteventid does not exist")
            })?;

        let event_id = EventId::parse_arc(
            utils::string_from_bytes(&bytes).map_err(|_| {
                Error::bad_database(
                    "EventID in shorteventid_eventid is invalid unicode.",
                )
            })?,
        )
        .map_err(|_| {
            Error::bad_database("EventId in shorteventid_eventid is invalid.")
        })?;

        Ok(event_id)
    }

    #[t::instrument(skip(self))]
    fn get_statekey_from_short(
        &self,
        shortstatekey: ShortStateKey,
    ) -> Result<(StateEventType, String)> {
        let bytes = self
            .shortstatekey_statekey
            .get(&shortstatekey.get().to_be_bytes())?
            .ok_or_else(|| {
                Error::bad_database("Shortstatekey does not exist")
            })?;

        let mut parts = bytes.splitn(2, |&b| b == 0xFF);
        let eventtype_bytes =
            parts.next().expect("split always returns one entry");
        let statekey_bytes = parts.next().ok_or_else(|| {
            Error::bad_database("Invalid statekey in shortstatekey_statekey.")
        })?;

        let event_type = StateEventType::from(
            utils::string_from_bytes(eventtype_bytes).map_err(|_| {
                Error::bad_database(
                    "Event type in shortstatekey_statekey is invalid unicode.",
                )
            })?,
        );

        let state_key =
            utils::string_from_bytes(statekey_bytes).map_err(|_| {
                Error::bad_database(
                    "Statekey in shortstatekey_statekey is invalid unicode.",
                )
            })?;

        Ok((event_type, state_key))
    }

    /// Returns `(shortstatehash, already_existed)`
    #[t::instrument(skip(self))]
    fn get_or_create_shortstatehash(
        &self,
        state_hash: &[u8],
    ) -> Result<(ShortStateHash, bool)> {
        let (short, existed) = if let Some(shortstatehash) =
            self.statehash_shortstatehash.get(state_hash)?
        {
            (
                utils::u64_from_bytes(&shortstatehash).map_err(|_| {
                    Error::bad_database("Invalid shortstatehash in db.")
                })?,
                true,
            )
        } else {
            let shortstatehash = services().globals.next_count()?;
            self.statehash_shortstatehash
                .insert(state_hash, &shortstatehash.to_be_bytes())?;
            (shortstatehash, false)
        };

        Ok((ShortStateHash::new(short), existed))
    }

    fn get_shortroomid(&self, room_id: &RoomId) -> Result<Option<ShortRoomId>> {
        self.roomid_shortroomid
            .get(room_id.as_bytes())?
            .map(|bytes| {
                utils::u64_from_bytes(&bytes)
                    .map_err(|_| {
                        Error::bad_database("Invalid shortroomid in db.")
                    })
                    .map(ShortRoomId::new)
            })
            .transpose()
    }

    fn get_or_create_shortroomid(
        &self,
        room_id: &RoomId,
    ) -> Result<ShortRoomId> {
        let short = if let Some(short) =
            self.roomid_shortroomid.get(room_id.as_bytes())?
        {
            utils::u64_from_bytes(&short).map_err(|_| {
                Error::bad_database("Invalid shortroomid in db.")
            })?
        } else {
            let short = services().globals.next_count()?;
            self.roomid_shortroomid
                .insert(room_id.as_bytes(), &short.to_be_bytes())?;
            short
        };

        Ok(ShortRoomId::new(short))
    }
}