use crate::error::MqError;
use crate::queue::{Message, MessageQueue};
use async_trait::async_trait;
use futures::StreamExt;
use lapin::{
options::*, types::FieldTable, BasicProperties, Channel, Connection, ConnectionProperties,
Consumer,
};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};
pub struct RealActivemqQueue {
broker_url: String,
connection: Option<Arc<Connection>>,
channel: Option<Channel>,
consumers: Arc<RwLock<HashMap<String, Arc<Mutex<Consumer>>>>>,
in_flight: Arc<RwLock<HashMap<String, u64>>>,
}
impl RealActivemqQueue {
pub fn new(broker_url: impl Into<String>) -> Self {
Self {
broker_url: broker_url.into(),
connection: None,
channel: None,
consumers: Arc::new(RwLock::new(HashMap::new())),
in_flight: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn connect(&mut self) -> Result<(), MqError> {
let conn = Connection::connect(&self.broker_url, ConnectionProperties::default())
.await
.map_err(|e| MqError::Connection(format!("ActiveMQ connect failed: {e}")))?;
let channel = conn
.create_channel()
.await
.map_err(|e| MqError::Connection(format!("ActiveMQ channel failed: {e}")))?;
self.connection = Some(Arc::new(conn));
self.channel = Some(channel);
Ok(())
}
pub async fn reconnect(&mut self) -> Result<(), MqError> {
self.channel = None;
self.connection = None;
self.consumers.write().await.clear();
self.in_flight.write().await.clear();
self.connect().await
}
}
impl Default for RealActivemqQueue {
fn default() -> Self {
Self::new("amqp://localhost:61616")
}
}
#[async_trait]
impl MessageQueue for RealActivemqQueue {
async fn publish(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
let channel = self
.channel
.as_ref()
.ok_or_else(|| MqError::Connection("ActiveMQ not connected".into()))?;
channel
.queue_declare(
topic.into(),
QueueDeclareOptions::default(),
FieldTable::default(),
)
.await
.map_err(|e| MqError::Publish(format!("ActiveMQ queue declare failed: {e}")))?;
channel
.basic_publish(
"".into(),
topic.into(),
BasicPublishOptions::default(),
message,
BasicProperties::default(),
)
.await
.map_err(|e| MqError::Publish(format!("ActiveMQ publish failed: {e}")))?
.await
.map_err(|e| MqError::Publish(format!("ActiveMQ confirm failed: {e}")))?;
Ok(())
}
async fn consume(&self, topic: &str) -> Result<Option<Message>, MqError> {
let channel = self
.channel
.as_ref()
.ok_or_else(|| MqError::Connection("ActiveMQ not connected".into()))?;
let consumer_arc = {
let subs = self.consumers.read().await;
subs.get(topic).cloned()
};
let consumer_arc = match consumer_arc {
Some(arc) => arc,
None => {
channel
.queue_declare(
topic.into(),
QueueDeclareOptions::default(),
FieldTable::default(),
)
.await
.map_err(|e| {
MqError::Subscribe(format!("ActiveMQ queue declare failed: {e}"))
})?;
let consumer = channel
.basic_consume(
topic.into(),
topic.into(),
BasicConsumeOptions::default(),
FieldTable::default(),
)
.await
.map_err(|e| MqError::Subscribe(format!("ActiveMQ consume failed: {e}")))?;
let arc = Arc::new(Mutex::new(consumer));
self.consumers
.write()
.await
.insert(topic.to_string(), arc.clone());
arc
}
};
let mut consumer = consumer_arc.lock().await;
match tokio::time::timeout(std::time::Duration::from_millis(100), consumer.next()).await {
Ok(Some(Ok(delivery))) => {
let delivery_tag = delivery.delivery_tag;
let msg_id = format!("activemq-{delivery_tag}");
let message = Message {
topic: topic.to_string(),
payload: delivery.data.clone(),
key: None,
timestamp: current_timestamp_millis(),
headers: HashMap::new(),
id: msg_id.clone(),
};
self.in_flight.write().await.insert(msg_id, delivery_tag);
Ok(Some(message))
}
Ok(Some(Err(_))) | Ok(None) => Ok(None),
Err(_) => Ok(None),
}
}
async fn ack(&self, message_id: &str) -> Result<(), MqError> {
let channel = self
.channel
.as_ref()
.ok_or_else(|| MqError::Connection("ActiveMQ not connected".into()))?;
let mut in_flight = self.in_flight.write().await;
let delivery_tag = in_flight
.remove(message_id)
.ok_or_else(|| MqError::Publish(format!("unknown message id: {message_id}")))?;
channel
.basic_ack(delivery_tag, BasicAckOptions::default())
.await
.map_err(|e| MqError::Publish(format!("ActiveMQ ack failed: {e}")))?;
Ok(())
}
async fn subscribe(&self, topic: &str) -> Result<(), MqError> {
let channel = self
.channel
.as_ref()
.ok_or_else(|| MqError::Connection("ActiveMQ not connected".into()))?;
channel
.queue_declare(
topic.into(),
QueueDeclareOptions::default(),
FieldTable::default(),
)
.await
.map_err(|e| MqError::Subscribe(format!("ActiveMQ queue declare failed: {e}")))?;
Ok(())
}
}
fn current_timestamp_millis() -> i64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as i64
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_real_activemq_queue_new() {
let queue = RealActivemqQueue::new("amqp://localhost:61616");
assert_eq!(queue.broker_url, "amqp://localhost:61616");
assert!(queue.connection.is_none());
}
#[test]
fn test_real_activemq_queue_default() {
let queue = RealActivemqQueue::default();
assert_eq!(queue.broker_url, "amqp://localhost:61616");
}
#[tokio::test]
async fn test_real_activemq_not_connected_publish() {
let queue = RealActivemqQueue::new("amqp://localhost:61616");
let result = queue.publish("topic", b"msg").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_real_activemq_not_connected_ack() {
let queue = RealActivemqQueue::new("amqp://localhost:61616");
let result = queue.ack("any-id").await;
assert!(result.is_err());
}
#[tokio::test]
#[ignore = "需真实 ActiveMQ Artemis 服务器"]
async fn test_real_activemq_publish_and_consume() {
let mut queue = RealActivemqQueue::new("amqp://localhost:61616");
queue.connect().await.unwrap();
queue.subscribe("test-queue").await.unwrap();
queue
.publish("test-queue", b"hello activemq")
.await
.unwrap();
let msg = queue
.consume("test-queue")
.await
.unwrap()
.expect("message should exist");
assert_eq!(msg.payload, b"hello activemq");
assert_eq!(msg.topic, "test-queue");
queue.ack(&msg.id).await.unwrap();
}
}