main
use rocketmq_client_rust::{ClientResult, DefaultMQPushConsumer, MQPushConsumer};
use rocketmq_model::common::message::message_ext::MessageExt;
use std::collections::{HashMap, VecDeque};
use std::sync::mpsc::SyncSender;
use std::time::Duration;
use tokio::sync::mpsc;
use crate::gateway::plan;
use config::event::consumer::RocketMqConsumerConfig;
use super::listener::{Forwarder, PendingEvent};
use super::result::ConsumeResult;
pub const QUEUE_CAPACITY: usize = 1024;
/// 有序消费线程等待单条事件完成的上限；必须远小于 MQ 侧 `consume_timeout`，
/// 否则 MQ 先判超时重投，同一条消息会被两个消费线程并发处理。
const MAX_LISTENER_WAIT: Duration = Duration::from_secs(60);
pub struct EventsPipe {
consumer: DefaultMQPushConsumer,
events_rx: mpsc::Receiver<PendingEvent>,
pending: HashMap<String, VecDeque<SyncSender<bool>>>,
prefixes: Vec<String>,
result_tx: mpsc::Sender<ConsumeResult>,
result_rx: mpsc::Receiver<ConsumeResult>,
}
impl EventsPipe {
pub fn new(
config: &RocketMqConsumerConfig,
runtime: &queue::QueueRuntime,
) -> Result<Self, String> {
config.validate()?;
let (events_tx, events_rx) = mpsc::channel(QUEUE_CAPACITY);
let (result_tx, result_rx) = mpsc::channel(QUEUE_CAPACITY);
let mut consumer = DefaultMQPushConsumer::builder(runtime.client())
.consumer_group(config.consumer_group.clone())
.name_server_addr(config.name_servers().join(";"))
.consume_timeout(config.consume_timeout_minutes())
.max_reconsume_times(config.max_reconsume_times as i32)
.pull_batch_size(config.pull_batch_size)
.pull_batch_size_in_bytes(config.pull_batch_size_in_bytes)
.pull_interval(config.pull_interval_ms)
.build();
consumer.register_message_listener_orderly(Forwarder::new(
events_tx,
Duration::from_millis(config.consume_timeout_ms).min(MAX_LISTENER_WAIT),
));
Ok(Self {
consumer,
events_rx,
pending: HashMap::new(),
prefixes: config.topic_prefixes.clone(),
result_tx,
result_rx,
})
}
pub async fn next_event(&mut self) -> Option<MessageExt> {
let pending = self.events_rx.recv().await?;
self.pending
.entry(pending.message.msg_id().to_string())
.or_default()
.push_back(pending.completion);
Some(pending.message)
}
pub fn result_sender(&self) -> mpsc::Sender<ConsumeResult> {
self.result_tx.clone()
}
pub async fn next_result(&mut self) -> Option<ConsumeResult> {
self.result_rx.recv().await
}
pub async fn apply_result(&mut self, result: ConsumeResult) -> ClientResult<()> {
let (message, success) = match result {
ConsumeResult::Ack { message } => (message, true),
ConsumeResult::Retry { message, .. } => (message, false),
};
let key = message.msg_id().to_string();
if let Some(queue) = self.pending.get_mut(&key) {
while let Some(completion) = queue.pop_front() {
let retry = completion.clone();
match completion.try_send(success) {
Ok(()) => break,
Err(std::sync::mpsc::TrySendError::Disconnected(_)) => continue,
Err(std::sync::mpsc::TrySendError::Full(_)) => {
queue.push_front(retry);
break;
}
}
}
if queue.is_empty() {
self.pending.remove(&key);
}
}
Ok(())
}
pub async fn subscribe(&mut self, topic: &str, expression: &str) -> ClientResult<()> {
self.consumer.subscribe(topic, expression).await
}
pub async fn subscribe_all(&mut self) -> ClientResult<()> {
for (topic, expression) in plan::plan(&self.prefixes) {
self.consumer.subscribe(topic, expression).await?;
}
Ok(())
}
pub async fn start(&mut self) -> ClientResult<()> {
self.consumer.start().await
}
pub async fn shutdown(&mut self) {
self.events_rx.close();
while self.events_rx.try_recv().is_ok() {}
self.pending.clear();
self.consumer.shutdown().await;
}
}