From e26cd5e296401418a3f533cd10bd9614b992d286 Mon Sep 17 00:00:00 2001 From: strawberry Date: Mon, 22 Apr 2024 17:27:54 -0400 Subject: [PATCH] rename OutgoingKind to Destination, add QueuedRequests and ActiveRequestsFor admin query commands Signed-off-by: strawberry --- src/database/key_value/sending.rs | 22 ++-- src/service/admin/query/mod.rs | 45 ++++++++ src/service/admin/query/sending.rs | 166 ++++++++++++++++++++++++++++- src/service/sending/data.rs | 14 +-- src/service/sending/mod.rs | 62 ++++++----- 5 files changed, 258 insertions(+), 51 deletions(-) diff --git a/src/database/key_value/sending.rs b/src/database/key_value/sending.rs index 9f871ee7..049c6236 100644 --- a/src/database/key_value/sending.rs +++ b/src/database/key_value/sending.rs @@ -4,7 +4,7 @@ use crate::{ database::KeyValueDatabase, service::{ self, - sending::{OutgoingKind, SendingEventType}, + sending::{Destination, SendingEventType}, }, services, utils, Error, Result, }; @@ -12,7 +12,7 @@ use crate::{ impl service::sending::Data for KeyValueDatabase { fn active_requests<'a>( &'a self, - ) -> Box, OutgoingKind, SendingEventType)>> + 'a> { + ) -> Box, Destination, SendingEventType)>> + 'a> { Box::new( self.servercurrentevent_data .iter() @@ -21,7 +21,7 @@ impl service::sending::Data for KeyValueDatabase { } fn active_requests_for<'a>( - &'a self, outgoing_kind: &OutgoingKind, + &'a self, outgoing_kind: &Destination, ) -> Box, SendingEventType)>> + 'a> { let prefix = outgoing_kind.get_prefix(); Box::new( @@ -33,7 +33,7 @@ impl service::sending::Data for KeyValueDatabase { fn delete_active_request(&self, key: Vec) -> Result<()> { self.servercurrentevent_data.remove(&key) } - fn delete_all_active_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()> { + fn delete_all_active_requests_for(&self, outgoing_kind: &Destination) -> Result<()> { let prefix = outgoing_kind.get_prefix(); for (key, _) in self.servercurrentevent_data.scan_prefix(prefix) { self.servercurrentevent_data.remove(&key)?; @@ -42,7 +42,7 @@ impl service::sending::Data for KeyValueDatabase { Ok(()) } - fn delete_all_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()> { + fn delete_all_requests_for(&self, outgoing_kind: &Destination) -> Result<()> { let prefix = outgoing_kind.get_prefix(); for (key, _) in self.servercurrentevent_data.scan_prefix(prefix.clone()) { self.servercurrentevent_data.remove(&key).unwrap(); @@ -55,7 +55,7 @@ impl service::sending::Data for KeyValueDatabase { Ok(()) } - fn queue_requests(&self, requests: &[(&OutgoingKind, SendingEventType)]) -> Result>> { + fn queue_requests(&self, requests: &[(&Destination, SendingEventType)]) -> Result>> { let mut batch = Vec::new(); let mut keys = Vec::new(); for (outgoing_kind, event) in requests { @@ -79,7 +79,7 @@ impl service::sending::Data for KeyValueDatabase { } fn queued_requests<'a>( - &'a self, outgoing_kind: &OutgoingKind, + &'a self, outgoing_kind: &Destination, ) -> Box)>> + 'a> { let prefix = outgoing_kind.get_prefix(); return Box::new( @@ -122,7 +122,7 @@ impl service::sending::Data for KeyValueDatabase { } #[tracing::instrument(skip(key))] -fn parse_servercurrentevent(key: &[u8], value: Vec) -> Result<(OutgoingKind, SendingEventType)> { +fn parse_servercurrentevent(key: &[u8], value: Vec) -> Result<(Destination, SendingEventType)> { // Appservices start with a plus Ok::<_, Error>(if key.starts_with(b"+") { let mut parts = key[1..].splitn(2, |&b| b == 0xFF); @@ -136,7 +136,7 @@ fn parse_servercurrentevent(key: &[u8], value: Vec) -> Result<(OutgoingKind, .map_err(|_| Error::bad_database("Invalid server bytes in server_currenttransaction"))?; ( - OutgoingKind::Appservice(server), + Destination::Appservice(server), if value.is_empty() { SendingEventType::Pdu(event.to_vec()) } else { @@ -163,7 +163,7 @@ fn parse_servercurrentevent(key: &[u8], value: Vec) -> Result<(OutgoingKind, .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?; ( - OutgoingKind::Push(user_id, pushkey_string), + Destination::Push(user_id, pushkey_string), if value.is_empty() { SendingEventType::Pdu(event.to_vec()) } else { @@ -183,7 +183,7 @@ fn parse_servercurrentevent(key: &[u8], value: Vec) -> Result<(OutgoingKind, .map_err(|_| Error::bad_database("Invalid server bytes in server_currenttransaction"))?; ( - OutgoingKind::Normal( + Destination::Normal( ServerName::parse(server) .map_err(|_| Error::bad_database("Invalid server string in server_currenttransaction"))?, ), diff --git a/src/service/admin/query/mod.rs b/src/service/admin/query/mod.rs index a1d7c53a..f3b4515d 100644 --- a/src/service/admin/query/mod.rs +++ b/src/service/admin/query/mod.rs @@ -147,8 +147,53 @@ pub(crate) enum Globals { #[derive(Subcommand)] /// All the getters and iterators from src/database/key_value/sending.rs pub(crate) enum Sending { + /// - Queries database for all `servercurrentevent_data` ActiveRequests, + /// - Queries database for `servercurrentevent_data` but for a specific + /// destination + /// + /// This command takes only *one* format of these arguments: + /// + /// appservice_id + /// server_name + /// user_id AND push_key + /// + /// See src/service/sending/mod.rs for the definition of the `Destination` + /// enum + ActiveRequestsFor { + #[arg(short, long)] + appservice_id: Option, + #[arg(short, long)] + server_name: Option>, + #[arg(short, long)] + user_id: Option>, + #[arg(short, long)] + push_key: Option, + }, + + /// - Queries database for `servernameevent_data` which are the queued up + /// requests that will eventually be sent + /// + /// This command takes only *one* format of these arguments: + /// + /// appservice_id + /// server_name + /// user_id AND push_key + /// + /// See src/service/sending/mod.rs for the definition of the `Destination` + /// enum + QueuedRequests { + #[arg(short, long)] + appservice_id: Option, + #[arg(short, long)] + server_name: Option>, + #[arg(short, long)] + user_id: Option>, + #[arg(short, long)] + push_key: Option, + }, + GetLatestEduCount { server_name: Box, }, diff --git a/src/service/admin/query/sending.rs b/src/service/admin/query/sending.rs index e980a6ee..da96c8fe 100644 --- a/src/service/admin/query/sending.rs +++ b/src/service/admin/query/sending.rs @@ -1,7 +1,7 @@ use ruma::events::room::message::RoomMessageEventContent; use super::Sending; -use crate::{services, Result}; +use crate::{service::sending::Destination, services, Result}; /// All the getters and iterators in key_value/sending.rs pub(super) async fn sending(subcommand: Sending) -> Result { @@ -21,6 +21,170 @@ pub(super) async fn sending(subcommand: Sending) -> Result { + if appservice_id.is_none() && server_name.is_none() && user_id.is_none() && push_key.is_none() { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via arguments. See \ + --help for more details.", + )); + } + + let (results, query_time) = match (appservice_id, server_name, user_id, push_key) { + (Some(appservice_id), None, None, None) => { + if appservice_id.is_empty() { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via \ + arguments. See --help for more details.", + )); + } + + let timer = tokio::time::Instant::now(); + let results = services() + .sending + .db + .queued_requests(&Destination::Appservice(appservice_id)); + let query_time = timer.elapsed(); + + (results, query_time) + }, + (None, Some(server_name), None, None) => { + let timer = tokio::time::Instant::now(); + let results = services() + .sending + .db + .queued_requests(&Destination::Normal(server_name.into())); + let query_time = timer.elapsed(); + + (results, query_time) + }, + (None, None, Some(user_id), Some(push_key)) => { + if push_key.is_empty() { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via \ + arguments. See --help for more details.", + )); + } + + let timer = tokio::time::Instant::now(); + let results = services() + .sending + .db + .queued_requests(&Destination::Push(user_id.into(), push_key)); + let query_time = timer.elapsed(); + + (results, query_time) + }, + (Some(_), Some(_), Some(_), Some(_)) => { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via arguments. \ + Not all of them See --help for more details.", + )); + }, + _ => { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via arguments. \ + See --help for more details.", + )); + }, + }; + + let queued_requests = results.collect::>>(); + + Ok(RoomMessageEventContent::text_html( + format!("Query completed in {query_time:?}:\n\n```\n{:?}```", queued_requests), + format!( + "

Query completed in {query_time:?}:

\n
{:?}\n
", + queued_requests + ), + )) + }, + Sending::ActiveRequestsFor { + appservice_id, + server_name, + user_id, + push_key, + } => { + if appservice_id.is_none() && server_name.is_none() && user_id.is_none() && push_key.is_none() { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via arguments. See \ + --help for more details.", + )); + } + + let (results, query_time) = match (appservice_id, server_name, user_id, push_key) { + (Some(appservice_id), None, None, None) => { + if appservice_id.is_empty() { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via \ + arguments. See --help for more details.", + )); + } + + let timer = tokio::time::Instant::now(); + let results = services() + .sending + .db + .active_requests_for(&Destination::Appservice(appservice_id)); + let query_time = timer.elapsed(); + + (results, query_time) + }, + (None, Some(server_name), None, None) => { + let timer = tokio::time::Instant::now(); + let results = services() + .sending + .db + .active_requests_for(&Destination::Normal(server_name.into())); + let query_time = timer.elapsed(); + + (results, query_time) + }, + (None, None, Some(user_id), Some(push_key)) => { + if push_key.is_empty() { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via \ + arguments. See --help for more details.", + )); + } + + let timer = tokio::time::Instant::now(); + let results = services() + .sending + .db + .active_requests_for(&Destination::Push(user_id.into(), push_key)); + let query_time = timer.elapsed(); + + (results, query_time) + }, + (Some(_), Some(_), Some(_), Some(_)) => { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via arguments. \ + Not all of them See --help for more details.", + )); + }, + _ => { + return Ok(RoomMessageEventContent::text_plain( + "An appservice ID, server name, or a user ID with push key must be specified via arguments. \ + See --help for more details.", + )); + }, + }; + + let active_requests = results.collect::>>(); + + Ok(RoomMessageEventContent::text_html( + format!("Query completed in {query_time:?}:\n\n```\n{:?}```", active_requests), + format!( + "

Query completed in {query_time:?}:

\n
{:?}\n
", + active_requests + ), + )) + }, Sending::GetLatestEduCount { server_name, } => { diff --git a/src/service/sending/data.rs b/src/service/sending/data.rs index 46f3cd71..33066a67 100644 --- a/src/service/sending/data.rs +++ b/src/service/sending/data.rs @@ -1,20 +1,20 @@ use ruma::ServerName; -use super::{OutgoingKind, SendingEventType}; +use super::{Destination, SendingEventType}; use crate::Result; -type OutgoingSendingIter<'a> = Box, OutgoingKind, SendingEventType)>> + 'a>; +type OutgoingSendingIter<'a> = Box, Destination, SendingEventType)>> + 'a>; type SendingEventTypeIter<'a> = Box, SendingEventType)>> + 'a>; pub trait Data: Send + Sync { fn active_requests(&self) -> OutgoingSendingIter<'_>; - fn active_requests_for(&self, outgoing_kind: &OutgoingKind) -> SendingEventTypeIter<'_>; + fn active_requests_for(&self, destination: &Destination) -> SendingEventTypeIter<'_>; fn delete_active_request(&self, key: Vec) -> Result<()>; - fn delete_all_active_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()>; - fn delete_all_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()>; - fn queue_requests(&self, requests: &[(&OutgoingKind, SendingEventType)]) -> Result>>; + fn delete_all_active_requests_for(&self, destination: &Destination) -> Result<()>; + fn delete_all_requests_for(&self, destination: &Destination) -> Result<()>; + fn queue_requests(&self, requests: &[(&Destination, SendingEventType)]) -> Result>>; fn queued_requests<'a>( - &'a self, outgoing_kind: &OutgoingKind, + &'a self, destination: &Destination, ) -> Box)>> + 'a>; fn mark_as_active(&self, events: &[(SendingEventType, Vec)]) -> Result<()>; fn set_latest_educount(&self, server_name: &ServerName, educount: u64) -> Result<()>; diff --git a/src/service/sending/mod.rs b/src/service/sending/mod.rs index 32a7409a..040422dd 100644 --- a/src/service/sending/mod.rs +++ b/src/service/sending/mod.rs @@ -42,15 +42,15 @@ pub struct Service { /// The state for a given state hash. pub(super) maximum_requests: Arc, - pub sender: loole::Sender<(OutgoingKind, SendingEventType, Vec)>, - receiver: Mutex)>>, + pub sender: loole::Sender<(Destination, SendingEventType, Vec)>, + receiver: Mutex)>>, startup_netburst: bool, startup_netburst_keep: i64, timeout: u64, } #[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub enum OutgoingKind { +pub enum Destination { Appservice(String), Push(OwnedUserId, String), // user and pushkey Normal(OwnedServerName), @@ -86,7 +86,7 @@ impl Service { #[tracing::instrument(skip(self, pdu_id, user, pushkey))] pub fn send_pdu_push(&self, pdu_id: &[u8], user: &UserId, pushkey: String) -> Result<()> { - let outgoing_kind = OutgoingKind::Push(user.to_owned(), pushkey); + let outgoing_kind = Destination::Push(user.to_owned(), pushkey); let event = SendingEventType::Pdu(pdu_id.to_owned()); let _cork = services().globals.db.cork()?; let keys = self.db.queue_requests(&[(&outgoing_kind, event.clone())])?; @@ -99,7 +99,7 @@ impl Service { #[tracing::instrument(skip(self))] pub fn send_pdu_appservice(&self, appservice_id: String, pdu_id: Vec) -> Result<()> { - let outgoing_kind = OutgoingKind::Appservice(appservice_id); + let outgoing_kind = Destination::Appservice(appservice_id); let event = SendingEventType::Pdu(pdu_id); let _cork = services().globals.db.cork()?; let keys = self.db.queue_requests(&[(&outgoing_kind, event.clone())])?; @@ -126,7 +126,7 @@ impl Service { pub fn send_pdu_servers>(&self, servers: I, pdu_id: &[u8]) -> Result<()> { let requests = servers .into_iter() - .map(|server| (OutgoingKind::Normal(server), SendingEventType::Pdu(pdu_id.to_owned()))) + .map(|server| (Destination::Normal(server), SendingEventType::Pdu(pdu_id.to_owned()))) .collect::>(); let _cork = services().globals.db.cork()?; let keys = self.db.queue_requests( @@ -146,7 +146,7 @@ impl Service { #[tracing::instrument(skip(self, server, serialized))] pub fn send_edu_server(&self, server: &ServerName, serialized: Vec) -> Result<()> { - let outgoing_kind = OutgoingKind::Normal(server.to_owned()); + let outgoing_kind = Destination::Normal(server.to_owned()); let event = SendingEventType::Edu(serialized); let _cork = services().globals.db.cork()?; let keys = self.db.queue_requests(&[(&outgoing_kind, event.clone())])?; @@ -173,7 +173,7 @@ impl Service { pub fn send_edu_servers>(&self, servers: I, serialized: Vec) -> Result<()> { let requests = servers .into_iter() - .map(|server| (OutgoingKind::Normal(server), SendingEventType::Edu(serialized.clone()))) + .map(|server| (Destination::Normal(server), SendingEventType::Edu(serialized.clone()))) .collect::>(); let _cork = services().globals.db.cork()?; let keys = self.db.queue_requests( @@ -205,7 +205,7 @@ impl Service { #[tracing::instrument(skip(self, servers))] pub fn flush_servers>(&self, servers: I) -> Result<()> { - let requests = servers.into_iter().map(OutgoingKind::Normal); + let requests = servers.into_iter().map(Destination::Normal); for outgoing_kind in requests { self.sender @@ -221,7 +221,7 @@ impl Service { #[tracing::instrument(skip(self))] pub fn cleanup_events(&self, appservice_id: String) -> Result<()> { self.db - .delete_all_requests_for(&OutgoingKind::Appservice(appservice_id))?; + .delete_all_requests_for(&Destination::Appservice(appservice_id))?; Ok(()) } @@ -277,11 +277,11 @@ impl Service { let receiver = self.receiver.lock().await; let mut futures = FuturesUnordered::new(); - let mut current_transaction_status = HashMap::::new(); + let mut current_transaction_status = HashMap::::new(); // Retry requests we could not finish yet if self.startup_netburst { - let mut initial_transactions = HashMap::>::new(); + let mut initial_transactions = HashMap::>::new(); for (key, outgoing_kind, event) in self.db.active_requests().filter_map(Result::ok) { let entry = initial_transactions .entry(outgoing_kind.clone()) @@ -361,9 +361,9 @@ impl Service { #[tracing::instrument(skip(self, outgoing_kind, new_events, current_transaction_status))] fn select_events( &self, - outgoing_kind: &OutgoingKind, + outgoing_kind: &Destination, new_events: Vec<(SendingEventType, Vec)>, // Events we want to send: event and full key - current_transaction_status: &mut HashMap, + current_transaction_status: &mut HashMap, ) -> Result>> { let (allow, retry) = self.select_events_current(outgoing_kind.clone(), current_transaction_status)?; @@ -395,7 +395,7 @@ impl Service { } // Add EDU's into the transaction - if let OutgoingKind::Normal(server_name) = outgoing_kind { + if let Destination::Normal(server_name) = outgoing_kind { if let Ok((select_edus, last_count)) = self.select_edus(server_name) { events.extend(select_edus.into_iter().map(SendingEventType::Edu)); self.db.set_latest_educount(server_name, last_count)?; @@ -407,7 +407,7 @@ impl Service { #[tracing::instrument(skip(self, outgoing_kind, current_transaction_status))] fn select_events_current( - &self, outgoing_kind: OutgoingKind, current_transaction_status: &mut HashMap, + &self, outgoing_kind: Destination, current_transaction_status: &mut HashMap, ) -> Result<(bool, bool)> { let (mut allow, mut retry) = (true, false); current_transaction_status @@ -596,20 +596,18 @@ pub fn select_edus_receipts( Ok(true) } -async fn handle_events( - kind: OutgoingKind, events: Vec, -) -> Result { +async fn handle_events(kind: Destination, events: Vec) -> Result { match kind { - OutgoingKind::Appservice(ref id) => handle_events_kind_appservice(&kind, id, events).await, - OutgoingKind::Push(ref userid, ref pushkey) => handle_events_kind_push(&kind, userid, pushkey, events).await, - OutgoingKind::Normal(ref server) => handle_events_kind_normal(&kind, server, events).await, + Destination::Appservice(ref id) => handle_events_kind_appservice(&kind, id, events).await, + Destination::Push(ref userid, ref pushkey) => handle_events_kind_push(&kind, userid, pushkey, events).await, + Destination::Normal(ref server) => handle_events_kind_normal(&kind, server, events).await, } } #[tracing::instrument(skip(kind, events))] async fn handle_events_kind_appservice( - kind: &OutgoingKind, id: &String, events: Vec, -) -> Result { + kind: &Destination, id: &String, events: Vec, +) -> Result { let mut pdu_jsons = Vec::new(); for event in &events { @@ -677,8 +675,8 @@ async fn handle_events_kind_appservice( #[tracing::instrument(skip(kind, events))] async fn handle_events_kind_push( - kind: &OutgoingKind, userid: &OwnedUserId, pushkey: &String, events: Vec, -) -> Result { + kind: &Destination, userid: &OwnedUserId, pushkey: &String, events: Vec, +) -> Result { let mut pdus = Vec::new(); for event in &events { @@ -755,8 +753,8 @@ async fn handle_events_kind_push( #[tracing::instrument(skip(kind, events), name = "")] async fn handle_events_kind_normal( - kind: &OutgoingKind, dest: &OwnedServerName, events: Vec, -) -> Result { + kind: &Destination, dest: &OwnedServerName, events: Vec, +) -> Result { let mut edu_jsons = Vec::new(); let mut pdu_jsons = Vec::new(); @@ -829,23 +827,23 @@ async fn handle_events_kind_normal( response } -impl OutgoingKind { +impl Destination { #[tracing::instrument(skip(self))] pub fn get_prefix(&self) -> Vec { let mut prefix = match self { - OutgoingKind::Appservice(server) => { + Destination::Appservice(server) => { let mut p = b"+".to_vec(); p.extend_from_slice(server.as_bytes()); p }, - OutgoingKind::Push(user, pushkey) => { + Destination::Push(user, pushkey) => { let mut p = b"$".to_vec(); p.extend_from_slice(user.as_bytes()); p.push(0xFF); p.extend_from_slice(pushkey.as_bytes()); p }, - OutgoingKind::Normal(server) => { + Destination::Normal(server) => { let mut p = Vec::new(); p.extend_from_slice(server.as_bytes()); p