main
//! 消息路径的扇出有界化：固定 worker 池替代裸 `tokio::spawn`。
//!
//! 三条 lane：`dm` 承接 DM 推送，`email` 承接 @mention 邮件（单 job 有界），
//! `bulk` 承接 @everyone（按 `EVERYONE_CHUNK` 切块）。队列满时转后台等待
//! 投递，只有 channel 关闭时才丢弃，绝不把背压施加到消息写入路径。
use std::sync::{Arc, OnceLock};
use tokio::sync::{Mutex, Semaphore, mpsc};
use tokio_util::sync::CancellationToken;
use tracing::Instrument as _;
use uuid::Uuid;
use crate::services::messages::DirectMessageKind;
use crate::services::service::ChatService;
use crate::share::notify::{
EveryoneScope, everyone_candidates, fanout_everyone_chunk, mention_email_allowed,
};
const DM_WORKERS: usize = 2;
const DM_CAPACITY: usize = 256;
const EMAIL_WORKERS: usize = 2;
const EMAIL_CAPACITY: usize = 64;
const BULK_WORKERS: usize = 2;
const BULK_CAPACITY: usize = 64;
/// @everyone 单块人数：一块的数据库往返次数有界，块与块之间让出 worker。
const EVERYONE_CHUNK: usize = 100;
/// 满队列时最多允许多少个后台等待任务：这是 deferral 的硬上限，
/// 保证持续过载时内存有界（超出即丢弃并打点，而不是无界堆积）。
const MAX_DEFERRED: usize = 256;
pub(crate) enum Job {
MentionEmails {
room_id: Uuid,
space_id: Option<Uuid>,
room_name: String,
body: String,
recipients: Vec<i64>,
request_id: Option<String>,
trace: Option<String>,
},
DirectMessagePush {
kind: DirectMessageKind,
room_id: Uuid,
space_id: Option<Uuid>,
message_id: Uuid,
actor: i64,
body: String,
recipients: Vec<i64>,
request_id: Option<String>,
trace: Option<String>,
},
EveryoneStart {
scope: EveryoneScope,
},
EveryoneChunk {
scope: Arc<EveryoneScope>,
recipients: Vec<i64>,
},
}
impl Job {
fn kind(&self) -> &'static str {
match self {
Self::MentionEmails { .. } => "mention_emails",
Self::DirectMessagePush { .. } => "direct_message_push",
Self::EveryoneStart { .. } => "everyone_start",
Self::EveryoneChunk { .. } => "everyone_chunk",
}
}
fn request_id(&self) -> Option<&str> {
match self {
Self::MentionEmails { request_id, .. } | Self::DirectMessagePush { request_id, .. } => {
request_id.as_deref()
}
Self::EveryoneStart { scope } => scope.request_id.as_deref(),
Self::EveryoneChunk { scope, .. } => scope.request_id.as_deref(),
}
}
fn trace(&self) -> Option<&str> {
match self {
Self::MentionEmails { trace, .. } | Self::DirectMessagePush { trace, .. } => {
trace.as_deref()
}
Self::EveryoneStart { scope } => scope.trace.as_deref(),
Self::EveryoneChunk { scope, .. } => scope.trace.as_deref(),
}
}
}
#[derive(Clone, Copy)]
enum Lane {
Dm,
Email,
Bulk,
}
impl Lane {
fn of(job: &Job) -> Self {
match job {
Job::DirectMessagePush { .. } => Self::Dm,
Job::MentionEmails { .. } => Self::Email,
Job::EveryoneStart { .. } | Job::EveryoneChunk { .. } => Self::Bulk,
}
}
fn name(self) -> &'static str {
match self {
Self::Dm => "dm",
Self::Email => "email",
Self::Bulk => "bulk",
}
}
}
struct Fanout {
dm: mpsc::Sender<Job>,
email: mpsc::Sender<Job>,
bulk: mpsc::Sender<Job>,
}
static FANOUT: OnceLock<Fanout> = OnceLock::new();
/// 满队列转后台等待时占用的名额；名额耗尽说明已经过载，直接丢弃并打点。
static DEFERRED: OnceLock<Arc<Semaphore>> = OnceLock::new();
fn deferred_slots() -> &'static Arc<Semaphore> {
DEFERRED.get_or_init(|| Arc::new(Semaphore::new(MAX_DEFERRED)))
}
/// 启动固定 worker 池；必须在 tokio 运行时内调用一次。
pub fn start(service: ChatService, cancel: CancellationToken) {
let (dm_tx, dm_rx) = mpsc::channel(DM_CAPACITY);
let (email_tx, email_rx) = mpsc::channel(EMAIL_CAPACITY);
let (bulk_tx, bulk_rx) = mpsc::channel(BULK_CAPACITY);
let dm_rx = Arc::new(Mutex::new(dm_rx));
let email_rx = Arc::new(Mutex::new(email_rx));
let bulk_rx = Arc::new(Mutex::new(bulk_rx));
for _ in 0..DM_WORKERS {
tokio::spawn(worker_loop(
service.clone(),
dm_rx.clone(),
bulk_tx.clone(),
cancel.clone(),
));
}
for _ in 0..EMAIL_WORKERS {
tokio::spawn(worker_loop(
service.clone(),
email_rx.clone(),
bulk_tx.clone(),
cancel.clone(),
));
}
for _ in 0..BULK_WORKERS {
tokio::spawn(worker_loop(
service.clone(),
bulk_rx.clone(),
bulk_tx.clone(),
cancel.clone(),
));
}
if FANOUT
.set(Fanout {
dm: dm_tx,
email: email_tx,
bulk: bulk_tx,
})
.is_err()
{
tracing::warn!("chat fanout already started");
} else {
tracing::info!(
dm = DM_WORKERS,
email = EMAIL_WORKERS,
bulk = BULK_WORKERS,
"chat fanout pool started"
);
}
}
/// 投递一个扇出任务；队列满时转后台等待投递，不阻塞调用方，只有 channel
/// 关闭（进程退出）或未启动时才会丢弃。
pub(crate) fn submit(job: Job) {
let lane = Lane::of(&job);
let Some(fanout) = FANOUT.get() else {
tracing::warn!(
kind = job.kind(),
lane = lane.name(),
"chat fanout not started; dropping job"
);
oltp::metrics::chat_fanout(lane.name(), "unstarted");
return;
};
let tx = match lane {
Lane::Dm => &fanout.dm,
Lane::Email => &fanout.email,
Lane::Bulk => &fanout.bulk,
};
match tx.try_send(job) {
Ok(()) => oltp::metrics::chat_fanout(lane.name(), "queued"),
Err(mpsc::error::TrySendError::Full(job)) => {
let Ok(permit) = deferred_slots().clone().try_acquire_owned() else {
tracing::warn!(
kind = job.kind(),
lane = lane.name(),
"chat fanout deferral saturated; dropping job"
);
oltp::metrics::chat_fanout(lane.name(), "dropped");
return;
};
tracing::warn!(
kind = job.kind(),
lane = lane.name(),
"chat fanout queue full; deferring job"
);
oltp::metrics::chat_fanout(lane.name(), "deferred");
let tx = tx.clone();
tokio::spawn(async move {
let _permit = permit;
if tx.send(job).await.is_err() {
tracing::warn!(lane = lane.name(), "chat fanout closed while deferring job");
oltp::metrics::chat_fanout(lane.name(), "closed");
}
});
}
Err(mpsc::error::TrySendError::Closed(job)) => {
tracing::warn!(
kind = job.kind(),
lane = lane.name(),
"chat fanout queue closed; dropping job"
);
oltp::metrics::chat_fanout(lane.name(), "closed");
}
}
}
async fn worker_loop(
service: ChatService,
rx: Arc<Mutex<mpsc::Receiver<Job>>>,
bulk: mpsc::Sender<Job>,
cancel: CancellationToken,
) {
loop {
let job = {
let mut guard = rx.lock().await;
tokio::select! {
_ = cancel.cancelled() => return,
job = guard.recv() => job,
}
};
match job {
Some(job) => run(&service, job, &bulk).await,
None => return,
}
}
}
async fn run(service: &ChatService, job: Job, bulk: &mpsc::Sender<Job>) {
let lane = Lane::of(&job).name();
let request_id = job.request_id().map(str::to_owned);
let trace = job.trace().map(str::to_owned);
let span = tracing::info_span!(
"chat.fanout",
kind = job.kind(),
request_id = tracing::field::Empty
);
if let Some(request_id) = request_id.as_deref() {
span.record("request_id", request_id);
}
oltp::propagation::set_parent_from_traceparent(&span, trace.as_deref());
api::middleware::with_request_id(
request_id.as_deref(),
run_job(service, job, bulk).instrument(span),
)
.await;
oltp::metrics::chat_fanout(lane, "done");
}
async fn run_job(service: &ChatService, job: Job, bulk: &mpsc::Sender<Job>) {
match job {
Job::MentionEmails {
room_id,
space_id,
room_name,
body,
recipients,
..
} => mention_emails(service, room_id, space_id, room_name, body, recipients).await,
Job::DirectMessagePush {
kind,
room_id,
space_id,
message_id,
actor,
body,
recipients,
..
} => {
crate::services::messages::broadcast_direct_message_push(
service,
kind,
room_id,
space_id,
message_id,
actor,
body,
&recipients,
)
.await;
}
Job::EveryoneStart { scope } => everyone_start(service, scope, bulk).await,
Job::EveryoneChunk { scope, recipients } => {
fanout_everyone_chunk(service, &scope, &recipients).await;
}
}
}
async fn everyone_start(service: &ChatService, scope: EveryoneScope, bulk: &mpsc::Sender<Job>) {
let candidates = everyone_candidates(service, &scope).await;
let scope = Arc::new(scope);
for chunk in candidates.chunks(EVERYONE_CHUNK) {
let job = Job::EveryoneChunk {
scope: scope.clone(),
recipients: chunk.to_vec(),
};
match bulk.try_send(job) {
Ok(()) => {}
// 队列已满时自己消化这一块：不丢通知，也不会互相等待而死锁。
Err(mpsc::error::TrySendError::Full(job)) => {
if let Job::EveryoneChunk { scope, recipients } = job {
fanout_everyone_chunk(service, &scope, &recipients).await;
}
}
Err(mpsc::error::TrySendError::Closed(_)) => break,
}
}
}
async fn mention_emails(
service: &ChatService,
room_id: Uuid,
space_id: Option<Uuid>,
room_name: String,
body: String,
recipients: Vec<i64>,
) {
let Ok(mut emails) = service.audience.primary_emails(&recipients).await else {
return;
};
let fallback = service.audience.default_locale;
let locales = i18n::resolve_user_locales(service.database.read(), &recipients, fallback).await;
let base_url = service.settings.web.web_base_url.clone();
for uid in recipients {
if !mention_email_allowed(service, room_id, space_id, uid)
.await
.unwrap_or(false)
{
continue;
}
let Some(addr) = emails.remove(&uid) else {
continue;
};
let locale = locales.get(&uid).copied().unwrap_or(fallback);
let subject = i18n::catalog::mention_email_subject(locale, &room_name);
let link = format!("{base_url}/chat/rooms/{room_id}");
let text = format!(
"{subject}\n\n{body}\n\n{}: {link}",
i18n::catalog::view_channel(locale)
);
let html = format!(
r#"<p>{subject}</p><p>{body}</p><p><a href="{link}">{}</a></p>"#,
i18n::catalog::view_channel(locale)
);
if let Err(error) = service
.audience
.send_email(&addr, &subject, &text, &html)
.await
{
tracing::warn!(%error, "failed to send mention email");
}
}
}