e40c667
use std::collections::HashMap;
use std::time::Duration;
use config::rocketmq::RocketMqConfig;
use rocketmq_client_rust::DefaultMQProducer;
use rocketmq_model::common::message::message_single::Message;
use crate::error::{QueueError, QueueResult};
use crate::runtime::QueueRuntime;
pub const DEFAULT_MAX_MESSAGE_BYTES: usize = 4 * 1024 * 1024;
/// 事件信封的字节预算：为消息属性与协议头预留 4 KiB。
pub const MAX_ENVELOPE_BYTES: usize = DEFAULT_MAX_MESSAGE_BYTES - 4096;
/// 建 topic 时向 broker 提交的模板 topic，与 `TopicValidator::AUTO_CREATE_TOPIC_KEY_TOPIC` 一致。
const CREATE_TOPIC_KEY: &str = "TBW102";
const DEFAULT_TOPIC_QUEUE_NUMS: i32 = 4;
const ENSURE_ROUTE_TIMEOUT: Duration = Duration::from_secs(3);
const ENSURE_CREATE_TIMEOUT: Duration = Duration::from_secs(5);
/// `ensure_topics` 的结果：补齐的 topic、已存在的 topic 与失败明细。
#[derive(Debug, Default, Clone)]
pub struct TopicEnsureReport {
pub created: Vec<String>,
pub existing: Vec<String>,
pub failed: Vec<(String, String)>,
}
#[derive(Clone)]
pub struct RocketMqConn {
producer: DefaultMQProducer,
config: RocketMqConfig,
}
impl std::fmt::Debug for RocketMqConn {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("RocketMqConn")
.field("name_server", &self.config.name_server)
.field("producer_group", &self.config.producer_group)
.field("topic_prefix", &self.config.topic_prefix)
.finish()
}
}
impl RocketMqConn {
pub async fn init(config: &RocketMqConfig, runtime: &QueueRuntime) -> QueueResult<Self> {
config.validate().map_err(QueueError::config)?;
if config.max_message_bytes < MAX_ENVELOPE_BYTES {
return Err(QueueError::config(format!(
"rocketmq.max_message_bytes ({}) is smaller than the event envelope budget ({MAX_ENVELOPE_BYTES})",
config.max_message_bytes
)));
}
probe_nameservers(config).await?;
let mut producer = DefaultMQProducer::builder(runtime.client())
.producer_group(config.producer_group.clone())
.name_server_addr(config.name_server.replace(',', ";"))
.build();
producer.start().await.map_err(|error| {
QueueError::connect(format!("failed to start the rocketmq producer: {error}"))
})?;
Ok(Self {
producer,
config: config.clone(),
})
}
pub fn config(&self) -> &RocketMqConfig {
&self.config
}
pub fn max_payload(&self) -> usize {
self.config.max_message_bytes
}
/// 全部 topic 的路由都能查到才算健康：任一 topic 缺失说明事件通道不完整。
pub async fn health(&self, topics: &[String]) -> QueueResult<()> {
let mut producer = self.producer.clone();
let mut failures = Vec::new();
for topic in topics {
if let Err(error) = producer.fetch_publish_message_queues(topic).await {
failures.push((topic, error));
}
}
match failures.first() {
None => Ok(()),
Some((topic, error)) => Err(QueueError::connect(format!(
"rocketmq route lookup failed for {}/{} topics (first: `{topic}`: {error})",
failures.len(),
topics.len()
))),
}
}
/// 幂等补齐事件 topic：路由存在即跳过，缺失才向 broker 提交创建请求。
pub async fn ensure_topics(&self, topics: &[String]) -> TopicEnsureReport {
let mut producer = self.producer.clone();
let mut report = TopicEnsureReport::default();
for topic in topics {
let route = tokio::time::timeout(
ENSURE_ROUTE_TIMEOUT,
producer.fetch_publish_message_queues(topic),
)
.await;
if matches!(route, Ok(Ok(_))) {
report.existing.push(topic.clone());
continue;
}
let created = tokio::time::timeout(
ENSURE_CREATE_TIMEOUT,
producer.create_topic(
CREATE_TOPIC_KEY,
topic,
DEFAULT_TOPIC_QUEUE_NUMS,
HashMap::new(),
),
)
.await;
match created {
Ok(Ok(())) => report.created.push(topic.clone()),
Ok(Err(error)) => report.failed.push((topic.clone(), error.to_string())),
Err(_) => report
.failed
.push((topic.clone(), "create topic timed out".to_owned())),
}
}
report
}
pub async fn send(&self, topic: &str, message: Message) -> QueueResult<()> {
let result = self
.producer
.send_with_timeout(message, self.config.send_timeout_ms)
.await
.map_err(|error| {
QueueError::publish(format!("failed to publish to `{topic}`: {error}"))
})?;
if result.is_none() {
return Err(QueueError::publish(format!(
"producer returned no send result for `{topic}`"
)));
}
Ok(())
}
pub async fn shutdown(&self) {
let mut producer = self.producer.clone();
producer.shutdown().await;
}
}
/// The RocketMQ client starts lazily and keeps retrying an unreachable
/// nameserver in the background. KenApi now treats RocketMQ as mandatory, so a
/// dead nameserver must fail startup instead of silently running without an
/// event channel.
async fn probe_nameservers(config: &RocketMqConfig) -> QueueResult<()> {
probe_nameservers_with(config, std::time::Duration::from_secs(3)).await
}
async fn probe_nameservers_with(
config: &RocketMqConfig,
timeout: std::time::Duration,
) -> QueueResult<()> {
for server in config.name_servers() {
match tokio::time::timeout(timeout, tokio::net::TcpStream::connect(server)).await {
Ok(Ok(_)) => return Ok(()),
Ok(Err(error)) => {
tracing::warn!(%server, %error, "RocketMQ nameserver is not reachable");
}
Err(_) => {
tracing::warn!(%server, "RocketMQ nameserver probe timed out");
}
}
}
Err(QueueError::connect(format!(
"no RocketMQ nameserver is reachable ({})",
config.name_servers().join(",")
)))
}