feat(consumers): rust consumers quantized rebalance

rust_snuba/Cargo.toml

@@ -54,6 +54,7 @@ data-encoding = "2.5.0"
 zstd = "0.12.3"
 serde_with = "3.8.1"
 seq-macro = "0.3"
+redis = "0.27.5"
 
 
 [patch.crates-io]

rust_snuba/src/consumer.rs

@@ -23,6 +23,7 @@ use crate::metrics::global_tags::set_global_tag;
 use crate::metrics::statsd::StatsDBackend;
 use crate::mutations::factory::MutConsumerStrategyFactory;
 use crate::processors;
+use crate::rebalancing;
 use crate::types::{InsertOrReplacement, KafkaMessageMetadata};
 
 #[pyfunction]
@@ -234,6 +235,11 @@ pub fn consumer_impl(
 
     let topic = Topic::new(&consumer_config.raw_topic.physical_topic_name);
 
+    let rebalance_delay_secs = rebalancing::get_rebalance_delay_secs(consumer_group);
+    if let Some(secs) = rebalance_delay_secs {
+        rebalancing::delay_kafka_rebalance(secs)
+    }
+
     let processor = if mutations_mode {
         let mut_factory = MutConsumerStrategyFactory {
             storage_config: first_storage,
@@ -286,10 +292,21 @@ pub fn consumer_impl(
 
     let mut handle = processor.get_handle();
 
-    ctrlc::set_handler(move || {
-        handle.signal_shutdown();
-    })
-    .expect("Error setting Ctrl-C handler");
+    match rebalance_delay_secs {
+        Some(secs) => {
+            ctrlc::set_handler(move || {
+                rebalancing::delay_kafka_rebalance(secs);
+                handle.signal_shutdown();
+            })
+            .expect("Error setting Ctrl-C handler");
+        }
+        None => {
+            ctrlc::set_handler(move || {
+                handle.signal_shutdown();
+            })
+            .expect("Error setting Ctrl-C handler");
+        }
+    }
 
     if let Err(error) = processor.run() {
         let error: &dyn std::error::Error = &error;

rust_snuba/src/lib.rs

@@ -6,6 +6,7 @@ mod logging;
 mod metrics;
 mod mutations;
 mod processors;
+mod rebalancing;
 mod runtime_config;
 mod strategies;
 mod types;

rust_snuba/src/rebalancing.rs

@@ -0,0 +1,81 @@
+use crate::runtime_config;
+use std::thread;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+pub fn delay_kafka_rebalance(configured_delay_secs: u64) {
+    /*
+     *  Introduces a configurable delay to the consumer topic
+     * subscription and consumer shutdown steps (handled by the
+     * StreamProcessor). The idea behind is that by forcing
+     * these steps to occur at certain time "ticks" (for example, at
+     * every 15 second tick in a minute), we can reduce the number of
+     * rebalances that are triggered during a deploy. This means
+     * fewer "stop the world rebalancing" occurrences and more time
+     * for the consumer group to stabilize and make progress.
+     */
+    let current_time = SystemTime::now();
+    let time_elapsed_in_slot = match current_time.duration_since(UNIX_EPOCH) {
+        Ok(duration) => duration.as_secs(),
+        Err(_) => 0,
+    } % configured_delay_secs;
+    tracing::info!(
+        "Delaying rebalance by {} seconds",
+        configured_delay_secs - time_elapsed_in_slot
+    );
+
+    thread::sleep(Duration::from_secs(
+        configured_delay_secs - time_elapsed_in_slot,
+    ));
+}
+
+pub fn get_rebalance_delay_secs(consumer_group: &str) -> Option<u64> {
+    match runtime_config::get_str_config(
+        format!(
+            "quantized_rebalance_consumer_group_delay_secs__{}",
+            consumer_group
+        )
+        .as_str(),
+    ) {
+        Ok(delay_secs) => match delay_secs {
+            Some(secs) => match secs.parse() {
+                Ok(v) => Some(v),
+                Err(_) => None,
+            },
+            None => None,
+        },
+        Err(_) => None,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_delay_config() {
+        runtime_config::del_str_config_direct(
+            "quantized_rebalance_consumer_group_delay_secs__spans",
+        )
+        .unwrap();
+        let delay_secs = get_rebalance_delay_secs("spans");
+        assert_eq!(delay_secs, None);
+        runtime_config::set_str_config_direct(
+            "quantized_rebalance_consumer_group_delay_secs__spans",
+            "420",
+        )
+        .unwrap();
+        let delay_secs = get_rebalance_delay_secs("spans");
+        assert_eq!(delay_secs, Some(420));
+        runtime_config::set_str_config_direct(
+            "quantized_rebalance_consumer_group_delay_secs__spans",
+            "garbage",
+        )
+        .unwrap();
+        let delay_secs = get_rebalance_delay_secs("spans");
+        assert_eq!(delay_secs, None);
+        runtime_config::del_str_config_direct(
+            "quantized_rebalance_consumer_group_delay_secs__spans",
+        )
+        .unwrap();
+    }
+}

rust_snuba/src/runtime_config.rs

@@ -1,18 +1,77 @@
 use anyhow::Error;
 use parking_lot::RwLock;
 use pyo3::prelude::{PyModule, Python};
+use redis::Commands;
 use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::time::Duration;
 
 use rust_arroyo::timer;
 use rust_arroyo::utils::timing::Deadline;
 
 static CONFIG: RwLock<BTreeMap<String, (Option<String>, Deadline)>> = RwLock::new(BTreeMap::new());
 
+static CONFIG_HASHSET_KEY: &str = "snuba-config";
+
+fn get_redis_client() -> Result<redis::Client, redis::RedisError> {
+    let redis_host = std::env::var("REDIS_HOST").unwrap_or(String::from("127.0.0.1"));
+    let redis_port = std::env::var("REDIS_PORT").unwrap_or(String::from("6379"));
+    let redis_password = std::env::var("REDIS_PASSWORD").unwrap_or(String::from(""));
+    let redis_db = std::env::var("REDIS_DB").unwrap_or(String::from("1"));
+    // TODO: handle SSL?
+    let url = format!(
+        "redis://{}:{}@{}:{}/{}",
+        "default", redis_password, redis_host, redis_port, redis_db
+    );
+    redis::Client::open(url)
+}
+
+fn get_str_config_direct(key: &str) -> Result<Option<String>, Error> {
+    let deadline = Deadline::new(Duration::from_secs(10));
+
+    let client = get_redis_client()?;
+    let mut con = client.get_connection()?;
+
+    let configmap: HashMap<String, String> = con.hgetall(CONFIG_HASHSET_KEY)?;
+    let val = match configmap.get(key) {
+        Some(val) => Some(val.clone()),
+        None => return Ok(None),
+    };
+
+    CONFIG
+        .write()
+        .insert(key.to_string(), (val.clone(), deadline));
+    Ok(CONFIG.read().get(key).unwrap().0.clone())
+}
+
+#[allow(dead_code)]
+pub fn set_str_config_direct(key: &str, val: &str) -> Result<(), Error> {
+    let client = get_redis_client()?;
+    let mut con = client.get_connection()?;
+    con.hset(CONFIG_HASHSET_KEY, key, val)?;
+    Ok(())
+}
+
+#[allow(dead_code)]
+pub fn del_str_config_direct(key: &str) -> Result<(), Error> {
+    let client = get_redis_client()?;
+    let mut con = client.get_connection()?;
+    con.hdel(CONFIG_HASHSET_KEY, key)?;
+    Ok(())
+}
+
 /// Runtime config is cached for 10 seconds
 pub fn get_str_config(key: &str) -> Result<Option<String>, Error> {
     let deadline = Deadline::new(Duration::from_secs(10));
 
+    match get_str_config_direct(key) {
+        Ok(val) => return Ok(val),
+        Err(error) => tracing::error!(
+            "Could not get config from redis directly, falling back to python {}",
+            error
+        ),
+    }
+
     if let Some(value) = CONFIG.read().get(key) {
         let (config, deadline) = value;
         if !deadline.has_elapsed() {