magmaus3/continuwuity
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: improved state store

Browse source
This commit is contained in:
Timo Kösters 2020-12-19 16:00:11 +01:00
parent 6e5b35ea92
commit 6606e41dde
No known key found for this signature in database
GPG key ID: 24DA7517711A2BA4
13 changed files with 405 additions and 251 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/database/globals.rs
View file

@ -4,6 +4,7 @@ use ruma::ServerName;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::RwLock;
use std::time::Duration;
use trust_dns_resolver::TokioAsyncResolver;
pub const COUNTER: &str = "c";
@ -54,11 +55,18 @@ impl Globals {
}
};
let reqwest_client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(30))
.timeout(Duration::from_secs(60 * 3))
.pool_max_idle_per_host(1)
.build()
.unwrap();
Ok(Self {
globals,
config,
keypair: Arc::new(keypair),
reqwest_client: reqwest::Client::new(),
reqwest_client,
dns_resolver: TokioAsyncResolver::tokio_from_system_conf()
.await
.map_err(|_| {

178
src/database/rooms.rs
View file

@ -62,7 +62,8 @@ pub struct Rooms {
/// Remember the state hash at events in the past.
pub(super) pduid_statehash: sled::Tree,
/// The state for a given state hash.
pub(super) stateid_pduid: sled::Tree, // StateId = StateHash + EventType + StateKey
pub(super) statekey_short: sled::Tree, // StateKey = EventType + StateKey, Short = Count
pub(super) stateid_pduid: sled::Tree, // StateId = StateHash + Short, PduId = Count (without roomid)
}
impl StateStore for Rooms {
@ -106,21 +107,28 @@ impl StateStore for Rooms {
impl Rooms {
/// Builds a StateMap by iterating over all keys that start
/// with state_hash, this gives the full state for the given state_hash.
pub fn state_full(&self, state_hash: &StateHashId) -> Result<StateMap<PduEvent>> {
pub fn state_full(
&self,
room_id: &RoomId,
state_hash: &StateHashId,
) -> Result<StateMap<PduEvent>> {
self.stateid_pduid
.scan_prefix(&state_hash)
.values()
.map(|pduid| {
self.pduid_pdu.get(&pduid?)?.map_or_else(
|| Err(Error::bad_database("Failed to find StateMap.")),
.map(|pduid_short| {
let mut pduid = room_id.as_bytes().to_vec();
pduid.push(0xff);
pduid.extend_from_slice(&pduid_short?);
self.pduid_pdu.get(&pduid)?.map_or_else(
|| Err(Error::bad_database("Failed to find PDU in state snapshot.")),
|b| {
serde_json::from_slice::<PduEvent>(&b)
.map_err(|_| Error::bad_database("Invalid PDU in db."))
},
)
})
.filter_map(|r| r.ok())
.map(|pdu| {
let pdu = pdu?;
Ok((
(
pdu.kind.clone(),
@ -135,64 +143,45 @@ impl Rooms {
.collect::<Result<StateMap<_>>>()
}
/// Returns all state entries for this type.
pub fn state_type(
&self,
state_hash: &StateHashId,
event_type: &EventType,
) -> Result<HashMap<String, PduEvent>> {
let mut prefix = state_hash.to_vec();
prefix.push(0xff);
prefix.extend_from_slice(&event_type.to_string().as_bytes());
prefix.push(0xff);
let mut hashmap = HashMap::new();
for pdu in self
.stateid_pduid
.scan_prefix(&prefix)
.values()
.map(|pdu_id| {
Ok::<_, Error>(
serde_json::from_slice::<PduEvent>(&self.pduid_pdu.get(pdu_id?)?.ok_or_else(
|| Error::bad_database("PDU in state not found in database."),
)?)
.map_err(|_| Error::bad_database("Invalid PDU bytes in room state."))?,
)
})
{
let pdu = pdu?;
let state_key = pdu.state_key.clone().ok_or_else(|| {
Error::bad_database("Room state contains event without state_key.")
})?;
hashmap.insert(state_key, pdu);
}
Ok(hashmap)
}
/// Returns a single PDU from `room_id` with key (`event_type`, `state_key`).
pub fn state_get(
&self,
room_id: &RoomId,
state_hash: &StateHashId,
event_type: &EventType,
state_key: &str,
) -> Result<Option<(IVec, PduEvent)>> {
let mut key = state_hash.to_vec();
key.push(0xff);
key.extend_from_slice(&event_type.to_string().as_bytes());
let mut key = event_type.to_string().as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(&state_key.as_bytes());
self.stateid_pduid.get(&key)?.map_or(Ok(None), |pdu_id| {
Ok::<_, Error>(Some((
pdu_id.clone(),
serde_json::from_slice::<PduEvent>(
&self.pduid_pdu.get(&pdu_id)?.ok_or_else(|| {
Error::bad_database("PDU in state not found in database.")
})?,
)
.map_err(|_| Error::bad_database("Invalid PDU bytes in room state."))?,
)))
})
let short = self.statekey_short.get(&key)?;
if let Some(short) = short {
let mut stateid = state_hash.to_vec();
stateid.push(0xff);
stateid.extend_from_slice(&short);
self.stateid_pduid
.get(&stateid)?
.map_or(Ok(None), |pdu_id_short| {
let mut pdu_id = room_id.as_bytes().to_vec();
pdu_id.push(0xff);
pdu_id.extend_from_slice(&pdu_id_short);
Ok::<_, Error>(Some((
pdu_id.clone().into(),
serde_json::from_slice::<PduEvent>(
&self.pduid_pdu.get(&pdu_id)?.ok_or_else(|| {
Error::bad_database("PDU in state not found in database.")
})?,
)
.map_err(|_| Error::bad_database("Invalid PDU bytes in room state."))?,
)))
})
} else {
return Ok(None);
}
}
/// Returns the last state hash key added to the db.
@ -260,6 +249,7 @@ impl Rooms {
&self,
room_id: &RoomId,
state: HashMap<(EventType, String), Vec<u8>>,
globals: &super::globals::Globals,
) -> Result<()> {
let state_hash =
self.calculate_hash(&state.values().map(|pdu_id| &**pdu_id).collect::<Vec<_>>())?;
@ -267,11 +257,29 @@ impl Rooms {
prefix.push(0xff);
for ((event_type, state_key), pdu_id) in state {
let mut statekey = event_type.as_ref().as_bytes().to_vec();
statekey.push(0xff);
statekey.extend_from_slice(&state_key.as_bytes());
let short = match self.statekey_short.get(&statekey)? {
Some(short) => utils::u64_from_bytes(&short)
.map_err(|_| Error::bad_database("Invalid short bytes in statekey_short."))?,
None => {
let short = globals.next_count()?;
self.statekey_short
.insert(&statekey, &short.to_be_bytes())?;
short
}
};
let pdu_id_short = pdu_id
.splitn(2, |&b| b == 0xff)
.nth(1)
.ok_or_else(|| Error::bad_database("Invalid pduid in state."))?;
let mut state_id = prefix.clone();
state_id.extend_from_slice(&event_type.as_ref().as_bytes());
state_id.push(0xff);
state_id.extend_from_slice(&state_key.as_bytes());
self.stateid_pduid.insert(state_id, pdu_id)?;
state_id.extend_from_slice(&short.to_be_bytes());
self.stateid_pduid.insert(state_id, pdu_id_short)?;
}
self.roomid_statehash
@ -283,25 +291,12 @@ impl Rooms {
/// Returns the full room state.
pub fn room_state_full(&self, room_id: &RoomId) -> Result<StateMap<PduEvent>> {
if let Some(current_state_hash) = self.current_state_hash(room_id)? {
self.state_full(&current_state_hash)
self.state_full(&room_id, &current_state_hash)
} else {
Ok(BTreeMap::new())
}
}
/// Returns all state entries for this type.
pub fn room_state_type(
&self,
room_id: &RoomId,
event_type: &EventType,
) -> Result<HashMap<String, PduEvent>> {
if let Some(current_state_hash) = self.current_state_hash(room_id)? {
self.state_type(&current_state_hash, event_type)
} else {
Ok(HashMap::new())
}
}
/// Returns a single PDU from `room_id` with key (`event_type`, `state_key`).
pub fn room_state_get(
&self,
@ -310,7 +305,7 @@ impl Rooms {
state_key: &str,
) -> Result<Option<(IVec, PduEvent)>> {
if let Some(current_state_hash) = self.current_state_hash(room_id)? {
self.state_get(&current_state_hash, event_type, state_key)
self.state_get(&room_id, &current_state_hash, event_type, state_key)
} else {
Ok(None)
}
@ -593,7 +588,12 @@ impl Rooms {
/// This adds all current state events (not including the incoming event)
/// to `stateid_pduid` and adds the incoming event to `pduid_statehash`.
/// The incoming event is the `pdu_id` passed to this method.
pub fn append_to_state(&self, new_pdu_id: &[u8], new_pdu: &PduEvent) -> Result<StateHashId> {
pub fn append_to_state(
&self,
new_pdu_id: &[u8],
new_pdu: &PduEvent,
globals: &super::globals::Globals,
) -> Result<StateHashId> {
let old_state =
if let Some(old_state_hash) = self.roomid_statehash.get(new_pdu.room_id.as_bytes())? {
// Store state for event. The state does not include the event itself.
@ -608,7 +608,7 @@ impl Rooms {
self.stateid_pduid
.scan_prefix(&prefix)
.filter_map(|pdu| pdu.map_err(|e| error!("{}", e)).ok())
// Chop the old state_hash out leaving behind the (EventType, StateKey)
// Chop the old state_hash out leaving behind the short key (u64)
.map(|(k, v)| (k.subslice(prefix.len(), k.len() - prefix.len()), v))
.collect::<HashMap<IVec, IVec>>()
} else {
@ -620,7 +620,23 @@ impl Rooms {
let mut pdu_key = new_pdu.kind.as_ref().as_bytes().to_vec();
pdu_key.push(0xff);
pdu_key.extend_from_slice(state_key.as_bytes());
new_state.insert(pdu_key.into(), new_pdu_id.into());
let short = match self.statekey_short.get(&pdu_key)? {
Some(short) => utils::u64_from_bytes(&short)
.map_err(|_| Error::bad_database("Invalid short bytes in statekey_short."))?,
None => {
let short = globals.next_count()?;
self.statekey_short.insert(&pdu_key, &short.to_be_bytes())?;
short
}
};
let new_pdu_id_short = new_pdu_id
.splitn(2, |&b| b == 0xff)
.nth(1)
.ok_or_else(|| Error::bad_database("Invalid pduid in state."))?;
new_state.insert((&short.to_be_bytes()).into(), new_pdu_id_short.into());
let new_state_hash =
self.calculate_hash(&new_state.values().map(|b| &**b).collect::<Vec<_>>())?;
@ -628,12 +644,10 @@ impl Rooms {
let mut key = new_state_hash.to_vec();
key.push(0xff);
// TODO: we could avoid writing to the DB on every state event by keeping
// track of the delta and write that every so often
for (key_without_prefix, pdu_id) in new_state {
for (short, short_pdu_id) in new_state {
let mut state_id = key.clone();
state_id.extend_from_slice(&key_without_prefix);
self.stateid_pduid.insert(&state_id, &pdu_id)?;
state_id.extend_from_slice(&short);
self.stateid_pduid.insert(&state_id, &short_pdu_id)?;
}
self.roomid_statehash
@ -887,7 +901,7 @@ impl Rooms {
// We append to state before appending the pdu, so we don't have a moment in time with the
// pdu without it's state. This is okay because append_pdu can't fail.
self.append_to_state(&pdu_id, &pdu)?;
self.append_to_state(&pdu_id, &pdu, &globals)?;
self.append_pdu(
&pdu,

137
src/database/sending.rs
View file

@ -1,21 +1,29 @@
use std::{collections::HashMap, convert::TryFrom, time::SystemTime};
use std::{
collections::HashMap,
convert::TryFrom,
fmt::Debug,
sync::Arc,
time::{Duration, Instant, SystemTime},
};
use crate::{appservice_server, server_server, utils, Error, PduEvent, Result};
use federation::transactions::send_transaction_message;
use log::warn;
use rocket::futures::stream::{FuturesUnordered, StreamExt};
use ruma::{
api::{appservice, federation},
api::{appservice, federation, OutgoingRequest},
ServerName,
};
use sled::IVec;
use tokio::select;
use tokio::sync::Semaphore;
#[derive(Clone)]
pub struct Sending {
/// The state for a given state hash.
pub(super) servernamepduids: sled::Tree, // ServernamePduId = (+)ServerName + PduId
pub(super) servercurrentpdus: sled::Tree, // ServerCurrentPdus = (+)ServerName + PduId (pduid can be empty for reservation)
pub(super) maximum_requests: Arc<Semaphore>,
}
impl Sending {
@ -40,35 +48,7 @@ impl Sending {
for (server, pdu, is_appservice) in servercurrentpdus
.iter()
.filter_map(|r| r.ok())
.map(|(key, _)| {
let mut parts = key.splitn(2, |&b| b == 0xff);
let server = parts.next().expect("splitn always returns one element");
let pdu = parts.next().ok_or_else(|| {
Error::bad_database("Invalid bytes in servercurrentpdus.")
})?;
let server = utils::string_from_bytes(&server).map_err(|_| {
Error::bad_database("Invalid server bytes in server_currenttransaction")
})?;
// Appservices start with a plus
let (server, is_appservice) = if server.starts_with("+") {
(&server[1..], true)
} else {
(&*server, false)
};
Ok::<_, Error>((
Box::<ServerName>::try_from(server).map_err(|_| {
Error::bad_database(
"Invalid server string in server_currenttransaction",
)
})?,
IVec::from(pdu),
is_appservice,
))
})
.filter_map(|r| r.ok())
.filter_map(|(key, _)| Self::parse_servercurrentpdus(key).ok())
.filter(|(_, pdu, _)| !pdu.is_empty()) // Skip reservation key
.take(50)
// This should not contain more than 50 anyway
@ -90,6 +70,8 @@ impl Sending {
));
}
let mut last_failed_try: HashMap<Box<ServerName>, (u32, Instant)> = HashMap::new();
let mut subscriber = servernamepduids.watch_prefix(b"");
loop {
select! {
@ -140,9 +122,24 @@ impl Sending {
// servercurrentpdus with the prefix should be empty now
}
}
Err((server, _is_appservice, e)) => {
warn!("Couldn't send transaction to {}: {}", server, e)
// TODO: exponential backoff
Err((server, is_appservice, e)) => {
warn!("Couldn't send transaction to {}: {}", server, e);
let mut prefix = if is_appservice {
"+".as_bytes().to_vec()
} else {
Vec::new()
};
prefix.extend_from_slice(server.as_bytes());
prefix.push(0xff);
last_failed_try.insert(server.clone(), match last_failed_try.get(&server) {
Some(last_failed) => {
(last_failed.0+1, Instant::now())
},
None => {
(1, Instant::now())
}
});
servercurrentpdus.remove(&prefix).unwrap();
}
};
},
@ -174,8 +171,19 @@ impl Sending {
.ok()
.map(|pdu_id| (server, is_appservice, pdu_id))
)
// TODO: exponential backoff
.filter(|(server, is_appservice, _)| {
if last_failed_try.get(server).map_or(false, |(tries, instant)| {
// Fail if a request has failed recently (exponential backoff)
let mut min_elapsed_duration = Duration::from_secs(60) * *tries * *tries;
if min_elapsed_duration > Duration::from_secs(60*60*24) {
min_elapsed_duration = Duration::from_secs(60*60*24);
}
instant.elapsed() < min_elapsed_duration
}) {
return false;
}
let mut prefix = if *is_appservice {
"+".as_bytes().to_vec()
} else {
@ -308,4 +316,63 @@ impl Sending {
.map_err(|e| (server, is_appservice, e))
}
}
fn parse_servercurrentpdus(key: IVec) -> Result<(Box<ServerName>, IVec, bool)> {
let mut parts = key.splitn(2, |&b| b == 0xff);
let server = parts.next().expect("splitn always returns one element");
let pdu = parts
.next()
.ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
let server = utils::string_from_bytes(&server).map_err(|_| {
Error::bad_database("Invalid server bytes in server_currenttransaction")
})?;
// Appservices start with a plus
let (server, is_appservice) = if server.starts_with("+") {
(&server[1..], true)
} else {
(&*server, false)
};
Ok::<_, Error>((
Box::<ServerName>::try_from(server).map_err(|_| {
Error::bad_database("Invalid server string in server_currenttransaction")
})?,
IVec::from(pdu),
is_appservice,
))
}
pub async fn send_federation_request<T: OutgoingRequest>(
&self,
globals: &crate::database::globals::Globals,
destination: Box<ServerName>,
request: T,
) -> Result<T::IncomingResponse>
where
T: Debug,
{
let permit = self.maximum_requests.acquire().await;
let response = server_server::send_request(globals, destination, request).await;
drop(permit);
response
}
pub async fn send_appservice_request<T: OutgoingRequest>(
&self,
globals: &crate::database::globals::Globals,
registration: serde_yaml::Value,
request: T,
) -> Result<T::IncomingResponse>
where
T: Debug,
{
let permit = self.maximum_requests.acquire().await;
let response = appservice_server::send_request(globals, registration, request).await;
drop(permit);
response
}
}