From 861016ce0f331ef6b48020b3159e84e8ebaf0482 Mon Sep 17 00:00:00 2001 From: Charles Hall Date: Thu, 10 Oct 2024 22:13:24 -0700 Subject: [PATCH] inline state accessor service code from data trait These were all calling into services() and not actually directly accessing the database at all anyway. --- .../key_value/rooms/state_accessor.rs | 181 +----------------- src/service/rooms/state_accessor.rs | 116 ++++++++++- src/service/rooms/state_accessor/data.rs | 65 +------ 3 files changed, 114 insertions(+), 248 deletions(-) diff --git a/src/database/key_value/rooms/state_accessor.rs b/src/database/key_value/rooms/state_accessor.rs index b2265ba2..1e9e1b64 100644 --- a/src/database/key_value/rooms/state_accessor.rs +++ b/src/database/key_value/rooms/state_accessor.rs @@ -1,139 +1,12 @@ -use std::{collections::HashMap, sync::Arc}; - -use async_trait::async_trait; -use ruma::{events::StateEventType, EventId, RoomId}; +use ruma::EventId; use crate::{ database::KeyValueDatabase, - service::{ - self, - rooms::short::{ShortStateHash, ShortStateKey}, - }, - services, utils, Error, PduEvent, Result, + service::{self, rooms::short::ShortStateHash}, + utils, Error, Result, }; -#[async_trait] impl service::rooms::state_accessor::Data for KeyValueDatabase { - async fn state_full_ids( - &self, - shortstatehash: ShortStateHash, - ) -> Result>> { - let full_state = services() - .rooms - .state_compressor - .load_shortstatehash_info(shortstatehash)? - .pop() - .expect("there is always one layer") - .full_state; - let mut result = HashMap::new(); - let mut i = 0; - for compressed in full_state.iter() { - let parsed = services() - .rooms - .state_compressor - .parse_compressed_state_event(compressed)?; - result.insert(parsed.0, parsed.1); - - i += 1; - if i % 100 == 0 { - tokio::task::yield_now().await; - } - } - Ok(result) - } - - async fn state_full( - &self, - shortstatehash: ShortStateHash, - ) -> Result>> { - let full_state = services() - .rooms - .state_compressor - .load_shortstatehash_info(shortstatehash)? - .pop() - .expect("there is always one layer") - .full_state; - - let mut result = HashMap::new(); - let mut i = 0; - for compressed in full_state.iter() { - let (_, eventid) = services() - .rooms - .state_compressor - .parse_compressed_state_event(compressed)?; - if let Some(pdu) = services().rooms.timeline.get_pdu(&eventid)? { - result.insert( - ( - pdu.kind.to_string().into(), - pdu.state_key - .as_ref() - .ok_or_else(|| { - Error::bad_database( - "State event has no state key.", - ) - })? - .clone(), - ), - pdu, - ); - } - - i += 1; - if i % 100 == 0 { - tokio::task::yield_now().await; - } - } - - Ok(result) - } - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn state_get_id( - &self, - shortstatehash: ShortStateHash, - event_type: &StateEventType, - state_key: &str, - ) -> Result>> { - let Some(shortstatekey) = - services().rooms.short.get_shortstatekey(event_type, state_key)? - else { - return Ok(None); - }; - let full_state = services() - .rooms - .state_compressor - .load_shortstatehash_info(shortstatehash)? - .pop() - .expect("there is always one layer") - .full_state; - Ok(full_state - .iter() - .find(|compressed| compressed.state == shortstatekey) - .and_then(|compressed| { - services() - .rooms - .state_compressor - .parse_compressed_state_event(compressed) - .ok() - .map(|(_, id)| id) - })) - } - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn state_get( - &self, - shortstatehash: ShortStateHash, - event_type: &StateEventType, - state_key: &str, - ) -> Result>> { - self.state_get_id(shortstatehash, event_type, state_key)? - .map_or(Ok(None), |event_id| { - services().rooms.timeline.get_pdu(&event_id) - }) - } - /// Returns the state hash for this pdu. fn pdu_shortstatehash( &self, @@ -158,52 +31,4 @@ impl service::rooms::state_accessor::Data for KeyValueDatabase { }, ) } - - /// Returns the full room state. - async fn room_state_full( - &self, - room_id: &RoomId, - ) -> Result>> { - if let Some(current_shortstatehash) = - services().rooms.state.get_room_shortstatehash(room_id)? - { - self.state_full(current_shortstatehash).await - } else { - Ok(HashMap::new()) - } - } - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn room_state_get_id( - &self, - room_id: &RoomId, - event_type: &StateEventType, - state_key: &str, - ) -> Result>> { - if let Some(current_shortstatehash) = - services().rooms.state.get_room_shortstatehash(room_id)? - { - self.state_get_id(current_shortstatehash, event_type, state_key) - } else { - Ok(None) - } - } - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn room_state_get( - &self, - room_id: &RoomId, - event_type: &StateEventType, - state_key: &str, - ) -> Result>> { - if let Some(current_shortstatehash) = - services().rooms.state.get_room_shortstatehash(room_id)? - { - self.state_get(current_shortstatehash, event_type, state_key) - } else { - Ok(None) - } - } } diff --git a/src/service/rooms/state_accessor.rs b/src/service/rooms/state_accessor.rs index 108e510d..e9794e53 100644 --- a/src/service/rooms/state_accessor.rs +++ b/src/service/rooms/state_accessor.rs @@ -68,7 +68,28 @@ impl Service { &self, shortstatehash: ShortStateHash, ) -> Result>> { - self.db.state_full_ids(shortstatehash).await + let full_state = services() + .rooms + .state_compressor + .load_shortstatehash_info(shortstatehash)? + .pop() + .expect("there is always one layer") + .full_state; + let mut result = HashMap::new(); + let mut i = 0; + for compressed in full_state.iter() { + let parsed = services() + .rooms + .state_compressor + .parse_compressed_state_event(compressed)?; + result.insert(parsed.0, parsed.1); + + i += 1; + if i % 100 == 0 { + tokio::task::yield_now().await; + } + } + Ok(result) } #[tracing::instrument(skip(self))] @@ -76,7 +97,45 @@ impl Service { &self, shortstatehash: ShortStateHash, ) -> Result>> { - self.db.state_full(shortstatehash).await + let full_state = services() + .rooms + .state_compressor + .load_shortstatehash_info(shortstatehash)? + .pop() + .expect("there is always one layer") + .full_state; + + let mut result = HashMap::new(); + let mut i = 0; + for compressed in full_state.iter() { + let (_, eventid) = services() + .rooms + .state_compressor + .parse_compressed_state_event(compressed)?; + if let Some(pdu) = services().rooms.timeline.get_pdu(&eventid)? { + result.insert( + ( + pdu.kind.to_string().into(), + pdu.state_key + .as_ref() + .ok_or_else(|| { + Error::bad_database( + "State event has no state key.", + ) + })? + .clone(), + ), + pdu, + ); + } + + i += 1; + if i % 100 == 0 { + tokio::task::yield_now().await; + } + } + + Ok(result) } /// Returns a single PDU from `room_id` with key (`event_type`, @@ -88,7 +147,29 @@ impl Service { event_type: &StateEventType, state_key: &str, ) -> Result>> { - self.db.state_get_id(shortstatehash, event_type, state_key) + let Some(shortstatekey) = + services().rooms.short.get_shortstatekey(event_type, state_key)? + else { + return Ok(None); + }; + let full_state = services() + .rooms + .state_compressor + .load_shortstatehash_info(shortstatehash)? + .pop() + .expect("there is always one layer") + .full_state; + Ok(full_state + .iter() + .find(|compressed| compressed.state == shortstatekey) + .and_then(|compressed| { + services() + .rooms + .state_compressor + .parse_compressed_state_event(compressed) + .ok() + .map(|(_, id)| id) + })) } /// Returns a single PDU from `room_id` with key (`event_type`, @@ -100,7 +181,10 @@ impl Service { event_type: &StateEventType, state_key: &str, ) -> Result>> { - self.db.state_get(shortstatehash, event_type, state_key) + self.state_get_id(shortstatehash, event_type, state_key)? + .map_or(Ok(None), |event_id| { + services().rooms.timeline.get_pdu(&event_id) + }) } /// Get membership for given user in state @@ -360,7 +444,13 @@ impl Service { &self, room_id: &RoomId, ) -> Result>> { - self.db.room_state_full(room_id).await + if let Some(current_shortstatehash) = + services().rooms.state.get_room_shortstatehash(room_id)? + { + self.state_full(current_shortstatehash).await + } else { + Ok(HashMap::new()) + } } /// Returns a single PDU from `room_id` with key (`event_type`, @@ -372,7 +462,13 @@ impl Service { event_type: &StateEventType, state_key: &str, ) -> Result>> { - self.db.room_state_get_id(room_id, event_type, state_key) + if let Some(current_shortstatehash) = + services().rooms.state.get_room_shortstatehash(room_id)? + { + self.state_get_id(current_shortstatehash, event_type, state_key) + } else { + Ok(None) + } } /// Returns a single PDU from `room_id` with key (`event_type`, @@ -384,7 +480,13 @@ impl Service { event_type: &StateEventType, state_key: &str, ) -> Result>> { - self.db.room_state_get(room_id, event_type, state_key) + if let Some(current_shortstatehash) = + services().rooms.state.get_room_shortstatehash(room_id)? + { + self.state_get(current_shortstatehash, event_type, state_key) + } else { + Ok(None) + } } #[tracing::instrument(skip(self))] diff --git a/src/service/rooms/state_accessor/data.rs b/src/service/rooms/state_accessor/data.rs index 4fd1b2af..4399dc66 100644 --- a/src/service/rooms/state_accessor/data.rs +++ b/src/service/rooms/state_accessor/data.rs @@ -1,72 +1,11 @@ -use std::{collections::HashMap, sync::Arc}; +use ruma::EventId; -use async_trait::async_trait; -use ruma::{events::StateEventType, EventId, RoomId}; +use crate::{service::rooms::short::ShortStateHash, Result}; -use crate::{ - service::rooms::short::{ShortStateHash, ShortStateKey}, - PduEvent, Result, -}; - -#[async_trait] pub(crate) trait Data: Send + Sync { - /// Builds a StateMap by iterating over all keys that start - /// with state_hash, this gives the full state for the given state_hash. - async fn state_full_ids( - &self, - shortstatehash: ShortStateHash, - ) -> Result>>; - - async fn state_full( - &self, - shortstatehash: ShortStateHash, - ) -> Result>>; - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn state_get_id( - &self, - shortstatehash: ShortStateHash, - event_type: &StateEventType, - state_key: &str, - ) -> Result>>; - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn state_get( - &self, - shortstatehash: ShortStateHash, - event_type: &StateEventType, - state_key: &str, - ) -> Result>>; - /// Returns the state hash for this pdu. fn pdu_shortstatehash( &self, event_id: &EventId, ) -> Result>; - - /// Returns the full room state. - async fn room_state_full( - &self, - room_id: &RoomId, - ) -> Result>>; - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn room_state_get_id( - &self, - room_id: &RoomId, - event_type: &StateEventType, - state_key: &str, - ) -> Result>>; - - /// Returns a single PDU from `room_id` with key (`event_type`, - /// `state_key`). - fn room_state_get( - &self, - room_id: &RoomId, - event_type: &StateEventType, - state_key: &str, - ) -> Result>>; }