sending.rs: add RequestKey

Much easier to reason about than with a bunch of Vec<u8> everywhere.
This commit is contained in:
Lambda 2024-05-22 18:28:26 +00:00
parent 507de063f5
commit e294543ddb
3 changed files with 50 additions and 27 deletions

View file

@ -4,7 +4,7 @@ use crate::{
database::KeyValueDatabase, database::KeyValueDatabase,
service::{ service::{
self, self,
sending::{OutgoingKind, SendingEventType}, sending::{OutgoingKind, RequestKey, SendingEventType},
}, },
services, utils, Error, Result, services, utils, Error, Result,
}; };
@ -13,10 +13,12 @@ impl service::sending::Data for KeyValueDatabase {
fn active_requests<'a>( fn active_requests<'a>(
&'a self, &'a self,
) -> Box< ) -> Box<
dyn Iterator<Item = Result<(Vec<u8>, OutgoingKind, SendingEventType)>> dyn Iterator<
+ 'a, Item = Result<(RequestKey, OutgoingKind, SendingEventType)>,
> + 'a,
> { > {
Box::new(self.servercurrentevent_data.iter().map(|(key, v)| { Box::new(self.servercurrentevent_data.iter().map(|(key, v)| {
let key = RequestKey::new(key);
parse_servercurrentevent(&key, v).map(|(k, e)| (key, k, e)) parse_servercurrentevent(&key, v).map(|(k, e)| (key, k, e))
})) }))
} }
@ -24,16 +26,19 @@ impl service::sending::Data for KeyValueDatabase {
fn active_requests_for<'a>( fn active_requests_for<'a>(
&'a self, &'a self,
outgoing_kind: &OutgoingKind, outgoing_kind: &OutgoingKind,
) -> Box<dyn Iterator<Item = Result<(Vec<u8>, SendingEventType)>> + 'a> ) -> Box<dyn Iterator<Item = Result<(RequestKey, SendingEventType)>> + 'a>
{ {
let prefix = outgoing_kind.get_prefix(); let prefix = outgoing_kind.get_prefix();
Box::new(self.servercurrentevent_data.scan_prefix(prefix).map( Box::new(self.servercurrentevent_data.scan_prefix(prefix).map(
|(key, v)| parse_servercurrentevent(&key, v).map(|(_, e)| (key, e)), |(key, v)| {
let key = RequestKey::new(key);
parse_servercurrentevent(&key, v).map(|(_, e)| (key, e))
},
)) ))
} }
fn delete_active_request(&self, key: Vec<u8>) -> Result<()> { fn delete_active_request(&self, key: RequestKey) -> Result<()> {
self.servercurrentevent_data.remove(&key) self.servercurrentevent_data.remove(key.as_bytes())
} }
fn delete_all_active_requests_for( fn delete_all_active_requests_for(
@ -51,7 +56,7 @@ impl service::sending::Data for KeyValueDatabase {
fn queue_requests( fn queue_requests(
&self, &self,
requests: &[(&OutgoingKind, SendingEventType)], requests: &[(&OutgoingKind, SendingEventType)],
) -> Result<Vec<Vec<u8>>> { ) -> Result<Vec<RequestKey>> {
let mut batch = Vec::new(); let mut batch = Vec::new();
let mut keys = Vec::new(); let mut keys = Vec::new();
for (outgoing_kind, event) in requests { for (outgoing_kind, event) in requests {
@ -69,7 +74,7 @@ impl service::sending::Data for KeyValueDatabase {
&[] &[]
}; };
batch.push((key.clone(), value.to_owned())); batch.push((key.clone(), value.to_owned()));
keys.push(key); keys.push(RequestKey::new(key));
} }
self.servernameevent_data.insert_batch(&mut batch.into_iter())?; self.servernameevent_data.insert_batch(&mut batch.into_iter())?;
Ok(keys) Ok(keys)
@ -78,17 +83,20 @@ impl service::sending::Data for KeyValueDatabase {
fn queued_requests<'a>( fn queued_requests<'a>(
&'a self, &'a self,
outgoing_kind: &OutgoingKind, outgoing_kind: &OutgoingKind,
) -> Box<dyn Iterator<Item = Result<(SendingEventType, Vec<u8>)>> + 'a> ) -> Box<dyn Iterator<Item = Result<(SendingEventType, RequestKey)>> + 'a>
{ {
let prefix = outgoing_kind.get_prefix(); let prefix = outgoing_kind.get_prefix();
return Box::new(self.servernameevent_data.scan_prefix(prefix).map( return Box::new(self.servernameevent_data.scan_prefix(prefix).map(
|(k, v)| parse_servercurrentevent(&k, v).map(|(_, ev)| (ev, k)), |(k, v)| {
let k = RequestKey::new(k);
parse_servercurrentevent(&k, v).map(|(_, ev)| (ev, k))
},
)); ));
} }
fn mark_as_active( fn mark_as_active(
&self, &self,
events: &[(SendingEventType, Vec<u8>)], events: &[(SendingEventType, RequestKey)],
) -> Result<()> { ) -> Result<()> {
for (e, key) in events { for (e, key) in events {
let value = if let SendingEventType::Edu(value) = &e { let value = if let SendingEventType::Edu(value) = &e {
@ -96,8 +104,8 @@ impl service::sending::Data for KeyValueDatabase {
} else { } else {
&[] &[]
}; };
self.servercurrentevent_data.insert(key, value)?; self.servercurrentevent_data.insert(key.as_bytes(), value)?;
self.servernameevent_data.remove(key)?; self.servernameevent_data.remove(key.as_bytes())?;
} }
Ok(()) Ok(())
@ -126,9 +134,10 @@ impl service::sending::Data for KeyValueDatabase {
#[tracing::instrument(skip(key))] #[tracing::instrument(skip(key))]
fn parse_servercurrentevent( fn parse_servercurrentevent(
key: &[u8], key: &RequestKey,
value: Vec<u8>, value: Vec<u8>,
) -> Result<(OutgoingKind, SendingEventType)> { ) -> Result<(OutgoingKind, SendingEventType)> {
let key = key.as_bytes();
// Appservices start with a plus // Appservices start with a plus
Ok::<_, Error>(if key.starts_with(b"+") { Ok::<_, Error>(if key.starts_with(b"+") {
let mut parts = key[1..].splitn(2, |&b| b == 0xFF); let mut parts = key[1..].splitn(2, |&b| b == 0xFF);

View file

@ -88,15 +88,28 @@ pub(crate) enum SendingEventType {
Edu(Vec<u8>), Edu(Vec<u8>),
} }
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) struct RequestKey(Vec<u8>);
impl RequestKey {
pub(crate) fn new(key: Vec<u8>) -> Self {
Self(key)
}
pub(crate) fn as_bytes(&self) -> &[u8] {
&self.0
}
}
pub(crate) struct Service { pub(crate) struct Service {
db: &'static dyn Data, db: &'static dyn Data,
/// The state for a given state hash. /// The state for a given state hash.
pub(super) maximum_requests: Arc<Semaphore>, pub(super) maximum_requests: Arc<Semaphore>,
pub(crate) sender: pub(crate) sender:
mpsc::UnboundedSender<(OutgoingKind, SendingEventType, Vec<u8>)>, mpsc::UnboundedSender<(OutgoingKind, SendingEventType, RequestKey)>,
receiver: Mutex< receiver: Mutex<
mpsc::UnboundedReceiver<(OutgoingKind, SendingEventType, Vec<u8>)>, mpsc::UnboundedReceiver<(OutgoingKind, SendingEventType, RequestKey)>,
>, >,
} }
@ -284,7 +297,7 @@ impl Service {
&self, &self,
outgoing_kind: OutgoingKind, outgoing_kind: OutgoingKind,
event: SendingEventType, event: SendingEventType,
key: Vec<u8>, key: RequestKey,
current_transaction_status: &mut TransactionStatusMap, current_transaction_status: &mut TransactionStatusMap,
) -> Option<HandlerInputs> { ) -> Option<HandlerInputs> {
if let Ok(Some(events)) = self.select_events( if let Ok(Some(events)) = self.select_events(
@ -312,7 +325,7 @@ impl Service {
&self, &self,
outgoing_kind: &OutgoingKind, outgoing_kind: &OutgoingKind,
// Events we want to send: event and full key // Events we want to send: event and full key
new_events: Vec<(SendingEventType, Vec<u8>)>, new_events: Vec<(SendingEventType, RequestKey)>,
current_transaction_status: &mut HashMap< current_transaction_status: &mut HashMap<
OutgoingKind, OutgoingKind,
TransactionStatus, TransactionStatus,

View file

@ -1,6 +1,6 @@
use ruma::ServerName; use ruma::ServerName;
use super::{OutgoingKind, SendingEventType}; use super::{OutgoingKind, RequestKey, SendingEventType};
use crate::Result; use crate::Result;
pub(crate) trait Data: Send + Sync { pub(crate) trait Data: Send + Sync {
@ -8,14 +8,15 @@ pub(crate) trait Data: Send + Sync {
fn active_requests<'a>( fn active_requests<'a>(
&'a self, &'a self,
) -> Box< ) -> Box<
dyn Iterator<Item = Result<(Vec<u8>, OutgoingKind, SendingEventType)>> dyn Iterator<
+ 'a, Item = Result<(RequestKey, OutgoingKind, SendingEventType)>,
> + 'a,
>; >;
fn active_requests_for<'a>( fn active_requests_for<'a>(
&'a self, &'a self,
outgoing_kind: &OutgoingKind, outgoing_kind: &OutgoingKind,
) -> Box<dyn Iterator<Item = Result<(Vec<u8>, SendingEventType)>> + 'a>; ) -> Box<dyn Iterator<Item = Result<(RequestKey, SendingEventType)>> + 'a>;
fn delete_active_request(&self, key: Vec<u8>) -> Result<()>; fn delete_active_request(&self, key: RequestKey) -> Result<()>;
fn delete_all_active_requests_for( fn delete_all_active_requests_for(
&self, &self,
outgoing_kind: &OutgoingKind, outgoing_kind: &OutgoingKind,
@ -23,14 +24,14 @@ pub(crate) trait Data: Send + Sync {
fn queue_requests( fn queue_requests(
&self, &self,
requests: &[(&OutgoingKind, SendingEventType)], requests: &[(&OutgoingKind, SendingEventType)],
) -> Result<Vec<Vec<u8>>>; ) -> Result<Vec<RequestKey>>;
fn queued_requests<'a>( fn queued_requests<'a>(
&'a self, &'a self,
outgoing_kind: &OutgoingKind, outgoing_kind: &OutgoingKind,
) -> Box<dyn Iterator<Item = Result<(SendingEventType, Vec<u8>)>> + 'a>; ) -> Box<dyn Iterator<Item = Result<(SendingEventType, RequestKey)>> + 'a>;
fn mark_as_active( fn mark_as_active(
&self, &self,
events: &[(SendingEventType, Vec<u8>)], events: &[(SendingEventType, RequestKey)],
) -> Result<()>; ) -> Result<()>;
fn set_latest_educount( fn set_latest_educount(
&self, &self,