Update routing tables via Chitchat (#4241)

quickwit/Cargo.toml

@@ -49,7 +49,7 @@ base64 = "0.21"
 bytes = { version = "1", features = ["serde"] }
 bytesize = { version = "1.3.0", features = ["serde"] }
 bytestring = "1.3.0"
-chitchat = { git = "https://github.com/quickwit-oss/chitchat.git", rev = "72a994f" }
+chitchat = { git = "https://github.com/quickwit-oss/chitchat.git", rev = "a3e3f8b" }
 chrono = { version = "0.4.23", default-features = false, features = [
   "clock",
   "std",

quickwit/quickwit-cluster/src/cluster.rs

@@ -28,7 +28,7 @@ use anyhow::Context;
 use chitchat::transport::Transport;
 use chitchat::{
     spawn_chitchat, Chitchat, ChitchatConfig, ChitchatHandle, ChitchatId, ClusterStateSnapshot,
-    FailureDetectorConfig, ListenerHandle, NodeState,
+    FailureDetectorConfig, KeyChangeEvent, ListenerHandle, NodeState,
 };
 use futures::Stream;
 use itertools::Itertools;
@@ -266,7 +266,7 @@ impl Cluster {
     pub async fn subscribe(
         &self,
         key_prefix: &str,
-        callback: impl Fn(&str, &str) + Send + Sync + 'static,
+        callback: impl Fn(KeyChangeEvent) + Send + Sync + 'static,
     ) -> ListenerHandle {
         self.chitchat()
             .await

quickwit/quickwit-cluster/src/lib.rs

@@ -28,7 +28,7 @@ mod node;
 pub use chitchat::transport::ChannelTransport;
 use chitchat::transport::UdpTransport;
 use chitchat::FailureDetectorConfig;
-pub use chitchat::ListenerHandle;
+pub use chitchat::{KeyChangeEvent, ListenerHandle};
 use quickwit_config::service::QuickwitService;
 use quickwit_config::NodeConfig;
 use quickwit_proto::indexing::CpuCapacity;

quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs

@@ -1166,6 +1166,7 @@ mod tests {
             ingestion_rate: RateMibPerSec(1),
         }]);
         let local_shards_update = LocalShardsUpdate {
+            leader_id: "test-ingester".into(),
             source_uid: source_uid.clone(),
             shard_infos,
         };
@@ -1227,6 +1228,7 @@ mod tests {
             },
         ]);
         let local_shards_update = LocalShardsUpdate {
+            leader_id: "test-ingester".into(),
             source_uid: source_uid.clone(),
             shard_infos,
         };
@@ -1248,6 +1250,7 @@ mod tests {
             },
         ]);
         let local_shards_update = LocalShardsUpdate {
+            leader_id: "test-ingester".into(),
             source_uid: source_uid.clone(),
             shard_infos,
         };

quickwit/quickwit-indexing/Cargo.toml

@@ -20,7 +20,6 @@ async-trait = { workspace = true }
 backoff = { workspace = true, optional = true }
 bytes = { workspace = true }
 bytesize = { workspace = true }
-chitchat = { workspace = true }
 fail = { workspace = true }
 flume = { workspace = true }
 fnv = { workspace = true }

quickwit/quickwit-indexing/src/actors/indexing_service.rs

@@ -857,9 +857,8 @@ mod tests {
     use std::path::Path;
     use std::time::Duration;
 
-    use chitchat::transport::ChannelTransport;
     use quickwit_actors::{Health, ObservationType, Supervisable, Universe, HEARTBEAT};
-    use quickwit_cluster::create_cluster_for_test;
+    use quickwit_cluster::{create_cluster_for_test, ChannelTransport};
     use quickwit_common::rand::append_random_suffix;
     use quickwit_config::{
         IngestApiConfig, KafkaSourceParams, SourceConfig, SourceInputFormat, SourceParams,

quickwit/quickwit-indexing/src/models/shard_positions.rs

@@ -23,10 +23,9 @@ use std::fmt::Debug;
 
 use anyhow::Context;
 use async_trait::async_trait;
-use chitchat::ListenerHandle;
 use fnv::FnvHashMap;
 use quickwit_actors::{Actor, ActorContext, ActorExitStatus, Handler, SpawnContext};
-use quickwit_cluster::Cluster;
+use quickwit_cluster::{Cluster, ListenerHandle};
 use quickwit_common::pubsub::{Event, EventBroker};
 use quickwit_proto::indexing::ShardPositionsUpdate;
 use quickwit_proto::types::{IndexUid, Position, ShardId, SourceUid};
@@ -117,13 +116,13 @@ impl Actor for ShardPositionsService {
         let mailbox = ctx.mailbox().clone();
         self.cluster_listener_handle_opt = Some(
             self.cluster
-                .subscribe(SHARD_POSITIONS_PREFIX, move |key, value| {
-                    let shard_positions= match parse_shard_positions_from_kv(key, value) {
+                .subscribe(SHARD_POSITIONS_PREFIX, move |event| {
+                    let shard_positions= match parse_shard_positions_from_kv(event.key, event.value) {
                         Ok(shard_positions) => {
                             shard_positions
                         }
                         Err(error) => {
-                            error!(key=key, value=value, error=%error, "failed to parse shard positions from cluster kv");
+                            error!(key=event.key, value=event.value, error=%error, "failed to parse shard positions from cluster kv");
                             return;
                         }
                     };
@@ -276,9 +275,8 @@ impl ShardPositionsService {
 mod tests {
     use std::time::Duration;
 
-    use chitchat::transport::ChannelTransport;
     use quickwit_actors::Universe;
-    use quickwit_cluster::create_cluster_for_test;
+    use quickwit_cluster::{create_cluster_for_test, ChannelTransport};
     use quickwit_common::pubsub::EventBroker;
     use quickwit_proto::types::IndexUid;
 

quickwit/quickwit-indexing/src/test_utils.rs

@@ -23,9 +23,8 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
 use bytes::Bytes;
-use chitchat::transport::ChannelTransport;
 use quickwit_actors::{Mailbox, Universe};
-use quickwit_cluster::create_cluster_for_test;
+use quickwit_cluster::{create_cluster_for_test, ChannelTransport};
 use quickwit_common::pubsub::EventBroker;
 use quickwit_common::rand::append_random_suffix;
 use quickwit_common::uri::Uri;

quickwit/quickwit-ingest/src/ingest_v2/broadcast.rs

@@ -28,7 +28,7 @@ use quickwit_common::shared_consts::INGESTER_PRIMARY_SHARDS_PREFIX;
 use quickwit_common::sorted_iter::{KeyDiff, SortedByKeyIterator};
 use quickwit_common::tower::Rate;
 use quickwit_proto::ingest::ShardState;
-use quickwit_proto::types::{split_queue_id, QueueId, ShardId, SourceUid};
+use quickwit_proto::types::{split_queue_id, NodeId, QueueId, ShardId, SourceUid};
 use serde::{Deserialize, Serialize, Serializer};
 use tokio::sync::RwLock;
 use tokio::task::JoinHandle;
@@ -276,8 +276,7 @@ fn parse_key(key: &str) -> Option<SourceUid> {
 
 #[derive(Debug, Clone)]
 pub struct LocalShardsUpdate {
-    // TODO: add leader ID in order to update routing table.
-    // leader_id: NodeId,
+    pub leader_id: NodeId,
     pub source_uid: SourceUid,
     pub shard_infos: ShardInfos,
 }
@@ -289,16 +288,19 @@ pub async fn setup_local_shards_update_listener(
     event_broker: EventBroker,
 ) -> ListenerHandle {
     cluster
-        .subscribe(INGESTER_PRIMARY_SHARDS_PREFIX, move |key, value| {
-            let Some(source_uid) = parse_key(key) else {
-                warn!("failed to parse source UID `{key}`");
+        .subscribe(INGESTER_PRIMARY_SHARDS_PREFIX, move |event| {
+            let Some(source_uid) = parse_key(event.key) else {
+                warn!("failed to parse source UID `{}`", event.key);
                 return;
             };
-            let Ok(shard_infos) = serde_json::from_str::<ShardInfos>(value) else {
-                warn!("failed to parse shard infos `{value}`");
+            let Ok(shard_infos) = serde_json::from_str::<ShardInfos>(event.value) else {
+                warn!("failed to parse shard infos `{}`", event.value);
                 return;
             };
+            let leader_id: NodeId = event.node.node_id.clone().into();
+
             let local_shards_update = LocalShardsUpdate {
+                leader_id,
                 source_uid,
                 shard_infos,
             };