use async_trait::async_trait;
use bevy::log::{debug, warn};
use crossbeam_channel::{
Receiver, RecvTimeoutError, Sender, TryRecvError, TrySendError, bounded,
};
use rdkafka::{
ClientContext, Offset, TopicPartitionList,
admin::{AdminClient, AdminOptions, NewTopic, TopicReplication},
client::DefaultClientContext,
config::ClientConfig,
consumer::{BaseConsumer, CommitMode, Consumer},
error::{KafkaError, RDKafkaErrorCode},
message::{Header, Headers, Message, OwnedHeaders},
producer::{BaseProducer, BaseRecord, DeliveryResult, Producer, ProducerContext},
};
use std::{
collections::{BTreeMap, HashMap, VecDeque},
fmt::Debug,
sync::{
Arc, Mutex, RwLock,
atomic::{AtomicBool, AtomicUsize, Ordering},
},
thread,
time::{Duration, Instant},
};
use tokio::task::JoinHandle;
use bevy::prelude::*;
use bevy_event_bus::{
TopologyMode,
backends::event_bus_backend::{
BackendConfigError, BackendPluginSetup, DeliveryFailure, DeliveryFailureCallback,
EventBusBackend, LagReportingBackend, LagReportingDescriptor, LagReportingHandle,
ManualCommitController, ManualCommitDescriptor, ManualCommitHandle, ManualCommitStyle,
ReceiveOptions, SendOptions,
},
config::kafka::{
KafkaBackendConfig, KafkaConnectionConfig, KafkaConsumerGroupSpec, KafkaInitialOffset,
KafkaTopicSpec, KafkaTopologyConfig,
},
error::EventBusErrorType,
resources::{
ConsumerMetrics, IncomingMessage, KafkaCommitQueue, KafkaCommitResultChannel,
KafkaLagCacheResource, MessageMetadata, backend_metadata::KafkaMetadata,
},
runtime,
};
#[derive(Debug, Clone)]
pub struct KafkaCommitRequest {
pub topic: String,
pub partition: i32,
pub offset: i64,
pub consumer_group: String,
}
#[derive(Debug, Clone)]
pub struct KafkaCommitResult {
pub topic: String,
pub partition: i32,
pub offset: i64,
pub consumer_group: String,
pub error: Option<String>,
}
#[derive(Debug, Clone)]
pub struct KafkaLagMeasurement {
pub lag: i64,
pub last_updated: Instant,
}
#[derive(Clone, Default)]
pub struct KafkaLagCache {
inner: Arc<RwLock<HashMap<(String, String), KafkaLagMeasurement>>>,
}
impl KafkaLagCache {
pub fn update(&self, consumer_group: &str, topic: &str, measurement: KafkaLagMeasurement) {
let mut guard = self.inner.write().unwrap();
guard.insert((consumer_group.to_string(), topic.to_string()), measurement);
}
pub fn snapshot_for_group(&self, consumer_group: &str) -> HashMap<String, KafkaLagMeasurement> {
let guard = self.inner.read().unwrap();
guard
.iter()
.filter(|((group, _), _)| group == consumer_group)
.map(|((_, topic), measurement)| (topic.clone(), measurement.clone()))
.collect()
}
pub fn get(&self, consumer_group: &str, topic: &str) -> Option<KafkaLagMeasurement> {
let guard = self.inner.read().unwrap();
guard
.get(&(consumer_group.to_string(), topic.to_string()))
.cloned()
}
}
struct KafkaManualCommitHandle {
sender: Sender<KafkaCommitRequest>,
results: Mutex<Option<Receiver<KafkaCommitResult>>>,
}
impl KafkaManualCommitHandle {
fn new(
sender: Sender<KafkaCommitRequest>,
results: Option<Receiver<KafkaCommitResult>>,
) -> Self {
Self {
sender,
results: Mutex::new(results),
}
}
}
impl ManualCommitHandle for KafkaManualCommitHandle {
fn register_resources(&self, world: &mut World) {
world.insert_resource(KafkaCommitQueue {
sender: self.sender.clone(),
});
if let Some(receiver) = self.results.lock().unwrap().take() {
world.insert_resource(KafkaCommitResultChannel { receiver });
}
}
fn descriptor(&self) -> ManualCommitDescriptor {
ManualCommitDescriptor {
backend: "kafka",
style: ManualCommitStyle::OffsetQueue,
}
}
}
struct KafkaLagHandle {
cache: KafkaLagCache,
}
impl KafkaLagHandle {
fn new(cache: KafkaLagCache) -> Self {
Self { cache }
}
}
impl LagReportingHandle for KafkaLagHandle {
fn register_resources(&self, world: &mut World) {
world.insert_resource(KafkaLagCacheResource {
cache: self.cache.clone(),
});
}
fn descriptor(&self) -> LagReportingDescriptor {
LagReportingDescriptor {
backend: "kafka",
detail: "consumer_lag",
}
}
}
fn kafka_commit_result_dispatch_system(
mut commands: Commands,
maybe_channel: Option<Res<KafkaCommitResultChannel>>,
mut messages: MessageWriter<KafkaCommitResultMessage>,
) {
if let Some(channel) = maybe_channel {
loop {
match channel.receiver.try_recv() {
Ok(result) => {
messages.write(KafkaCommitResultMessage {
backend: "kafka".into(),
consumer_group: result.consumer_group,
topic: result.topic,
partition: result.partition,
offset: result.offset,
error: result.error,
});
}
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => {
commands.remove_resource::<KafkaCommitResultChannel>();
break;
}
}
}
}
}
fn kafka_commit_result_stats_system(
mut stats: ResMut<KafkaCommitResultStats>,
mut messages: MessageReader<KafkaCommitResultMessage>,
) {
for message in messages.read() {
let key = (
message.backend.clone(),
message.consumer_group.clone(),
message.topic.clone(),
);
let entry = stats.totals.entry(key).or_default();
entry.partition = message.partition;
entry.last_offset = message.offset;
match &message.error {
Some(err) => {
entry.failures += 1;
entry.last_error = Some(err.clone());
}
None => {
entry.successes += 1;
entry.last_error = None;
}
}
}
}
#[derive(Message, Debug, Clone)]
pub struct KafkaCommitResultMessage {
pub backend: String,
pub consumer_group: String,
pub topic: String,
pub partition: i32,
pub offset: i64,
pub error: Option<String>,
}
#[derive(Resource, Debug, Clone, Default)]
pub struct KafkaCommitResultStats {
pub totals: HashMap<(String, String, String), KafkaCommitOutcome>,
}
#[derive(Debug, Clone, Default)]
pub struct KafkaCommitOutcome {
pub partition: i32,
pub last_offset: i64,
pub successes: usize,
pub failures: usize,
pub last_error: Option<String>,
}
#[derive(Clone)]
struct ConsumerRuntime {
consumer: Arc<BaseConsumer>,
manual_commit: bool,
group_id: String,
topics: Vec<String>,
}
struct KafkaMessageIngress {
tx: Sender<IncomingMessage>,
rx: Mutex<Option<Receiver<IncomingMessage>>>,
pending: Mutex<VecDeque<IncomingMessage>>,
dropped: Arc<AtomicUsize>,
}
impl KafkaMessageIngress {
fn new(capacity: usize) -> Self {
let (tx, rx) = bounded(capacity);
Self {
tx,
rx: Mutex::new(Some(rx)),
pending: Mutex::new(VecDeque::new()),
dropped: Arc::new(AtomicUsize::new(0)),
}
}
fn take_receiver(&self) -> Option<Receiver<IncomingMessage>> {
self.rx.lock().unwrap().take()
}
fn restore_receiver(&self, receiver: Receiver<IncomingMessage>) {
*self.rx.lock().unwrap() = Some(receiver);
}
}
struct KafkaCommitPipeline {
request_tx: Sender<KafkaCommitRequest>,
request_rx: Mutex<Option<Receiver<KafkaCommitRequest>>>,
outcome_tx: Sender<KafkaCommitResult>,
outcome_rx: Mutex<Option<Receiver<KafkaCommitResult>>>,
}
impl KafkaCommitPipeline {
fn new(capacity: usize) -> Self {
let (request_tx, request_rx) = bounded(capacity);
let (outcome_tx, outcome_rx) = bounded(capacity);
Self {
request_tx,
request_rx: Mutex::new(Some(request_rx)),
outcome_tx,
outcome_rx: Mutex::new(Some(outcome_rx)),
}
}
fn request_sender(&self) -> Sender<KafkaCommitRequest> {
self.request_tx.clone()
}
fn take_outcome_receiver(&self) -> Option<Receiver<KafkaCommitResult>> {
self.outcome_rx.lock().unwrap().take()
}
fn take_request_receiver(&self) -> Receiver<KafkaCommitRequest> {
self.request_rx
.lock()
.unwrap()
.take()
.expect("commit request receiver already taken")
}
}
struct KafkaWorkerRuntime {
running: Arc<AtomicBool>,
consumer_handles: Mutex<Vec<JoinHandle<()>>>,
commit: Mutex<Option<JoinHandle<()>>>,
lag: Mutex<Option<JoinHandle<()>>>,
producer_poll: Mutex<Option<JoinHandle<()>>>,
}
impl KafkaWorkerRuntime {
fn new() -> Self {
Self {
running: Arc::new(AtomicBool::new(false)),
consumer_handles: Mutex::new(Vec::new()),
commit: Mutex::new(None),
lag: Mutex::new(None),
producer_poll: Mutex::new(None),
}
}
fn stop_all(&self) {
if let Ok(mut handles) = self.consumer_handles.lock() {
for handle in handles.drain(..) {
handle.abort();
}
}
for slot in [&self.commit, &self.lag, &self.producer_poll] {
if let Ok(mut guard) = slot.lock() {
if let Some(handle) = guard.take() {
handle.abort();
}
}
}
}
}
struct KafkaBackendState {
config: KafkaBackendConfig,
producer: Arc<BaseProducer<EventBusProducerContext>>,
consumers: Arc<Mutex<HashMap<String, ConsumerRuntime>>>,
ingress: KafkaMessageIngress,
commits: KafkaCommitPipeline,
lag_cache: KafkaLagCache,
workers: KafkaWorkerRuntime,
}
#[derive(Clone)]
pub struct KafkaEventBusBackend {
state: Arc<KafkaBackendState>,
}
impl Debug for KafkaEventBusBackend {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let consumers = self.state.consumers.lock().unwrap();
f.debug_struct("KafkaEventBusBackend")
.field(
"bootstrap",
&self.state.config.connection.bootstrap_servers(),
)
.field(
"consumer_groups",
&consumers.keys().cloned().collect::<Vec<_>>(),
)
.finish()
}
}
#[derive(Clone)]
struct EventBusProducerContext;
impl ClientContext for EventBusProducerContext {}
impl ProducerContext for EventBusProducerContext {
type DeliveryOpaque = Arc<DeliveryFailureCallback>;
fn delivery(&self, delivery_result: &DeliveryResult, delivery_opaque: Self::DeliveryOpaque) {
match delivery_result {
Err((kafka_error, owned_message)) => {
warn!(
topic = %owned_message.topic(),
partition = owned_message.partition(),
error = %kafka_error,
"Kafka message delivery failed"
);
let mut headers_map = HashMap::new();
if let Some(headers) = owned_message.headers() {
for header in headers.iter() {
if let Some(value) = header.value {
if let Ok(str_value) = String::from_utf8(value.to_vec()) {
headers_map.insert(header.key.to_string(), str_value);
}
}
}
}
let backend_metadata = KafkaMetadata {
topic: owned_message.topic().to_string(),
partition: owned_message.partition(),
offset: owned_message.offset(),
consumer_group: None,
manual_commit: false,
headers: headers_map,
};
let metadata = MessageMetadata::new(
owned_message.topic().to_string(),
Instant::now(),
owned_message
.key()
.and_then(|k| String::from_utf8(k.to_vec()).ok()),
Some(Box::new(backend_metadata)),
);
delivery_opaque.call(DeliveryFailure {
backend: "kafka",
kind: EventBusErrorType::DeliveryFailure,
topic: owned_message.topic().to_string(),
error: kafka_error.to_string(),
metadata: Some(metadata),
});
}
Ok(delivery) => {
debug!(
topic = %delivery.topic(),
partition = delivery.partition(),
offset = delivery.offset(),
"Kafka message delivered successfully"
);
}
}
}
}
impl KafkaEventBusBackend {
pub fn new(config: KafkaBackendConfig) -> Result<Self, BackendConfigError> {
prepare_kafka_topology(&config.connection, &config.topology)?;
let producer = Arc::new(
build_producer(&config.connection)
.map_err(|err| BackendConfigError::new("kafka", err.to_string()))?,
);
let consumers = Arc::new(Mutex::new(build_consumers(
&config.connection,
&config.topology,
)));
let capacities = config.channel_capacities.clone();
let state = KafkaBackendState {
config,
producer,
consumers: consumers.clone(),
ingress: KafkaMessageIngress::new(capacities.message),
commits: KafkaCommitPipeline::new(capacities.commit),
lag_cache: KafkaLagCache::default(),
workers: KafkaWorkerRuntime::new(),
};
Ok(Self {
state: Arc::new(state),
})
}
pub fn bootstrap_servers(&self) -> &str {
self.state.config.connection.bootstrap_servers()
}
pub fn configured_topics(&self) -> Vec<String> {
self.state
.config
.topology
.topics()
.iter()
.map(|spec| spec.name.clone())
.collect()
}
pub fn take_receiver(&self) -> Option<Receiver<IncomingMessage>> {
self.state.ingress.take_receiver()
}
pub fn commit_sender(&self) -> Sender<KafkaCommitRequest> {
self.state.commits.request_sender()
}
pub fn take_commit_results(&self) -> Option<Receiver<KafkaCommitResult>> {
self.state.commits.take_outcome_receiver()
}
pub fn lag_cache(&self) -> KafkaLagCache {
self.state.lag_cache.clone()
}
pub fn dropped_count(&self) -> usize {
self.state.ingress.dropped.load(Ordering::Relaxed)
}
pub fn poll_producer(&self) {
self.state.producer.poll(Duration::from_millis(0));
}
pub fn flush(&self, timeout: Duration) -> Result<(), String> {
let result = self.state.producer.flush(timeout);
for _ in 0..10 {
self.state.producer.poll(Duration::from_millis(0));
}
result.map_err(|err| err.to_string())
}
fn drain_matching_messages(&self, topic: &str, group: Option<&str>) -> Vec<Vec<u8>> {
let mut matches = Vec::new();
let mut deferred = VecDeque::new();
{
let mut pending = self.state.ingress.pending.lock().unwrap();
while let Some(msg) = pending.pop_front() {
if Self::message_matches(&msg, topic, group) {
let payload = msg.payload;
matches.push(payload);
} else {
deferred.push_back(msg);
}
}
}
if let Some(receiver) = self.state.ingress.take_receiver() {
loop {
match receiver.try_recv() {
Ok(msg) => {
let matches_topic = Self::message_matches(&msg, topic, group);
if matches_topic {
let payload = msg.payload;
matches.push(payload);
} else {
deferred.push_back(msg);
}
}
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => break,
}
}
self.state.ingress.restore_receiver(receiver);
}
if !deferred.is_empty() {
let mut pending = self.state.ingress.pending.lock().unwrap();
while let Some(msg) = deferred.pop_front() {
pending.push_back(msg);
}
}
matches
}
fn message_matches(message: &IncomingMessage, topic: &str, group: Option<&str>) -> bool {
if message.source != topic {
return false;
}
match group {
Some(group_id) => message
.backend_metadata
.as_ref()
.and_then(|meta| meta.as_any().downcast_ref::<KafkaMetadata>())
.and_then(|kafka| kafka.consumer_group.as_deref())
.map(|candidate| candidate == group_id)
.unwrap_or(false),
None => true,
}
}
fn spawn_runtime_tasks(&self) {
self.spawn_producer_poll_worker();
self.spawn_consumer_tasks();
self.spawn_commit_worker();
self.spawn_lag_worker();
}
fn spawn_producer_poll_worker(&self) {
let mut guard = self.state.workers.producer_poll.lock().unwrap();
if guard.is_some() {
return;
}
let running = self.state.workers.running.clone();
let producer = self.state.producer.clone();
let handle = runtime::runtime().spawn_blocking(move || {
while running.load(Ordering::Relaxed) {
producer.poll(Duration::from_millis(10));
std::thread::sleep(Duration::from_millis(10));
}
});
*guard = Some(handle);
}
fn spawn_consumer_tasks(&self) {
let consumers = self.state.consumers.lock().unwrap().clone();
let running = self.state.workers.running.clone();
let producer = self.state.producer.clone();
let tx = self.state.ingress.tx.clone();
let dropped = self.state.ingress.dropped.clone();
let mut handles = self.state.workers.consumer_handles.lock().unwrap();
for runtime in consumers.into_values() {
let tx_clone = tx.clone();
let running_clone = running.clone();
let producer_clone = producer.clone();
let dropped_clone = dropped.clone();
let handle = runtime::runtime().spawn_blocking(move || {
consumer_loop(runtime, tx_clone, running_clone, producer_clone, dropped_clone);
});
handles.push(handle);
}
}
fn spawn_commit_worker(&self) {
let mut guard = self.state.workers.commit.lock().unwrap();
if guard.is_some() {
return;
}
let running = self.state.workers.running.clone();
let consumers = self.state.consumers.clone();
let rx = self.state.commits.take_request_receiver();
let outcome_tx = self.state.commits.outcome_tx.clone();
let handle = runtime::runtime().spawn(async move {
while running.load(Ordering::Relaxed) {
match tokio::task::block_in_place(|| rx.recv_timeout(Duration::from_millis(50))) {
Ok(req) => {
let consumer_opt =
consumers.lock().unwrap().get(&req.consumer_group).cloned();
let result = if let Some(runtime) = consumer_opt {
commit_offset_sync(&runtime, &req)
} else {
Err(format!(
"Consumer group '{}' not found for commit",
req.consumer_group
))
};
let _ = outcome_tx.try_send(KafkaCommitResult {
topic: req.topic,
partition: req.partition,
offset: req.offset,
consumer_group: req.consumer_group,
error: result.err(),
});
}
Err(RecvTimeoutError::Timeout) => {}
Err(RecvTimeoutError::Disconnected) => break,
}
}
});
*guard = Some(handle);
}
fn spawn_lag_worker(&self) {
let mut guard = self.state.workers.lag.lock().unwrap();
if guard.is_some() {
return;
}
let running = self.state.workers.running.clone();
let consumers = self.state.consumers.clone();
let lag_cache = self.state.lag_cache.clone();
let poll_interval = self.state.config.consumer_lag_poll_interval;
let connection = self.state.config.connection.clone();
const TRANSPORT_WARN_THRESHOLD: u32 = 3;
let handle = runtime::runtime().spawn(async move {
let mut transport_failures: HashMap<(String, String), u32> = HashMap::new();
let mut ticker = tokio::time::interval(poll_interval);
while running.load(Ordering::Relaxed) {
ticker.tick().await;
let snapshot = consumers.lock().unwrap().clone();
for runtime in snapshot.values() {
let group_id = runtime.group_id.clone();
let topics = runtime.topics.clone();
let connection_for_lag = connection.clone();
let group_for_task = group_id.clone();
let group_measurement = tokio::task::spawn_blocking(move || {
compute_group_lag(&connection_for_lag, &group_for_task, &topics)
})
.await;
let per_topic = match group_measurement {
Ok(Ok(per_topic)) => per_topic,
Ok(Err(err)) => {
warn!(
consumer_group = %group_id,
error = %err,
"Failed to create consumer for lag measurement"
);
continue;
}
Err(err) => {
warn!(
consumer_group = %group_id,
error = %err,
"Consumer lag measurement task failed"
);
continue;
}
};
for (topic, measurement) in per_topic {
let key = (group_id.clone(), topic.clone());
match measurement {
Ok(lag) => {
transport_failures.remove(&key);
lag_cache.update(
&group_id,
&topic,
KafkaLagMeasurement {
lag,
last_updated: Instant::now(),
},
);
}
Err(KafkaError::MetadataFetch(RDKafkaErrorCode::BrokerTransportFailure)) => {
let count = transport_failures.entry(key).or_insert(0);
*count += 1;
if *count <= TRANSPORT_WARN_THRESHOLD {
debug!(
consumer_group = %group_id,
topic = %topic,
attempt = *count,
"Broker transport not yet ready; will retry consumer lag"
);
} else {
warn!(
consumer_group = %group_id,
topic = %topic,
consecutive_failures = *count,
"Failed to refresh consumer lag: broker transport failure"
);
}
}
Err(err) => {
transport_failures.remove(&key);
warn!(
consumer_group = %group_id,
topic = %topic,
error = %err,
"Failed to refresh consumer lag"
);
}
};
}
}
}
});
*guard = Some(handle);
}
fn stop_tasks(&self) {
self.state.workers.stop_all();
}
}
impl Drop for KafkaEventBusBackend {
fn drop(&mut self) {
if self.state.workers.running.swap(false, Ordering::SeqCst) {
self.stop_tasks();
}
for _ in 0..10 {
self.state.producer.poll(Duration::from_millis(10));
}
let _ = self.state.producer.flush(Duration::from_millis(250));
}
}
#[async_trait]
impl EventBusBackend for KafkaEventBusBackend {
fn clone_box(&self) -> Box<dyn EventBusBackend> {
Box::new(self.clone())
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
fn backend_name(&self) -> &'static str {
"kafka"
}
fn configure_plugin(&self, app: &mut App) {
app.add_message::<KafkaCommitResultMessage>();
app.init_resource::<KafkaCommitResultStats>();
app.add_systems(
PreUpdate,
(
kafka_commit_result_dispatch_system,
kafka_commit_result_stats_system,
)
.chain(),
);
}
fn setup_plugin(&self, _world: &mut World) -> BackendPluginSetup {
BackendPluginSetup {
ready_topics: self.configured_topics(),
message_stream: self.take_receiver(),
manual_commit: Some(Box::new(KafkaManualCommitHandle::new(
self.commit_sender(),
self.take_commit_results(),
))),
lag_reporting: Some(Box::new(KafkaLagHandle::new(self.lag_cache()))),
kafka_topology: Some(self.state.config.topology.clone()),
..BackendPluginSetup::default()
}
}
fn augment_metrics(&self, metrics: &mut ConsumerMetrics) {
metrics.dropped_messages = self.dropped_count();
}
fn apply_event_bindings(&self, app: &mut App) {
for binding in self.state.config.topology.event_bindings() {
binding.apply(app);
}
}
async fn connect(&mut self) -> bool {
if self
.state
.workers
.running
.compare_exchange(false, true, Ordering::SeqCst, Ordering::Relaxed)
.is_err()
{
return true;
}
self.spawn_runtime_tasks();
true
}
async fn disconnect(&mut self) -> bool {
if !self.state.workers.running.swap(false, Ordering::SeqCst) {
return true;
}
self.stop_tasks();
true
}
fn try_send_serialized(
&self,
event_json: &[u8],
topic: &str,
options: SendOptions<'_>,
failure_handler: Option<Arc<DeliveryFailureCallback>>,
) -> bool {
let handler = failure_handler
.unwrap_or_else(|| Arc::new(DeliveryFailureCallback::new(|_failure| {})));
let mut record = BaseRecord::<[u8], [u8], Arc<DeliveryFailureCallback>>::with_opaque_to(
topic,
handler.clone(),
)
.payload(event_json);
if let Some(key) = options.partition_key {
record = record.key(key.as_bytes());
}
if let Some(backend_data) = options.backend.as_any() {
if let Some(headers) = backend_data.downcast_ref::<HashMap<String, String>>() {
if !headers.is_empty() {
let mut kafka_headers = OwnedHeaders::new();
for (key, value) in headers {
kafka_headers = kafka_headers.insert(Header {
key,
value: Some(value.as_bytes()),
});
}
record = record.headers(kafka_headers);
}
}
}
match self.state.producer.send(record) {
Ok(_) => true,
Err((err, _)) => {
warn!(target = %topic, error = %err, "Failed to enqueue Kafka message");
handler.call(DeliveryFailure {
backend: "kafka",
kind: EventBusErrorType::DeliveryFailure,
topic: topic.to_string(),
error: err.to_string(),
metadata: None,
});
false
}
}
}
async fn receive_serialized(&self, topic: &str, options: ReceiveOptions<'_>) -> Vec<Vec<u8>> {
self.drain_matching_messages(topic, options.consumer_group)
}
async fn flush(&self) -> Result<(), String> {
match self.state.producer.flush(Duration::from_secs(30)) {
Ok(_) => Ok(()),
Err(err) => Err(format!("Flush failed: {err}")),
}
}
}
#[async_trait]
impl ManualCommitController for KafkaEventBusBackend {
async fn enable_manual_commits(&mut self, group_id: &str) -> Result<(), String> {
let consumers = self.state.consumers.lock().unwrap();
if let Some(runtime) = consumers.get(group_id) {
if runtime.manual_commit {
Ok(())
} else {
Err(format!(
"Consumer group '{}' is configured for auto commits",
group_id
))
}
} else {
Err(format!("Consumer group '{}' not found", group_id))
}
}
async fn commit_offset(&self, topic: &str, partition: i32, offset: i64) -> Result<(), String> {
let consumers = self.state.consumers.lock().unwrap();
let runtime = consumers
.values()
.find(|cg| cg.manual_commit && cg.topics.contains(&topic.to_string()))
.cloned()
.ok_or_else(|| format!("No manual commit consumer group for topic '{}'", topic))?;
drop(consumers);
let request = KafkaCommitRequest {
topic: topic.to_string(),
partition,
offset,
consumer_group: runtime.group_id.clone(),
};
match self.state.commits.request_tx.try_send(request.clone()) {
Ok(_) => Ok(()),
Err(crossbeam_channel::TrySendError::Full(_)) => commit_offset_sync(&runtime, &request),
Err(crossbeam_channel::TrySendError::Disconnected(_)) => {
Err("Kafka commit queue is not available".to_string())
}
}
}
}
#[async_trait]
impl LagReportingBackend for KafkaEventBusBackend {
async fn get_consumer_lag(&self, topic: &str, group_id: &str) -> Result<i64, String> {
let connection = self.state.config.connection.clone();
let group_id = group_id.to_string();
let topic = topic.to_string();
tokio::task::spawn_blocking(move || {
let consumer = build_lag_consumer(&connection, &group_id)?;
compute_topic_lag(&consumer, &topic)
})
.await
.map_err(|e| e.to_string())?
.map_err(|e| e.to_string())
}
}
fn apply_connection_settings(cfg: &mut ClientConfig, connection: &KafkaConnectionConfig) {
cfg.set("bootstrap.servers", connection.bootstrap_servers());
if let Some(client_id) = connection.client_id() {
cfg.set("client.id", client_id);
}
for (key, value) in connection.additional_config() {
cfg.set(key, value);
}
}
fn build_producer(
connection: &KafkaConnectionConfig,
) -> Result<BaseProducer<EventBusProducerContext>, KafkaError> {
let mut cfg = ClientConfig::new();
apply_connection_settings(&mut cfg, connection);
if !connection
.additional_config()
.contains_key("message.timeout.ms")
{
cfg.set("message.timeout.ms", connection.timeout_ms().to_string());
}
cfg.create_with_context(EventBusProducerContext)
}
fn prepare_kafka_topology(
connection: &KafkaConnectionConfig,
topology: &KafkaTopologyConfig,
) -> Result<(), BackendConfigError> {
prepare_topics(connection, topology)
.map_err(|reason| BackendConfigError::new("kafka", reason))?;
prepare_consumer_groups(connection, topology)
.map_err(|reason| BackendConfigError::new("kafka", reason))?;
Ok(())
}
fn prepare_topics(
connection: &KafkaConnectionConfig,
topology: &KafkaTopologyConfig,
) -> Result<(), String> {
if topology.topics().is_empty() {
return Ok(());
}
let unique_specs = deduplicate_topic_specs(topology.topics())?;
let mut cfg = ClientConfig::new();
apply_connection_settings(&mut cfg, connection);
cfg.set("allow.auto.create.topics", "false");
let admin = cfg
.create::<AdminClient<DefaultClientContext>>()
.map_err(|err| format!("Failed to create Kafka admin client: {err}"))?;
let mut provision_specs = Vec::new();
let mut validate_specs = Vec::new();
for spec in &unique_specs {
match spec.mode {
TopologyMode::Provision => provision_specs.push(spec),
TopologyMode::Validate => validate_specs.push(spec),
}
}
if !provision_specs.is_empty() {
let options = AdminOptions::new();
for spec in provision_specs {
let partitions = spec.partitions.unwrap_or(1);
let replication = spec
.replication
.map(|r| TopicReplication::Fixed(r as i32))
.unwrap_or(TopicReplication::Fixed(1));
let new_topic = NewTopic::new(&spec.name, partitions, replication);
match runtime::block_on(admin.create_topics(std::iter::once(&new_topic), &options)) {
Ok(results) => {
for result in results {
match result {
Ok(_topic) => {}
Err((_topic, RDKafkaErrorCode::TopicAlreadyExists)) => {}
Err((topic, code)) => {
return Err(format!(
"Topic creation failed for '{topic}': {code:?}"
));
}
}
}
}
Err(err) => {
let msg = err.to_string();
if !msg.contains("TopicAlreadyExists") {
return Err(format!("Topic creation failed for '{}': {msg}", spec.name));
}
}
}
}
}
if !validate_specs.is_empty() {
for spec in validate_specs {
let metadata = admin
.inner()
.fetch_metadata(Some(&spec.name), Duration::from_secs(5))
.map_err(|err| {
format!(
"Failed to fetch metadata for Kafka topic '{}': {err}",
spec.name
)
})?;
let topic_metadata = metadata
.topics()
.iter()
.find(|topic| topic.name() == spec.name)
.ok_or_else(|| {
format!(
"Kafka topic '{}' metadata missing during validation",
spec.name
)
})?;
if let Some(error) = topic_metadata.error() {
return Err(format!(
"Kafka topic '{}' reported error during validation: {:?}",
spec.name, error
));
}
}
}
Ok(())
}
fn deduplicate_topic_specs(specs: &[KafkaTopicSpec]) -> Result<Vec<KafkaTopicSpec>, String> {
let mut unique = BTreeMap::<String, KafkaTopicSpec>::new();
for spec in specs {
match unique.get(&spec.name) {
None => {
unique.insert(spec.name.clone(), spec.clone());
}
Some(existing) => {
if existing.partitions != spec.partitions
|| existing.replication != spec.replication
|| existing.mode != spec.mode
{
return Err(format!(
"Kafka topology contains duplicate topic '{}' with conflicting settings (existing: partitions={:?} replication={:?} mode={:?}; new: partitions={:?} replication={:?} mode={:?})",
spec.name,
existing.partitions,
existing.replication,
existing.mode,
spec.partitions,
spec.replication,
spec.mode
));
}
}
}
}
Ok(unique.into_values().collect())
}
fn prepare_consumer_groups(
connection: &KafkaConnectionConfig,
topology: &KafkaTopologyConfig,
) -> Result<(), String> {
if topology.consumer_groups().is_empty() {
return Ok(());
}
for (group_id, spec) in topology.consumer_groups() {
if spec.mode == TopologyMode::Provision {
provision_consumer_group(connection, group_id, spec)?;
}
}
let mut validator = build_validation_consumer(connection)?;
for (group_id, spec) in topology.consumer_groups() {
validate_consumer_group(&mut validator, group_id).map_err(|err| match spec.mode {
TopologyMode::Provision => format!(
"Kafka consumer group '{}' could not be provisioned: {err}",
group_id
),
TopologyMode::Validate => format!(
"Kafka consumer group '{}' is missing while configured with TopologyMode::Validate: {err}",
group_id
),
})?;
}
Ok(())
}
fn provision_consumer_group(
connection: &KafkaConnectionConfig,
group_id: &str,
spec: &KafkaConsumerGroupSpec,
) -> Result<(), String> {
let consumer = build_consumer(connection, group_id, spec).map_err(|err| {
format!(
"Failed to build Kafka consumer while provisioning group '{}': {err}",
group_id
)
})?;
if !spec.topics.is_empty() {
let topic_refs: Vec<&str> = spec.topics.iter().map(String::as_str).collect();
consumer.subscribe(&topic_refs).map_err(|err| {
format!(
"Failed to subscribe while provisioning Kafka consumer group '{}': {err}",
group_id
)
})?;
}
const REGISTRATION_ATTEMPTS: usize = 30;
for attempt in 0..REGISTRATION_ATTEMPTS {
if let Some(Err(err)) = consumer.poll(Duration::from_millis(100)) {
warn!(
group = %group_id,
error = %err,
"Kafka consumer poll error while provisioning group"
);
}
match consumer
.client()
.fetch_group_list(Some(group_id), Duration::from_secs(2))
{
Ok(metadata) => {
let found = metadata
.groups()
.iter()
.any(|group| group.name() == group_id);
if found {
return Ok(());
}
}
Err(err) => {
warn!(
group = %group_id,
error = %err,
"Kafka consumer group metadata fetch failed during provisioning"
);
}
}
if attempt + 1 < REGISTRATION_ATTEMPTS {
thread::sleep(Duration::from_millis(500));
}
}
Err(format!(
"Kafka consumer group '{}' not reported by broker after provisioning",
group_id
))
}
fn build_validation_consumer(connection: &KafkaConnectionConfig) -> Result<BaseConsumer, String> {
let mut cfg = ClientConfig::new();
apply_connection_settings(&mut cfg, connection);
cfg.set("group.id", "__bevy_event_bus.topology_validator");
cfg.set("enable.auto.commit", "false");
cfg.set("allow.auto.create.topics", "false");
cfg.set("socket.timeout.ms", "5000");
cfg.create()
.map_err(|err| format!("Failed to create Kafka validation consumer: {err}"))
}
fn validate_consumer_group(consumer: &mut BaseConsumer, group_id: &str) -> Result<(), String> {
const VALIDATION_RETRIES: usize = 10;
for attempt in 0..VALIDATION_RETRIES {
let metadata = consumer
.client()
.fetch_group_list(Some(group_id), Duration::from_secs(5))
.map_err(|err| {
format!(
"Failed to fetch metadata for Kafka consumer group '{}': {err}",
group_id
)
})?;
let found = metadata
.groups()
.iter()
.any(|group| group.name() == group_id);
if found {
return Ok(());
}
if attempt + 1 < VALIDATION_RETRIES {
thread::sleep(Duration::from_millis(100));
}
}
Err("group not reported by broker".to_string())
}
fn build_consumers(
connection: &KafkaConnectionConfig,
topology: &KafkaTopologyConfig,
) -> HashMap<String, ConsumerRuntime> {
let mut map = HashMap::new();
for (group_id, spec) in topology.consumer_groups() {
let consumer = build_consumer(connection, group_id, spec)
.unwrap_or_else(|e| panic!("Failed to build consumer for group {}: {e}", group_id));
if !spec.topics.is_empty() {
let topic_refs: Vec<&str> = spec.topics.iter().map(String::as_str).collect();
consumer
.subscribe(&topic_refs)
.unwrap_or_else(|e| panic!("Failed to subscribe group {}: {e}", group_id));
}
map.insert(
group_id.clone(),
ConsumerRuntime {
consumer: Arc::new(consumer),
manual_commit: spec.manual_commits,
group_id: group_id.clone(),
topics: spec.topics.clone(),
},
);
}
map
}
fn build_consumer(
connection: &KafkaConnectionConfig,
group_id: &str,
spec: &KafkaConsumerGroupSpec,
) -> Result<BaseConsumer, KafkaError> {
let mut cfg = ClientConfig::new();
apply_connection_settings(&mut cfg, connection);
cfg.set("group.id", group_id);
cfg.set(
"enable.auto.commit",
if spec.manual_commits { "false" } else { "true" },
);
cfg.set("session.timeout.ms", "6000");
cfg.set(
"auto.offset.reset",
match spec.initial_offset {
KafkaInitialOffset::Earliest => "earliest",
KafkaInitialOffset::Latest => "latest",
KafkaInitialOffset::None => "none",
},
);
if let Some(client_id) = connection.client_id() {
cfg.set("client.id", format!("{}_{}", client_id, group_id));
}
cfg.create()
}
fn consumer_loop(
runtime: ConsumerRuntime,
tx: Sender<IncomingMessage>,
running: Arc<AtomicBool>,
producer: Arc<BaseProducer<EventBusProducerContext>>,
dropped: Arc<AtomicUsize>,
) {
let mut buffered_messages: VecDeque<IncomingMessage> = VecDeque::new();
let pause_partitions = |consumer: &BaseConsumer, group_id: &str| {
match consumer.assignment() {
Ok(assignment) => {
if let Err(err) = consumer.pause(&assignment) {
warn!(error = %err, consumer_group = %group_id, "Failed to pause Kafka consumer");
}
}
Err(err) => {
warn!(error = %err, consumer_group = %group_id, "Failed to fetch Kafka assignment for pause");
}
}
};
let resume_partitions = |consumer: &BaseConsumer, group_id: &str| {
match consumer.assignment() {
Ok(assignment) => {
if let Err(err) = consumer.resume(&assignment) {
warn!(error = %err, consumer_group = %group_id, "Failed to resume Kafka consumer");
}
}
Err(err) => {
warn!(error = %err, consumer_group = %group_id, "Failed to fetch Kafka assignment for resume");
}
}
};
while running.load(Ordering::Relaxed) {
producer.poll(Duration::from_millis(0));
let mut sent_any = false;
while let Some(message) = buffered_messages.pop_front() {
match tx.try_send(message) {
Ok(()) => {
sent_any = true;
}
Err(TrySendError::Full(returned)) => {
buffered_messages.push_front(returned);
break;
}
Err(TrySendError::Disconnected(_returned)) => {
dropped.fetch_add(1, Ordering::Relaxed);
return;
}
}
}
let mut backpressured = !buffered_messages.is_empty();
if backpressured {
pause_partitions(&runtime.consumer, &runtime.group_id);
} else {
resume_partitions(&runtime.consumer, &runtime.group_id);
}
match runtime
.consumer
.poll(if backpressured {
Duration::from_millis(0)
} else {
Duration::from_millis(50)
})
{
Some(Ok(message)) => {
let payload = match message.payload() {
Some(payload) => payload.to_vec(),
None => Vec::new(),
};
let key = message
.key()
.map(|k| String::from_utf8_lossy(k).into_owned());
let mut headers_map = HashMap::new();
if let Some(headers) = message.headers() {
for header in headers.iter() {
if let Some(value) = header.value {
if let Ok(str_value) = String::from_utf8(value.to_vec()) {
headers_map.insert(header.key.to_string(), str_value);
}
}
}
}
let kafka_metadata = KafkaMetadata {
topic: message.topic().to_string(),
partition: message.partition(),
offset: message.offset(),
consumer_group: Some(runtime.group_id.clone()),
manual_commit: runtime.manual_commit,
headers: headers_map,
};
let msg = IncomingMessage::plain(
message.topic().to_string(),
payload,
key,
Some(Box::new(kafka_metadata)),
);
match tx.try_send(msg) {
Ok(()) => {}
Err(TrySendError::Full(returned)) => {
buffered_messages.push_back(returned);
backpressured = true;
pause_partitions(&runtime.consumer, &runtime.group_id);
}
Err(TrySendError::Disconnected(_returned)) => {
dropped.fetch_add(1, Ordering::Relaxed);
return;
}
}
}
Some(Err(err)) => {
warn!("Kafka consumer poll error: {err}");
}
None => {}
}
if backpressured && !sent_any {
std::thread::sleep(Duration::from_millis(5));
}
}
}
fn commit_offset_sync(runtime: &ConsumerRuntime, req: &KafkaCommitRequest) -> Result<(), String> {
if !runtime.manual_commit {
return Err(format!(
"Consumer group '{}' is not configured for manual commits",
runtime.group_id
));
}
let mut tpl = TopicPartitionList::new();
tpl.add_partition_offset(&req.topic, req.partition, Offset::Offset(req.offset + 1))
.map_err(|err| err.to_string())?;
runtime
.consumer
.commit(&tpl, CommitMode::Sync)
.map_err(|err| err.to_string())
}
fn build_lag_consumer(
connection: &KafkaConnectionConfig,
group_id: &str,
) -> Result<BaseConsumer, KafkaError> {
let mut cfg = ClientConfig::new();
apply_connection_settings(&mut cfg, connection);
cfg.set("group.id", group_id);
cfg.set("enable.auto.commit", "false");
cfg.set("session.timeout.ms", "6000");
cfg.create()
}
fn compute_group_lag(
connection: &KafkaConnectionConfig,
group_id: &str,
topics: &[String],
) -> Result<Vec<(String, Result<i64, KafkaError>)>, KafkaError> {
let consumer = build_lag_consumer(connection, group_id)?;
Ok(topics
.iter()
.map(|topic| (topic.clone(), compute_topic_lag(&consumer, topic)))
.collect())
}
fn compute_topic_lag(consumer: &BaseConsumer, topic: &str) -> Result<i64, KafkaError> {
let metadata = consumer
.fetch_metadata(Some(topic), Duration::from_secs(5))?;
let topic_metadata = metadata
.topics()
.iter()
.find(|t| t.name() == topic)
.ok_or(KafkaError::MetadataFetch(RDKafkaErrorCode::UnknownTopicOrPartition))?;
let mut total_lag = 0i64;
for partition in topic_metadata.partitions() {
let partition_id = partition.id();
let (low, high) = consumer
.fetch_watermarks(topic, partition_id, Duration::from_secs(5))?;
let mut tpl = TopicPartitionList::new();
tpl.add_partition_offset(topic, partition_id, Offset::Invalid)?;
let committed = consumer
.committed_offsets(tpl, Duration::from_secs(5))?;
if let Some(elem) = committed.elements().first() {
match elem.offset() {
Offset::Offset(committed_offset) => {
let lag = high - committed_offset;
total_lag += lag.max(0);
}
_ => {
let lag = high - low;
total_lag += lag.max(0);
}
}
}
}
Ok(total_lag)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_producer_context_creation() {
let context = EventBusProducerContext;
assert_eq!(std::mem::size_of_val(&context), 0);
}
#[test]
fn test_deduplicate_topic_specs_allows_exact_duplicates() {
let specs = vec![
KafkaTopicSpec::new("topic-a").partitions(3).replication(1),
KafkaTopicSpec::new("topic-a").partitions(3).replication(1),
];
let unique = deduplicate_topic_specs(&specs).unwrap();
assert_eq!(unique.len(), 1);
assert_eq!(unique[0].name, "topic-a");
}
#[test]
fn test_deduplicate_topic_specs_rejects_conflicting_duplicates() {
let specs = vec![
KafkaTopicSpec::new("topic-a").partitions(3).replication(1),
KafkaTopicSpec::new("topic-a").partitions(1).replication(1),
];
let err = deduplicate_topic_specs(&specs).unwrap_err();
assert!(err.contains("topic-a"));
assert!(err.contains("conflicting"));
}
#[test]
fn test_deduplicate_topic_specs_rejects_different_modes() {
let specs = vec![
KafkaTopicSpec::new("topic-a").mode(TopologyMode::Provision),
KafkaTopicSpec::new("topic-a").mode(TopologyMode::Validate),
];
let err = deduplicate_topic_specs(&specs).unwrap_err();
assert!(err.contains("topic-a"));
assert!(err.contains("conflicting"));
}
}