b7cd2b1
//! 邮件发送的本地有界队列。
//!
//! 邮件请求不再经过 RocketMQ：入队后停在进程内 channel 里，由内嵌 worker 直接走 SMTP。
//! 队列满时按反压阻塞，等满 `enqueue_timeout` 仍无空间才把错误抛回调用方。
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tokio::sync::{Mutex, mpsc};
use tokio_util::sync::CancellationToken;
use tracing::{Instrument, debug, error, info, warn};
use uuid::Uuid;
use crate::error::EmailError;
use crate::events::{RenderAndSendEmailRequest, SendBatchEmailRequest, SendEmailRequest};
use config::identity::email::EmailPoolConfig;
/// 队列里的一条工作项。
#[derive(Debug, Clone)]
pub enum Job {
Send(SendEmailRequest),
Batch(SendBatchEmailRequest),
Render(RenderAndSendEmailRequest),
}
impl Job {
pub fn kind(&self) -> &'static str {
match self {
Self::Send(_) => "send",
Self::Batch(_) => "batch",
Self::Render(_) => "render",
}
}
pub fn request_id(&self) -> Uuid {
match self {
Self::Send(req) => req.request_id,
Self::Batch(req) => req.request_id,
Self::Render(req) => req.request_id,
}
}
}
/// 实际执行投递的一方（由 `EmailService` 提供）。
#[async_trait::async_trait]
pub trait JobHandler: Send + Sync {
async fn handle(&self, job: &Job) -> Result<(), EmailError>;
/// 重试用尽后的收尾：打 ERROR 日志并广播失败事件。
async fn give_up(&self, job: &Job, error: &EmailError);
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueueSettings {
pub workers: usize,
pub capacity: usize,
pub enqueue_timeout: Duration,
pub retry_delay: Duration,
pub max_attempts: u32,
}
impl QueueSettings {
pub fn from_config(config: &EmailPoolConfig) -> Self {
Self {
workers: config.workers,
capacity: config.capacity,
enqueue_timeout: config.enqueue_timeout(),
retry_delay: config.retry_delay(),
max_attempts: config.max_attempts,
}
}
}
impl Default for QueueSettings {
fn default() -> Self {
Self::from_config(&EmailPoolConfig::default())
}
}
#[derive(Clone)]
pub struct EmailQueue {
tx: mpsc::Sender<Job>,
settings: Arc<QueueSettings>,
depth: Arc<AtomicUsize>,
workers: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
stop: CancellationToken,
}
impl EmailQueue {
/// 启动队列与 worker；返回的队列可以随意克隆，底层共享同一条 channel。
pub fn start(settings: QueueSettings, handler: Arc<dyn JobHandler>) -> Self {
let (tx, rx) = mpsc::channel(settings.capacity);
let rx = Arc::new(Mutex::new(rx));
let depth = Arc::new(AtomicUsize::new(0));
let stop = CancellationToken::new();
let settings = Arc::new(settings);
let mut workers = Vec::with_capacity(settings.workers);
for index in 0..settings.workers {
workers.push(tokio::spawn(run_worker(
index,
rx.clone(),
handler.clone(),
settings.clone(),
depth.clone(),
stop.clone(),
)));
}
info!(
workers = settings.workers,
capacity = settings.capacity,
"email queue started"
);
Self {
tx,
settings,
depth,
workers: Arc::new(Mutex::new(workers)),
stop,
}
}
pub fn settings(&self) -> &QueueSettings {
&self.settings
}
/// 当前积压条数。
pub fn depth(&self) -> usize {
self.depth.load(Ordering::Relaxed)
}
pub fn capacity(&self) -> usize {
self.settings.capacity
}
/// 入队：队列满时阻塞等待，超时返回 [`EmailError::QueueFull`]。
/// 深度计数在发送前预增，发送失败回滚，避免 worker 抢先扣减导致下溢。
#[tracing::instrument(skip_all, fields(request_id = %_request_id.unwrap_or_default()))]
pub async fn enqueue(&self, job: Job, _request_id: Option<&str>) -> Result<usize, EmailError> {
if self.stop.is_cancelled() {
return Err(EmailError::QueueClosed);
}
let kind = job.kind();
let depth = self.depth.fetch_add(1, Ordering::Relaxed) + 1;
macro_rules! rollback {
() => {{
let depth = self.depth.fetch_sub(1, Ordering::Relaxed) - 1;
crate::metrics::queue_depth(depth as i64);
}};
}
match self.tx.try_send(job) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(job)) => {
debug!(
kind,
depth = self.depth(),
capacity = self.settings.capacity,
"email queue is full; waiting for capacity"
);
match tokio::time::timeout(self.settings.enqueue_timeout, self.tx.send(job)).await {
Ok(Ok(())) => debug!(kind, "email queue accepted after backpressure"),
Ok(Err(_)) => {
rollback!();
return Err(EmailError::QueueClosed);
}
Err(_) => {
rollback!();
let depth = self.depth();
warn!(
kind,
depth,
capacity = self.settings.capacity,
timeout_ms = self.settings.enqueue_timeout.as_millis() as u64,
"email queue is full; rejecting request"
);
return Err(EmailError::QueueFull {
depth,
capacity: self.settings.capacity,
});
}
}
}
Err(mpsc::error::TrySendError::Closed(_)) => {
rollback!();
return Err(EmailError::QueueClosed);
}
}
crate::metrics::queue_depth(depth as i64);
Ok(depth)
}
/// 优雅退出：先等积压排空（最多 `grace`），再让 worker 停下。
pub async fn shutdown(&self, grace: Duration) -> Result<(), EmailError> {
let deadline = tokio::time::Instant::now() + grace;
while self.depth() > 0 && tokio::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(20)).await;
}
let pending = self.depth();
self.stop.cancel();
let handles: Vec<_> = {
let mut guard = self.workers.lock().await;
guard.drain(..).collect()
};
for mut handle in handles {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
handle.abort();
continue;
}
if tokio::time::timeout(remaining, &mut handle).await.is_err() {
// 超时的 worker 及其后所有 worker 一律中止，保证 shutdown 必然返回。
}
handle.abort();
}
crate::metrics::queue_depth(0);
if pending > 0 {
warn!(pending, "email queue shut down before draining every job");
} else {
info!("email queue shut down cleanly");
}
Ok(())
}
}
async fn run_worker(
index: usize,
rx: Arc<Mutex<mpsc::Receiver<Job>>>,
handler: Arc<dyn JobHandler>,
settings: Arc<QueueSettings>,
depth: Arc<AtomicUsize>,
stop: CancellationToken,
) {
loop {
let job = tokio::select! {
biased;
_ = stop.cancelled() => break,
job = async { rx.lock().await.recv().await } => job,
};
let Some(job) = job else {
break;
};
depth.fetch_sub(1, Ordering::Relaxed);
let request_id = job.request_id();
let kind = job.kind();
let span = tracing::info_span!("email.job", worker = index, kind, request_id = %request_id);
let mut attempt = 1;
loop {
match handler.handle(&job).instrument(span.clone()).await {
Ok(()) => break,
Err(error) => {
if attempt >= settings.max_attempts {
error!(
worker = index,
%request_id,
kind,
attempt,
%error,
"email job failed permanently"
);
handler.give_up(&job, &error).await;
break;
}
warn!(
worker = index,
%request_id,
kind,
attempt,
%error,
"email job failed; retrying"
);
tokio::time::sleep(settings.retry_delay).await;
attempt += 1;
}
}
}
}
debug!(worker = index, "email worker stopped");
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::sync::Notify;
use super::*;
/// 可阻塞、可失败的假 handler。
struct Probe {
calls: AtomicUsize,
give_ups: AtomicUsize,
gate: Notify,
open: AtomicBool,
fails: bool,
}
impl Probe {
fn new(fails: bool) -> Arc<Self> {
Arc::new(Self {
calls: AtomicUsize::new(0),
give_ups: AtomicUsize::new(0),
gate: Notify::new(),
open: AtomicBool::new(false),
fails,
})
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
fn give_ups(&self) -> usize {
self.give_ups.load(Ordering::SeqCst)
}
fn release(&self) {
self.open.store(true, Ordering::SeqCst);
self.gate.notify_waiters();
}
}
#[async_trait::async_trait]
impl JobHandler for Probe {
async fn handle(&self, _job: &Job) -> Result<(), EmailError> {
self.calls.fetch_add(1, Ordering::SeqCst);
while !self.open.load(Ordering::SeqCst) {
self.gate.notified().await;
}
if self.fails {
return Err(EmailError::Other("smtp refused".into()));
}
Ok(())
}
async fn give_up(&self, _job: &Job, _error: &EmailError) {
self.give_ups.fetch_add(1, Ordering::SeqCst);
}
}
fn job() -> Job {
Job::Send(SendEmailRequest {
request_id: Uuid::now_v7(),
to: vec![crate::events::EmailAddress::new("a@b.c")],
subject: "s".into(),
..Default::default()
})
}
fn settings(workers: usize, capacity: usize, timeout_ms: u64) -> QueueSettings {
QueueSettings {
workers,
capacity,
enqueue_timeout: Duration::from_millis(timeout_ms),
retry_delay: Duration::from_millis(5),
max_attempts: 3,
}
}
#[tokio::test]
async fn enqueue_fills_the_queue_and_then_rejects_after_the_timeout() {
let probe = Probe::new(false);
let queue = EmailQueue::start(settings(1, 1, 30), probe.clone());
let mut rejected = None;
for _ in 0..8 {
if let Err(error) = queue.enqueue(job(), None).await {
rejected = Some(error);
break;
}
}
let error = rejected.expect("queue must apply backpressure");
assert!(
matches!(error, EmailError::QueueFull { capacity: 1, .. }),
"unexpected error: {error}"
);
assert!(queue.depth() >= 1);
probe.release();
queue.shutdown(Duration::from_secs(2)).await.unwrap();
assert_eq!(queue.depth(), 0);
}
#[tokio::test]
async fn shutdown_drains_pending_jobs() {
let probe = Probe::new(false);
let queue = EmailQueue::start(settings(2, 8, 1_000), probe.clone());
for _ in 0..4 {
queue.enqueue(job(), None).await.unwrap();
}
probe.release();
queue.shutdown(Duration::from_secs(3)).await.unwrap();
assert_eq!(probe.calls(), 4);
assert_eq!(queue.depth(), 0);
}
#[tokio::test]
async fn failing_jobs_are_retried_then_reported_once() {
let probe = Probe::new(true);
let queue = EmailQueue::start(settings(1, 4, 1_000), probe.clone());
probe.release();
queue.enqueue(job(), None).await.unwrap();
for _ in 0..200 {
if probe.give_ups() > 0 {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
assert_eq!(probe.calls(), 3);
assert_eq!(probe.give_ups(), 1);
queue.shutdown(Duration::from_secs(1)).await.unwrap();
}
}