chore(producer): Add produce time by topic (#423)

untitaker · web-flow · commit 8ab18aa03935 · 2025-06-23T14:49:50.000+02:00
diff --git a/arroyo/backends/kafka/__init__.py b/arroyo/backends/kafka/__init__.py
@@ -1,9 +1,14 @@
-from .configuration import build_kafka_configuration, build_kafka_consumer_configuration
+from .configuration import (
+    build_kafka_configuration,
+    build_kafka_consumer_configuration,
+    build_kafka_producer_configuration,
+)
 from .consumer import KafkaConsumer, KafkaPayload, KafkaProducer
 
 __all__ = [
     "build_kafka_configuration",
     "build_kafka_consumer_configuration",
+    "build_kafka_producer_configuration",
     "KafkaConsumer",
     "KafkaPayload",
     "KafkaProducer",
diff --git a/arroyo/backends/kafka/configuration.py b/arroyo/backends/kafka/configuration.py
@@ -49,6 +49,77 @@ def stats_callback(stats_json: str) -> None:
     )
 
 
+def producer_stats_callback(stats_json: str) -> None:
+    stats = json.loads(stats_json)
+    metrics = get_metrics()
+
+    # Extract broker-level int_latency metrics
+    brokers = stats.get("brokers", {})
+    for broker_id, broker_stats in brokers.items():
+        int_latency = broker_stats.get("int_latency", {})
+        if int_latency:
+            p99_latency_ms = int_latency.get("p99", 0) / 1000.0
+            metrics.timing(
+                "arroyo.producer.librdkafka.p99_int_latency",
+                p99_latency_ms,
+                tags={"broker_id": str(broker_id)},
+            )
+            avg_latency_ms = int_latency.get("avg", 0) / 1000.0
+            metrics.timing(
+                "arroyo.producer.librdkafka.avg_int_latency",
+                avg_latency_ms,
+                tags={"broker_id": str(broker_id)},
+            )
+
+        outbuf_latency = broker_stats.get("outbuf_latency", {})
+        if outbuf_latency:
+            p99_latency_ms = outbuf_latency.get("p99", 0) / 1000.0
+            metrics.timing(
+                "arroyo.producer.librdkafka.p99_outbuf_latency",
+                p99_latency_ms,
+                tags={"broker_id": str(broker_id)},
+            )
+            avg_latency_ms = outbuf_latency.get("avg", 0) / 1000.0
+            metrics.timing(
+                "arroyo.producer.librdkafka.avg_outbuf_latency",
+                avg_latency_ms,
+                tags={"broker_id": str(broker_id)},
+            )
+
+        rtt = broker_stats.get("rtt", {})
+        if rtt:
+            p99_rtt_ms = rtt.get("p99", 0) / 1000.0
+            metrics.timing(
+                "arroyo.producer.librdkafka.p99_rtt",
+                p99_rtt_ms,
+                tags={"broker_id": str(broker_id)},
+            )
+            avg_rtt_ms = rtt.get("avg", 0) / 1000.0
+            metrics.timing(
+                "arroyo.producer.librdkafka.avg_rtt",
+                avg_rtt_ms,
+                tags={"broker_id": str(broker_id)},
+            )
+
+
+def build_kafka_producer_configuration(
+    default_config: Mapping[str, Any],
+    bootstrap_servers: Optional[Sequence[str]] = None,
+    override_params: Optional[Mapping[str, Any]] = None,
+) -> KafkaBrokerConfig:
+    broker_config = build_kafka_configuration(
+        default_config, bootstrap_servers, override_params
+    )
+
+    broker_config.update(
+        {
+            "statistics.interval.ms": STATS_COLLECTION_FREQ_MS,
+            "stats_cb": producer_stats_callback,
+        }
+    )
+    return broker_config
+
+
 def build_kafka_consumer_configuration(
     default_config: Mapping[str, Any],
     group_id: str,
diff --git a/arroyo/utils/metric_defs.py b/arroyo/utils/metric_defs.py
@@ -106,4 +106,22 @@
     "arroyo.consumer.dlq_buffer.exceeded",
     # Gauge: Number of partitions being tracked in the DLQ buffer
     "arroyo.consumer.dlq_buffer.assigned_partitions",
+    # Time: Internal producer queue latency from librdkafka statistics.
+    # Tagged by broker_id.
+    "arroyo.producer.librdkafka.p99_int_latency",
+    # Time: Output buffer latency from librdkafka statistics.
+    # Tagged by broker_id.
+    "arroyo.producer.librdkafka.p99_outbuf_latency",
+    # Time: Round-trip time to brokers from librdkafka statistics.
+    # Tagged by broker_id.
+    "arroyo.producer.librdkafka.p99_rtt",
+    # Time: Average internal producer queue latency from librdkafka statistics.
+    # Tagged by broker_id.
+    "arroyo.producer.librdkafka.avg_int_latency",
+    # Time: Average output buffer latency from librdkafka statistics.
+    # Tagged by broker_id.
+    "arroyo.producer.librdkafka.avg_outbuf_latency",
+    # Time: Average round-trip time to brokers from librdkafka statistics.
+    # Tagged by broker_id.
+    "arroyo.producer.librdkafka.avg_rtt",
 ]
diff --git a/rust-arroyo/src/backends/kafka/config.rs b/rust-arroyo/src/backends/kafka/config.rs
@@ -77,7 +77,11 @@ impl KafkaConfig {
         bootstrap_servers: Vec<String>,
         override_params: Option<HashMap<String, String>>,
     ) -> Self {
-        let config = KafkaConfig::new_config(bootstrap_servers, None);
+        let mut config = KafkaConfig::new_config(bootstrap_servers, None);
+        config.config_map.insert(
+            "statistics.interval.ms".to_string(),
+            STATS_COLLECTION_FREQ_MS.to_string(),
+        );
 
         apply_override_params(config, override_params)
     }
diff --git a/rust-arroyo/src/backends/kafka/producer.rs b/rust-arroyo/src/backends/kafka/producer.rs
@@ -2,18 +2,89 @@ use crate::backends::kafka::config::KafkaConfig;
 use crate::backends::kafka::types::KafkaPayload;
 use crate::backends::Producer as ArroyoProducer;
 use crate::backends::ProducerError;
+use crate::timer;
 use crate::types::TopicOrPartition;
+use rdkafka::client::ClientContext;
 use rdkafka::config::ClientConfig;
-use rdkafka::producer::{DefaultProducerContext, ThreadedProducer};
+use rdkafka::producer::{
+    DeliveryResult, ProducerContext as RdkafkaProducerContext, ThreadedProducer,
+};
+use rdkafka::Statistics;
+use std::time::Duration;
+
+pub struct ProducerContext;
+
+impl ClientContext for ProducerContext {
+    fn stats(&self, stats: Statistics) {
+        for (broker_id, broker_stats) in &stats.brokers {
+            if let Some(int_latency) = &broker_stats.int_latency {
+                let p99_latency_ms = int_latency.p99 as f64 / 1000.0;
+                timer!(
+                    "arroyo.producer.librdkafka.p99_int_latency",
+                    Duration::from_millis(p99_latency_ms as u64),
+                    "broker_id" => broker_id.to_string()
+                );
+                // Also record average latency
+                let avg_latency_ms = int_latency.avg as f64 / 1000.0;
+                timer!(
+                    "arroyo.producer.librdkafka.avg_int_latency",
+                    Duration::from_millis(avg_latency_ms as u64),
+                    "broker_id" => broker_id.to_string()
+                );
+            }
+            if let Some(outbuf_latency) = &broker_stats.outbuf_latency {
+                let p99_latency_ms = outbuf_latency.p99 as f64 / 1000.0;
+                timer!(
+                    "arroyo.producer.librdkafka.p99_outbuf_latency",
+                    Duration::from_millis(p99_latency_ms as u64),
+                    "broker_id" => broker_id.to_string()
+                );
+                let avg_latency_ms = outbuf_latency.avg as f64 / 1000.0;
+                timer!(
+                    "arroyo.producer.librdkafka.avg_outbuf_latency",
+                    Duration::from_millis(avg_latency_ms as u64),
+                    "broker_id" => broker_id.to_string()
+                );
+            }
+            if let Some(rtt) = &broker_stats.rtt {
+                let p99_rtt_ms = rtt.p99 as f64 / 1000.0;
+                timer!(
+                    "arroyo.producer.librdkafka.p99_rtt",
+                    Duration::from_millis(p99_rtt_ms as u64),
+                    "broker_id" => broker_id.to_string()
+                );
+                let avg_rtt_ms = rtt.avg as f64 / 1000.0;
+                timer!(
+                    "arroyo.producer.librdkafka.avg_rtt",
+                    Duration::from_millis(avg_rtt_ms as u64),
+                    "broker_id" => broker_id.to_string()
+                );
+            }
+        }
+    }
+}
+
+impl RdkafkaProducerContext for ProducerContext {
+    type DeliveryOpaque = ();
+
+    fn delivery(
+        &self,
+        _delivery_result: &DeliveryResult<'_>,
+        _delivery_opaque: Self::DeliveryOpaque,
+    ) {
+    }
+}
 
 pub struct KafkaProducer {
-    producer: ThreadedProducer<DefaultProducerContext>,
+    producer: ThreadedProducer<ProducerContext>,
 }
 
 impl KafkaProducer {
     pub fn new(config: KafkaConfig) -> Self {
+        let context = ProducerContext;
         let config_obj: ClientConfig = config.into();
-        let threaded_producer: ThreadedProducer<_> = config_obj.create().unwrap();
+        let threaded_producer: ThreadedProducer<_> =
+            config_obj.create_with_context(context).unwrap();
 
         Self {
             producer: threaded_producer,
@@ -39,11 +110,125 @@ impl ArroyoProducer<KafkaPayload> for KafkaProducer {
 
 #[cfg(test)]
 mod tests {
-    use super::KafkaProducer;
+    use super::{KafkaProducer, ProducerContext};
     use crate::backends::kafka::config::KafkaConfig;
     use crate::backends::kafka::types::KafkaPayload;
     use crate::backends::Producer;
     use crate::types::{Topic, TopicOrPartition};
+    use rdkafka::client::ClientContext;
+    use rdkafka::statistics::{Broker, Statistics, Window};
+    use std::collections::HashMap;
+
+    fn create_test_statistics_with_all_metrics() -> Statistics {
+        let mut brokers = HashMap::new();
+        brokers.insert(
+            "1".to_string(),
+            Broker {
+                int_latency: Some(Window {
+                    p99: 2000, // microseconds -> 2.0 ms
+                    avg: 1000, // microseconds -> 1.0 ms
+                    ..Default::default()
+                }),
+                outbuf_latency: Some(Window {
+                    p99: 4000, // microseconds -> 4.0 ms
+                    avg: 2000, // microseconds -> 2.0 ms
+                    ..Default::default()
+                }),
+                rtt: Some(Window {
+                    p99: 1500, // microseconds -> 1.5 ms
+                    avg: 750,  // microseconds -> 0.75 ms
+                    ..Default::default()
+                }),
+                ..Default::default()
+            },
+        );
+
+        Statistics {
+            brokers,
+            ..Default::default()
+        }
+    }
+
+    fn create_test_statistics_with_partial_metrics() -> Statistics {
+        let mut brokers = HashMap::new();
+        brokers.insert(
+            "1".to_string(),
+            Broker {
+                int_latency: Some(Window {
+                    p99: 2000,
+                    avg: 1000,
+                    ..Default::default()
+                }),
+                outbuf_latency: Some(Window {
+                    p99: 4000,
+                    avg: 2000,
+                    ..Default::default()
+                }),
+                // No RTT data
+                ..Default::default()
+            },
+        );
+
+        Statistics {
+            brokers,
+            ..Default::default()
+        }
+    }
+
+    fn create_test_statistics_empty_brokers() -> Statistics {
+        Statistics {
+            brokers: HashMap::new(),
+            ..Default::default()
+        }
+    }
+
+    fn create_test_statistics_empty_broker_stats() -> Statistics {
+        let mut brokers = HashMap::new();
+        brokers.insert("1".to_string(), Broker::default());
+
+        Statistics {
+            brokers,
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn test_producer_context_stats_with_all_metrics() {
+        let context = ProducerContext;
+        let stats = create_test_statistics_with_all_metrics();
+
+        // This test verifies that the stats callback processes all metrics correctly
+        // We can't easily mock the timer! macro, but we can verify the method runs without panicking
+        context.stats(stats);
+    }
+
+    #[test]
+    fn test_producer_context_stats_with_partial_metrics() {
+        let context = ProducerContext;
+        let stats = create_test_statistics_with_partial_metrics();
+
+        // This test verifies that the stats callback handles missing RTT data gracefully
+        context.stats(stats);
+    }
+
+    #[test]
+    fn test_producer_context_stats_no_brokers() {
+        let context = ProducerContext;
+        let stats = create_test_statistics_empty_brokers();
+
+        // This test verifies that the stats callback handles empty broker data gracefully
+        context.stats(stats);
+    }
+
+    #[test]
+    fn test_producer_context_stats_empty_broker_stats() {
+        let context = ProducerContext;
+        let stats = create_test_statistics_empty_broker_stats();
+
+        // This test verifies that the stats callback handles broker with no metrics gracefully
+        context.stats(stats);
+    }
+
     #[test]
     fn test_producer() {
         let topic = Topic::new("test");
diff --git a/tests/backends/test_kafka_producer.py b/tests/backends/test_kafka_producer.py
