gitdataai/lib/queue/producer.rs
2026-05-30 01:38:40 +08:00

91 lines
2.4 KiB
Rust

use std::time::Duration;
use async_nats::{HeaderMap, jetstream};
use config::AppConfig;
use serde::Serialize;
#[derive(Clone)]
pub struct NatsProducer {
jetstream: jetstream::Context,
}
impl NatsProducer {
pub async fn new(config: &AppConfig) -> anyhow::Result<Self> {
let jetstream = connect_jetstream(config).await?;
ensure_stream(config, &jetstream).await?;
Ok(Self { jetstream })
}
pub async fn send<T>(
&self,
subject: &str,
key: &str,
payload: &T,
headers: Option<HeaderMap>,
) -> anyhow::Result<()>
where
T: Serialize + ?Sized,
{
let payload_bytes = serde_json::to_vec(payload)?;
self.send_raw(subject, key, &payload_bytes, headers).await
}
pub async fn send_raw(
&self,
subject: &str,
key: &str,
payload: &[u8],
headers: Option<HeaderMap>,
) -> anyhow::Result<()> {
let mut headers = headers.unwrap_or_default();
if !key.is_empty() {
headers.append("x-message-key", key);
}
let subject = subject.to_string();
let publish = if headers.is_empty() {
self.jetstream
.publish(subject.clone(), payload.to_vec().into())
.await?
} else {
self.jetstream
.publish_with_headers(subject, headers, payload.to_vec().into())
.await?
};
tokio::time::timeout(Duration::from_secs(5), publish).await??;
Ok(())
}
}
pub(crate) async fn connect_jetstream(
config: &AppConfig,
) -> anyhow::Result<jetstream::Context> {
let client = match config.nats_token() {
Some(token) if !token.is_empty() => {
async_nats::ConnectOptions::with_token(token)
.connect(config.nats_url())
.await?
}
_ => async_nats::connect(config.nats_url()).await?,
};
Ok(jetstream::new(client))
}
pub(crate) async fn ensure_stream(
config: &AppConfig,
jetstream: &jetstream::Context,
) -> anyhow::Result<jetstream::stream::Stream> {
Ok(jetstream
.get_or_create_stream(jetstream::stream::Config {
name: config.nats_stream_name(),
subjects: config.nats_stream_subjects(),
max_age: Duration::from_secs(config.nats_max_age_secs()),
..Default::default()
})
.await?)
}