horizon-sdk 16.0.0

//! Kafka producer wrapper with SASL/SSL and idempotent delivery.

use std::fmt;
use std::time::Duration;

use rdkafka::config::ClientConfig;
use rdkafka::producer::{FutureProducer, FutureRecord, Producer as _};
use tracing::{debug, error, info, warn};

use crate::types::config::KafkaConfig;
use crate::types::constants::{
    DEFAULT_KAFKA_MESSAGE_MAX_BYTES, DEFAULT_KAFKA_PRODUCE_TIMEOUT_SECS,
};
use crate::types::error::{KafkaError, Result};

/// Async Kafka producer for writing Arrow IPC payloads to Horizon topics.
///
/// Configured for exactly-once semantics within a single producer session:
/// idempotence enabled, acks=all, LZ4 compression.
pub struct KafkaProducer {
    /// The underlying rdkafka async producer.
    producer: FutureProducer,
    /// Kafka topic name this producer writes to.
    topic: String,
}

impl KafkaProducer {
    /// Graceful shutdown: flush with 30-second timeout.
    #[allow(
        clippy::cognitive_complexity,
        reason = "3-arm match with tracing macros expands to more branches than the logic implies"
    )]
    pub fn close(&self) {
        info!("closing Kafka producer, flushing pending messages...");
        match self.flush(Duration::from_secs(30)) {
            Ok(0_i32) => info!("all messages delivered successfully"),
            Ok(remaining) => {
                warn!(remaining, "messages were not delivered before timeout");
            }
            Err(error) => {
                warn!(error = %error, "Kafka producer flush failed during close");
            }
        }
    }

    /// Flush all pending messages, waiting up to `timeout`.
    ///
    /// On success, returns the number of messages still in flight after
    /// the timeout (matching `FutureProducer::in_flight_count`'s `i32`
    /// return type); a value of 0 means all produced messages are durable
    /// in Kafka.
    ///
    /// # Errors
    ///
    /// Returns `KafkaError::Flush` if the underlying rdkafka flush fails
    /// (broker disconnect, client shutdown, etc.). The in-flight count is
    /// intentionally not reported in this case because rdkafka may have
    /// stopped tracking queued messages.
    pub fn flush(&self, timeout: Duration) -> Result<i32> {
        self.producer
            .flush(timeout)
            .map_err(|source| KafkaError::Flush { source })?;
        Ok(self.producer.in_flight_count())
    }

    /// Create a new Kafka producer from SDK configuration.
    ///
    /// # Errors
    ///
    /// Returns `KafkaError::Client` if the producer cannot be created
    /// (e.g. invalid broker address or TLS configuration).
    pub fn new(config: &KafkaConfig, topic: &str) -> Result<Self> {
        let mut client_config = ClientConfig::new();
        client_config
            .set("bootstrap.servers", &config.bootstrap_servers)
            .set("security.protocol", &config.security_protocol)
            .set("sasl.mechanism", &config.sasl_mechanism)
            .set("compression.type", &config.compression_type)
            .set("acks", "all")
            .set("enable.idempotence", "true")
            .set("max.in.flight.requests.per.connection", "5")
            .set("message.max.bytes", DEFAULT_KAFKA_MESSAGE_MAX_BYTES)
            .set("retries", "5")
            .set("retry.backoff.ms", "1000");

        if let Some(username) = &config.sasl_username {
            client_config.set("sasl.username", username);
        }
        if let Some(password) = &config.sasl_password {
            client_config.set("sasl.password", password);
        }

        let producer: FutureProducer = client_config
            .create()
            .map_err(|source| KafkaError::Client { source })?;

        info!(topic = topic, "initialized Kafka producer");

        Ok(Self {
            producer,
            topic: topic.to_owned(),
        })
    }

    /// Produce a message to the configured Kafka topic.
    ///
    /// Returns when the message is acknowledged by the broker (acks=all).
    /// The per-message delivery wait is bounded by
    /// `DEFAULT_KAFKA_PRODUCE_TIMEOUT_SECS`.
    ///
    /// # Errors
    ///
    /// Returns `KafkaError::Delivery` if the message cannot be delivered.
    pub async fn produce(&self, key: &[u8], value: &[u8]) -> Result<()> {
        let record = FutureRecord::to(&self.topic).key(key).payload(value);

        let produce_timeout = Duration::from_secs(DEFAULT_KAFKA_PRODUCE_TIMEOUT_SECS);
        match self.producer.send(record, produce_timeout).await {
            Ok(delivery) => {
                debug!(
                    topic = self.topic.as_str(),
                    partition = delivery.partition,
                    offset = delivery.offset,
                    "message delivered"
                );
                Ok(())
            }
            Err((source, _)) => {
                error!(
                    topic = self.topic.as_str(),
                    payload_bytes = value.len(),
                    error = %source,
                    "message delivery failed"
                );
                Err(KafkaError::Delivery { source }.into())
            }
        }
    }

    /// The topic this producer writes to.
    #[must_use]
    pub fn topic(&self) -> &str {
        &self.topic
    }
}

impl fmt::Debug for KafkaProducer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KafkaProducer")
            .field("topic", &self.topic)
            .finish_non_exhaustive()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, reason = "test code uses unwrap for clarity")]
mod tests {
    use super::*;

    #[test]
    fn kafka_config_defaults_produce_valid_settings() {
        let config = KafkaConfig {
            bootstrap_servers: "localhost:9092".to_owned(),
            security_protocol: "PLAINTEXT".to_owned(),
            sasl_mechanism: "PLAIN".to_owned(),
            sasl_username: None,
            sasl_password: None,
            compression_type: "lz4".to_owned(),
            data_row_topic: "horizon.data_row".to_owned(),
            metadata_row_topic: "horizon.metadata_row".to_owned(),
        };

        let producer = KafkaProducer::new(&config, &config.data_row_topic);
        assert!(producer.is_ok());
        let producer = producer.unwrap();
        assert_eq!(producer.topic(), "horizon.data_row");
    }

    #[test]
    fn debug_impl_does_not_leak_internals() {
        let config = KafkaConfig {
            bootstrap_servers: "localhost:9092".to_owned(),
            security_protocol: "PLAINTEXT".to_owned(),
            sasl_mechanism: "PLAIN".to_owned(),
            sasl_username: None,
            sasl_password: None,
            compression_type: "lz4".to_owned(),
            data_row_topic: "horizon.data_row".to_owned(),
            metadata_row_topic: "horizon.metadata_row".to_owned(),
        };
        let producer = KafkaProducer::new(&config, "test.topic").unwrap();
        let debug = format!("{producer:?}");
        assert!(debug.contains("test.topic"));
        assert!(debug.contains("KafkaProducer"));
    }
}