ec0bdb8
use std::sync::Arc;
use config::rocketmq::RocketMqConfig;
use rocketmq_client_rust::producer::default_mq_producer::DefaultMQProducer;
use rocketmq_common::common::message::message_single::Message;
use tokio::sync::Mutex;
use crate::error::{QueueError, QueueResult};
pub const DEFAULT_MAX_MESSAGE_BYTES: usize = 4 * 1024 * 1024;
/// 事件信封的字节预算：为消息属性与协议头预留 4 KiB。
pub const MAX_ENVELOPE_BYTES: usize = DEFAULT_MAX_MESSAGE_BYTES - 4096;
#[derive(Clone)]
pub struct RocketMqConn {
producer: Arc<dyn Producer>,
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) -> 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()
.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: Arc::new(RocketProducer(Mutex::new(producer))),
config: config.clone(),
})
}
pub fn config(&self) -> &RocketMqConfig {
&self.config
}
pub fn max_payload(&self) -> usize {
self.config.max_message_bytes
}
pub async fn send(&self, topic: &str, message: Message) -> QueueResult<()> {
self.producer
.send(topic, message, self.config.send_timeout_ms)
.await
}
pub async fn health(&self, topic: &str) -> QueueResult<()> {
self.producer.health(topic).await
}
pub async fn shutdown(&self) {
self.producer.shutdown().await;
}
}
#[async_trait::async_trait]
trait Producer: Send + Sync {
async fn send(&self, topic: &str, message: Message, timeout_ms: u64) -> QueueResult<()>;
async fn health(&self, topic: &str) -> QueueResult<()>;
async fn shutdown(&self);
}
struct RocketProducer(Mutex<DefaultMQProducer>);
#[async_trait::async_trait]
impl Producer for RocketProducer {
async fn send(&self, topic: &str, message: Message, timeout_ms: u64) -> QueueResult<()> {
let mut producer = self.0.lock().await;
let result = producer
.send_with_timeout(message, 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(())
}
async fn health(&self, topic: &str) -> QueueResult<()> {
let mut producer = self.0.lock().await;
producer
.fetch_publish_message_queues(topic)
.await
.map(|_| ())
.map_err(|error| {
QueueError::connect(format!(
"rocketmq route lookup for `{topic}` failed: {error}"
))
})
}
async fn shutdown(&self) {
self.0.lock().await.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
}
pub 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(",")
)))
}