feat(redis): Redis Backed Taskbroker

Cargo.toml

@@ -10,10 +10,14 @@ debug = 1
 
 [dependencies]
 anyhow = "1.0.92"
+async-backtrace = "0.2.7"
+base64 = "0.21.0"
 bytes = "1.10.0"
 chrono = { version = "0.4.26" }
 clap = { version = "4.5.20", features = ["derive"] }
+deadpool-redis = { version = "0.22.0", features = ["cluster"] }
 elegant-departure = { version = "0.3.1", features = ["tokio"] }
+cityhasher = "0.1.0"
 figment = { version = "0.10.19", features = ["env", "yaml", "test"] }
 futures = "0.3.31"
 futures-util = "0.3.31"
@@ -28,6 +32,7 @@ prost = "0.13"
 prost-types = "0.13.3"
 rand = "0.8.5"
 rdkafka = { version = "0.37.0", features = ["cmake-build", "ssl"] }
+redis = "0.32.7"
 sentry = { version = "0.41.0", default-features = false, features = [
     # default features, except `release-health` is disabled
     "backtrace",
@@ -44,6 +49,7 @@ serde = "1.0.214"
 serde_yaml = "0.9.34"
 sha2 = "0.10.8"
 sqlx = { version = "0.8.3", features = ["sqlite", "runtime-tokio", "chrono"] }
+thiserror = "2.0.17"
 tokio = { version = "1.43.1", features = ["full"] }
 tokio-stream = { version = "0.1.16", features = ["full"] }
 tokio-util = "0.7.12"

src/config.rs

@@ -214,6 +214,15 @@ pub struct Config {
 
     /// Enable additional metrics for the sqlite.
     pub enable_sqlite_status_metrics: bool,
+
+    /// Redis cluster information
+    pub redis_cluster_urls: Vec<String>,
+
+    pub namespaces: Vec<String>,
+
+    pub num_redis_buckets: usize,
+
+    pub payload_ttl_seconds: u64,
 }
 
 impl Default for Config {
@@ -279,6 +288,11 @@ impl Default for Config {
             full_vacuum_on_upkeep: true,
             vacuum_interval_ms: 30000,
             enable_sqlite_status_metrics: true,
+            // Redis information
+            redis_cluster_urls: vec!["redis://127.0.0.1:6379".to_owned()],
+            namespaces: vec!["default".to_owned()],
+            num_redis_buckets: 256,
+            payload_ttl_seconds: 60 * 60 * 24,
         }
     }
 }

src/grpc/server.rs

@@ -9,11 +9,12 @@ use std::sync::Arc;
 use std::time::Instant;
 use tonic::{Request, Response, Status};
 
-use crate::store::inflight_activation::{InflightActivationStatus, InflightActivationStore};
+use crate::store::inflight_activation::InflightActivationStatus;
+use crate::store::inflight_redis_activation::RedisActivationStore;
 use tracing::{error, instrument};
 
 pub struct TaskbrokerServer {
-    pub store: Arc<InflightActivationStore>,
+    pub store: Arc<RedisActivationStore>,
 }
 
 #[tonic::async_trait]
@@ -29,7 +30,6 @@ impl ConsumerService for TaskbrokerServer {
             .store
             .get_pending_activation(namespace.as_deref())
             .await;
-
         match inflight {
             Ok(Some(inflight)) => {
                 let now = Utc::now();

src/grpc/server_tests.rs

@@ -4,11 +4,12 @@ use tonic::{Code, Request};
 
 use crate::grpc::server::TaskbrokerServer;
 
-use crate::test_utils::{create_test_store, make_activations};
+use crate::test_utils::{create_redis_test_store, make_activations};
 
 #[tokio::test]
 async fn test_get_task() {
-    let store = create_test_store().await;
+    let store = create_redis_test_store().await;
+    store.delete_all_keys().await.unwrap();
     let service = TaskbrokerServer { store };
     let request = GetTaskRequest { namespace: None };
     let response = service.get_task(Request::new(request)).await;
@@ -21,7 +22,8 @@ async fn test_get_task() {
 #[tokio::test]
 #[allow(deprecated)]
 async fn test_set_task_status() {
-    let store = create_test_store().await;
+    let store = create_redis_test_store().await;
+    store.delete_all_keys().await.unwrap();
     let service = TaskbrokerServer { store };
     let request = SetTaskStatusRequest {
         id: "test_task".to_string(),
@@ -37,7 +39,8 @@ async fn test_set_task_status() {
 #[tokio::test]
 #[allow(deprecated)]
 async fn test_set_task_status_invalid() {
-    let store = create_test_store().await;
+    let store = create_redis_test_store().await;
+    store.delete_all_keys().await.unwrap();
     let service = TaskbrokerServer { store };
     let request = SetTaskStatusRequest {
         id: "test_task".to_string(),
@@ -57,46 +60,68 @@ async fn test_set_task_status_invalid() {
 #[tokio::test]
 #[allow(deprecated)]
 async fn test_get_task_success() {
-    let store = create_test_store().await;
+    let store = create_redis_test_store().await;
+    store.delete_all_keys().await.unwrap();
     let activations = make_activations(1);
     store.store(activations).await.unwrap();
 
-    let service = TaskbrokerServer { store };
+    let service = TaskbrokerServer {
+        store: store.clone(),
+    };
     let request = GetTaskRequest { namespace: None };
     let response = service.get_task(Request::new(request)).await;
     assert!(response.is_ok());
     let resp = response.unwrap();
     assert!(resp.get_ref().task.is_some());
     let task = resp.get_ref().task.as_ref().unwrap();
     assert!(task.id == "id_0");
+    assert!(store.count_pending_activations().await.unwrap() == 0);
+    assert!(store.count_processing_activations().await.unwrap() == 1);
 }
 
 #[tokio::test]
 #[allow(deprecated)]
 async fn test_set_task_status_success() {
-    let store = create_test_store().await;
+    let store = create_redis_test_store().await;
+    store.delete_all_keys().await.unwrap();
     let activations = make_activations(2);
     store.store(activations).await.unwrap();
 
-    let service = TaskbrokerServer { store };
+    let service = TaskbrokerServer {
+        store: store.clone(),
+    };
 
     let request = GetTaskRequest { namespace: None };
     let response = service.get_task(Request::new(request)).await;
     assert!(response.is_ok());
     let resp = response.unwrap();
     assert!(resp.get_ref().task.is_some());
     let task = resp.get_ref().task.as_ref().unwrap();
-    assert!(task.id == "id_0");
-
+    assert!(task.id == "id_0" || task.id == "id_1");
+    let first_task_id = task.id.clone();
     let request = SetTaskStatusRequest {
-        id: "id_0".to_string(),
+        id: first_task_id.clone(),
         status: 5, // Complete
         fetch_next_task: Some(FetchNextTask { namespace: None }),
     };
     let response = service.set_task_status(Request::new(request)).await;
-    assert!(response.is_ok());
+    println!("response: {:?}", response);
+    assert!(response.is_ok(), "response: {:?}", response);
     let resp = response.unwrap();
     assert!(resp.get_ref().task.is_some());
     let task = resp.get_ref().task.as_ref().unwrap();
-    assert_eq!(task.id, "id_1");
+    let second_task_id = if first_task_id == "id_0" {
+        "id_1"
+    } else {
+        "id_0"
+    };
+    assert_eq!(task.id, second_task_id);
+    let pending_count = store.count_pending_activations().await.unwrap();
+    let processing_count = store.count_processing_activations().await.unwrap();
+    assert!(pending_count == 0, "pending_count: {:?}", pending_count);
+    assert!(
+        processing_count == 1,
+        "processing_count: {:?}",
+        processing_count
+    );
 }

src/kafka/consumer.rs

@@ -1,3 +1,4 @@
+use crate::store::inflight_redis_activation::RedisActivationStore;
 use anyhow::{Error, anyhow};
 use futures::{
     Stream, StreamExt,
@@ -21,10 +22,8 @@ use std::{
     future::Future,
     iter,
     mem::take,
-    sync::{
-        Arc,
-        mpsc::{SyncSender, sync_channel},
-    },
+    sync::Arc,
+    sync::mpsc::{SyncSender, sync_channel},
     time::Duration,
 };
 use tokio::{
@@ -44,14 +43,15 @@ use tracing::{debug, error, info, instrument, warn};
 pub async fn start_consumer(
     topics: &[&str],
     kafka_client_config: &ClientConfig,
+    redis_store: Arc<RedisActivationStore>,
     spawn_actors: impl FnMut(
         Arc<StreamConsumer<KafkaContext>>,
         &BTreeSet<(String, i32)>,
     ) -> ActorHandles,
 ) -> Result<(), Error> {
     let (client_shutdown_sender, client_shutdown_receiver) = oneshot::channel();
     let (event_sender, event_receiver) = unbounded_channel();
-    let context = KafkaContext::new(event_sender.clone());
+    let context = KafkaContext::new(event_sender.clone(), redis_store.clone());
     let consumer: Arc<StreamConsumer<KafkaContext>> = Arc::new(
         kafka_client_config
             .create_with_context(context)
@@ -67,6 +67,7 @@ pub async fn start_consumer(
     metrics::gauge!("arroyo.consumer.current_partitions").set(0);
     handle_events(
         consumer,
+        redis_store,
         event_receiver,
         client_shutdown_sender,
         spawn_actors,
@@ -118,11 +119,18 @@ pub fn poll_consumer_client(
 #[derive(Debug)]
 pub struct KafkaContext {
     event_sender: UnboundedSender<(Event, SyncSender<()>)>,
+    redis_store: Arc<RedisActivationStore>,
 }
 
 impl KafkaContext {
-    pub fn new(event_sender: UnboundedSender<(Event, SyncSender<()>)>) -> Self {
-        Self { event_sender }
+    pub fn new(
+        event_sender: UnboundedSender<(Event, SyncSender<()>)>,
+        redis_store: Arc<RedisActivationStore>,
+    ) -> Self {
+        Self {
+            event_sender,
+            redis_store,
+        }
     }
 }
 
@@ -339,6 +347,7 @@ enum ConsumerState {
 #[instrument(skip_all)]
 pub async fn handle_events(
     consumer: Arc<StreamConsumer<KafkaContext>>,
+    redis_store: Arc<RedisActivationStore>,
     events: UnboundedReceiver<(Event, SyncSender<()>)>,
     shutdown_client: oneshot::Sender<()>,
     mut spawn_actors: impl FnMut(
@@ -372,6 +381,17 @@ pub async fn handle_events(
                 state = match (state, event) {
                     (ConsumerState::Ready, Event::Assign(tpl)) => {
                         metrics::gauge!("arroyo.consumer.current_partitions").set(tpl.len() as f64);
+                        let mut topics = HashMap::<String, Vec<i32>>::new();
+                        for (topic, partition) in tpl.iter() {
+                            if !topics.contains_key(topic) {
+                                topics.insert(topic.clone(), vec![*partition]);
+                            } else {
+                                topics.get_mut(topic).unwrap().push(*partition);
+                            }
+                        }
+                        for (topic, partitions) in topics.iter() {
+                            redis_store.rebalance_partitions(topic.clone(), partitions.clone());
+                        }
                         ConsumerState::Consuming(spawn_actors(consumer.clone(), &tpl), tpl)
                     }
                     (ConsumerState::Ready, Event::Revoke(_)) => {

src/kafka/deserialize_activation.rs

@@ -81,6 +81,7 @@ pub fn new(
             id: activation.id.clone(),
             activation: payload.to_vec(),
             status,
+            topic: msg.topic().to_string(),
             partition: msg.partition(),
             offset: msg.offset(),
             added_at: Utc::now(),

src/kafka/inflight_activation_batcher.rs

@@ -260,6 +260,7 @@ demoted_namespaces:
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),
@@ -306,6 +307,7 @@ demoted_namespaces:
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),
@@ -355,6 +357,7 @@ demoted_namespaces:
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),
@@ -406,6 +409,7 @@ demoted_namespaces:
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),
@@ -437,6 +441,7 @@ demoted_namespaces:
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),
@@ -497,6 +502,7 @@ demoted_topic: taskworker-demoted"#;
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),
@@ -528,6 +534,7 @@ demoted_topic: taskworker-demoted"#;
             }
             .encode_to_vec(),
             status: InflightActivationStatus::Pending,
+            topic: "test_topic".to_string(),
             partition: 0,
             offset: 0,
             added_at: Utc::now(),