TudbuT/conduit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

rename OutgoingKind to Destination, add QueuedRequests and ActiveRequestsFor admin query commands

Browse source 
Signed-off-by: strawberry <strawberry@puppygock.gay>
This commit is contained in:
strawberry 2024-04-22 17:27:54 -04:00 • committed by June
parent 77d73583f6
commit e26cd5e296
5 changed files with 258 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
src/database/key_value/sending.rs
View file

@ -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<dyn Iterator<Item = Result<(Vec<u8>, OutgoingKind, SendingEventType)>> + 'a> {
) -> Box<dyn Iterator<Item = Result<(Vec<u8>, 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<dyn Iterator<Item = Result<(Vec<u8>, 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<u8>) -> 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<Vec<Vec<u8>>> {
fn queue_requests(&self, requests: &[(&Destination, SendingEventType)]) -> Result<Vec<Vec<u8>>> {
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<dyn Iterator<Item = Result<(SendingEventType, Vec<u8>)>> + '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<u8>) -> Result<(OutgoingKind, SendingEventType)> {
fn parse_servercurrentevent(key: &[u8], value: Vec<u8>) -> 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<u8>) -> 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<u8>) -> 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<u8>) -> 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"))?,
),

45
src/service/admin/query/mod.rs
View file

@ -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<String>,
#[arg(short, long)]
server_name: Option<Box<ServerName>>,
#[arg(short, long)]
user_id: Option<Box<UserId>>,
#[arg(short, long)]
push_key: Option<String>,
},
/// - 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<String>,
#[arg(short, long)]
server_name: Option<Box<ServerName>>,
#[arg(short, long)]
user_id: Option<Box<UserId>>,
#[arg(short, long)]
push_key: Option<String>,
},
GetLatestEduCount {
server_name: Box<ServerName>,
},

166
src/service/admin/query/sending.rs
View file

@ -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<RoomMessageEventContent> {
@ -21,6 +21,170 @@ pub(super) async fn sending(subcommand: Sending) -> Result<RoomMessageEventConte
),
))
},
Sending::QueuedRequests {
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
.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::<Result<Vec<(_, _)>>>();
Ok(RoomMessageEventContent::text_html(
format!("Query completed in {query_time:?}:\n\n```\n{:?}```", queued_requests),
format!(
"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
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::<Result<Vec<(_, _)>>>();
Ok(RoomMessageEventContent::text_html(
format!("Query completed in {query_time:?}:\n\n```\n{:?}```", active_requests),
format!(
"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
active_requests
),
))
},
Sending::GetLatestEduCount {
server_name,
} => {

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

@ -1,20 +1,20 @@
use ruma::ServerName;
use super::{OutgoingKind, SendingEventType};
use super::{Destination, SendingEventType};
use crate::Result;
type OutgoingSendingIter<'a> = Box<dyn Iterator<Item = Result<(Vec<u8>, OutgoingKind, SendingEventType)>> + 'a>;
type OutgoingSendingIter<'a> = Box<dyn Iterator<Item = Result<(Vec<u8>, Destination, SendingEventType)>> + 'a>;
type SendingEventTypeIter<'a> = Box<dyn Iterator<Item = Result<(Vec<u8>, 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<u8>) -> 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<Vec<Vec<u8>>>;
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<Vec<Vec<u8>>>;
fn queued_requests<'a>(
&'a self, outgoing_kind: &OutgoingKind,
&'a self, destination: &Destination,
) -> Box<dyn Iterator<Item = Result<(SendingEventType, Vec<u8>)>> + 'a>;
fn mark_as_active(&self, events: &[(SendingEventType, Vec<u8>)]) -> Result<()>;
fn set_latest_educount(&self, server_name: &ServerName, educount: u64) -> Result<()>;

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

@ -42,15 +42,15 @@ pub struct Service {
/// The state for a given state hash.
pub(super) maximum_requests: Arc<Semaphore>,
pub sender: loole::Sender<(OutgoingKind, SendingEventType, Vec<u8>)>,
receiver: Mutex<loole::Receiver<(OutgoingKind, SendingEventType, Vec<u8>)>>,
pub sender: loole::Sender<(Destination, SendingEventType, Vec<u8>)>,
receiver: Mutex<loole::Receiver<(Destination, SendingEventType, Vec<u8>)>>,
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<u8>) -> 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<I: Iterator<Item = OwnedServerName>>(&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::<Vec<_>>();
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<u8>) -> 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<I: Iterator<Item = OwnedServerName>>(&self, servers: I, serialized: Vec<u8>) -> 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::<Vec<_>>();
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<I: Iterator<Item = OwnedServerName>>(&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::<OutgoingKind, TransactionStatus>::new();
let mut current_transaction_status = HashMap::<Destination, TransactionStatus>::new();
// Retry requests we could not finish yet
if self.startup_netburst {
let mut initial_transactions = HashMap::<OutgoingKind, Vec<SendingEventType>>::new();
let mut initial_transactions = HashMap::<Destination, Vec<SendingEventType>>::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<u8>)>, // Events we want to send: event and full key
current_transaction_status: &mut HashMap<OutgoingKind, TransactionStatus>,
current_transaction_status: &mut HashMap<Destination, TransactionStatus>,
) -> Result<Option<Vec<SendingEventType>>> {
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<OutgoingKind, TransactionStatus>,
&self, outgoing_kind: Destination, current_transaction_status: &mut HashMap<Destination, TransactionStatus>,
) -> 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<SendingEventType>,
) -> Result<OutgoingKind, (OutgoingKind, Error)> {
async fn handle_events(kind: Destination, events: Vec<SendingEventType>) -> Result<Destination, (Destination, Error)> {
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<SendingEventType>,
) -> Result<OutgoingKind, (OutgoingKind, Error)> {
kind: &Destination, id: &String, events: Vec<SendingEventType>,
) -> Result<Destination, (Destination, Error)> {
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<SendingEventType>,
) -> Result<OutgoingKind, (OutgoingKind, Error)> {
kind: &Destination, userid: &OwnedUserId, pushkey: &String, events: Vec<SendingEventType>,
) -> Result<Destination, (Destination, Error)> {
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<SendingEventType>,
) -> Result<OutgoingKind, (OutgoingKind, Error)> {
kind: &Destination, dest: &OwnedServerName, events: Vec<SendingEventType>,
) -> Result<Destination, (Destination, Error)> {
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<u8> {
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