d6d42eb
use std::sync::Arc;
use std::time::Duration;
use futures::future::BoxFuture;
use rdkafka::ClientConfig;
use rdkafka::config::RDKafkaLogLevel;
use rdkafka::consumer::{Consumer, StreamConsumer};
use rdkafka::error::KafkaError;
use rdkafka::message::Message;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, warn};
use uuid::Uuid;
use crate::config::AppQueueSettings;
use crate::error::QueueError;
pub type RawHandler =
Arc<dyn Fn(RawMessage) -> BoxFuture<'static, Result<(), QueueError>> + Send + Sync>;
#[derive(Debug, Clone)]
pub enum AutoOffsetReset {
Earliest,
Latest,
None,
}
impl AutoOffsetReset {
fn as_str(&self) -> &'static str {
match self {
AutoOffsetReset::Earliest => "earliest",
AutoOffsetReset::Latest => "latest",
AutoOffsetReset::None => "none",
}
}
}
#[derive(Debug, Clone)]
pub enum SubscriptionMode {
Group {
group_id: String,
auto_offset_reset: AutoOffsetReset,
},
Broadcast {
subscriber_id: String,
},
}
#[derive(Clone)]
pub struct RawMessage {
pub topic: String,
pub partition: i32,
pub offset: i64,
pub key: Option<Vec<u8>>,
pub payload: Vec<u8>,
}
pub struct SubscribeHandle {
cancel: CancellationToken,
join: Option<JoinHandle<()>>,
graceful_timeout: Duration,
}
impl std::fmt::Debug for SubscribeHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SubscribeHandle")
.field("cancelled", &self.cancel.is_cancelled())
.finish()
}
}
impl SubscribeHandle {
pub fn cancel(&self) {
self.cancel.cancel();
}
pub fn is_cancelled(&self) -> bool {
self.cancel.is_cancelled()
}
pub async fn shutdown(mut self) -> Result<(), QueueError> {
self.cancel.cancel();
if let Some(join) = self.join.take() {
match tokio::time::timeout(self.graceful_timeout, join).await {
Ok(Ok(())) => Ok(()),
Ok(Err(e)) => Err(QueueError::Join(e)),
Err(_) => Err(QueueError::Cancelled),
}
} else {
Ok(())
}
}
}
impl Drop for SubscribeHandle {
fn drop(&mut self) {
self.cancel.cancel();
if let Some(join) = self.join.take() {
tokio::spawn(async move {
let _ = join.await;
});
}
}
}
pub(super) fn build_consumer(
cfg: &AppQueueSettings,
mode: &SubscriptionMode,
topics: &[&str],
) -> Result<(StreamConsumer, String), QueueError> {
let mut client_cfg = ClientConfig::new();
client_cfg
.set("bootstrap.servers", &cfg.brokers)
.set("client.id", &cfg.client_id)
.set("session.timeout.ms", cfg.session_timeout_ms.to_string())
.set("enable.auto.commit", "true")
.set("auto.commit.interval.ms", "1000")
.set_log_level(RDKafkaLogLevel::Warning);
let group_id = match mode {
SubscriptionMode::Group {
group_id,
auto_offset_reset,
} => {
client_cfg
.set("group.id", group_id.as_str())
.set("auto.offset.reset", auto_offset_reset.as_str());
group_id.clone()
}
SubscriptionMode::Broadcast { subscriber_id } => {
let unique = format!("{subscriber_id}.{}", Uuid::new_v4());
client_cfg
.set("group.id", unique.as_str())
.set("auto.offset.reset", "latest")
.set("enable.auto.commit", "false");
unique
}
};
for (k, v) in &cfg.additional_properties {
client_cfg.set(k.as_str(), v.as_str());
}
let consumer: StreamConsumer = client_cfg.create().map_err(QueueError::Kafka)?;
consumer
.subscribe(topics)
.map_err(|e: KafkaError| QueueError::Kafka(e))?;
Ok((consumer, group_id))
}
pub(super) fn spawn_raw_loop(
cfg: Arc<AppQueueSettings>,
consumer: StreamConsumer,
group_id: String,
handler: RawHandler,
) -> SubscribeHandle {
let cancel = CancellationToken::new();
let cancel_inner = cancel.clone();
let max_retries = cfg.subscribe_max_retries as usize;
let graceful = Duration::from_millis(cfg.subscribe_graceful_shutdown_ms);
let join = tokio::spawn(async move {
loop {
tokio::select! {
_ = cancel_inner.cancelled() => {
debug!(group = %group_id, "subscriber cancelled");
break;
}
recv = consumer.recv() => {
match recv {
Err(e) => {
warn!(group = %group_id, error = %e, "consumer recv error");
tokio::time::sleep(Duration::from_millis(500)).await;
continue;
}
Ok(msg) => {
let raw = RawMessage {
topic: msg.topic().to_string(),
partition: msg.partition(),
offset: msg.offset(),
key: msg.key().map(|k| k.to_vec()),
payload: msg.payload().map(|p| p.to_vec()).unwrap_or_default(),
};
let mut attempts = 0usize;
loop {
attempts += 1;
match handler(raw.clone()).await {
Ok(()) => break,
Err(e) => {
let delay = backoff_delay(attempts);
warn!(
group = %group_id,
topic = %raw.topic,
attempt = attempts,
error = %e,
"handler failed"
);
if attempts > max_retries {
error!(
group = %group_id,
topic = %raw.topic,
offset = raw.offset,
"handler gave up after retries, dropping message"
);
break;
}
tokio::select! {
_ = cancel_inner.cancelled() => return,
_ = tokio::time::sleep(delay) => {}
}
}
}
}
}
}
}
}
}
});
SubscribeHandle {
cancel,
join: Some(join),
graceful_timeout: graceful,
}
}
fn backoff_delay(attempt: usize) -> Duration {
match attempt {
1 => Duration::from_millis(100),
2 => Duration::from_millis(500),
_ => Duration::from_secs(2),
}
}