test: add HA chaos harness that drives the scheduler's fault-tolerance paths on a real cluster

Cargo.toml

@@ -24,6 +24,7 @@ members = [
     "ballista/executor",
     "ballista/scheduler",
     "benchmarks",
+    "chaos-testing",
     "examples",
 ]
 resolver = "3"

chaos-testing/Cargo.toml

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "ballista-chaos"
+description = "Fault-injection harness for testing Ballista high availability"
+license = "Apache-2.0"
+version = "0.1.0"
+edition = { workspace = true }
+rust-version = { workspace = true }
+publish = false
+
+[dependencies]
+arrow = { workspace = true }
+ballista = { path = "../ballista/client" }
+ballista-core = { path = "../ballista/core" }
+ballista-executor = { path = "../ballista/executor" }
+ballista-scheduler = { path = "../ballista/scheduler" }
+datafusion = { workspace = true }
+env_logger = { workspace = true }
+log = { workspace = true }
+nix = { version = "0.29", features = ["fs", "signal"] }
+reqwest = { version = "0.12", default-features = false, features = ["json"] }
+serde_json = "1.0"
+tempfile = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread", "process", "sync", "time"] }
+
+[dev-dependencies]
+rstest = { workspace = true }
+
+[lib]
+name = "chaos_testing"
+path = "src/lib.rs"
+
+[[bin]]
+name = "chaos-executor"
+path = "src/bin/chaos-executor.rs"
+
+[[bin]]
+name = "chaos-scheduler"
+path = "src/bin/chaos-scheduler.rs"

chaos-testing/src/bin/chaos-executor.rs

@@ -0,0 +1,62 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A Ballista executor whose function registry includes the chaos UDFs.
+//!
+//! Configured entirely from the environment, because `TestCluster` spawns it as
+//! a child process. Mirrors `examples/examples/custom-executor.rs`.
+
+use ballista_executor::executor_process::{
+    ExecutorProcessConfig, start_executor_process,
+};
+use chaos_testing::registry::chaos_function_registry;
+use std::sync::Arc;
+
+fn env_u16(key: &str) -> u16 {
+    std::env::var(key)
+        .unwrap_or_else(|_| panic!("{key} must be set"))
+        .parse()
+        .unwrap_or_else(|e| panic!("{key} must be a u16: {e}"))
+}
+
+#[tokio::main]
+async fn main() -> ballista_core::error::Result<()> {
+    env_logger::init();
+
+    let config = ExecutorProcessConfig {
+        bind_host: "127.0.0.1".to_string(),
+        port: env_u16("CHAOS_EXECUTOR_PORT"),
+        grpc_port: env_u16("CHAOS_EXECUTOR_GRPC_PORT"),
+        scheduler_host: "127.0.0.1".to_string(),
+        scheduler_port: env_u16("CHAOS_SCHEDULER_PORT"),
+        vcores: std::env::var("CHAOS_CONCURRENT_TASKS")
+            .ok()
+            .and_then(|v| v.parse().ok())
+            .unwrap_or(4),
+        work_dir: std::env::var("CHAOS_WORK_DIR").ok(),
+        // The default is 60s. Executor-loss scenarios need the scheduler to see a
+        // missing heartbeat within seconds, not minutes.
+        executor_heartbeat_interval_seconds: std::env::var("CHAOS_HEARTBEAT_SECONDS")
+            .ok()
+            .and_then(|v| v.parse().ok())
+            .unwrap_or(1),
+        override_function_registry: Some(chaos_function_registry()),
+        ..Default::default()
+    };
+
+    start_executor_process(Arc::new(config)).await
+}

chaos-testing/src/bin/chaos-scheduler.rs

@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A Ballista scheduler whose session builder includes the chaos UDFs, with
+//! executor-loss detection tuned down from minutes to seconds.
+//!
+//! Mirrors `examples/examples/custom-scheduler.rs`.
+
+use ballista_core::error::BallistaError;
+use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::config::SchedulerConfig;
+use ballista_scheduler::scheduler_process::start_server;
+use chaos_testing::registry::chaos_session_state;
+use std::net::AddrParseError;
+use std::sync::Arc;
+
+fn env_parsed<T: std::str::FromStr>(key: &str, default: T) -> T {
+    std::env::var(key)
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(default)
+}
+
+#[tokio::main]
+async fn main() -> ballista_core::error::Result<()> {
+    env_logger::init();
+
+    let config = SchedulerConfig {
+        bind_host: "127.0.0.1".to_string(),
+        bind_port: std::env::var("CHAOS_SCHEDULER_PORT")
+            .expect("CHAOS_SCHEDULER_PORT must be set")
+            .parse()
+            .expect("CHAOS_SCHEDULER_PORT must be a u16"),
+        // Defaults are 180s / 15s, which would make every executor-kill scenario
+        // take three minutes. Tests override these per scenario.
+        executor_timeout_seconds: env_parsed("CHAOS_EXECUTOR_TIMEOUT_SECONDS", 5),
+        expire_dead_executor_interval_seconds: env_parsed(
+            "CHAOS_EXPIRE_INTERVAL_SECONDS",
+            1,
+        ),
+        task_max_failures: env_parsed("CHAOS_TASK_MAX_FAILURES", 4),
+        stage_max_failures: env_parsed("CHAOS_STAGE_MAX_FAILURES", 4),
+        override_session_builder: Some(Arc::new(chaos_session_state)),
+        ..Default::default()
+    };
+
+    let addr = format!("{}:{}", config.bind_host, config.bind_port);
+    let addr = addr
+        .parse()
+        .map_err(|e: AddrParseError| BallistaError::Configuration(e.to_string()))?;
+
+    let cluster = BallistaCluster::new_from_config(&config).await?;
+    start_server(cluster, addr, Arc::new(config)).await
+}

chaos-testing/src/budget.rs

@@ -0,0 +1,126 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::path::{Path, PathBuf};
+
+/// A filesystem-backed budget of injectable faults, shared across processes.
+///
+/// A budget of `n` is a directory containing `n` token files. Consuming a token
+/// is `fs::remove_file`, which is atomic across processes: exactly one caller
+/// can succeed for a given token, whichever executor it runs in. The directory
+/// outlives task retries and executor restarts, so the budget bounds the total
+/// number of injected faults for the whole run rather than per attempt.
+#[derive(Debug, Clone)]
+pub struct FaultBudget {
+    dir: PathBuf,
+}
+
+impl FaultBudget {
+    /// Create the budget directory with `tokens` tokens, replacing any existing one.
+    pub fn create(dir: &Path, tokens: usize) -> std::io::Result<Self> {
+        let _ = std::fs::remove_dir_all(dir);
+        std::fs::create_dir_all(dir)?;
+        for i in 0..tokens {
+            std::fs::write(dir.join(format!("token-{i}")), b"")?;
+        }
+        Ok(Self {
+            dir: dir.to_path_buf(),
+        })
+    }
+
+    /// Open an existing budget directory. Used by executor processes, which only
+    /// ever consume; a missing directory simply means no faults are available.
+    pub fn open(dir: &Path) -> Self {
+        Self {
+            dir: dir.to_path_buf(),
+        }
+    }
+
+    pub fn dir(&self) -> &Path {
+        &self.dir
+    }
+
+    pub fn remaining(&self) -> usize {
+        std::fs::read_dir(&self.dir)
+            .map(|entries| entries.flatten().count())
+            .unwrap_or(0)
+    }
+
+    /// Attempt to consume one token. Returns true iff this caller won the token.
+    ///
+    /// `remove_file` is the atomicity primitive: if two executors race for the
+    /// last token, exactly one `remove_file` returns Ok and the other errors.
+    pub fn try_consume(&self) -> bool {
+        let Ok(entries) = std::fs::read_dir(&self.dir) else {
+            return false;
+        };
+        for entry in entries.flatten() {
+            if std::fs::remove_file(entry.path()).is_ok() {
+                return true;
+            }
+        }
+        false
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::budget::FaultBudget;
+    use std::path::PathBuf;
+
+    fn temp_dir(name: &str) -> PathBuf {
+        let dir = std::env::temp_dir().join(format!("ballista-chaos-{name}"));
+        let _ = std::fs::remove_dir_all(&dir);
+        dir
+    }
+
+    #[test]
+    fn consumes_exactly_the_budget() {
+        let dir = temp_dir("budget-exact");
+        let budget = FaultBudget::create(&dir, 2).unwrap();
+
+        assert_eq!(budget.remaining(), 2);
+        assert!(budget.try_consume());
+        assert!(budget.try_consume());
+        assert!(
+            !budget.try_consume(),
+            "third consume must fail: budget was 2"
+        );
+        assert_eq!(budget.remaining(), 0);
+    }
+
+    #[test]
+    fn zero_budget_never_consumes() {
+        let dir = temp_dir("budget-zero");
+        let budget = FaultBudget::create(&dir, 0).unwrap();
+        assert!(!budget.try_consume());
+    }
+
+    #[test]
+    fn open_sees_tokens_created_by_another_handle() {
+        // This models a separate executor process reading the same budget dir.
+        let dir = temp_dir("budget-shared");
+        let creator = FaultBudget::create(&dir, 1).unwrap();
+        let other = FaultBudget::open(&dir);
+
+        assert!(other.try_consume(), "second handle must see the token");
+        assert!(
+            !creator.try_consume(),
+            "token already consumed by the other handle"
+        );
+    }
+}