feat: Rewrite acks (#741)

* feat: rewrite ack system

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* feat: rewrite acks to crond + rabbit task

* fix: review changes

---------

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>
This commit is contained in:
Tom
2026-05-08 14:39:16 -07:00
committed by GitHub
parent 9fd7128f80
commit ab5bd47a39
19 changed files with 993 additions and 63 deletions

View File

@@ -30,6 +30,10 @@ host = "rabbit"
port = 5672
username = "rabbituser"
password = "rabbitpass"
default_exchange = "revolt"
[rabbit.queues]
acks = "internal.ack"
[api]

View File

@@ -122,12 +122,19 @@ pub struct Database {
pub redis_pubsub: Option<String>,
}
#[derive(Deserialize, Debug, Clone)]
pub struct RabbitQueues {
pub acks: String,
}
#[derive(Deserialize, Debug, Clone)]
pub struct Rabbit {
pub host: String,
pub port: u16,
pub username: String,
pub password: String,
pub default_exchange: String,
pub queues: RabbitQueues,
}
#[derive(Deserialize, Debug, Clone)]

View File

@@ -38,6 +38,7 @@ revolt-models = { workspace = true, features = ["validator"] }
revolt-presence = { workspace = true }
revolt-permissions = { workspace = true, features = ["serde", "bson"] }
revolt-parser = { workspace = true }
revolt-coalesced = { workspace = true }
# Utility
log = { workspace = true }

View File

@@ -2,7 +2,10 @@ use std::collections::HashSet;
use crate::events::rabbit::*;
use crate::User;
use amqprs::channel::{BasicPublishArguments, ExchangeDeclareArguments};
use amqprs::channel::{
BasicPublishArguments, ExchangeDeclareArguments, ExchangeType, QueueBindArguments,
QueueDeclareArguments,
};
use amqprs::connection::OpenConnectionArguments;
use amqprs::{channel::Channel, connection::Connection, error::Error as AMQPError};
use amqprs::{BasicProperties, FieldTable};
@@ -55,6 +58,43 @@ impl AMQP {
AMQP::new(connection, channel)
}
pub async fn configure_channels(&self) -> revolt_result::Result<()> {
let config = revolt_config::config().await;
self.channel
.exchange_declare(
ExchangeDeclareArguments::new(
&config.rabbit.default_exchange,
&ExchangeType::Topic.to_string(),
)
.durable(true)
.finish(),
)
.await
.expect("Failed to declare exchange");
// Configure acks channel & routing
self.channel
.queue_declare(
QueueDeclareArguments::new(&config.rabbit.queues.acks)
.durable(true)
.no_wait(true)
.finish(),
)
.await
.expect("Failed to bind queue");
self.channel
.queue_bind(QueueBindArguments::new(
&config.rabbit.queues.acks,
&config.rabbit.default_exchange,
&config.rabbit.queues.acks,
))
.await
.expect("Failed to bind channel");
Ok(())
}
pub async fn friend_request_accepted(
&self,
accepted_request_user: &User,
@@ -232,7 +272,9 @@ impl AMQP {
.await
}
pub async fn ack_message(
/// # Sends an ack to pushd to update badges on iPhones.
/// Not to be confused with the process_ack function, which handles sending all acks to crond for processing.
pub async fn ack_notification_message(
&self,
user_id: String,
channel_id: String,
@@ -316,4 +358,41 @@ impl AMQP {
)
.await
}
/// # Send an ack to crond for processing
pub async fn process_ack(
&self,
user_id: &str,
channel_id: Option<&str>,
server_id: Option<&str>,
) -> Result<(), AMQPError> {
let config = revolt_config::config().await;
let payload = AckEventPayload {
user_id: user_id.to_string(),
channel_id: channel_id.map(|value| value.to_string()),
server_id: server_id.map(|value| value.to_string()),
};
let payload = to_string(&payload).unwrap();
info!(
"Sending ack processor event on exchange {}, channel {}: {}",
config.rabbit.default_exchange, config.rabbit.queues.acks, payload
);
self.channel
.basic_publish(
BasicProperties::default()
.with_content_type("application/json")
.with_persistence(true)
//.with_headers(headers)
.finish(),
payload.into(),
BasicPublishArguments::new(
&config.rabbit.default_exchange,
&config.rabbit.queues.acks,
),
)
.await
}
}

View File

@@ -78,3 +78,11 @@ pub struct AckPayload {
pub channel_id: String,
pub message_id: String,
}
/// This is not the same as the AckPayload above, as the state for this event is stored in redis to allow for state updates while the event is queued.
#[derive(Serialize, Deserialize, Debug)]
pub struct AckEventPayload {
pub user_id: String,
pub channel_id: Option<String>,
pub server_id: Option<String>,
}

View File

@@ -1,6 +1,7 @@
#![allow(deprecated)]
use std::{borrow::Cow, collections::HashMap};
use redis_kiss::get_connection;
use revolt_config::config;
use revolt_models::v0::{self, MessageAuthor};
use revolt_permissions::OverrideField;
@@ -212,7 +213,7 @@ impl Channel {
role_permissions: HashMap::new(),
nsfw: data.nsfw.unwrap_or(false),
voice: data.voice.map(|voice| voice.into()),
slowmode: None
slowmode: None,
},
v0::LegacyServerChannelType::Voice => Channel::TextChannel {
id: id.clone(),
@@ -225,7 +226,7 @@ impl Channel {
role_permissions: HashMap::new(),
nsfw: data.nsfw.unwrap_or(false),
voice: Some(data.voice.unwrap_or_default().into()),
slowmode: None
slowmode: None,
},
};
@@ -643,7 +644,7 @@ impl Channel {
}
/// Acknowledge a message
pub async fn ack(&self, user: &str, message: &str) -> Result<()> {
pub async fn ack(&self, user: &str, message: &str, amqp: &AMQP) -> Result<()> {
EventV1::ChannelAck {
id: self.id().to_string(),
user: user.to_string(),
@@ -652,17 +653,7 @@ impl Channel {
.private(user.to_string())
.await;
#[cfg(feature = "tasks")]
crate::tasks::ack::queue_ack(
self.id().to_string(),
user.to_string(),
crate::tasks::ack::AckEvent::AckMessage {
id: message.to_string(),
},
)
.await;
Ok(())
crate::util::acker::ack_channel(user, self.id(), message, amqp).await
}
/// Remove user from a group

View File

@@ -105,7 +105,11 @@ pub async fn handle_ack_event(
if mentions_acked > 0 {
if let Err(err) = amqp
.ack_message(user.to_string(), channel.to_string(), id.to_owned())
.ack_notification_message(
user.to_string(),
channel.to_string(),
id.to_owned(),
)
.await
{
revolt_config::capture_error(&err);
@@ -192,9 +196,7 @@ pub async fn handle_ack_event(
.expect("Failed to fetch channel from db");
if let TextChannel { server, .. } = channel {
if let Err(err) =
amqp.mass_mention_message_sent(server, mass_mentions).await
{
if let Err(err) = amqp.mass_mention_message_sent(server, mass_mentions).await {
revolt_config::capture_error(&err);
}
} else {

View File

@@ -0,0 +1,77 @@
use redis_kiss::{get_connection, AsyncCommands};
use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
use revolt_result::{Result, ToRevoltError};
use crate::{events::client::EventV1, Channel, Database, Server, User, AMQP};
pub async fn ack_channel(user: &str, channel: &str, message: &str, amqp: &AMQP) -> Result<()> {
let mut redis = get_connection()
.await
.map_err(|_| create_error!(InternalError))?;
let old: Option<String> = redis
.getset(format!("acker:{user}+{channel}"), message)
.await
.to_internal_error()?;
if old.is_none() || old.unwrap() == message {
amqp.process_ack(user, Some(channel), None)
.await
.to_internal_error()?;
}
Ok(())
}
pub async fn ack_server(user: &User, server: &Server, db: &Database, amqp: &AMQP) -> Result<()> {
let mut redis = get_connection()
.await
.map_err(|_| create_error!(InternalError))?;
let channels = db.fetch_channels(&server.channels).await?;
let query = crate::util::permissions::DatabasePermissionQuery::new(db, user).server(server);
for channel in channels {
let channel_id = channel.id();
let mut q = query.clone().channel(&channel);
if calculate_channel_permissions(&mut q)
.await
.has_channel_permission(ChannelPermission::ViewChannel)
{
let channel_last_msg = match &channel {
Channel::TextChannel {
last_message_id, ..
} => last_message_id,
_ => unreachable!(),
}
.clone();
if let Some(channel_last_msg) = channel_last_msg {
let old: Option<String> = redis
.getset(
format!("acker:{}+{}", user.id, channel_id),
&channel_last_msg,
)
.await
.to_internal_error()?;
if old.is_none() || old.unwrap() == channel_last_msg {
amqp.process_ack(&user.id, Some(channel_id), Some(&server.id))
.await
.to_internal_error()?;
EventV1::ChannelAck {
id: channel_id.to_string(),
user: user.id.clone(),
message_id: channel_last_msg,
}
.private(user.id.clone())
.await;
}
}
}
}
Ok(())
}

View File

@@ -1,3 +1,4 @@
pub mod acker;
pub mod bridge;
pub mod bulk_permissions;
mod funcs;

View File

@@ -16,8 +16,22 @@ log = { workspace = true }
# Async
tokio = { workspace = true }
# Redis
redis-kiss = { workspace = true }
# RabbitMQ
lapin = { workspace = true }
futures-lite = { workspace = true }
# Processing
serde_json = { workspace = true }
revolt_optional_struct = { workspace = true }
serde = { workspace = true }
iso8601-timestamp = { workspace = true, features = ["serde", "bson"] }
# Core
revolt-database = { workspace = true }
revolt-result = { workspace = true }
revolt-config = { workspace = true }
revolt-files = { workspace = true }
revolt-permissions = { workspace = true }

View File

@@ -1,7 +1,7 @@
use revolt_config::configure;
use revolt_database::DatabaseInfo;
use revolt_result::Result;
use tasks::{file_deletion, prune_dangling_files, prune_members};
use tasks::{acks, file_deletion, prune_dangling_files, prune_members};
use tokio::try_join;
pub mod tasks;
@@ -14,7 +14,8 @@ async fn main() -> Result<()> {
try_join!(
file_deletion::task(db.clone()),
prune_dangling_files::task(db.clone()),
prune_members::task(db.clone())
prune_members::task(db.clone()),
acks::task(db.clone())
)
.map(|_| ())
}

View File

@@ -0,0 +1,129 @@
use futures_lite::stream::StreamExt;
use lapin::{
options::*,
types::FieldTable,
uri::{AMQPAuthority, AMQPQueryString, AMQPUri, AMQPUserInfo},
ConnectionBuilder, ConnectionProperties,
};
use log::info;
use redis_kiss::{get_connection, AsyncCommands, Conn as RedisConnection};
use revolt_config::config;
use revolt_database::{events::rabbit::AckEventPayload, Database};
use revolt_result::{Result, ToRevoltError};
use serde_json;
pub async fn task(db: Database) -> Result<()> {
let config = config().await;
let mut redis = get_connection()
.await
.expect("Failed to get redis connection");
let uri = AMQPUri {
scheme: lapin::uri::AMQPScheme::AMQP,
authority: AMQPAuthority {
userinfo: AMQPUserInfo {
username: config.rabbit.username,
password: config.rabbit.password,
},
host: config.rabbit.host,
port: config.rabbit.port,
},
vhost: "/".to_string(),
query: AMQPQueryString::default(),
};
let connection = ConnectionBuilder::new()
.expect("Builder")
.with_uri(uri)
.with_properties(ConnectionProperties::default())
.connect()
.await
.expect("Failed to connect to rabbitmq");
let reader_channel = connection
.create_channel()
.await
.expect("Failed to create channel");
let mut consumer = reader_channel
.basic_consume(
config.rabbit.queues.acks.into(),
"crond-ack-consumer".into(),
BasicConsumeOptions::default(),
FieldTable::default(),
)
.await
.expect("Failed to create consumer");
while let Some(delivery) = consumer.next().await {
if let Ok(delivery) = delivery {
let payload: std::result::Result<AckEventPayload, _> =
serde_json::from_slice(&delivery.data);
if let Ok(payload) = payload {
info!("{:?}", payload);
if let Err(e) = process_channel_ack(
&db,
payload.user_id,
payload.channel_id.unwrap(),
&mut redis,
)
.await
{
revolt_config::capture_error(&e);
_ = delivery.reject(BasicRejectOptions { requeue: false }).await;
} else {
_ = delivery.ack(BasicAckOptions { multiple: false }).await;
}
} else {
revolt_config::capture_message(
format!("Failed to decode ack data: {:?}", delivery.data).as_str(),
revolt_config::Level::Error,
);
}
}
}
Ok(())
}
#[allow(clippy::disallowed_methods)]
async fn process_channel_ack(
db: &Database,
user: String,
channel: String,
redis: &mut RedisConnection,
) -> Result<()> {
let message_id: Option<String> = redis
.get_del(format!("acker:{user}+{channel}"))
.await
.to_internal_error()?;
if let Some(message_id) = message_id {
// This will be uncommented eventually, but we need to sort out the transition to lapin first. For now we'll simply disable the badge update logic.
// We also drop a db request as a bonus.
//let unread = db.fetch_unread(&user, &channel).await?;
let _updated = db.acknowledge_message(&channel, &user, &message_id).await?;
info!("Set new state for ack: {}:{}:{}", channel, user, message_id);
// if let (Some(before), Some(after)) = (unread, updated) {
// let before_mentions = before.mentions.unwrap_or_default().len();
// let after_mentions = after.mentions.unwrap_or_default().len();
// let mentions_acked = before_mentions - after_mentions;
// if mentions_acked > 0 {
// if let Err(err) = amqp
// .ack_message(user.to_string(), channel.to_string(), payload.message_id)
// .await
// {
// revolt_config::capture_error(&err);
// }
// };
// }
Ok(())
} else {
Err(message_id.to_internal_error().expect_err("no err"))
}
}

View File

@@ -1,3 +1,4 @@
pub mod acks;
pub mod file_deletion;
pub mod prune_dangling_files;
pub mod prune_members;

View File

@@ -24,8 +24,8 @@ use amqprs::{
};
use async_std::channel::unbounded;
use authifier::AuthifierEvent;
use rocket::data::ToByteUnit;
use revolt_database::voice::VoiceClient;
use rocket::data::ToByteUnit;
pub async fn web() -> Rocket<Build> {
// Get settings
@@ -93,22 +93,6 @@ pub async fn web() -> Rocket<Build> {
)
.into();
let swagger_0_8 = revolt_rocket_okapi::swagger_ui::make_swagger_ui(
&revolt_rocket_okapi::swagger_ui::SwaggerUIConfig {
url: "/0.8/openapi.json".to_owned(),
..Default::default()
},
)
.into();
let swagger_0_8 = revolt_rocket_okapi::swagger_ui::make_swagger_ui(
&revolt_rocket_okapi::swagger_ui::SwaggerUIConfig {
url: "/0.8/openapi.json".to_owned(),
..Default::default()
},
)
.into();
// Voice handler
let voice_client = VoiceClient::new(config.api.livekit.nodes.clone());
// Configure Rabbit
@@ -136,6 +120,9 @@ pub async fn web() -> Rocket<Build> {
.expect("Failed to declare exchange");
let amqp = AMQP::new(connection, channel);
amqp.configure_channels()
.await
.expect("Failed to configure channels");
// Launch background task workers
revolt_database::tasks::start_workers(db.clone(), amqp.clone());
@@ -153,7 +140,6 @@ pub async fn web() -> Rocket<Build> {
.mount("/", rocket_cors::catch_all_options_routes())
.mount("/", ratelimiter::routes())
.mount("/swagger/", swagger)
.mount("/0.8/swagger/", swagger_0_8)
.manage(authifier)
.manage(db)
.manage(amqp)

View File

@@ -1,6 +1,6 @@
use revolt_database::{
util::{permissions::DatabasePermissionQuery, reference::Reference},
Database, User,
Database, User, AMQP,
};
use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
use revolt_result::{create_error, Result};
@@ -14,6 +14,7 @@ use rocket_empty::EmptyResponse;
#[put("/<target>/ack/<message>")]
pub async fn ack(
db: &State<Database>,
amqp: &State<AMQP>,
user: User,
target: Reference<'_>,
message: Reference<'_>,
@@ -29,7 +30,7 @@ pub async fn ack(
.throw_if_lacking_channel_permission(ChannelPermission::ViewChannel)?;
channel
.ack(&user.id, message.id)
.ack(&user.id, message.id, amqp)
.await
.map(|_| EmptyResponse)
}

View File

@@ -1,6 +1,6 @@
use revolt_database::{
util::{permissions::DatabasePermissionQuery, reference::Reference},
Database, User,
util::{acker, permissions::DatabasePermissionQuery, reference::Reference},
Database, User, AMQP,
};
use revolt_permissions::PermissionQuery;
use revolt_result::{create_error, Result};
@@ -12,7 +12,12 @@ use rocket_empty::EmptyResponse;
/// Mark all channels in a server as read.
#[openapi(tag = "Server Information")]
#[put("/<target>/ack")]
pub async fn ack(db: &State<Database>, user: User, target: Reference<'_>) -> Result<EmptyResponse> {
pub async fn ack(
db: &State<Database>,
amqp: &State<AMQP>,
user: User,
target: Reference<'_>,
) -> Result<EmptyResponse> {
if user.bot.is_some() {
return Err(create_error!(IsBot));
}
@@ -23,7 +28,6 @@ pub async fn ack(db: &State<Database>, user: User, target: Reference<'_>) -> Res
return Err(create_error!(NotFound));
}
db.acknowledge_channels(&user.id, &server.channels)
.await
.map(|_| EmptyResponse)
acker::ack_server(&user, &server, db, amqp).await?;
Ok(EmptyResponse)
}