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

dissolve key_value/*

Browse source 
Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk 2024-05-26 21:29:19 +00:00 • committed by June 🍓🦴
parent 3122648767
commit b94045a468
88 changed files with 4556 additions and 4751 deletions
Download patch file Download diff file Expand all files Collapse all files

189
src/service/sending/data.rs
View file

@ -1,7 +1,7 @@
use ruma::ServerName;
use ruma::{ServerName, UserId};
use super::{Destination, SendingEvent};
use crate::Result;
use crate::{services, utils, Error, KeyValueDatabase, Result};
type OutgoingSendingIter<'a> = Box<dyn Iterator<Item = Result<(Vec<u8>, Destination, SendingEvent)>> + 'a>;
type SendingEventIter<'a> = Box<dyn Iterator<Item = Result<(Vec<u8>, SendingEvent)>> + 'a>;
@ -23,3 +23,188 @@ pub trait Data: Send + Sync {
fn set_latest_educount(&self, server_name: &ServerName, educount: u64) -> Result<()>;
fn get_latest_educount(&self, server_name: &ServerName) -> Result<u64>;
}
impl Data for KeyValueDatabase {
fn active_requests<'a>(&'a self) -> Box<dyn Iterator<Item = Result<(Vec<u8>, Destination, SendingEvent)>> + 'a> {
Box::new(
self.servercurrentevent_data
.iter()
.map(|(key, v)| parse_servercurrentevent(&key, v).map(|(k, e)| (key, k, e))),
)
}
fn active_requests_for<'a>(
&'a self, destination: &Destination,
) -> Box<dyn Iterator<Item = Result<(Vec<u8>, SendingEvent)>> + 'a> {
let prefix = destination.get_prefix();
Box::new(
self.servercurrentevent_data
.scan_prefix(prefix)
.map(|(key, v)| parse_servercurrentevent(&key, v).map(|(_, e)| (key, e))),
)
}
fn delete_active_request(&self, key: Vec<u8>) -> Result<()> { self.servercurrentevent_data.remove(&key) }
fn delete_all_active_requests_for(&self, destination: &Destination) -> Result<()> {
let prefix = destination.get_prefix();
for (key, _) in self.servercurrentevent_data.scan_prefix(prefix) {
self.servercurrentevent_data.remove(&key)?;
}
Ok(())
}
fn delete_all_requests_for(&self, destination: &Destination) -> Result<()> {
let prefix = destination.get_prefix();
for (key, _) in self.servercurrentevent_data.scan_prefix(prefix.clone()) {
self.servercurrentevent_data.remove(&key).unwrap();
}
for (key, _) in self.servernameevent_data.scan_prefix(prefix) {
self.servernameevent_data.remove(&key).unwrap();
}
Ok(())
}
fn queue_requests(&self, requests: &[(&Destination, SendingEvent)]) -> Result<Vec<Vec<u8>>> {
let mut batch = Vec::new();
let mut keys = Vec::new();
for (destination, event) in requests {
let mut key = destination.get_prefix();
if let SendingEvent::Pdu(value) = &event {
key.extend_from_slice(value);
} else {
key.extend_from_slice(&services().globals.next_count()?.to_be_bytes());
}
let value = if let SendingEvent::Edu(value) = &event {
&**value
} else {
&[]
};
batch.push((key.clone(), value.to_owned()));
keys.push(key);
}
self.servernameevent_data
.insert_batch(&mut batch.into_iter())?;
Ok(keys)
}
fn queued_requests<'a>(
&'a self, destination: &Destination,
) -> Box<dyn Iterator<Item = Result<(SendingEvent, Vec<u8>)>> + 'a> {
let prefix = destination.get_prefix();
return Box::new(
self.servernameevent_data
.scan_prefix(prefix)
.map(|(k, v)| parse_servercurrentevent(&k, v).map(|(_, ev)| (ev, k))),
);
}
fn mark_as_active(&self, events: &[(SendingEvent, Vec<u8>)]) -> Result<()> {
for (e, key) in events {
if key.is_empty() {
continue;
}
let value = if let SendingEvent::Edu(value) = &e {
&**value
} else {
&[]
};
self.servercurrentevent_data.insert(key, value)?;
self.servernameevent_data.remove(key)?;
}
Ok(())
}
fn set_latest_educount(&self, server_name: &ServerName, last_count: u64) -> Result<()> {
self.servername_educount
.insert(server_name.as_bytes(), &last_count.to_be_bytes())
}
fn get_latest_educount(&self, server_name: &ServerName) -> Result<u64> {
self.servername_educount
.get(server_name.as_bytes())?
.map_or(Ok(0), |bytes| {
utils::u64_from_bytes(&bytes).map_err(|_| Error::bad_database("Invalid u64 in servername_educount."))
})
}
}
#[tracing::instrument(skip(key))]
fn parse_servercurrentevent(key: &[u8], value: Vec<u8>) -> Result<(Destination, SendingEvent)> {
// Appservices start with a plus
Ok::<_, Error>(if key.starts_with(b"+") {
let mut parts = key[1..].splitn(2, |&b| b == 0xFF);
let server = parts.next().expect("splitn always returns one element");
let event = 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"))?;
(
Destination::Appservice(server),
if value.is_empty() {
SendingEvent::Pdu(event.to_vec())
} else {
SendingEvent::Edu(value)
},
)
} else if key.starts_with(b"$") {
let mut parts = key[1..].splitn(3, |&b| b == 0xFF);
let user = parts.next().expect("splitn always returns one element");
let user_string = utils::string_from_bytes(user)
.map_err(|_| Error::bad_database("Invalid user string in servercurrentevent"))?;
let user_id =
UserId::parse(user_string).map_err(|_| Error::bad_database("Invalid user id in servercurrentevent"))?;
let pushkey = parts
.next()
.ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
let pushkey_string = utils::string_from_bytes(pushkey)
.map_err(|_| Error::bad_database("Invalid pushkey in servercurrentevent"))?;
let event = parts
.next()
.ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
(
Destination::Push(user_id, pushkey_string),
if value.is_empty() {
SendingEvent::Pdu(event.to_vec())
} else {
// I'm pretty sure this should never be called
SendingEvent::Edu(value)
},
)
} else {
let mut parts = key.splitn(2, |&b| b == 0xFF);
let server = parts.next().expect("splitn always returns one element");
let event = 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"))?;
(
Destination::Normal(
ServerName::parse(server)
.map_err(|_| Error::bad_database("Invalid server string in server_currenttransaction"))?,
),
if value.is_empty() {
SendingEvent::Pdu(event.to_vec())
} else {
SendingEvent::Edu(value)
},
)
})
}

12
src/service/sending/mod.rs
View file

@ -1,6 +1,6 @@
use std::{fmt::Debug, sync::Arc};
pub use data::Data;
use data::Data;
use ruma::{
api::{appservice::Registration, OutgoingRequest},
OwnedServerName, OwnedUserId, RoomId, ServerName, UserId,
@ -81,7 +81,7 @@ impl Service {
pub fn send_pdu_push(&self, pdu_id: &[u8], user: &UserId, pushkey: String) -> Result<()> {
let dest = Destination::Push(user.to_owned(), pushkey);
let event = SendingEvent::Pdu(pdu_id.to_owned());
let _cork = services().globals.db.cork()?;
let _cork = services().globals.cork()?;
let keys = self.db.queue_requests(&[(&dest, event.clone())])?;
self.dispatch(Msg {
dest,
@ -94,7 +94,7 @@ impl Service {
pub fn send_pdu_appservice(&self, appservice_id: String, pdu_id: Vec<u8>) -> Result<()> {
let dest = Destination::Appservice(appservice_id);
let event = SendingEvent::Pdu(pdu_id);
let _cork = services().globals.db.cork()?;
let _cork = services().globals.cork()?;
let keys = self.db.queue_requests(&[(&dest, event.clone())])?;
self.dispatch(Msg {
dest,
@ -121,7 +121,7 @@ impl Service {
.into_iter()
.map(|server| (Destination::Normal(server), SendingEvent::Pdu(pdu_id.to_owned())))
.collect::<Vec<_>>();
let _cork = services().globals.db.cork()?;
let _cork = services().globals.cork()?;
let keys = self.db.queue_requests(
&requests
.iter()
@ -143,7 +143,7 @@ impl Service {
pub fn send_edu_server(&self, server: &ServerName, serialized: Vec<u8>) -> Result<()> {
let dest = Destination::Normal(server.to_owned());
let event = SendingEvent::Edu(serialized);
let _cork = services().globals.db.cork()?;
let _cork = services().globals.cork()?;
let keys = self.db.queue_requests(&[(&dest, event.clone())])?;
self.dispatch(Msg {
dest,
@ -170,7 +170,7 @@ impl Service {
.into_iter()
.map(|server| (Destination::Normal(server), SendingEvent::Edu(serialized.clone())))
.collect::<Vec<_>>();
let _cork = services().globals.db.cork()?;
let _cork = services().globals.cork()?;
let keys = self.db.queue_requests(
&requests
.iter()

6
src/service/sending/sender.rs
View file

@ -100,7 +100,7 @@ impl Service {
fn handle_response_ok(
&self, dest: &Destination, futures: &mut SendingFutures<'_>, statuses: &mut CurTransactionStatus,
) {
let _cork = services().globals.db.cork();
let _cork = services().globals.cork();
self.db
.delete_all_active_requests_for(dest)
.expect("all active requests deleted");
@ -173,7 +173,7 @@ impl Service {
return Ok(None);
}
let _cork = services().globals.db.cork();
let _cork = services().globals.cork();
let mut events = Vec::new();
// Must retry any previous transaction for this remote.
@ -187,7 +187,7 @@ impl Service {
}
// Compose the next transaction
let _cork = services().globals.db.cork();
let _cork = services().globals.cork();
if !new_events.is_empty() {
self.db.mark_as_active(&new_events)?;
for (e, _) in new_events {