From 56d361a3e22b224bfd7de418b209ba324b7f974e Mon Sep 17 00:00:00 2001
From: Pierre Massat <pierre.massat@sentry.io>
Date: Tue, 5 Mar 2024 14:30:44 -0500
Subject: [PATCH] feat: Ingest a metric for rdkafka queue size (#342)

---
 rust-arroyo/src/backends/kafka/mod.rs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/rust-arroyo/src/backends/kafka/mod.rs b/rust-arroyo/src/backends/kafka/mod.rs
index 805ad45c..22e6991b 100644
--- a/rust-arroyo/src/backends/kafka/mod.rs
+++ b/rust-arroyo/src/backends/kafka/mod.rs
@@ -4,6 +4,7 @@ use super::CommitOffsets;
 use super::Consumer as ArroyoConsumer;
 use super::ConsumerError;
 use crate::backends::kafka::types::KafkaPayload;
+use crate::gauge;
 use crate::types::{BrokerMessage, Partition, Topic};
 use chrono::{DateTime, NaiveDateTime, Utc};
 use parking_lot::Mutex;
@@ -15,6 +16,7 @@ use rdkafka::error::KafkaError;
 use rdkafka::message::{BorrowedMessage, Message};
 use rdkafka::topic_partition_list::{Offset, TopicPartitionList};
 use rdkafka::types::{RDKafkaErrorCode, RDKafkaRespErr};
+use rdkafka::Statistics;
 use sentry::Hub;
 use std::collections::HashMap;
 use std::collections::HashSet;
@@ -167,6 +169,13 @@ impl<C: AssignmentCallbacks + Send + Sync> ClientContext for CustomContext<C> {
             tracing::error!(error, "librdkafka: {error}: {reason}");
         })
     }
+
+    fn stats(&self, stats: Statistics) {
+        gauge!(
+            "arroyo.consumer.librdkafka.total_queue_size",
+            stats.replyq as u64,
+        );
+    }
 }
 
 impl<C: AssignmentCallbacks> ConsumerContext for CustomContext<C> {