d5939db
use std::sync::Arc;
use std::time::Duration;
use rdkafka::ClientConfig;
use rdkafka::error::KafkaError;
use rdkafka::message::{Header, OwnedHeaders};
use rdkafka::producer::{FutureProducer, FutureRecord};
use tracing::warn;
use uuid::Uuid;
use crate::config::AppQueueSettings;
use crate::error::QueueError;
#[derive(Clone)]
pub struct KafkaPublisher {
inner: Arc<FutureProducer>,
cfg: Arc<AppQueueSettings>,
}
impl KafkaPublisher {
pub(super) fn from_config(cfg: &AppQueueSettings) -> Result<Self, QueueError> {
cfg.is_valid().map_err(QueueError::Config)?;
let mut client_cfg = ClientConfig::new();
client_cfg
.set("bootstrap.servers", &cfg.brokers)
.set("client.id", &cfg.client_id)
.set("compression.type", &cfg.producer_compression)
.set("acks", &cfg.producer_acks)
.set("message.timeout.ms", cfg.message_timeout_ms.to_string())
.set("queue.buffering.max.ms", "10")
.set("enable.idempotence", "true");
for (k, v) in &cfg.additional_properties {
client_cfg.set(k.as_str(), v.as_str());
}
let producer: FutureProducer = client_cfg
.create()
.map_err(|e: KafkaError| QueueError::Kafka(e))?;
Ok(Self {
inner: Arc::new(producer),
cfg: Arc::new(cfg.clone()),
})
}
pub async fn publish(
&self,
topic: &str,
key: Option<&[u8]>,
payload: &[u8],
) -> Result<(), QueueError> {
let timeout = Duration::from_millis(self.cfg.publish_queue_timeout_ms as u64);
let message_id = Uuid::new_v4().to_string();
let headers = OwnedHeaders::new()
.insert(Header {
key: "message-id",
value: Some(&message_id),
})
.insert(Header {
key: "content-type",
value: Some("application/octet-stream"),
});
let mut record = FutureRecord::to(topic).payload(payload).headers(headers);
if let Some(k) = key {
record = record.key(k);
}
match self.inner.send(record, timeout).await {
Ok(_) => Ok(()),
Err((e, _msg)) => {
warn!(topic, error = %e, "kafka publish failed");
Err(QueueError::Kafka(e))
}
}
}
pub fn producer(&self) -> &FutureProducer {
&self.inner
}
pub fn config(&self) -> &AppQueueSettings {
&self.cfg
}
}