zhenyi
821537186e
- Reorganized import statements in adapter tests for better readability - Replaced or_insert_with(Vec::new) with or_default() in test closures - Updated Cargo.lock with new dependency versions and checksums - Added TLS features to tonic dependency configuration - Included sqlx, chrono, and uuid dependencies with specific features - Added jsonwebtoken and arc-swap as project dependencies - Reformatted assertion statements to comply with line length limits - Adjusted base64 import order in engine codec module - Updated protobuf include statement formatting
90 lines
2.7 KiB
Rust
90 lines
2.7 KiB
Rust
use async_trait::async_trait;
|
|
use dashmap::DashMap;
|
|
use tokio::sync::{mpsc, watch};
|
|
|
|
use crate::socket::message_bus::{MessageBus, MessageBusError};
|
|
|
|
pub struct NatsMessageBus {
|
|
client: async_nats::Client,
|
|
shutdowns: DashMap<String, watch::Sender<bool>>,
|
|
}
|
|
|
|
impl NatsMessageBus {
|
|
pub async fn new(nats_url: &str) -> Result<Self, MessageBusError> {
|
|
let client = async_nats::connect(nats_url)
|
|
.await
|
|
.map_err(|e| MessageBusError::Nats(e.to_string()))?;
|
|
Ok(Self {
|
|
client,
|
|
shutdowns: DashMap::new(),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl MessageBus for NatsMessageBus {
|
|
async fn publish(&self, channel: &str, message: &[u8]) -> Result<(), MessageBusError> {
|
|
self.client
|
|
.publish(channel.to_string(), message.to_vec().into())
|
|
.await
|
|
.map_err(|e| MessageBusError::Nats(e.to_string()))?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn subscribe(&self, channel: &str) -> Result<mpsc::Receiver<Vec<u8>>, MessageBusError> {
|
|
let (tx, rx) = mpsc::channel::<Vec<u8>>(256);
|
|
|
|
let mut subscriber = self
|
|
.client
|
|
.subscribe(channel.to_string())
|
|
.await
|
|
.map_err(|e| MessageBusError::Nats(e.to_string()))?;
|
|
|
|
let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
|
|
self.shutdowns.insert(channel.to_string(), shutdown_tx);
|
|
|
|
tokio::spawn(async move {
|
|
use futures_util::StreamExt;
|
|
loop {
|
|
tokio::select! {
|
|
_ = shutdown_rx.changed() => {
|
|
break;
|
|
}
|
|
message = subscriber.next() => {
|
|
match message {
|
|
Some(msg) => {
|
|
let data = msg.payload.to_vec();
|
|
if tx.send(data).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
None => break,
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if let Err(e) = subscriber.unsubscribe().await {
|
|
tracing::warn!("NATS unsubscribe error: {}", e);
|
|
}
|
|
});
|
|
|
|
Ok(rx)
|
|
}
|
|
|
|
async fn unsubscribe(&self, channel: &str) -> Result<(), MessageBusError> {
|
|
if let Some((_, tx)) = self.shutdowns.remove(channel) {
|
|
let _ = tx.send(true);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn close(&self) -> Result<(), MessageBusError> {
|
|
// Signal all subscribers to shutdown
|
|
self.shutdowns.iter().for_each(|entry| {
|
|
let _ = entry.value().send(true);
|
|
});
|
|
self.shutdowns.clear();
|
|
Ok(())
|
|
}
|
|
}
|