128 lines
3.5 KiB
Rust
128 lines
3.5 KiB
Rust
use std::time::Duration;
|
|
|
|
use async_nats::{HeaderMap, jetstream};
|
|
use config::AppConfig;
|
|
use serde::Serialize;
|
|
use track::CounterVec;
|
|
|
|
#[derive(Clone)]
|
|
pub struct NatsProducer {
|
|
jetstream: jetstream::Context,
|
|
metrics: Option<track::MetricsRegistry>,
|
|
}
|
|
|
|
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,
|
|
metrics: None,
|
|
})
|
|
}
|
|
|
|
pub fn set_metrics(&mut self, registry: track::MetricsRegistry) {
|
|
self.metrics = Some(registry);
|
|
}
|
|
|
|
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_result: anyhow::Result<()> = async {
|
|
let publish = if headers.is_empty() {
|
|
self.jetstream
|
|
.publish(subject.clone(), payload.to_vec().into())
|
|
.await?
|
|
} else {
|
|
self.jetstream
|
|
.publish_with_headers(
|
|
subject.clone(),
|
|
headers,
|
|
payload.to_vec().into(),
|
|
)
|
|
.await?
|
|
};
|
|
|
|
tokio::time::timeout(Duration::from_secs(5), publish).await??;
|
|
Ok(())
|
|
}
|
|
.await;
|
|
|
|
self.record_published(&subject, publish_result.is_ok());
|
|
publish_result
|
|
}
|
|
|
|
fn record_published(&self, topic: &str, success: bool) {
|
|
if let Some(reg) = &self.metrics {
|
|
let status = if success { "published" } else { "error" };
|
|
queue_messages_vec(reg)
|
|
.with_label_values(&[topic, status])
|
|
.inc();
|
|
}
|
|
}
|
|
}
|
|
|
|
pub 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 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?)
|
|
}
|
|
|
|
fn queue_messages_vec(registry: &track::MetricsRegistry) -> CounterVec {
|
|
registry
|
|
.register_counter_vec(
|
|
"queue_messages_total",
|
|
"Total queue messages",
|
|
&["topic", "status"],
|
|
)
|
|
.expect("failed to register queue_messages_total")
|
|
}
|