Add AZ awareness to node state (quickwit-oss#6092)

* Initial stab at AZ awareness

Signed-off-by: Nadav Gov-Ari <nadav.govari@datadoghq.com>

* Comment

Signed-off-by: Nadav Gov-Ari <nadav.govari@datadoghq.com>

* Add tests, PR comments

Signed-off-by: Nadav Gov-Ari <nadav.govari@datadoghq.com>

---------

Signed-off-by: Nadav Gov-Ari <nadav.govari@datadoghq.com>

Author: nadav-govari

quickwit/quickwit-cli/src/tool.rs

@@ -938,6 +938,7 @@ async fn create_empty_cluster(config: &NodeConfig) -> anyhow::Result<Cluster> {
         grpc_advertise_addr: config.grpc_advertise_addr,
         indexing_cpu_capacity: CpuCapacity::zero(),
         indexing_tasks: Vec::new(),
+        availability_zone: None,
     };
     let client_grpc_config = make_client_grpc_config(&config.grpc_config)?;
     let cluster = Cluster::join(

quickwit/quickwit-cluster/src/cluster.rs

@@ -39,9 +39,9 @@ use tracing::{info, warn};
 use crate::change::{ClusterChange, ClusterChangeStreamFactory, compute_cluster_change_events};
 use crate::grpc_gossip::spawn_catchup_callback_task;
 use crate::member::{
-    ClusterMember, ENABLED_SERVICES_KEY, GRPC_ADVERTISE_ADDR_KEY, NodeStateExt,
-    PIPELINE_METRICS_PREFIX, READINESS_KEY, READINESS_VALUE_NOT_READY, READINESS_VALUE_READY,
-    build_cluster_member,
+    AVAILABILITY_ZONE_KEY, ClusterMember, ENABLED_SERVICES_KEY, GRPC_ADVERTISE_ADDR_KEY,
+    NodeStateExt, PIPELINE_METRICS_PREFIX, READINESS_KEY, READINESS_VALUE_NOT_READY,
+    READINESS_VALUE_READY, build_cluster_member,
 };
 use crate::metrics::spawn_metrics_task;
 use crate::{ClusterChangeStream, ClusterNode};
@@ -150,25 +150,26 @@ impl Cluster {
             catchup_callback: Some(Box::new(catchup_callback)),
             extra_liveness_predicate: Some(Box::new(extra_liveness_predicate)),
         };
-        let chitchat_handle = spawn_chitchat(
-            chitchat_config,
-            vec![
-                (
-                    ENABLED_SERVICES_KEY.to_string(),
-                    self_node.enabled_services.iter().join(","),
-                ),
-                (
-                    GRPC_ADVERTISE_ADDR_KEY.to_string(),
-                    self_node.grpc_advertise_addr.to_string(),
-                ),
-                (
-                    READINESS_KEY.to_string(),
-                    READINESS_VALUE_NOT_READY.to_string(),
-                ),
-            ],
-            transport,
-        )
-        .await?;
+        let mut initial_key_values = vec![
+            (
+                ENABLED_SERVICES_KEY.to_string(),
+                self_node.enabled_services.iter().join(","),
+            ),
+            (
+                GRPC_ADVERTISE_ADDR_KEY.to_string(),
+                self_node.grpc_advertise_addr.to_string(),
+            ),
+            (
+                READINESS_KEY.to_string(),
+                READINESS_VALUE_NOT_READY.to_string(),
+            ),
+        ];
+
+        if let Some(az) = &self_node.availability_zone {
+            initial_key_values.push((AVAILABILITY_ZONE_KEY.to_string(), az.clone()));
+        }
+        let chitchat_handle =
+            spawn_chitchat(chitchat_config, initial_key_values, transport).await?;
 
         let chitchat = chitchat_handle.chitchat();
         let chitchat_guard = chitchat.lock().await;
@@ -706,6 +707,7 @@ pub async fn create_cluster_for_test_with_id(
         grpc_advertise_addr: grpc_addr_from_listen_addr_for_test(gossip_advertise_addr),
         indexing_tasks: Vec::new(),
         indexing_cpu_capacity: PIPELINE_FULL_CAPACITY,
+        availability_zone: None,
     };
     let failure_detector_config = create_failure_detector_config_for_test();
     let cluster = Cluster::join(

quickwit/quickwit-cluster/src/lib.rs

@@ -143,6 +143,7 @@ pub async fn start_cluster_service(node_config: &NodeConfig) -> anyhow::Result<C
         grpc_advertise_addr: node_config.grpc_advertise_addr,
         indexing_tasks,
         indexing_cpu_capacity,
+        availability_zone: node_config.availability_zone.clone(),
     };
     let failure_detector_config = FailureDetectorConfig {
         dead_node_grace_period: Duration::from_secs(2 * 60 * 60), // 2 hours

quickwit/quickwit-cluster/src/member.rs

@@ -36,6 +36,8 @@ pub(crate) const READINESS_KEY: &str = "readiness";
 pub(crate) const READINESS_VALUE_READY: &str = "READY";
 pub(crate) const READINESS_VALUE_NOT_READY: &str = "NOT_READY";
 
+pub(crate) const AVAILABILITY_ZONE_KEY: &str = "availability_zone";
+
 pub const INDEXING_CPU_CAPACITY_KEY: &str = "indexing_cpu_capacity";
 
 pub(crate) trait NodeStateExt {
@@ -44,6 +46,8 @@ pub(crate) trait NodeStateExt {
     fn is_ready(&self) -> bool;
 
     fn size_bytes(&self) -> usize;
+
+    fn availability_zone(&self) -> Option<String>;
 }
 
 impl NodeStateExt for NodeState {
@@ -74,6 +78,10 @@ impl NodeStateExt for NodeState {
             .map(|(key, value)| key.len() + value.value.len() + SIZE_OF_VERSION + SIZE_OF_TOMBSTONE)
             .sum()
     }
+
+    fn availability_zone(&self) -> Option<String> {
+        self.get(AVAILABILITY_ZONE_KEY).map(|az| az.to_string())
+    }
 }
 
 /// Cluster member.
@@ -101,6 +109,8 @@ pub struct ClusterMember {
     /// Indexing cpu capacity of the node expressed in milli cpu.
     pub indexing_cpu_capacity: CpuCapacity,
     pub is_ready: bool,
+    /// Availability zone the node is running in, if enabled.
+    pub availability_zone: Option<String>,
 }
 
 impl ClusterMember {
@@ -149,6 +159,7 @@ pub(crate) fn build_cluster_member(
         .map(|enabled_services_str| {
             parse_enabled_services_str(enabled_services_str, &chitchat_id.node_id)
         })?;
+    let availability_zone = node_state.availability_zone();
     let grpc_advertise_addr = node_state.grpc_advertise_addr()?;
     let indexing_tasks = parse_indexing_tasks(node_state);
     let indexing_cpu_capacity = parse_indexing_cpu_capacity(node_state);
@@ -161,6 +172,7 @@ pub(crate) fn build_cluster_member(
         grpc_advertise_addr,
         indexing_tasks,
         indexing_cpu_capacity,
+        availability_zone,
     };
     Ok(member)
 }

quickwit/quickwit-config/resources/tests/node_config/quickwit.json

@@ -3,6 +3,7 @@
     "version": "0.7",
     "cluster_id": "quickwit-cluster",
     "node_id": "my-unique-node-id",
+    "availability_zone": "az-1",
     "enabled_services": [
         "janitor",
         "metastore"

quickwit/quickwit-config/resources/tests/node_config/quickwit.toml

@@ -2,6 +2,7 @@ version = "0.7"
 
 cluster_id = "quickwit-cluster"
 node_id = "my-unique-node-id"
+availability_zone = "az-1"
 enabled_services = [ "janitor", "metastore" ]
 listen_address = "0.0.0.0"
 advertise_address = "172.0.0.12"

quickwit/quickwit-config/resources/tests/node_config/quickwit.yaml

@@ -2,6 +2,7 @@ version: 0.8
 
 cluster_id: quickwit-cluster
 node_id: my-unique-node-id
+availability_zone: az-1
 enabled_services:
   - janitor
   - metastore

quickwit/quickwit-config/src/config_value.rs

@@ -41,18 +41,17 @@ where
         }
     }
 
-    #[cfg(test)]
-    pub(crate) fn for_test(value: T) -> Self {
+    pub(crate) fn none() -> Self {
         Self {
-            provided: Some(value),
+            provided: None,
             default: None,
         }
     }
 
     #[cfg(test)]
-    pub(crate) fn none() -> Self {
+    pub(crate) fn for_test(value: T) -> Self {
         Self {
-            provided: None,
+            provided: Some(value),
             default: None,
         }
     }

quickwit/quickwit-config/src/node_config/mod.rs

@@ -566,6 +566,7 @@ impl Default for JaegerConfig {
 pub struct NodeConfig {
     pub cluster_id: String,
     pub node_id: NodeId,
+    pub availability_zone: Option<String>,
     pub enabled_services: HashSet<QuickwitService>,
     pub gossip_listen_addr: SocketAddr,
     pub grpc_listen_addr: SocketAddr,

quickwit/quickwit-config/src/node_config/serialize.rs

@@ -66,6 +66,10 @@ fn default_node_id() -> ConfigValue<String, QW_NODE_ID> {
     ConfigValue::with_default(node_id)
 }
 
+fn default_availability_zone() -> ConfigValue<String, QW_AVAILABILITY_ZONE> {
+    ConfigValue::none()
+}
+
 #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
 struct List(Vec<String>);
 
@@ -169,6 +173,8 @@ struct NodeConfigBuilder {
     cluster_id: ConfigValue<String, QW_CLUSTER_ID>,
     #[serde(default = "default_node_id")]
     node_id: ConfigValue<String, QW_NODE_ID>,
+    #[serde(default = "default_availability_zone")]
+    availability_zone: ConfigValue<String, QW_AVAILABILITY_ZONE>,
     #[serde(default = "default_enabled_services")]
     enabled_services: ConfigValue<List, QW_ENABLED_SERVICES>,
     #[serde(default = "default_listen_address")]
@@ -218,6 +224,7 @@ impl NodeConfigBuilder {
         env_vars: &HashMap<String, String>,
     ) -> anyhow::Result<NodeConfig> {
         let node_id = self.node_id.resolve(env_vars).map(NodeId::new)?;
+        let availability_zone = self.availability_zone.resolve_optional(env_vars)?;
 
         let enabled_services = self
             .enabled_services
@@ -305,6 +312,7 @@ impl NodeConfigBuilder {
         let node_config = NodeConfig {
             cluster_id: self.cluster_id.resolve(env_vars)?,
             node_id,
+            availability_zone,
             enabled_services,
             gossip_listen_addr,
             grpc_listen_addr,
@@ -401,6 +409,7 @@ impl Default for NodeConfigBuilder {
         Self {
             cluster_id: default_cluster_id(),
             node_id: default_node_id(),
+            availability_zone: ConfigValue::none(),
             enabled_services: default_enabled_services(),
             listen_address: default_listen_address(),
             rest_listen_port: None,
@@ -469,6 +478,7 @@ pub fn node_config_for_tests_from_ports(
 ) -> NodeConfig {
     let node_id = NodeId::new(default_node_id().unwrap());
     let enabled_services = QuickwitService::supported_services();
+    let availability_zone = Some(String::from("az-1"));
     let listen_address = Host::default();
     let rest_listen_addr = listen_address
         .with_port(rest_listen_port)
@@ -499,6 +509,7 @@ pub fn node_config_for_tests_from_ports(
     NodeConfig {
         cluster_id: default_cluster_id().unwrap(),
         node_id,
+        availability_zone,
         enabled_services,
         gossip_advertise_addr: gossip_listen_addr,
         grpc_advertise_addr: grpc_listen_addr,
@@ -553,6 +564,7 @@ mod tests {
         assert!(config.is_service_enabled(QuickwitService::Janitor));
         assert!(config.is_service_enabled(QuickwitService::Metastore));
 
+        assert_eq!(config.availability_zone.unwrap(), "az-1");
         assert_eq!(
             config.rest_config.listen_addr,
             SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 1111)
@@ -738,6 +750,7 @@ mod tests {
         .unwrap();
         assert_eq!(config.cluster_id, DEFAULT_CLUSTER_ID);
         assert_eq!(config.node_id, get_short_hostname().unwrap());
+        assert_eq!(config.availability_zone, None);
         assert_eq!(
             config.enabled_services,
             QuickwitService::supported_services()