feat: init

queue/mod.rs

@@ -0,0 +1,183 @@
+pub mod publisher;
+pub mod subscriber;
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use async_nats::jetstream::Context;
+use async_nats::jetstream::stream::{Config as StreamConfig, RetentionPolicy};
+
+use crate::config::AppConfig;
+use crate::error::{AppError, AppResult};
+
+#[derive(Clone)]
+pub struct NatsQueue {
+    inner: Arc<NatsQueueInner>,
+}
+
+struct NatsQueueInner {
+    js: Context,
+    client: async_nats::Client,
+    stream_prefix: String,
+    ack_wait: Duration,
+    max_deliver: i64,
+}
+
+pub struct StreamHandle {
+    pub name: String,
+    pub subjects: Vec<String>,
+}
+
+impl NatsQueue {
+    pub async fn connect(config: &AppConfig) -> AppResult<Self> {
+        let url = config.nats_url()?;
+        let timeout = config.nats_connection_timeout_secs()?;
+        let ping_interval = config.nats_ping_interval_secs()?;
+        let reconnect_delay = config.nats_reconnect_delay_secs()?;
+        let max_reconnects = config.nats_max_reconnects()?;
+
+        let mut opts = async_nats::ConnectOptions::new()
+            .connection_timeout(Duration::from_secs(timeout))
+            .ping_interval(Duration::from_secs(ping_interval))
+            .retry_on_initial_connect()
+            .max_reconnects(Some(max_reconnects))
+            .reconnect_delay_callback(move |_| Duration::from_secs(reconnect_delay))
+            .event_callback(|event| async move {
+                match event {
+                    async_nats::Event::Disconnected => {
+                        tracing::warn!("nats disconnected");
+                    }
+                    async_nats::Event::Connected => {
+                        tracing::info!("nats connected");
+                    }
+                    async_nats::Event::LameDuckMode => {
+                        tracing::warn!("nats server entering lame duck mode");
+                    }
+                    _ => {}
+                }
+            });
+
+        if let (Some(user), Some(pass)) = (config.nats_username()?, config.nats_password()?) {
+            opts = opts.user_and_password(user, pass);
+        } else if let Some(token) = config.nats_token()? {
+            opts = opts.token(token);
+        }
+
+        let client = async_nats::connect_with_options(&url, opts)
+            .await
+            .map_err(|e| AppError::Config(format!("nats connect failed: {e}")))?;
+
+        let js = async_nats::jetstream::new(client.clone());
+        let stream_prefix = config.nats_stream_prefix()?;
+        let ack_wait = config.nats_default_ack_wait_secs()?;
+        let max_deliver = config.nats_default_max_deliver()?;
+
+        tracing::info!(url = %url, "nats connected");
+
+        Ok(Self {
+            inner: Arc::new(NatsQueueInner {
+                js,
+                client,
+                stream_prefix,
+                ack_wait: Duration::from_secs(ack_wait),
+                max_deliver,
+            }),
+        })
+    }
+
+    pub fn client(&self) -> &async_nats::Client {
+        &self.inner.client
+    }
+
+    pub fn js(&self) -> &Context {
+        &self.inner.js
+    }
+
+    pub fn stream_name(&self, name: &str) -> String {
+        format!("{}_{}", self.inner.stream_prefix, name)
+    }
+
+    pub async fn ensure_stream(
+        &self,
+        name: &str,
+        subjects: Vec<String>,
+    ) -> AppResult<StreamHandle> {
+        let stream_name = self.stream_name(name);
+        let config = StreamConfig {
+            name: stream_name.clone(),
+            subjects,
+            retention: RetentionPolicy::Limits,
+            max_messages: 100_000,
+            max_age: Duration::from_secs(7 * 24 * 3600),
+            duplicate_window: Duration::from_secs(120),
+            storage: async_nats::jetstream::stream::StorageType::File,
+            ..Default::default()
+        };
+
+        self.inner
+            .js
+            .get_or_create_stream(config)
+            .await
+            .map_err(|e| AppError::Config(format!("ensure stream {stream_name} failed: {e}")))?;
+
+        let subjects = self
+            .inner
+            .js
+            .get_stream(&stream_name)
+            .await
+            .map_err(|e| AppError::Config(format!("get stream {stream_name} failed: {e}")))?
+            .info()
+            .await
+            .map_err(|e| AppError::Config(format!("stream info {stream_name} failed: {e}")))?
+            .config
+            .subjects
+            .clone();
+
+        tracing::info!(stream = %stream_name, subjects = ?subjects, "stream ready");
+
+        Ok(StreamHandle {
+            name: stream_name,
+            subjects,
+        })
+    }
+
+    pub async fn ensure_ephemeral_stream(
+        &self,
+        name: &str,
+        subjects: Vec<String>,
+        max_age_secs: u64,
+    ) -> AppResult<StreamHandle> {
+        let stream_name = self.stream_name(name);
+        let config = StreamConfig {
+            name: stream_name.clone(),
+            subjects,
+            retention: RetentionPolicy::Limits,
+            max_messages: 10_000,
+            max_age: Duration::from_secs(max_age_secs),
+            duplicate_window: Duration::from_secs(60),
+            storage: async_nats::jetstream::stream::StorageType::Memory,
+            ..Default::default()
+        };
+
+        self.inner
+            .js
+            .get_or_create_stream(config)
+            .await
+            .map_err(|e| AppError::Config(format!("ensure stream {stream_name} failed: {e}")))?;
+
+        Ok(StreamHandle {
+            name: stream_name,
+            subjects: vec![],
+        })
+    }
+
+    pub async fn delete_stream(&self, name: &str) -> AppResult<()> {
+        let stream_name = self.stream_name(name);
+        self.inner
+            .js
+            .delete_stream(&stream_name)
+            .await
+            .map_err(|e| AppError::Config(format!("delete stream {stream_name} failed: {e}")))?;
+        Ok(())
+    }
+}

queue/publisher.rs

@@ -0,0 +1,65 @@
+use serde::Serialize;
+
+use crate::error::{AppError, AppResult};
+
+use super::NatsQueue;
+
+pub struct PublishResult {
+    pub stream: String,
+    pub sequence: u64,
+}
+
+impl NatsQueue {
+    pub async fn publish(&self, subject: &str, payload: &[u8]) -> AppResult<PublishResult> {
+        let subject = subject.to_string();
+        let ack = self
+            .inner
+            .js
+            .publish(subject.clone(), payload.to_vec().into())
+            .await
+            .map_err(|e| AppError::Config(format!("publish to {subject} failed: {e}")))?
+            .await
+            .map_err(|e| AppError::Config(format!("publish ack for {subject} failed: {e}")))?;
+
+        Ok(PublishResult {
+            stream: ack.stream,
+            sequence: ack.sequence,
+        })
+    }
+
+    pub async fn publish_json<T: Serialize>(
+        &self,
+        subject: &str,
+        payload: &T,
+    ) -> AppResult<PublishResult> {
+        let data = serde_json::to_vec(payload)?;
+        self.publish(subject, &data).await
+    }
+
+    pub async fn publish_with_headers(
+        &self,
+        subject: &str,
+        payload: &[u8],
+        headers: Vec<(String, String)>,
+    ) -> AppResult<PublishResult> {
+        let subject = subject.to_string();
+        let mut nats_headers = async_nats::HeaderMap::new();
+        for (k, v) in headers {
+            nats_headers.insert(k, v);
+        }
+
+        let ack = self
+            .inner
+            .js
+            .publish_with_headers(subject.clone(), nats_headers, payload.to_vec().into())
+            .await
+            .map_err(|e| AppError::Config(format!("publish to {subject} failed: {e}")))?
+            .await
+            .map_err(|e| AppError::Config(format!("publish ack for {subject} failed: {e}")))?;
+
+        Ok(PublishResult {
+            stream: ack.stream,
+            sequence: ack.sequence,
+        })
+    }
+}

queue/subscriber.rs

@@ -0,0 +1,202 @@
+use std::time::Duration;
+
+use async_nats::jetstream::consumer::AckPolicy;
+use async_nats::jetstream::consumer::pull::Config as PullConfig;
+use async_nats::jetstream::consumer::push::Config as PushConfig;
+use async_nats::jetstream::consumer::push::Messages;
+use futures_util::StreamExt;
+use serde::de::DeserializeOwned;
+
+use crate::error::{AppError, AppResult};
+
+use super::NatsQueue;
+
+pub struct NatsMessage {
+    inner: async_nats::jetstream::Message,
+}
+
+impl NatsMessage {
+    pub fn subject(&self) -> &str {
+        self.inner.message.subject.as_str()
+    }
+
+    pub fn payload(&self) -> &[u8] {
+        &self.inner.message.payload
+    }
+
+    pub fn payload_json<T: DeserializeOwned>(&self) -> AppResult<T> {
+        serde_json::from_slice(&self.inner.message.payload).map_err(AppError::from)
+    }
+
+    pub fn headers(&self) -> Option<&async_nats::HeaderMap> {
+        self.inner.message.headers.as_ref()
+    }
+
+    pub async fn ack(self) -> AppResult<()> {
+        self.inner
+            .ack()
+            .await
+            .map_err(|e| AppError::Config(format!("ack failed: {e}")))
+    }
+
+    pub async fn nack(self) -> AppResult<()> {
+        self.inner
+            .ack_with(async_nats::jetstream::AckKind::Nak(None))
+            .await
+            .map_err(|e| AppError::Config(format!("nack failed: {e}")))
+    }
+
+    pub async fn nack_with_delay(self, delay: Duration) -> AppResult<()> {
+        self.inner
+            .ack_with(async_nats::jetstream::AckKind::Nak(Some(delay)))
+            .await
+            .map_err(|e| AppError::Config(format!("nack with delay failed: {e}")))
+    }
+
+    pub async fn ack_in_progress(&self) -> AppResult<()> {
+        self.inner
+            .ack_with(async_nats::jetstream::AckKind::Progress)
+            .await
+            .map_err(|e| AppError::Config(format!("ack in progress failed: {e}")))
+    }
+
+    pub async fn term(self) -> AppResult<()> {
+        self.inner
+            .ack_with(async_nats::jetstream::AckKind::Term)
+            .await
+            .map_err(|e| AppError::Config(format!("term failed: {e}")))
+    }
+}
+
+pub struct PullSubscription {
+    stream: async_nats::jetstream::stream::Stream,
+    consumer_name: String,
+}
+
+impl NatsQueue {
+    pub async fn subscribe_broadcast(
+        &self,
+        stream_name: &str,
+        consumer_name: &str,
+        filter_subject: Option<&str>,
+    ) -> AppResult<Messages> {
+        let full_stream = self.stream_name(stream_name);
+        let stream = self
+            .inner
+            .js
+            .get_stream(&full_stream)
+            .await
+            .map_err(|e| AppError::Config(format!("get stream {full_stream} failed: {e}")))?;
+
+        let deliver_subject = format!(
+            "deliver.{}.{}.{}",
+            self.inner.stream_prefix, stream_name, consumer_name
+        );
+
+        let config = PushConfig {
+            durable_name: Some(consumer_name.to_string()),
+            deliver_subject,
+            deliver_group: Some(consumer_name.to_string()),
+            ack_policy: AckPolicy::Explicit,
+            ack_wait: self.inner.ack_wait,
+            max_deliver: self.inner.max_deliver,
+            filter_subject: filter_subject.unwrap_or(">").to_string(),
+            ..Default::default()
+        };
+
+        let consumer = stream
+            .get_or_create_consumer(consumer_name, config)
+            .await
+            .map_err(|e| {
+                AppError::Config(format!("create push consumer {consumer_name} failed: {e}"))
+            })?;
+
+        let messages = consumer.messages().await.map_err(|e| {
+            AppError::Config(format!("subscribe broadcast {consumer_name} failed: {e}"))
+        })?;
+
+        Ok(messages)
+    }
+
+    pub async fn create_pull_consumer(
+        &self,
+        stream_name: &str,
+        consumer_name: &str,
+        filter_subject: Option<&str>,
+        max_batch: usize,
+    ) -> AppResult<PullSubscription> {
+        let full_stream = self.stream_name(stream_name);
+        let stream = self
+            .inner
+            .js
+            .get_stream(&full_stream)
+            .await
+            .map_err(|e| AppError::Config(format!("get stream {full_stream} failed: {e}")))?;
+
+        let mut config = PullConfig {
+            durable_name: Some(consumer_name.to_string()),
+            ack_policy: AckPolicy::Explicit,
+            ack_wait: self.inner.ack_wait,
+            max_deliver: self.inner.max_deliver,
+            max_ack_pending: max_batch as i64,
+            ..Default::default()
+        };
+
+        if let Some(subject) = filter_subject {
+            config.filter_subject = subject.to_string();
+        }
+
+        stream
+            .get_or_create_consumer(consumer_name, config)
+            .await
+            .map_err(|e| {
+                AppError::Config(format!("create pull consumer {consumer_name} failed: {e}"))
+            })?;
+
+        Ok(PullSubscription {
+            stream,
+            consumer_name: consumer_name.to_string(),
+        })
+    }
+
+    pub async fn fetch(
+        &self,
+        subscription: &PullSubscription,
+        batch_size: usize,
+    ) -> AppResult<Vec<NatsMessage>> {
+        let consumer = subscription
+            .stream
+            .get_consumer(&subscription.consumer_name)
+            .await
+            .map_err(|e| AppError::Config(format!("get consumer failed: {e}")))?;
+
+        let mut messages = consumer
+            .fetch()
+            .max_messages(batch_size)
+            .messages()
+            .await
+            .map_err(|e| AppError::Config(format!("fetch failed: {e}")))?;
+
+        let mut result = Vec::with_capacity(batch_size);
+        while let Some(msg) = messages.next().await {
+            match msg {
+                Ok(m) => result.push(NatsMessage { inner: m }),
+                Err(e) => {
+                    tracing::warn!(error = %e, "fetch message error");
+                    break;
+                }
+            }
+        }
+        Ok(result)
+    }
+
+    pub async fn subscribe_ephemeral(&self, subject: String) -> AppResult<async_nats::Subscriber> {
+        let sub = self
+            .inner
+            .client
+            .subscribe(subject.clone())
+            .await
+            .map_err(|e| AppError::Config(format!("subscribe ephemeral {subject} failed: {e}")))?;
+        Ok(sub)
+    }
+}