feat(consumer): periodically restart the Rust consumer to reconnect to Kafka

rust_snuba/src/auto_restart.rs

@@ -0,0 +1,138 @@
+use std::time::Duration;
+
+use crate::runtime_config;
+
+/// Default interval between consumer restarts when periodic restart is enabled
+/// but no explicit interval is configured: 15 minutes.
+pub const DEFAULT_RESTART_INTERVAL_SECS: u64 = 900;
+
+fn is_truthy(value: &str) -> bool {
+    matches!(value.trim().to_ascii_lowercase().as_str(), "1" | "true")
+}
+
+/// Returns the interval after which the consumer should gracefully shut down so
+/// that it can be restarted and reconnect to Kafka, or `None` when periodic
+/// restart is not enabled for `consumer_group`.
+///
+/// This is gated behind a feature flag and controlled by two runtime config
+/// keys (read from `snuba.state`, so they can be toggled live without a
+/// redeploy):
+///
+/// * `kafka_consumer_periodic_restart__<consumer_group>` — the feature flag.
+///   Periodic restart is only enabled when this is set to a truthy value
+///   (`"1"` or `"true"`). Defaults to disabled.
+/// * `kafka_consumer_periodic_restart_interval_secs__<consumer_group>` — the
+///   number of seconds between restarts. Must be a positive integer; defaults
+///   to [`DEFAULT_RESTART_INTERVAL_SECS`] (15 minutes) when unset or invalid.
+pub fn get_restart_interval(consumer_group: &str) -> Option<Duration> {
+    let enabled = runtime_config::get_str_config(
+        format!("kafka_consumer_periodic_restart__{consumer_group}").as_str(),
+    )
+    .ok()
+    .flatten()
+    .map(|value| is_truthy(&value))
+    .unwrap_or(false);
+
+    if !enabled {
+        return None;
+    }
+
+    let interval_secs = runtime_config::get_str_config(
+        format!("kafka_consumer_periodic_restart_interval_secs__{consumer_group}").as_str(),
+    )
+    .ok()
+    .flatten()
+    .and_then(|value| value.parse::<u64>().ok())
+    .filter(|&secs| secs > 0)
+    .unwrap_or(DEFAULT_RESTART_INTERVAL_SECS);
+
+    Some(Duration::from_secs(interval_secs))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn clear(consumer_group: &str) {
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart__{consumer_group}"),
+            None,
+        );
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart_interval_secs__{consumer_group}"),
+            None,
+        );
+    }
+
+    #[test]
+    fn test_disabled_by_default() {
+        crate::testutils::initialize_python();
+        assert_eq!(get_restart_interval("periodic_restart_default_test"), None);
+    }
+
+    #[test]
+    fn test_enabled_uses_default_interval() {
+        crate::testutils::initialize_python();
+        let group = "periodic_restart_enabled_test";
+        let _guard = scopeguard::guard((), |_| clear(group));
+
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart__{group}"),
+            Some("1"),
+        );
+        assert_eq!(
+            get_restart_interval(group),
+            Some(Duration::from_secs(DEFAULT_RESTART_INTERVAL_SECS))
+        );
+    }
+
+    #[test]
+    fn test_enabled_with_custom_interval() {
+        crate::testutils::initialize_python();
+        let group = "periodic_restart_custom_test";
+        let _guard = scopeguard::guard((), |_| clear(group));
+
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart__{group}"),
+            Some("true"),
+        );
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart_interval_secs__{group}"),
+            Some("60"),
+        );
+        assert_eq!(get_restart_interval(group), Some(Duration::from_secs(60)));
+    }
+
+    #[test]
+    fn test_invalid_interval_falls_back_to_default() {
+        crate::testutils::initialize_python();
+        let group = "periodic_restart_invalid_test";
+        let _guard = scopeguard::guard((), |_| clear(group));
+
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart__{group}"),
+            Some("1"),
+        );
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart_interval_secs__{group}"),
+            Some("garbage"),
+        );
+        assert_eq!(
+            get_restart_interval(group),
+            Some(Duration::from_secs(DEFAULT_RESTART_INTERVAL_SECS))
+        );
+    }
+
+    #[test]
+    fn test_interval_without_flag_is_ignored() {
+        crate::testutils::initialize_python();
+        let group = "periodic_restart_no_flag_test";
+        let _guard = scopeguard::guard((), |_| clear(group));
+
+        runtime_config::patch_str_config_for_test(
+            &format!("kafka_consumer_periodic_restart_interval_secs__{group}"),
+            Some("60"),
+        );
+        assert_eq!(get_restart_interval(group), None);
+    }
+}

rust_snuba/src/consumer.rs

@@ -1,11 +1,13 @@
 use std::sync::Arc;
+use std::thread;
 use std::time::Duration;
 
 use chrono::{DateTime, Utc};
 
 use sentry_arroyo::backends::kafka::config::KafkaConfig;
 use sentry_arroyo::backends::kafka::producer::KafkaProducer;
 use sentry_arroyo::backends::kafka::types::KafkaPayload;
+use sentry_arroyo::counter;
 use sentry_arroyo::metrics;
 use sentry_arroyo::processing::dlq::{DlqLimit, DlqPolicy, KafkaDlqProducer};
 
@@ -18,6 +20,7 @@ use pyo3::types::PyBytes;
 
 use sentry_options::init_with_schemas;
 
+use crate::auto_restart;
 use crate::config;
 use crate::factory_v2::ConsumerStrategyFactoryV2;
 use crate::logging::{setup_logging, setup_sentry};
@@ -297,6 +300,32 @@ pub fn consumer_impl(
 
     let mut handle = processor.get_handle();
 
+    // Periodically restart the consumer to reconnect to Kafka when the
+    // periodic-restart feature flag is enabled for this consumer group. After
+    // the configured interval (default 15 minutes) we signal a graceful
+    // shutdown; the process orchestrator then restarts the consumer, which
+    // re-establishes its Kafka connection. See `crate::auto_restart` for the
+    // controlling runtime config keys.
+    if let Some(restart_interval) = auto_restart::get_restart_interval(consumer_group) {
+        tracing::info!(
+            consumer_group,
+            interval_secs = restart_interval.as_secs(),
+            "Periodic Kafka reconnect enabled; consumer will restart after the configured interval",
+        );
+        let mut restart_handle = handle.clone();
+        let consumer_group = consumer_group.to_owned();
+        thread::spawn(move || {
+            thread::sleep(restart_interval);
+            tracing::info!(
+                consumer_group,
+                interval_secs = restart_interval.as_secs(),
+                "Periodic restart interval elapsed; signaling shutdown to reconnect to Kafka",
+            );
+            counter!("periodic_restart");
+            restart_handle.signal_shutdown();
+        });
+    }
+
     match rebalance_delay_secs {
         Some(secs) => {
             ctrlc::set_handler(move || {

rust_snuba/src/lib.rs

@@ -2,6 +2,7 @@ pub(crate) const SNUBA_SCHEMA: &str =
     include_str!("../../sentry-options/schemas/snuba/schema.json");
 
 mod accepted_outcomes_consumer;
+mod auto_restart;
 mod config;
 mod consumer;
 mod factory_v2;