Make KafkaSinkCluster rack aware

shotover · Mar 14, 2024 · efab1ba · efab1ba
1 parent 1d23452
commit efab1ba
Show file tree

Hide file tree

Showing 16 changed files with 357 additions and 220 deletions.
diff --git a/shotover-proxy/benches/windsock/kafka/bench.rs b/shotover-proxy/benches/windsock/kafka/bench.rs
@@ -13,7 +13,7 @@ use itertools::Itertools;
 use shotover::config::chain::TransformChainConfig;
 use shotover::sources::SourceConfig;
 use shotover::transforms::debug::force_parse::DebugForceEncodeConfig;
-use shotover::transforms::kafka::sink_cluster::KafkaSinkClusterConfig;
+use shotover::transforms::kafka::sink_cluster::{KafkaSinkClusterConfig, ShotoverNodeConfig};
 use shotover::transforms::kafka::sink_single::KafkaSinkSingleConfig;
 use shotover::transforms::TransformConfig;
 use std::sync::Arc;
@@ -92,7 +92,11 @@ impl KafkaBench {
                 connect_timeout_ms: 3000,
                 read_timeout: None,
                 first_contact_points: vec![kafka_address],
-                shotover_nodes: vec![host_address.clone()],
+                shotover_nodes: vec![ShotoverNodeConfig {
+                    address: host_address.parse().unwrap(),
+                    rack: "rack1".into(),
+                    broker_id: 0,
+                }],
                 tls: None,
                 sasl_enabled: Some(false),
             }),

diff --git a/shotover-proxy/tests/kafka_int_tests/mod.rs b/shotover-proxy/tests/kafka_int_tests/mod.rs
@@ -192,30 +192,6 @@ async fn cluster_1_rack_multi_shotover(#[case] driver: KafkaDriver) {
     }
 }
 
-#[cfg(feature = "rdkafka-driver-tests")] // temporarily needed to avoid a warning
-#[rstest]
-#[cfg_attr(feature = "rdkafka-driver-tests", case::cpp(KafkaDriver::Cpp))]
-// #[case::java(KafkaDriver::Java)]
-#[tokio::test(flavor = "multi_thread")] // multi_thread is needed since java driver will block when consuming, causing shotover logs to not appear
-async fn cluster_2_racks_single_shotover(#[case] driver: KafkaDriver) {
-    let _docker_compose =
-        docker_compose("tests/test-configs/kafka/cluster-2-racks/docker-compose.yaml");
-    let shotover =
-        shotover_process("tests/test-configs/kafka/cluster-2-racks/topology-single.yaml")
-            .start()
-            .await;
-
-    let connection_builder = KafkaConnectionBuilder::new(driver, "127.0.0.1:9192");
-    test_cases::basic(connection_builder).await;
-
-    tokio::time::timeout(
-        Duration::from_secs(10),
-        shotover.shutdown_and_then_consume_events(&[]),
-    )
-    .await
-    .expect("Shotover did not shutdown within 10s");
-}
-
 #[cfg(feature = "rdkafka-driver-tests")] // temporarily needed to avoid a warning
 #[rstest]
 #[cfg_attr(feature = "rdkafka-driver-tests", case::cpp(KafkaDriver::Cpp))]

diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology-single.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology-single.yaml
@@ -5,6 +5,9 @@ sources:
       listen_addr: "127.0.0.1:9192"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: ["127.0.0.1:9192"]
+            shotover_nodes:
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 0
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology1.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology1.yaml
@@ -5,9 +5,15 @@ sources:
       listen_addr: "127.0.0.1:9191"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
-              - "127.0.0.1:9193"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack0"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 1
+              - address: "127.0.0.1:9193"
+                rack: "rack0"
+                broker_id: 2
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology2.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology2.yaml
@@ -5,9 +5,15 @@ sources:
       listen_addr: "127.0.0.1:9192"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
-              - "127.0.0.1:9193"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack0"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 1
+              - address: "127.0.0.1:9193"
+                rack: "rack0"
+                broker_id: 2
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology3.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-1-rack/topology3.yaml
@@ -5,9 +5,15 @@ sources:
       listen_addr: "127.0.0.1:9193"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
-              - "127.0.0.1:9193"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack0"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 1
+              - address: "127.0.0.1:9193"
+                rack: "rack0"
+                broker_id: 2
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-2-racks/topology-rack1.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-2-racks/topology-rack1.yaml
@@ -5,8 +5,12 @@ sources:
       listen_addr: "127.0.0.1:9191"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack1"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack2"
+                broker_id: 1
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-2-racks/topology-rack2.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-2-racks/topology-rack2.yaml
@@ -5,8 +5,12 @@ sources:
       listen_addr: "127.0.0.1:9192"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack1"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack2"
+                broker_id: 1
             first_contact_points: ["172.16.1.5:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-2-racks/topology-single.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-2-racks/topology-single.yaml
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology-single.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology-single.yaml
@@ -5,7 +5,10 @@ sources:
       listen_addr: "127.0.0.1:9192"
       chain:
         - KafkaSinkCluster:
+            shotover_nodes:
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 0
             sasl_enabled: true
-            shotover_nodes: ["127.0.0.1:9192"]
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology1.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology1.yaml
@@ -6,9 +6,15 @@ sources:
       chain:
         - KafkaSinkCluster:
             sasl_enabled: true
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
-              - "127.0.0.1:9193"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack0"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 1
+              - address: "127.0.0.1:9193"
+                rack: "rack0"
+                broker_id: 2
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology2.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology2.yaml
@@ -6,9 +6,15 @@ sources:
       chain:
         - KafkaSinkCluster:
             sasl_enabled: true
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
-              - "127.0.0.1:9193"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack0"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 1
+              - address: "127.0.0.1:9193"
+                rack: "rack0"
+                broker_id: 2
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology3.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-sasl/topology3.yaml
@@ -6,9 +6,15 @@ sources:
       chain:
         - KafkaSinkCluster:
             sasl_enabled: true
-            shotover_nodes: 
-              - "127.0.0.1:9191"
-              - "127.0.0.1:9192"
-              - "127.0.0.1:9193"
+            shotover_nodes:
+              - address: "127.0.0.1:9191"
+                rack: "rack0"
+                broker_id: 0
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 1
+              - address: "127.0.0.1:9193"
+                rack: "rack0"
+                broker_id: 2
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
diff --git a/shotover-proxy/tests/test-configs/kafka/cluster-tls/topology.yaml b/shotover-proxy/tests/test-configs/kafka/cluster-tls/topology.yaml
@@ -5,7 +5,10 @@ sources:
       listen_addr: "127.0.0.1:9192"
       chain:
         - KafkaSinkCluster:
-            shotover_nodes: ["127.0.0.1:9192"]
+            shotover_nodes:
+              - address: "127.0.0.1:9192"
+                rack: "rack0"
+                broker_id: 0
             first_contact_points: ["172.16.1.2:9092"]
             connect_timeout_ms: 3000
             tls: