diff --git a/Cargo.lock b/Cargo.lock index 82c2528884..4551bf4cbd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1045,6 +1045,26 @@ dependencies = [ "tokio", ] +[[package]] +name = "ballista-chaos" +version = "0.1.0" +dependencies = [ + "arrow", + "ballista", + "ballista-core", + "ballista-executor", + "ballista-scheduler", + "datafusion", + "env_logger", + "log", + "nix 0.29.0", + "reqwest 0.12.28", + "rstest", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "ballista-cli" version = "54.0.0" diff --git a/Cargo.toml b/Cargo.toml index f529c82df2..b29fc92dc4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ members = [ "ballista/executor", "ballista/scheduler", "benchmarks", + "chaos-testing", "examples", ] resolver = "3" diff --git a/chaos-testing/Cargo.toml b/chaos-testing/Cargo.toml new file mode 100644 index 0000000000..e2647e90c9 --- /dev/null +++ b/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" diff --git a/chaos-testing/README.md b/chaos-testing/README.md new file mode 100644 index 0000000000..cc2af73b03 --- /dev/null +++ b/chaos-testing/README.md @@ -0,0 +1,350 @@ + + +# ballista-chaos + +A fault-injection harness that runs real, multi-process Ballista clusters and +injects faults into real queries, to exercise Ballista's high-availability +(HA) machinery end to end. + +**This is a bug-hunting harness, not a regression suite in the usual sense.** +Its job is to surface real defects in Ballista's HA behavior. Where it finds +one, the corresponding test reproduces the bug rather than working around it. +Such a test is marked `#[ignore]` with the issue it reproduces, so that it +does not hold CI red on a bug it did not introduce, and is un-ignored — not +rewritten — when that issue is fixed, at which point it becomes the regression +test for the fix. Run them with `cargo test -p ballista-chaos -- --ignored`. +See [Findings](#findings) below for the confirmed bugs this harness has +found so far, each with the test that reproduces it. + +## Why this crate exists + +Ballista's HA state machine — stage/task retry, executor-loss recovery, +map-stage resubmission — lives in +`ballista/scheduler/src/state/execution_graph.rs`. Before this crate, that code +was exercised only by unit tests that hand-construct `TaskStatus` protobufs and +feed them directly into `ExecutionGraph` methods. Those tests are useful for +pinning the state machine's transition logic, but nothing drove it end to end: +no test ran a real query against a real multi-process cluster, killed a real +executor process, and checked that the _result_ was still correct. That gap is +exactly where the bugs in [Findings](#findings) were hiding — they only show up +when a real executor process dies mid-task, a real gRPC connection is refused, +or a real DataFusion error is really propagated through the real serialization +path, none of which a hand-built `TaskStatus` reproduces. + +`ballista-chaos` closes that gap: it spawns a real `ballista-scheduler` and +one or more real `ballista-executor` processes, runs a real multi-stage query +against them through the `ballista` client, and injects faults or kills +processes while the query is in flight. + +## Why fault injection uses UDFs, not `ChaosExec` + +Ballista's AQE planner already has a fault-injection mechanism: +`ChaosCreatingRule` (`ballista/scheduler/src/state/aqe/planner.rs:542`), which +wraps a plan node in `ChaosExec` when `chaos_execution_enabled` is set. It was +deliberately not reused here, because it cannot do what this harness needs: + +- It is wired into the **AQE physical-optimizer pipeline only** + (`plan_preparation_optimizers` in `planner.rs`). It does not run at all when + AQE is off, and every scenario in this crate must run under _both_ AQE + settings — the two planners have materially different join and retry + behavior, and a bug that only reproduces on one side is easy to miss if you + only test the other. +- It picks a **uniformly random plan node** to wrap + (`ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs`), not a node + the test chooses. A scenario that wants to fault "the scan of `facts`" or + "the shared join build side" specifically has no way to target it. +- It fires **probabilistically** (`chaos_execution_probability`), not + deterministically. A test built on it would need to loop-and-retry until the + fault happened to fire the right number of times, which is exactly the kind + of flakiness this harness is trying to avoid introducing. + +A SQL-level UDF (`chaos_fail`, `chaos_delay`, in `src/udf.rs`) sidesteps all +three problems: it lives in the query text itself, so it plans identically +(modulo AQE's own re-planning) whether AQE is on or off; its `guard` argument +lets a scenario target specific rows (and therefore specific partitions/tasks) +by writing an ordinary predicate; and it fires on every row where the guard is +true, subject only to the fault budget below — no probability, no retries of +the test itself. + +## How determinism works + +Every chaos scenario needs two things to be true: which rows/tasks fault must +be controlled, and how many attempts fault (across the whole cluster, across +retries and executor restarts) must be bounded. Two mechanisms provide these: + +- **The `guard` predicate.** The fixture (`src/fixture.rs`) is a small, fully + deterministic dataset: `facts(key, value)` joined to `dims(key, name)`, with + a known key distribution. A scenario passes a boolean expression over that + data as `chaos_fail`'s/`chaos_delay`'s first argument (e.g. `f.key = 7`); + since the data is fixed, this expression deterministically selects which + partitions the fault can fire in. +- **The filesystem fault budget** (`src/budget.rs`). A budget is a directory + of token files, created with a fixed token count. Consuming a token is + `fs::remove_file`, which is atomic across processes, so a budget of `n` + bounds the fault to firing at most `n` times _cluster-wide_ — across every + executor process, every task attempt, and every retry or restart — not `n` + times per process or per attempt. This is what makes "exactly one retryable + fault, then it must succeed" (Scenario A) and "faults never stop, so retries + must exhaust" (Scenario B) both expressible and deterministic. + +## The `OR TRUE` trap + +`Fixture::chaos_query` splices the injection expression into the query as +`WHERE {injection} IS NOT NULL`, not the more obvious `WHERE {injection} OR +TRUE`. This is deliberate and load-bearing: DataFusion's optimizer +constant-folds `expr OR TRUE` to the literal `TRUE` during logical +optimization, and once the predicate is a literal, the plan no longer +references the UDF call at all — it is dropped, not merely skipped. Every +fault-injection scenario built on `OR TRUE` would silently become a no-op: the +budget would never be consumed, the fault would never fire, and the suite +would report green while testing nothing. + +`chaos_fail`/`chaos_delay` always return `Some(guard)` (never `NULL`), so +`... IS NOT NULL` is always true but is not foldable to a constant without +evaluating the call — the optimizer has no way to know the result is always +non-null without invoking the (volatile) UDF. Two regression tests in +`src/fixture.rs` pin this: + +- `or_true_predicate_is_optimized_away_and_never_fires` proves the bad form + is eliminated from the plan and never consumes a budget token — pinning the + trap so it cannot silently return if someone "simplifies" the predicate back + to `OR TRUE`. +- `chaos_query_predicate_survives_optimization_and_fires` proves the + `IS NOT NULL` form the harness actually uses survives into the physical plan + and does fire. + +## How to run + +```sh +cargo test -p ballista-chaos # everything except the known-bug scenarios +cargo test -p ballista-chaos -- --ignored # the known-bug scenarios; these fail, on purpose +``` + +Every test that spawns a cluster does so through `TestCluster`, which holds +two locks for the cluster's lifetime: a process-wide mutex, and a +machine-wide `flock` (on a file in the system temp dir) for cluster-spawning +processes no in-process lock can see — a second cargo invocation in another +shell, or a runner like nextest that parallelizes test binaries. So scenarios +serialize themselves no matter how the test harness is invoked. This is not +cosmetic: each one starts a whole scheduler-plus-executors cluster, and +concurrent clusters exhaust ports and CPU and fail for reasons unrelated to +the scenario under test. `--test-threads=1` is therefore no longer required +(it will simply make the run marginally less confusing to read). + +CI runners are slow enough that cluster startup alone has blown a 30-second +registration deadline ("timed out waiting for 2 executors to register"); the +deadline is now 120s, and on expiry the error message carries the tail of +every child process log, so a recurrence in CI is diagnosable from the test +output alone. + +The `chaos-scheduler`/`chaos-executor` binaries are spawned as real child +processes rather than run in-process, but `cargo test` builds this crate's bin +targets along with its tests, so no separate build step is needed. The +harness locates them next to the running test executable, which is what makes +it work under any cargo profile (CI uses `--profile ci`, not `dev` or +`release`). + +Unit tests only (fast; the ones that do not spawn a cluster): + +```sh +cargo test -p ballista-chaos --lib +``` + +Each cluster's child-process logs (`scheduler.log`, `executor-0.log`, ...) are +written under that cluster's own temp directory, in a `logs/` subdirectory +(`TestCluster::log_dir()`). When a scenario fails, those logs are the first +place to look for what the scheduler and executors were actually doing. + +## Scenarios + +Every scenario runs under both `ballista.planner.adaptive.enabled=false` (AQE +off, the default, static `DefaultDistributedPlanner`) and `=true` (AQE on, the +experimental dynamic-join-selection planner) — 14 test cases total across the +7 scenarios below, plus a non-lettered `baseline_matches_local_datafusion` +sanity check that every other scenario's assertions depend on. + +| Scenario | Test | What it does | Expected result | +| -------- | ------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------- | +| A | `retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}` | Injects one retryable IO fault (budget 1); the retry must succeed and match baseline. | **Ignored (both), reproduces the error flattening tracked by [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 2.** | +| B | `exhausted_retries_fail_the_job_and_leave_the_cluster_healthy` | Injects an inexhaustible IO fault (budget 99 ≫ `task_max_failures`); job must fail, cluster must stay usable after. | Pass (both). | +| C | `panicking_task_fails_the_job_but_the_executor_survives` | Injects a task panic; job must fail non-retryably, both executor processes must survive, cluster must stay usable after. | Pass (both). | +| D | `executor_killed_mid_stage_is_recovered` | SIGKILLs an executor while its tasks are genuinely running (held open by `chaos_delay`); scheduler must reschedule onto the survivor and return the correct result. | **Ignored (both), reproduces [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 1.** | +| E | `executor_killed_after_shuffle_write_is_recovered` | SIGKILLs the map-side executor _after_ it wrote shuffle output, with a long executor timeout to bias toward the fetch-failure path rather than heartbeat expiry; downstream stage must re-run the map stage. | **Ignored (both), reproduces [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 1.** | +| F | `restarted_executor_rejoins_and_serves_queries` | Kills an executor, waits for the scheduler to reap it, restarts it, asserts the registered count returns to 2 and the cluster still serves the baseline query. | Pass (both), after the race fix in this crate (see below). | +| G | `killing_every_executor_terminates_the_job` | SIGKILLs every executor mid-query; the only requirement is that the job _terminates_ within 120s rather than hanging. | **Ignored (both), reproduces [#2029](https://github.com/apache/datafusion-ballista/issues/2029) — Finding 3.** | + +An ignored scenario above is not a defect in this harness, and its assertions +have not been weakened to make it pass — it reproduces a real Ballista bug and +is ignored only so that CI is not red on a bug this crate did not introduce. +Run the ignored scenarios with `-- --ignored --test-threads=1` to see the +failures; see [Findings](#findings) for what each one proves. + +### A note on Scenario F: the harness race that was fixed here + +`restarted_executor_rejoins_and_serves_queries` used to kill executor 0 and +restart it immediately, then assert `registered_executors() == 2`. That is a +harness bug, not a Ballista bug: SIGKILL does not deregister the executor, so +the scheduler keeps listing it until its heartbeat times out +(`executor_timeout_seconds`, 5s in this harness's defaults). Restarting +immediately races the scheduler's own reap: the assertion could observe three +executors (the dead one, the untouched survivor, and the freshly restarted +one) depending on timing, and failed intermittently with `left: 3, right: 2`. +Ballista was behaving correctly; the test just hadn't waited for the state it +was asserting about. The fix adds `TestCluster::await_executor_count(n)` (an +exact-count analogue of the existing `await_executors(n)`, which only waits for +_at least_ `n` — the right primitive for growing a cluster, but not for +observing a shrink) and has the scenario wait for the count to drop to 1 +before restarting, so the final assertion tests what the scenario name +actually promises. + +## Findings + +Four scenarios above fail because they have found real bugs in Ballista, not +because the harness is broken. Each is `#[ignore]`d against the issue it +reproduces so CI stays green on a tree whose bugs predate this crate, and each +keeps its original assertions: nothing is relaxed to manufacture a pass. When +the issue is fixed, delete the `#[ignore]` — the scenario is then the +regression test for it. One of the findings (#2028) has already gone through +that cycle in reverse: it was fixed on main, but a concurrent refactor +re-broke the same scenario through a different mechanism — see Finding 2. + +### Finding 1 — Shuffle-fetch failures lose their type, so the map-stage resubmit never fires + +Tracked by [#2027](https://github.com/apache/datafusion-ballista/issues/2027). + +**Proven by:** Scenario D (`executor_killed_mid_stage_is_recovered`) and +Scenario E (`executor_killed_after_shuffle_write_is_recovered`), both AQE +settings, `#[ignore]`d against that issue. + +Scenarios D and E are **races**, not deterministic reproducers, and this is +the one place in the crate where that is true. Killing an executor can be +noticed by the scheduler in either of two ways, and they are in a footrace: if +the heartbeat expires first, the `ExecutorLost` path recovers the job +correctly and the scenario passes in a few seconds; if a downstream task tries +to fetch shuffle output from the dead executor first, the bug below bites and +the job fails (or hangs until the scenario's timeout). Scenario D failed +locally on two of three runs; Scenario E passed locally but failed in CI +(`aqe_off`, with the reduce stage reporting a `FetchFailed` flattened inside +`DataFusionError::Shared` — see the log excerpt in #2027). The CI failure also +settles which way `executor_timeout_seconds` biases the race: raising it to +60s (as Scenario E does) _delays_ heartbeat expiry, so the downstream fetch +hits the dead executor first and the scenario exercises the broken +fetch-failure path; it only passes when the kill happens to land after the +reduce tasks already fetched their input. Un-ignoring these scenarios once +#2027 is fixed therefore also means pinning which of the two paths each +exercises — `executor_timeout_seconds` is the knob that decides the race, and +Scenario D currently leaves it at the harness default — otherwise they will be +flaky regression tests. + +The shuffle reader (`ballista/core/src/execution_plans/shuffle_reader.rs`) +correctly produces a typed `BallistaError::FetchFailed(executor_id, +map_stage_id, map_partition_id, desc)` when it cannot reach a dead executor, +and `ballista/core/src/error.rs`'s `impl From for FailedTask` +has a dedicated arm for exactly that variant (around line 205) which produces +`FailedReason::FetchPartitionError` — the signal +`ballista/scheduler/src/state/execution_graph.rs` (around line 826) uses to +resubmit the lost map stage rather than simply failing the job. + +The type does not survive to that point, however. Two real, non-test code +paths erase it before the executor reports its `TaskStatus`: + +- `ballista/executor/src/executor.rs:238-239`, in + `Executor::execute_query_stage`, converts the stage's result with + `result.map_err(|e| BallistaError::DataFusionError(Box::new(e)))` instead of + `BallistaError::from(e)` / `e.into()`. That bypasses the very unwrapping + logic `error.rs`'s `impl From for BallistaError` exists to + provide (`DataFusionError::ArrowError(e, _) => Self::from(*e)`, which would + otherwise recover a `FetchFailed` wrapped inside an `ArrowError::ExternalError`). +- `ballista/core/src/execution_plans/shuffle_writer.rs:245`, in + `ShuffleWriterExec`'s unpartitioned write branch, Debug-formats a + `BallistaError` into an opaque `DataFusionError::Execution(format!("{e:?}"))` + — a conversion that can never be undone by any later `.into()`, because the + original variant no longer exists, only its printed form. + +Either path leaves the executor reporting something like +`BallistaError::DataFusionError(Execution("FetchFailed(\"\", ..., +\"...Connection refused...\")"))` — the `FetchFailed` information is present +only as inert text inside a string. `error.rs`'s `FetchFailed` arm cannot match +a `DataFusionError::Execution`, so the task falls to the catch-all arm +(around line 248) and is marked `retryable: false`, `FailedReason::ExecutionError`. + +**Net effect:** when an executor dies after producing shuffle output that a +downstream stage still needs, Ballista fails the whole query instead of +re-running the map stage that produced it. + +### Finding 2 — Retryable IO errors are misclassified because the shuffle writer flattens them + +Originally tracked by +[#2028](https://github.com/apache/datafusion-ballista/issues/2028) (now +fixed); the surviving flattening mechanism is the one +[#2027](https://github.com/apache/datafusion-ballista/issues/2027) tracks. + +**Proven by:** Scenario A, both cases +(`retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}`), +`#[ignore]`d against #2027. + +The history matters here because the failure mode moved underneath the +harness. As first found, only the `aqe_on` case failed: an `IoError` raised on +a join's shared broadcast build side arrived wrapped as +`DataFusionError::Shared(Arc)`, and the classifier in +`ballista/core/src/error.rs` matched only a _direct_ +`DataFusionError::IoError`. That was #2028, and it was fixed by classifying on +`find_root()` instead (#2119). + +The sort-shuffle writer refactor (#2038, #2106) then re-broke both cases at an +earlier point in the pipeline: the shuffle write coordinator's error arm +(`ballista/core/src/execution_plans/shuffle_writer.rs`, in the coordinator +fan-out that distributes results to output-partition streams) converts any +task error with `DataFusionError::Execution(format!("{e:?}"))`. The injected +fault now reaches the classifier as +`Execution("IoError(Custom { .. })")` (`aqe_off`) or +`Execution("Shared(IoError(Custom { .. }))")` (`aqe_on`) — the variant exists +only as printed text, so `find_root()` has nothing to unwrap and the task is +marked non-retryable. This is the same type-erasing conversion Finding 1 +describes for `FetchFailed`, which is why both scenarios are ignored against +#2027 rather than the fixed #2028. + +**Net effect:** any genuine transient IO error inside a stage that writes +shuffle output — which is every non-final stage — is classified non-retryable, +turning what should be a retried task into an immediate job failure. + +### Finding 3 — Killing every executor hangs the job instead of failing it + +Tracked by [#2029](https://github.com/apache/datafusion-ballista/issues/2029). + +**Proven by:** Scenario G (`killing_every_executor_terminates_the_job`), both +AQE settings, `#[ignore]`d against that issue. + +With every executor dead mid-query, there is nothing left to schedule tasks +onto. The job does not terminate within the scenario's 120s timeout — the +scheduler waits rather than failing the query once it can determine no +executor can ever satisfy the remaining tasks. Note that this scenario +deliberately asserts only _termination_, not success or a particular error; it +is a hang detector, and what it detects is the hang itself. + +### For comparison: the heartbeat-expiry path does recover + +Ballista's HA recovery is not uniformly broken. Whenever the kill in Scenario +D or E happens to be noticed by heartbeat expiry (`ExecutorLost`) before any +downstream task fetches from the dead executor, the job recovers and matches +the baseline — that is why both scenarios pass on some runs. What Finding 1 +breaks is specifically the fetch-failure path, and both scenarios are ignored +against #2027 because whether a given run exercises that path is a timing +race, not something the harness currently controls. diff --git a/chaos-testing/src/bin/chaos-executor.rs b/chaos-testing/src/bin/chaos-executor.rs new file mode 100644 index 0000000000..a2efe82a72 --- /dev/null +++ b/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 +} diff --git a/chaos-testing/src/bin/chaos-scheduler.rs b/chaos-testing/src/bin/chaos-scheduler.rs new file mode 100644 index 0000000000..1fd4616c89 --- /dev/null +++ b/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(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 +} diff --git a/chaos-testing/src/budget.rs b/chaos-testing/src/budget.rs new file mode 100644 index 0000000000..4b5e3ace21 --- /dev/null +++ b/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 { + 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" + ); + } +} diff --git a/chaos-testing/src/cluster.rs b/chaos-testing/src/cluster.rs new file mode 100644 index 0000000000..7cd201adf6 --- /dev/null +++ b/chaos-testing/src/cluster.rs @@ -0,0 +1,673 @@ +// 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 nix::fcntl::{Flock, FlockArg}; +use std::fs::{File, OpenOptions}; +use std::net::TcpListener; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::sync::{Arc, OnceLock}; +use std::time::{Duration, Instant}; +use tokio::sync::{Mutex, OwnedMutexGuard}; + +/// Reserve a free TCP port by binding to :0 and immediately releasing it. +/// +/// Inherently racy, but adequate here: the child binds within milliseconds and +/// the tests are the only thing running. +fn free_port() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port"); + listener.local_addr().expect("local addr").port() +} + +/// Open a child process log file for append. +/// +/// Appending (rather than truncating) matters for Task 6's kill/restart +/// scenarios: a restarted executor reuses the same log path, and the prior +/// process's output is the evidence of why it died. It must not be wiped out +/// by the replacement process starting up. +fn open_log(path: &Path) -> std::io::Result { + OpenOptions::new().create(true).append(true).open(path) +} + +/// Locate a binary built by this crate. +/// +/// The profile directory is taken from the running test executable +/// (`//deps/`) rather than inferred from the profile's +/// settings. Inferring it from `cfg!(debug_assertions)` only works for the +/// stock `dev` and `release` profiles: CI builds under `--profile ci`, which +/// inherits `dev` but turns debug assertions off, so the binaries land in +/// `target/ci/` while the inference points at `target/release/`. Deriving the +/// directory from `current_exe` is correct for any profile and honours +/// `CARGO_TARGET_DIR` for free. +fn binary(name: &str) -> PathBuf { + let mut path = std::env::current_exe().expect("locate the test executable"); + path.pop(); // deps/ + path.pop(); // // + path.push(name); + assert!( + path.exists(), + "{} not found at {}. Run `cargo build -p ballista-chaos --bins` first.", + name, + path.display() + ); + path +} + +/// One supervised executor process. +/// +/// `child` is used by this task's `kill_executor`/`executor_is_alive`. `port`, +/// `grpc_port`, and `work_dir` are still not read anywhere yet — no code in +/// this task needed to target an executor by its network address or inspect +/// its working directory — so they keep a narrower `#[allow(dead_code)]` than +/// the struct-wide one Task 5 left; a later scenario that needs to address a +/// specific executor's port or inspect its shuffle files can drop it then. +pub(crate) struct ExecutorHandle { + pub(crate) child: Child, + #[allow(dead_code)] + pub(crate) port: u16, + #[allow(dead_code)] + pub(crate) grpc_port: u16, + #[allow(dead_code)] + pub(crate) work_dir: PathBuf, +} + +pub struct TestClusterBuilder { + executors: usize, + executor_timeout_seconds: u64, + expire_interval_seconds: u64, + task_max_failures: usize, + stage_max_failures: usize, + concurrent_tasks: usize, +} + +impl Default for TestClusterBuilder { + fn default() -> Self { + Self { + executors: 2, + // Ballista's defaults are 180s/15s, which would make an executor-kill + // scenario take three minutes to even notice the death. + executor_timeout_seconds: 5, + expire_interval_seconds: 1, + task_max_failures: 4, + stage_max_failures: 4, + concurrent_tasks: 4, + } + } +} + +impl TestClusterBuilder { + pub fn executors(mut self, n: usize) -> Self { + self.executors = n; + self + } + + /// How long the scheduler waits on a missing heartbeat before declaring the + /// executor lost. Scenario E raises this deliberately to isolate the + /// FetchPartitionError path from the ExecutorLost path. + pub fn executor_timeout_seconds(mut self, seconds: u64) -> Self { + self.executor_timeout_seconds = seconds; + self + } + + pub fn task_max_failures(mut self, n: usize) -> Self { + self.task_max_failures = n; + self + } + + pub async fn start(self) -> Result { + // Held for the cluster's whole lifetime, so only one cluster exists in + // this process at a time. `--test-threads=1` gives the same guarantee, + // but nothing forces a caller (CI runs a plain `cargo test` over the + // whole workspace) to pass it, and a dozen concurrent clusters exhaust + // ports and CPU and fail for reasons that have nothing to do with the + // scenario under test. Making the harness enforce its own requirement + // is more robust than documenting a flag. + let cluster_lock = cluster_lock().lock_owned().await; + + // The mutex above only serializes within one process. `cargo test` + // runs test binaries sequentially, but nothing else does: a second + // cargo invocation in another shell, an IDE's background test run, or + // a runner like nextest (which parallelizes binaries) would put two + // multi-process clusters on the machine at once, and on a two-core CI + // runner that is enough for executors to miss the registration + // deadline. flock is advisory and machine-wide, so the second process + // blocks here until the first one's cluster is gone. + let machine_lock = tokio::task::spawn_blocking(machine_lock) + .await + .map_err(|e| format!("acquire machine-wide cluster lock: {e}"))??; + + let temp = tempfile::tempdir().map_err(|e| e.to_string())?; + let log_dir = temp.path().join("logs"); + std::fs::create_dir_all(&log_dir).map_err(|e| e.to_string())?; + let scheduler_port = free_port(); + + let scheduler_log = log_dir.join("scheduler.log"); + let scheduler_stdout = open_log(&scheduler_log).map_err(|e| { + format!("open scheduler log {}: {e}", scheduler_log.display()) + })?; + let scheduler_stderr = open_log(&scheduler_log).map_err(|e| { + format!("open scheduler log {}: {e}", scheduler_log.display()) + })?; + + let mut scheduler = Command::new(binary("chaos-scheduler")); + scheduler + .env("CHAOS_SCHEDULER_PORT", scheduler_port.to_string()) + .env( + "CHAOS_EXECUTOR_TIMEOUT_SECONDS", + self.executor_timeout_seconds.to_string(), + ) + .env( + "CHAOS_EXPIRE_INTERVAL_SECONDS", + self.expire_interval_seconds.to_string(), + ) + .env( + "CHAOS_TASK_MAX_FAILURES", + self.task_max_failures.to_string(), + ) + .env( + "CHAOS_STAGE_MAX_FAILURES", + self.stage_max_failures.to_string(), + ) + .env( + "RUST_LOG", + std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()), + ) + .stdout(Stdio::from(scheduler_stdout)) + .stderr(Stdio::from(scheduler_stderr)); + let scheduler = scheduler + .spawn() + .map_err(|e| format!("spawn scheduler: {e}"))?; + + let mut cluster = TestCluster { + scheduler, + scheduler_port, + executors: Vec::new(), + temp, + log_dir, + builder: self, + _cluster_lock: cluster_lock, + _machine_lock: machine_lock, + }; + + for i in 0..cluster.builder.executors { + cluster.spawn_executor(i)?; + } + + let n = cluster.builder.executors; + cluster.await_executors(n).await?; + Ok(cluster) + } +} + +/// A multi-process Ballista cluster under test. Every child is killed on drop. +pub struct TestCluster { + scheduler: Child, + scheduler_port: u16, + pub(crate) executors: Vec, + temp: tempfile::TempDir, + log_dir: PathBuf, + builder: TestClusterBuilder, + /// Serializes clusters across the whole test process; see + /// [`TestClusterBuilder::start`]. `Drop for TestCluster` reaps every child + /// before this field is dropped, so the next cluster never starts until the + /// previous one's processes are gone. + _cluster_lock: OwnedMutexGuard<()>, + /// Serializes clusters across test *processes* on the same machine; see + /// [`TestClusterBuilder::start`]. + _machine_lock: Flock, +} + +/// The process-wide lock guaranteeing one cluster at a time. +fn cluster_lock() -> Arc> { + static LOCK: OnceLock>> = OnceLock::new(); + LOCK.get_or_init(|| Arc::new(Mutex::new(()))).clone() +} + +/// The machine-wide lock guaranteeing one cluster at a time across processes. +/// +/// Blocks until acquired, so call it from a blocking-friendly context. The +/// kernel releases a flock when its file handle closes, so a SIGKILLed test +/// process cannot leave the lock stuck. +fn machine_lock() -> Result, String> { + let path = std::env::temp_dir().join("ballista-chaos-cluster.lock"); + let file = OpenOptions::new() + .create(true) + .write(true) + .truncate(false) + .open(&path) + .map_err(|e| format!("open lock file {}: {e}", path.display()))?; + Flock::lock(file, FlockArg::LockExclusive) + .map_err(|(_, e)| format!("flock {}: {e}", path.display())) +} + +impl TestCluster { + pub fn builder() -> TestClusterBuilder { + TestClusterBuilder::default() + } + + /// The Ballista client URL. + pub fn scheduler_url(&self) -> String { + format!("df://127.0.0.1:{}", self.scheduler_port) + } + + /// The scheduler REST base URL. gRPC and REST share one port. + pub fn rest_url(&self) -> String { + format!("http://127.0.0.1:{}", self.scheduler_port) + } + + /// The shared directory for fixtures and fault budgets. Every executor can + /// read it, which is what makes the fault budget cluster-wide. + pub fn shared_dir(&self) -> &std::path::Path { + self.temp.path() + } + + /// Directory containing each child process's stdout/stderr log + /// (`scheduler.log`, `executor-{index}.log`). When a scenario fails, these + /// logs are the evidence of what the scheduler and executors were doing. + pub fn log_dir(&self) -> &Path { + &self.log_dir + } + + pub(crate) fn spawn_executor(&mut self, index: usize) -> Result<(), String> { + let port = free_port(); + let grpc_port = free_port(); + let work_dir = self.temp.path().join(format!("executor-{index}")); + std::fs::create_dir_all(&work_dir).map_err(|e| e.to_string())?; + + // Appends rather than truncates: a respawn at this same index (Task 6's + // kill/restart scenarios) must not erase the log of the process that + // just died. + let executor_log = self.log_dir.join(format!("executor-{index}.log")); + let executor_stdout = open_log(&executor_log).map_err(|e| { + format!("open executor {index} log {}: {e}", executor_log.display()) + })?; + let executor_stderr = open_log(&executor_log).map_err(|e| { + format!("open executor {index} log {}: {e}", executor_log.display()) + })?; + + let child = Command::new(binary("chaos-executor")) + .env("CHAOS_EXECUTOR_PORT", port.to_string()) + .env("CHAOS_EXECUTOR_GRPC_PORT", grpc_port.to_string()) + .env("CHAOS_SCHEDULER_PORT", self.scheduler_port.to_string()) + .env( + "CHAOS_CONCURRENT_TASKS", + self.builder.concurrent_tasks.to_string(), + ) + .env("CHAOS_WORK_DIR", work_dir.display().to_string()) + .env("CHAOS_HEARTBEAT_SECONDS", "1") + .env( + "RUST_LOG", + std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()), + ) + .stdout(Stdio::from(executor_stdout)) + .stderr(Stdio::from(executor_stderr)) + .spawn() + .map_err(|e| format!("spawn executor {index}: {e}"))?; + + if self.executors.len() > index { + self.executors[index] = ExecutorHandle { + child, + port, + grpc_port, + work_dir, + }; + } else { + self.executors.push(ExecutorHandle { + child, + port, + grpc_port, + work_dir, + }); + } + Ok(()) + } + + /// Block until `n` executors have registered with the scheduler. + /// + /// The deadline is generous because a loaded CI runner starts these child + /// processes slowly; a healthy cluster returns in a couple of seconds + /// regardless. On timeout the error carries every child's log tail — the + /// only evidence of a startup failure CI would otherwise throw away. + pub async fn await_executors(&self, n: usize) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(120); + loop { + if let Ok(count) = self.registered_executors().await + && count >= n + { + return Ok(()); + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for {n} executors to register\n{}", + self.log_tails() + )); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + } + + /// The last lines of every child process log, for timeout diagnostics. + fn log_tails(&self) -> String { + let mut out = String::new(); + let mut paths: Vec = std::fs::read_dir(&self.log_dir) + .map(|d| d.filter_map(|e| e.ok().map(|e| e.path())).collect()) + .unwrap_or_default(); + paths.sort(); + for path in paths { + let content = std::fs::read_to_string(&path).unwrap_or_default(); + let tail: Vec<&str> = content.lines().rev().take(20).collect(); + out.push_str(&format!( + "--- {} (last {} lines) ---\n", + path.display(), + tail.len() + )); + for line in tail.into_iter().rev() { + out.push_str(line); + out.push('\n'); + } + } + out + } + + /// Block until the scheduler considers exactly `n` executors registered. + /// + /// Unlike `await_executors` (which waits for *at least* `n`, the right + /// condition when growing a cluster), a SIGKILLed executor is not dropped + /// from `/api/executors` the instant it dies — the scheduler keeps + /// listing it until its heartbeat times out (`executor_timeout_seconds`). + /// A scenario that kills an executor and wants to observe the scheduler + /// actually reaping it (rather than just transiently over-counting) needs + /// to wait for the count to come down to `n` exactly, not merely reach it. + pub async fn await_executor_count(&self, n: usize) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(120); + loop { + if let Ok(count) = self.registered_executors().await + && count == n + { + return Ok(()); + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for exactly {n} registered executors" + )); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + } + + /// How many executors the scheduler currently considers registered. + pub async fn registered_executors(&self) -> Result { + let body: serde_json::Value = + reqwest::get(format!("{}/api/executors", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + Ok(body.as_array().map(|a| a.len()).unwrap_or(0)) + } + + /// The id of the single job the scheduler currently knows about. + /// + /// The harness runs one query at a time, so "the running job" is unambiguous. + pub async fn running_job_id(&self) -> Result { + let deadline = Instant::now() + Duration::from_secs(30); + loop { + let body: serde_json::Value = + reqwest::get(format!("{}/api/jobs", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + + if let Some(job) = body.as_array().and_then(|jobs| jobs.first()) + && let Some(id) = job.get("job_id").and_then(|v| v.as_str()) + { + return Ok(id.to_string()); + } + if Instant::now() > deadline { + return Err("timed out waiting for a job to appear".to_string()); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + async fn stages(&self, job_id: &str) -> Result { + reqwest::get(format!("{}/api/job/{job_id}/stages", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string()) + } + + /// Block until `stage_id` has at least one task in state Running. + /// + /// This is what lets a kill land *while the stage is genuinely executing*, + /// rather than after an arbitrary sleep that may fire too early or too late. + pub async fn await_stage_running( + &self, + job_id: &str, + stage_id: usize, + ) -> Result<(), String> { + self.await_stage_task_state(job_id, stage_id, "Running") + .await + } + + /// Block until every task in `stage_id` is Successful. + pub async fn await_stage_successful( + &self, + job_id: &str, + stage_id: usize, + ) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = self.stages(job_id).await?; + if let Some(stage) = find_stage(&stages, stage_id) + && let Some(tasks) = stage.get("tasks").and_then(|t| t.as_array()) + { + let all_ok = !tasks.is_empty() + && tasks.iter().all(|t| { + t.get("status").and_then(|s| s.as_str()) == Some("Successful") + }); + if all_ok { + return Ok(()); + } + } + if Instant::now() > deadline { + return Err(format!("timed out waiting for stage {stage_id} to succeed")); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + async fn await_stage_task_state( + &self, + job_id: &str, + stage_id: usize, + state: &str, + ) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = self.stages(job_id).await?; + if let Some(stage) = find_stage(&stages, stage_id) + && let Some(tasks) = stage.get("tasks").and_then(|t| t.as_array()) + { + let hit = tasks + .iter() + .any(|t| t.get("status").and_then(|s| s.as_str()) == Some(state)); + if hit { + return Ok(()); + } + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for a {state} task in stage {stage_id}" + )); + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + } + + /// The scheduler's short categorical job status: "Queued", "Running", + /// "Completed", "Failed", or "Invalid". + /// + /// The brief's original version read the JSON field literally named + /// `job_status`, but that field (`JobResponse::job_status` / + /// `handlers::format_job_status`'s second return value) is actually a long + /// human-readable sentence, e.g. "Completed. Produced 1 partition + /// containing 50 rows. Elapsed time: 49 ms." — never `"Successful"` as the + /// original doc comment claimed (a completed job reports `"Completed"`, + /// not `"Successful"`), and not stable/matchable for assertions. The short + /// categorical value the doc comment actually promises lives in the + /// sibling `status` field, so this reads that one instead. Verified against + /// a live cluster: `{"status": "Completed", "job_status": "Completed. + /// Produced 1 partition containing 50 rows. Elapsed time: 49 ms.", ...}`. + pub async fn job_status(&self, job_id: &str) -> Result { + let body: serde_json::Value = + reqwest::get(format!("{}/api/job/{job_id}", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + Ok(body + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string()) + } + + /// SIGKILL an executor. Not SIGTERM: a graceful shutdown would let the + /// executor deregister, which is a different (and much easier) code path + /// than the crash we are trying to test. + pub fn kill_executor(&mut self, index: usize) -> Result<(), String> { + use nix::sys::signal::{Signal, kill}; + use nix::unistd::Pid; + + let pid = self.executors[index].child.id(); + kill(Pid::from_raw(pid as i32), Signal::SIGKILL).map_err(|e| e.to_string())?; + let _ = self.executors[index].child.wait(); + Ok(()) + } + + /// Start a fresh executor process in the given slot and wait for it to register. + pub async fn restart_executor(&mut self, index: usize) -> Result<(), String> { + let expected = self.executors.len(); + self.spawn_executor(index)?; + self.await_executors(expected).await + } + + /// Whether an executor process is still alive. + pub fn executor_is_alive(&mut self, index: usize) -> bool { + matches!(self.executors[index].child.try_wait(), Ok(None)) + } +} + +/// Stage ids come back from the REST API as strings. +fn find_stage(stages: &serde_json::Value, stage_id: usize) -> Option<&serde_json::Value> { + stages.get("stages")?.as_array()?.iter().find(|s| { + s.get("stage_id").and_then(|v| v.as_str()) == Some(stage_id.to_string().as_str()) + }) +} + +/// If `child` already exited with a non-zero status, log the path of its +/// output so a human investigating a failed scenario knows where to look. +/// Then make sure it is actually gone. +fn reap(child: &mut Child, log_path: &Path) { + if let Ok(Some(status)) = child.try_wait() + && !status.success() + { + log::warn!( + "process exited with {status}; see log at {}", + log_path.display() + ); + } + let _ = child.kill(); + let _ = child.wait(); +} + +impl Drop for TestCluster { + fn drop(&mut self) { + for (index, executor) in self.executors.iter_mut().enumerate() { + let log_path = self.log_dir.join(format!("executor-{index}.log")); + reap(&mut executor.child, &log_path); + } + let scheduler_log = self.log_dir.join("scheduler.log"); + reap(&mut self.scheduler, &scheduler_log); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn cluster_starts_with_the_requested_executors_registered() { + let cluster = TestCluster::builder() + .executors(2) + .start() + .await + .expect("cluster must start"); + + // The scheduler's own view is the source of truth: if the executors did + // not register, every later scenario would silently run single-executor. + let executors: serde_json::Value = + reqwest::get(format!("{}/api/executors", cluster.rest_url())) + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!( + executors.as_array().map(|a| a.len()), + Some(2), + "expected 2 registered executors, got {executors:?}" + ); + } + + #[tokio::test] + async fn killed_executor_is_reaped_and_can_be_restarted() { + let mut cluster = TestCluster::builder() + .executors(2) + .executor_timeout_seconds(5) + .start() + .await + .unwrap(); + + assert_eq!(cluster.registered_executors().await.unwrap(), 2); + + cluster.kill_executor(0).unwrap(); + + // The scheduler must notice the missing heartbeat and drop the executor. + // With the defaults (180s timeout, 60s heartbeat) this would never happen + // inside a test; it works only because the harness turns both down. + cluster + .await_executor_count(1) + .await + .expect("scheduler never reaped the killed executor"); + + cluster.restart_executor(0).await.unwrap(); + assert_eq!( + cluster.registered_executors().await.unwrap(), + 2, + "restarted executor must re-register" + ); + } +} diff --git a/chaos-testing/src/fixture.rs b/chaos-testing/src/fixture.rs new file mode 100644 index 0000000000..3871fce4ef --- /dev/null +++ b/chaos-testing/src/fixture.rs @@ -0,0 +1,373 @@ +// 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 datafusion::arrow::array::{Int64Array, StringArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +/// Rows per `facts` partition file. +const FACT_ROWS: i64 = 2_000; +/// Number of `facts` partition files. Multiple files give the scan enough +/// parallelism to spread across executors. +const FACT_FILES: i64 = 8; +/// Number of distinct keys in `dims`. The join on `key` forces a shuffle, which +/// is what creates the multi-stage plan the HA paths need. +const DIM_KEYS: i64 = 50; + +/// A small deterministic Parquet dataset: `facts(key, value)` joined to +/// `dims(key, name)`. +pub struct Fixture { + facts_dir: PathBuf, + dims_dir: PathBuf, +} + +impl Fixture { + pub async fn write(dir: &Path) -> Result { + let ctx = SessionContext::new(); + let facts_dir = dir.join("facts"); + let dims_dir = dir.join("dims"); + std::fs::create_dir_all(&facts_dir)?; + std::fs::create_dir_all(&dims_dir)?; + + let fact_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ])); + + for file in 0..FACT_FILES { + let keys: Vec = (0..FACT_ROWS) + .map(|r| (r + file * FACT_ROWS) % DIM_KEYS) + .collect(); + let values: Vec = (0..FACT_ROWS).map(|r| r + file * FACT_ROWS).collect(); + let batch = RecordBatch::try_new( + fact_schema.clone(), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from(values)), + ], + )?; + let df = ctx.read_batch(batch)?; + df.write_parquet( + facts_dir + .join(format!("part-{file}.parquet")) + .to_str() + .unwrap(), + datafusion::dataframe::DataFrameWriteOptions::new(), + None, + ) + .await?; + } + + let dim_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + dim_schema, + vec![ + Arc::new(Int64Array::from((0..DIM_KEYS).collect::>())), + Arc::new(StringArray::from( + (0..DIM_KEYS) + .map(|k| format!("dim-{k}")) + .collect::>(), + )), + ], + )?; + ctx.read_batch(batch)? + .write_parquet( + dims_dir.join("part-0.parquet").to_str().unwrap(), + datafusion::dataframe::DataFrameWriteOptions::new(), + None, + ) + .await?; + + Ok(Self { + facts_dir, + dims_dir, + }) + } + + /// `CREATE EXTERNAL TABLE` statements, to run against any SessionContext + /// (local for the baseline, Ballista for the cluster run). + pub fn register_sql(&self) -> Vec { + vec![ + format!( + "CREATE EXTERNAL TABLE facts STORED AS PARQUET LOCATION '{}'", + self.facts_dir.display() + ), + format!( + "CREATE EXTERNAL TABLE dims STORED AS PARQUET LOCATION '{}'", + self.dims_dir.display() + ), + ] + } + + /// The chaos-free query. A join plus a grouped aggregate: at least two + /// stages, with a shuffle between them. + pub fn baseline_query() -> &'static str { + "SELECT d.name, COUNT(*) AS n, SUM(f.value) AS total \ + FROM facts f JOIN dims d ON f.key = d.key \ + GROUP BY d.name ORDER BY d.name" + } + + /// The same query with a chaos UDF spliced into the WHERE clause. + /// + /// `injection` is a complete `chaos_*(...)` call returning BOOLEAN. Its own + /// guard argument selects which rows fault (and hence which partitions and + /// tasks). The predicate is written so every row passes through regardless: + /// the guard decides only *where the fault fires*, never which rows survive, + /// so a chaos run must return exactly the baseline result. + /// + /// The predicate uses `IS NOT NULL` rather than `OR TRUE`: DataFusion's + /// optimizer constant-folds `expr OR TRUE` to the literal `TRUE` and drops + /// the volatile UDF call entirely (verified empirically in the + /// `or_true_predicate_is_optimized_away_and_never_fires` and + /// `chaos_query_predicate_survives_optimization_and_fires` tests below), + /// which would silently disarm every fault injection. `chaos_fail`/ + /// `chaos_delay` always return `Some(guard)` (see udf.rs), so + /// `... IS NOT NULL` is always true without being foldable to a constant, + /// and the optimizer must still evaluate the call to determine nullness. + /// + /// Example: `chaos_query("chaos_fail(f.key = 7, 'io', '/tmp/b')")` + pub fn chaos_query(injection: &str) -> String { + format!( + "SELECT d.name, COUNT(*) AS n, SUM(f.value) AS total \ + FROM facts f JOIN dims d ON f.key = d.key \ + WHERE {injection} IS NOT NULL \ + GROUP BY d.name ORDER BY d.name" + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + #[tokio::test] + async fn baseline_query_is_deterministic_and_non_empty() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + + let ctx = SessionContext::new(); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + let rows = ctx + .sql(Fixture::baseline_query()) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total: usize = rows.iter().map(|b| b.num_rows()).sum(); + assert!(total > 0, "baseline query must return rows"); + + // Determinism: the same query over the same data must give the same answer. + let rows2 = ctx + .sql(Fixture::baseline_query()) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + datafusion::arrow::util::pretty::pretty_format_batches(&rows) + .unwrap() + .to_string(), + datafusion::arrow::util::pretty::pretty_format_batches(&rows2) + .unwrap() + .to_string(), + ); + } + + /// Empirically confirms the risk called out on `chaos_query`: DataFusion's + /// constant-folding simplifies `expr OR TRUE` to the literal `TRUE` during + /// logical optimization, and once the predicate is a literal the join/filter + /// no longer references the UDF call at all, so it is never invoked. If this + /// test ever starts failing (budget still consumed), the optimizer's folding + /// behavior changed and `chaos_query`'s `IS NOT NULL` predicate should be + /// re-verified instead of reverting to `OR TRUE`. + #[tokio::test] + async fn or_true_predicate_is_optimized_away_and_never_fires() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + let budget_dir = dir.path().join("budget"); + let budget = crate::budget::FaultBudget::create(&budget_dir, 1).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_udf(crate::udf::chaos_fail_udf().as_ref().clone()); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + // key = 7 genuinely exists in the generated data (keys cycle 0..DIM_KEYS), + // so an empty-partition short-circuit is not what's suppressing the call. + let sql = format!( + "SELECT d.name, COUNT(*) AS n, SUM(f.value) AS total \ + FROM facts f JOIN dims d ON f.key = d.key \ + WHERE chaos_fail(f.key = 7, 'io', '{}') OR TRUE \ + GROUP BY d.name ORDER BY d.name", + budget_dir.display() + ); + + let explain = ctx + .sql(&format!("EXPLAIN {sql}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + let explain_str = + datafusion::arrow::util::pretty::pretty_format_batches(&explain) + .unwrap() + .to_string(); + println!("EXPLAIN (OR TRUE form):\n{explain_str}"); + assert!( + !explain_str.contains("chaos_fail"), + "expected the optimizer to eliminate the chaos_fail call from the plan, but it is still present:\n{explain_str}" + ); + + let result = ctx.sql(&sql).await.unwrap().collect().await; + assert!( + result.is_ok(), + "expected the query to succeed because the fault never fires, got {result:?}" + ); + assert_eq!( + budget.remaining(), + 1, + "the token must be untouched: `OR TRUE` is constant-folded away, so chaos_fail is never invoked" + ); + } + + /// The predicate form `Fixture::chaos_query` actually uses. Unlike `OR TRUE`, + /// `chaos_fail(...) IS NOT NULL` is not foldable to a constant (the UDF's + /// return type/value is not known to the optimizer without invoking it), so + /// the call must survive into the physical plan and the fault must fire. + #[tokio::test] + async fn chaos_query_predicate_survives_optimization_and_fires() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + let budget_dir = dir.path().join("budget"); + let budget = crate::budget::FaultBudget::create(&budget_dir, 1).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_udf(crate::udf::chaos_fail_udf().as_ref().clone()); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + let injection = + format!("chaos_fail(f.key = 7, 'io', '{}')", budget_dir.display()); + let sql = Fixture::chaos_query(&injection); + + let explain = ctx + .sql(&format!("EXPLAIN {sql}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + let explain_str = + datafusion::arrow::util::pretty::pretty_format_batches(&explain) + .unwrap() + .to_string(); + println!("EXPLAIN (IS NOT NULL form):\n{explain_str}"); + assert!( + explain_str.contains("chaos_fail"), + "expected the chaos_fail call to survive into the plan, but it is missing:\n{explain_str}" + ); + + let err = ctx.sql(&sql).await.unwrap().collect().await.unwrap_err(); + // The join fans the same filtered stream out to multiple consumers + // (CollectLeft build side plus the probe side), so DataFusion wraps the + // propagated error in `DataFusionError::Shared` rather than surfacing the + // bare `IoError` directly; `find_root` unwraps that layer. + assert!( + matches!( + err.find_root(), + datafusion::error::DataFusionError::IoError(_) + ), + "expected an IoError (possibly Shared-wrapped) from the injected fault, got {err:?}" + ); + assert_eq!( + budget.remaining(), + 0, + "the token must be consumed: the predicate must force chaos_fail to be evaluated" + ); + } + + /// CRITICAL INVARIANT: When a chaos fault cannot fire (budget exhausted), + /// `chaos_query()` must return EXACTLY the same result as `baseline_query()`. + /// + /// This invariant underpins every HA scenario: we detect when a re-run stage + /// duplicates or drops partitions by asserting that the result after a fault + /// is identical to the baseline. The chaos_query predicate `WHERE ... IS NOT + /// NULL` preserves the row set only because chaos_fail/chaos_delay always + /// return `Some(guard)` (never NULL). If a future change ever returned NULL + /// from either UDF, `IS NOT NULL` would silently drop rows and no test would + /// catch it — until a real distributed run failed mysteriously. This test + /// pins that invariant before any such refactor happens. + #[tokio::test] + async fn chaos_query_without_a_firing_fault_equals_baseline() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + let budget_dir = dir.path().join("budget"); + // 0 tokens: the fault cannot fire, no matter how many times it is called. + let _budget = crate::budget::FaultBudget::create(&budget_dir, 0).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_udf(crate::udf::chaos_fail_udf().as_ref().clone()); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + let injection = + format!("chaos_fail(f.key = 7, 'io', '{}')", budget_dir.display()); + let chaos_sql = Fixture::chaos_query(&injection); + let baseline_sql = Fixture::baseline_query(); + + let chaos_rows = ctx.sql(&chaos_sql).await.unwrap().collect().await.unwrap(); + let baseline_rows = ctx + .sql(baseline_sql) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let chaos_str = + datafusion::arrow::util::pretty::pretty_format_batches(&chaos_rows) + .unwrap() + .to_string(); + let baseline_str = + datafusion::arrow::util::pretty::pretty_format_batches(&baseline_rows) + .unwrap() + .to_string(); + + assert_eq!( + chaos_str, baseline_str, + "chaos_query with a non-firing fault must return the same rows as baseline_query" + ); + } +} diff --git a/chaos-testing/src/lib.rs b/chaos-testing/src/lib.rs new file mode 100644 index 0000000000..5c1c527643 --- /dev/null +++ b/chaos-testing/src/lib.rs @@ -0,0 +1,27 @@ +// 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. + +//! Fault-injection harness for testing Ballista high availability. +//! +//! See `specs/2026-07-13-ha-chaos-integration-test.md` in the ballista-ai-dev +//! repository for the design. + +pub mod budget; +pub mod cluster; +pub mod fixture; +pub mod registry; +pub mod udf; diff --git a/chaos-testing/src/registry.rs b/chaos-testing/src/registry.rs new file mode 100644 index 0000000000..8f4df2946f --- /dev/null +++ b/chaos-testing/src/registry.rs @@ -0,0 +1,103 @@ +// 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 crate::udf::{chaos_delay_udf, chaos_fail_udf}; +use ballista_core::registry::BallistaFunctionRegistry; +use datafusion::execution::FunctionRegistry; +use datafusion::execution::session_state::{SessionState, SessionStateBuilder}; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +/// The executor's function registry, extended with the chaos UDFs. +/// +/// Starts from `BallistaFunctionRegistry::default()` (the DataFusion built-ins) +/// so overriding the registry does not silently drop the standard functions. +pub fn chaos_function_registry() -> Arc { + let mut registry = BallistaFunctionRegistry::default(); + for udf in [chaos_fail_udf(), chaos_delay_udf()] { + registry + .scalar_functions + .insert(udf.name().to_string(), udf); + } + Arc::new(registry) +} + +/// The scheduler's session state, extended with the chaos UDFs. +/// +/// The scheduler needs them to plan a logical plan that references them; the +/// executor needs them to decode the physical plan. Both must agree. +/// +/// Note: `SessionStateBuilder::with_scalar_functions` *replaces* the builder's +/// scalar function list rather than appending to it, so calling it after +/// `with_default_features()` would silently drop every built-in (including +/// `abs`). To avoid that footgun, the chaos UDFs are registered on the already +/// built `SessionState` via `FunctionRegistry::register_udf`, which inserts +/// into the existing map instead of replacing it. +pub fn chaos_session_state( + config: SessionConfig, +) -> datafusion::error::Result { + let mut state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + + for udf in [chaos_fail_udf(), chaos_delay_udf()] { + state.register_udf(udf)?; + } + + Ok(state) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registry_resolves_the_chaos_udfs_by_name() { + // The executor resolves UDFs by name at plan-decode time. If these names + // are absent, every chaos query fails to deserialize on the executor. + let registry = chaos_function_registry(); + assert!(registry.udf("chaos_fail").is_ok()); + assert!(registry.udf("chaos_delay").is_ok()); + } + + #[test] + fn registry_retains_the_standard_functions() { + // Regression guard: overriding the registry must not drop the built-ins, + // or ordinary query operators stop working on the executor. + let registry = chaos_function_registry(); + assert!(registry.udf("abs").is_ok()); + } + + #[test] + fn session_state_resolves_the_chaos_udfs_by_name() { + // Mirrors registry_resolves_the_chaos_udfs_by_name, but for the + // scheduler's session state rather than the executor's registry. + let state = chaos_session_state(SessionConfig::new()).unwrap(); + assert!(state.udf("chaos_fail").is_ok()); + assert!(state.udf("chaos_delay").is_ok()); + } + + #[test] + fn session_state_retains_the_standard_functions() { + // Regression guard for the with_scalar_functions replace-vs-append + // footgun described on chaos_session_state: the scheduler must still be + // able to plan ordinary queries that use built-in functions. + let state = chaos_session_state(SessionConfig::new()).unwrap(); + assert!(state.udf("abs").is_ok()); + } +} diff --git a/chaos-testing/src/udf.rs b/chaos-testing/src/udf.rs new file mode 100644 index 0000000000..afd7987997 --- /dev/null +++ b/chaos-testing/src/udf.rs @@ -0,0 +1,328 @@ +// 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 crate::budget::FaultBudget; +use arrow::array::{Array, BooleanArray}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::cast::as_boolean_array; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + TypeSignature, Volatility, +}; +use std::path::Path; +use std::sync::Arc; + +/// `chaos_fail(guard BOOLEAN, mode UTF8, budget_dir UTF8) -> BOOLEAN` +/// +/// Pass-through: returns `guard` unchanged. If any row in the batch has +/// `guard = true` and a fault token is available in `budget_dir`, consumes one +/// token and injects a fault: +/// +/// - `io` -> `DataFusionError::IoError` (Ballista: retryable, counts to failures) +/// - `exec` -> `DataFusionError::Execution` (Ballista: non-retryable) +/// - `panic` -> `panic!` (Ballista: caught, becomes non-retryable Internal) +/// +/// Volatile so DataFusion neither constant-folds nor CSEs it away. +#[derive(Debug, PartialEq, Eq, Hash)] +struct ChaosFail { + signature: Signature, +} + +impl Default for ChaosFail { + fn default() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![ + DataType::Boolean, + DataType::Utf8, + DataType::Utf8, + ]), + Volatility::Volatile, + ), + } + } +} + +impl ScalarUDFImpl for ChaosFail { + fn name(&self) -> &str { + "chaos_fail" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let guard = args.args[0].clone().into_array(args.number_rows)?; + let mode = scalar_utf8(&args.args[1], "mode")?; + let budget_dir = scalar_utf8(&args.args[2], "budget_dir")?; + + let guard = as_boolean_array(&guard)?; + let selected = (0..guard.len()).any(|i| !guard.is_null(i) && guard.value(i)); + + if selected && FaultBudget::open(Path::new(&budget_dir)).try_consume() { + match mode.as_str() { + "io" => { + let msg = "chaos_fail: injected retryable IO fault"; + log::error!("{msg}"); + return Err(DataFusionError::IoError(std::io::Error::other(msg))); + } + "exec" => { + let msg = + "chaos_fail: injected non-retryable execution fault".to_string(); + log::error!("{msg}"); + return Err(DataFusionError::Execution(msg)); + } + "panic" => { + log::error!("chaos_fail: injected panic"); + panic!("chaos_fail: injected panic"); + } + other => { + return Err(DataFusionError::Configuration(format!( + "chaos_fail: unknown mode {other:?} (expected io, exec, or panic)" + ))); + } + } + } + + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from( + (0..guard.len()) + .map(|i| (!guard.is_null(i)).then(|| guard.value(i))) + .collect::>>(), + )))) + } +} + +/// `chaos_delay(guard BOOLEAN, ms INT64) -> BOOLEAN` +/// +/// Pass-through: returns `guard` unchanged, sleeping `ms` milliseconds per batch +/// in which any row has `guard = true`. Used to hold a stage open long enough for +/// the harness to kill an executor while the stage is genuinely running. +#[derive(Debug, PartialEq, Eq, Hash)] +struct ChaosDelay { + signature: Signature, +} + +impl Default for ChaosDelay { + fn default() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![DataType::Boolean, DataType::Int64]), + Volatility::Volatile, + ), + } + } +} + +impl ScalarUDFImpl for ChaosDelay { + fn name(&self) -> &str { + "chaos_delay" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let guard = args.args[0].clone().into_array(args.number_rows)?; + let ms = match &args.args[1] { + ColumnarValue::Scalar(s) => match s { + datafusion::scalar::ScalarValue::Int64(Some(v)) => *v as u64, + other => { + return Err(DataFusionError::Configuration(format!( + "chaos_delay: ms must be a non-null INT64 literal, got {other:?}" + ))); + } + }, + ColumnarValue::Array(_) => { + return Err(DataFusionError::Configuration( + "chaos_delay: ms must be a literal, not a column".to_string(), + )); + } + }; + + let guard = as_boolean_array(&guard)?; + let selected = (0..guard.len()).any(|i| !guard.is_null(i) && guard.value(i)); + if selected { + std::thread::sleep(std::time::Duration::from_millis(ms)); + } + + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from( + (0..guard.len()) + .map(|i| (!guard.is_null(i)).then(|| guard.value(i))) + .collect::>>(), + )))) + } +} + +fn scalar_utf8(value: &ColumnarValue, arg: &str) -> Result { + match value { + ColumnarValue::Scalar(datafusion::scalar::ScalarValue::Utf8(Some(s))) => { + Ok(s.clone()) + } + other => Err(DataFusionError::Configuration(format!( + "chaos_fail: {arg} must be a non-null UTF8 literal, got {other:?}" + ))), + } +} + +/// The `chaos_fail` UDF. +pub fn chaos_fail_udf() -> Arc { + Arc::new(ScalarUDF::from(ChaosFail::default())) +} + +/// The `chaos_delay` UDF. +pub fn chaos_delay_udf() -> Arc { + Arc::new(ScalarUDF::from(ChaosDelay::default())) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::budget::FaultBudget; + use arrow::array::BooleanArray; + use datafusion::arrow::array::RecordBatch; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + + fn temp_dir(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join(format!("ballista-chaos-udf-{name}")); + let _ = std::fs::remove_dir_all(&dir); + dir + } + + /// Build a one-column table `t(guard BOOLEAN)` with the given values. + async fn ctx_with_guards(guards: Vec) -> SessionContext { + let ctx = SessionContext::new(); + ctx.register_udf(chaos_fail_udf().as_ref().clone()); + ctx.register_udf(chaos_delay_udf().as_ref().clone()); + + let schema = Arc::new(Schema::new(vec![Field::new( + "guard", + DataType::Boolean, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(BooleanArray::from(guards))], + ) + .unwrap(); + ctx.register_batch("t", batch).unwrap(); + ctx + } + + #[tokio::test] + async fn io_mode_errors_when_a_token_is_available() { + let dir = temp_dir("io-fires"); + FaultBudget::create(&dir, 1).unwrap(); + let ctx = ctx_with_guards(vec![true]).await; + + let sql = format!("SELECT chaos_fail(guard, 'io', '{}') FROM t", dir.display()); + let err = ctx.sql(&sql).await.unwrap().collect().await.unwrap_err(); + + // Must be an IoError: that is the only variant Ballista treats as retryable. + assert!( + matches!(err, datafusion::error::DataFusionError::IoError(_)), + "expected IoError, got {err:?}" + ); + } + + #[tokio::test] + async fn passes_through_when_budget_is_exhausted() { + let dir = temp_dir("io-exhausted"); + FaultBudget::create(&dir, 0).unwrap(); + let ctx = ctx_with_guards(vec![true, false]).await; + + let sql = format!( + "SELECT chaos_fail(guard, 'io', '{}') AS g FROM t", + dir.display() + ); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + + // Pass-through: output must equal the input guard column. + let col = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(col.value(0)); + assert!(!col.value(1)); + } + + #[tokio::test] + async fn does_not_fire_when_no_guard_row_is_true() { + let dir = temp_dir("io-no-guard"); + FaultBudget::create(&dir, 1).unwrap(); + let ctx = ctx_with_guards(vec![false, false]).await; + + let sql = format!("SELECT chaos_fail(guard, 'io', '{}') FROM t", dir.display()); + ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + + // The token must be untouched: the guard never selected this data. + assert_eq!(FaultBudget::open(&dir).remaining(), 1); + } + + /// Tests that panic mode injects a panic when a guard row is true. + /// The panic crosses the async stream boundary. The authoritative end-to-end + /// panic test is Scenario C in Task 8, where the executor's `catch_unwind` + /// converts the panic into a `FailedTask`. + #[should_panic(expected = "chaos_fail: injected panic")] + #[tokio::test] + async fn panic_mode_panics() { + let dir = temp_dir("panic-fires"); + FaultBudget::create(&dir, 1).unwrap(); + let ctx = ctx_with_guards(vec![true]).await; + + let sql = format!( + "SELECT chaos_fail(guard, 'panic', '{}') FROM t", + dir.display() + ); + let _ = ctx.sql(&sql).await.unwrap().collect().await; + } + + #[tokio::test] + async fn delay_sleeps_and_passes_through() { + let ctx = ctx_with_guards(vec![true]).await; + let start = std::time::Instant::now(); + + let batches = ctx + .sql("SELECT chaos_delay(guard, 150) AS g FROM t") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert!(start.elapsed() >= std::time::Duration::from_millis(150)); + let col = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(col.value(0)); + } +} diff --git a/chaos-testing/tests/common/mod.rs b/chaos-testing/tests/common/mod.rs new file mode 100644 index 0000000000..1a2755d3d2 --- /dev/null +++ b/chaos-testing/tests/common/mod.rs @@ -0,0 +1,131 @@ +// 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 ballista::prelude::{SessionConfigExt, SessionContextExt}; +use ballista_core::config::BALLISTA_ADAPTIVE_PLANNER_ENABLED; +use chaos_testing::budget::FaultBudget; +use chaos_testing::cluster::TestCluster; +use chaos_testing::fixture::Fixture; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::{SessionConfig, SessionContext}; + +/// One cluster plus its fixture, wired for a single scenario. +pub struct ChaosRun { + pub cluster: TestCluster, + pub fixture: Fixture, + ctx: SessionContext, +} + +impl ChaosRun { + pub async fn start(aqe: bool, executors: usize) -> Self { + Self::start_with(aqe, executors, 5).await + } + + /// `executor_timeout_seconds` selects which HA mechanism a kill surfaces + /// through: short biases toward ExecutorLost (heartbeat expiry), long biases + /// toward FetchPartitionError (a downstream fetch from a dead executor). + pub async fn start_with( + aqe: bool, + executors: usize, + executor_timeout_seconds: u64, + ) -> Self { + let _ = env_logger::builder().is_test(true).try_init(); + + let cluster = TestCluster::builder() + .executors(executors) + .executor_timeout_seconds(executor_timeout_seconds) + .start() + .await + .expect("cluster must start"); + + let fixture = Fixture::write(cluster.shared_dir()) + .await + .expect("fixture must be written"); + + let config = SessionConfig::new_with_ballista() + .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, aqe); + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + + let ctx = SessionContext::remote_with_state(&cluster.scheduler_url(), state) + .await + .expect("client must connect to the scheduler"); + + // Registered *after* `remote_with_state`, not baked into the initial + // SessionState: `SessionStateExt::upgrade_for_ballista` (ballista/core) + // rebuilds the state via `with_scalar_functions(ballista_scalar_functions())`, + // which replaces rather than merges the scalar-function map, silently + // dropping any custom UDF registered beforehand. `register_udf` mutates + // the already-upgraded state's registry directly, so it survives. + ctx.register_udf(chaos_testing::udf::chaos_fail_udf().as_ref().clone()); + ctx.register_udf(chaos_testing::udf::chaos_delay_udf().as_ref().clone()); + + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + Self { + cluster, + fixture, + ctx, + } + } + + /// A fault budget in the cluster's shared directory, visible to every executor. + pub fn budget(&self, name: &str, tokens: usize) -> FaultBudget { + let dir = self.cluster.shared_dir().join(format!("budget-{name}")); + FaultBudget::create(&dir, tokens).expect("budget must be created") + } + + /// Run a query on the cluster, returning the formatted result. + pub async fn sql(&self, sql: &str) -> Result { + let df = self.ctx.sql(sql).await.map_err(|e| e.to_string())?; + let batches = df.collect().await.map_err(|e| e.to_string())?; + Ok(pretty_format_batches(&batches) + .map_err(|e| e.to_string())? + .to_string()) + } + + /// The expected answer, computed by plain local DataFusion. + pub async fn local_baseline(&self) -> String { + let ctx = SessionContext::new(); + for stmt in self.fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + let batches = ctx + .sql(Fixture::baseline_query()) + .await + .unwrap() + .collect() + .await + .unwrap(); + pretty_format_batches(&batches).unwrap().to_string() + } + + /// A clone of the session context, for running a query concurrently with a fault. + /// + /// Not used until Task 9's executor-loss scenarios, which submit a query in + /// a spawned task while the main task polls the scheduler and kills an + /// executor. + #[allow(dead_code)] + pub fn clone_ctx(&self) -> SessionContext { + self.ctx.clone() + } +} diff --git a/chaos-testing/tests/ha.rs b/chaos-testing/tests/ha.rs new file mode 100644 index 0000000000..104d4a6e7b --- /dev/null +++ b/chaos-testing/tests/ha.rs @@ -0,0 +1,417 @@ +// 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. + +//! End-to-end high-availability scenarios against a real multi-process cluster. +//! +//! Every scenario runs under both AQE settings. The AQE-on axis is where bugs are +//! expected: a resubmitted stage under AQE is re-planned against runtime +//! statistics, so a re-run map stage can come back with a different plan than the +//! one whose output was lost. +//! +//! Every test in this file spawns a whole multi-process cluster, so this file +//! must always be run with `--test-threads=1` or ports and CPU will be +//! exhausted by concurrent clusters. + +mod common; + +use common::ChaosRun; +use rstest::rstest; + +/// The cluster must agree with local DataFusion before any fault is injected. +/// Every recovery scenario asserts against this baseline, so if it is wrong, +/// every other assertion is meaningless. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn baseline_matches_local_datafusion(#[case] aqe: bool) { + let run = ChaosRun::start(aqe, 2).await; + + let expected = run.local_baseline().await; + let actual = run + .sql(ballista_chaos_query_baseline()) + .await + .expect("baseline query must succeed on the cluster"); + + assert_eq!( + actual, expected, + "cluster result must match local DataFusion" + ); +} + +fn ballista_chaos_query_baseline() -> &'static str { + chaos_testing::fixture::Fixture::baseline_query() +} + +use chaos_testing::fixture::Fixture; + +/// Scenario A: one retryable (IO) fault, budget 1. +/// +/// A single task attempt anywhere in the cluster faults; the budget is then +/// exhausted, so the retry must succeed. The load-bearing assertion is that the +/// result still equals the baseline: a retried stage is exactly where duplicated +/// or dropped partitions would show up. +/// +/// This scenario is split into one test per AQE setting rather than being an +/// `rstest` over both, because historically only the AQE-on case failed and +/// `rstest` cannot ignore an individual case. +/// +/// Both cases originally passed or reproduced #2028 (the AQE-on case: +/// `Shared(IoError)` misses the classifier's shallow match). #2028 was fixed +/// by classifying on `find_root()` (#2119), but the sort-shuffle writer +/// refactor (#2038/#2106) then made both cases fail the same way: the shuffle +/// write coordinator flattens any task error into +/// `DataFusionError::Execution(format!("{e:?}"))` +/// (`ballista/core/src/execution_plans/shuffle_writer.rs`, error arm of the +/// coordinator fan-out), so the injected `IoError` reaches the classifier as +/// inert text inside an `Execution` string, `find_root()` has nothing to see +/// through, and the task is marked non-retryable. That is the same +/// type-erasing mechanism tracked for fetch failures by #2027. +/// +/// Ignored, not deleted or weakened. Un-ignore both as the regression tests +/// when #2027's error-flattening is fixed; see chaos-testing/README.md, +/// Finding 2. +#[tokio::test] +#[ignore = "reproduces #2027's error flattening: the shuffle writer Debug-formats the IoError into an opaque Execution string, so it is misclassified as non-retryable"] +async fn retryable_fault_is_retried_and_result_is_correct_aqe_off() { + retryable_fault_is_retried_and_result_is_correct(false).await; +} + +/// See `retryable_fault_is_retried_and_result_is_correct_aqe_off` above; the +/// AQE-on case fails identically (the error arrives as +/// `Execution("Shared(IoError(..))")` — stringified before the classifier, +/// which is what distinguishes this from the fixed #2028). +#[tokio::test] +#[ignore = "reproduces #2027's error flattening: the shuffle writer Debug-formats the IoError into an opaque Execution string, so it is misclassified as non-retryable"] +async fn retryable_fault_is_retried_and_result_is_correct_aqe_on() { + retryable_fault_is_retried_and_result_is_correct(true).await; +} + +async fn retryable_fault_is_retried_and_result_is_correct(aqe: bool) { + let run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + let budget = run.budget("scenario-a", 1); + let sql = Fixture::chaos_query(&format!( + "chaos_fail(f.key = 7, 'io', '{}')", + budget.dir().display() + )); + + let actual = run + .sql(&sql) + .await + .expect("query must recover from one IO fault"); + + assert_eq!( + actual, expected, + "result after retry must match the baseline" + ); + assert_eq!(budget.remaining(), 0, "the fault must actually have fired"); +} + +/// Scenario B: an inexhaustible retryable fault. +/// +/// The budget far exceeds task_max_failures (4), so every attempt faults. The job +/// must fail rather than retry forever, and the cluster must remain usable +/// afterwards — a scheduler that wedges after a failed job is an HA bug. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn exhausted_retries_fail_the_job_and_leave_the_cluster_healthy(#[case] aqe: bool) { + let run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + let budget = run.budget("scenario-b", 99); + let sql = Fixture::chaos_query(&format!( + "chaos_fail(f.key = 7, 'io', '{}')", + budget.dir().display() + )); + + let err = run + .sql(&sql) + .await + .expect_err("the job must fail once retries are exhausted"); + assert!(!err.is_empty(), "the failure must carry an error message"); + + // The cluster must still serve queries. A chaos-free query proves the + // scheduler and both executors survived the failed job. + let after = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must still be healthy after a failed job"); + assert_eq!(after, expected); +} + +/// Scenario C: a panicking task. +/// +/// The executor catches the panic (executor.rs:237) and turns it into a +/// non-retryable Internal error, so the job fails immediately with no retry. This +/// test encodes *today's* behaviour, not necessarily the desired behaviour: if we +/// later decide panics should be retryable, this is the test that changes. +/// +/// The second assertion is the important one: the executor process must survive. +/// A panic in one task must not take down the whole executor and every other task +/// running on it. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn panicking_task_fails_the_job_but_the_executor_survives(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + let budget = run.budget("scenario-c", 1); + let sql = Fixture::chaos_query(&format!( + "chaos_fail(f.key = 7, 'panic', '{}')", + budget.dir().display() + )); + + let err = run + .sql(&sql) + .await + .expect_err("a panicking task must fail the job"); + assert!(!err.is_empty()); + assert_eq!(budget.remaining(), 0, "the panic must actually have fired"); + + // Both executor processes must still be alive. + assert!( + run.cluster.executor_is_alive(0), + "executor 0 died on a task panic" + ); + assert!( + run.cluster.executor_is_alive(1), + "executor 1 died on a task panic" + ); + + // And the cluster must still serve queries. + let after = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must still be healthy after a panicking task"); + assert_eq!(after, expected); +} + +use std::time::Duration; + +/// Scenario D: SIGKILL an executor while it is running tasks. +/// +/// `chaos_delay` holds stage 1 open so the kill lands while tasks are genuinely +/// in flight. The scheduler must detect the loss, reschedule the dead executor's +/// tasks onto the survivor, and still return the correct result. +/// +/// Both cases reproduce #2027: the shuffle reader's typed +/// `BallistaError::FetchFailed` is flattened into an opaque +/// `DataFusionError::Execution` before the executor reports its `TaskStatus`, +/// so the scheduler never sees the `FetchPartitionError` that would make it +/// resubmit the lost map stage, and fails the query instead of recovering it. +/// +/// Ignored, not deleted or weakened. Un-ignore it as the regression test when +/// #2027 is fixed; see chaos-testing/README.md, Finding 1. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +#[ignore = "reproduces #2027: FetchFailed loses its type, so the map stage is never resubmitted"] +async fn executor_killed_mid_stage_is_recovered(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + // Delay every scan task by 300ms per batch so the stage stays running long + // enough to kill an executor inside it. + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 300)"); + + // Submit the query concurrently, then kill executor 0 once stage 1 is running. + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + let sql = sql.clone(); + async move { ctx.sql(&sql).await?.collect().await } + }); + + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_stage_running(&job_id, 1) + .await + .expect("stage 1 must start running"); + run.cluster.kill_executor(0).expect("kill executor 0"); + + let batches = tokio::time::timeout(Duration::from_secs(120), query) + .await + .expect("query must not hang after an executor is killed") + .expect("query task must not panic") + .expect("query must recover from the lost executor"); + + let actual = datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert_eq!( + actual, expected, + "result after executor loss must match the baseline" + ); +} + +/// Scenario E: SIGKILL a map-side executor after it wrote shuffle output. +/// +/// The downstream stage must fetch shuffle partitions from an executor that no +/// longer exists. Recovery requires re-running the map stage. The executor +/// timeout is raised to 60s to bias the failure toward the FetchPartitionError +/// path rather than the heartbeat-expiry ExecutorLost path; both are valid +/// recoveries, so the assertion is on correctness, and the path that actually +/// fired is only recorded. +/// +/// Reproduces #2027: whenever the reduce stage genuinely has to fetch from the +/// dead executor, the typed `FetchFailed` arrives flattened inside +/// `DataFusionError::Shared`, the scheduler never resubmits the map stage, and +/// the job fails. The scenario only passes when the kill happens to land after +/// the reduce tasks have already fetched their partitions, which makes it +/// timing-dependent (it failed in CI, aqe_off case). +/// +/// Ignored, not deleted or weakened. Un-ignore it as the regression test when +/// #2027 is fixed; see chaos-testing/README.md, Finding 1. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +#[ignore = "reproduces #2027: FetchFailed loses its type, so the map stage is never resubmitted"] +async fn executor_killed_after_shuffle_write_is_recovered(#[case] aqe: bool) { + let mut run = ChaosRun::start_with(aqe, 2, 60).await; + let expected = run.local_baseline().await; + + // Delay the *aggregate* side so the reduce stage is slow, giving us a window + // between "stage 1 succeeded" and "stage 2 has finished fetching". + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 50)"); + + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + let sql = sql.clone(); + async move { ctx.sql(&sql).await?.collect().await } + }); + + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_stage_successful(&job_id, 1) + .await + .expect("stage 1 must complete before we kill its executor"); + run.cluster.kill_executor(0).expect("kill executor 0"); + + let batches = tokio::time::timeout(Duration::from_secs(180), query) + .await + .expect("query must not hang after shuffle output is lost") + .expect("query task must not panic") + .expect("query must recover by re-running the map stage"); + + let actual = datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert_eq!( + actual, expected, + "result after shuffle-output loss must match the baseline" + ); +} + +/// Scenario F: an executor is killed and restarted; the cluster must reabsorb it. +/// +/// The kill and the restart are separated by a wait for the scheduler to +/// actually reap the dead executor (`registered_executors` dropping to 1), +/// rather than restarting immediately. SIGKILL does not deregister the +/// executor: the scheduler keeps listing it until its heartbeat expires +/// (`executor_timeout_seconds`), so a restart fired immediately after the +/// kill lands while the scheduler still counts *three* executors (the dead +/// one, the survivor, and the freshly restarted one) — a harness race, not a +/// Ballista bug, that used to make this scenario fail with `left: 3, right: +/// 2`. Waiting for the reap first means the final assertion is actually +/// testing what the scenario name promises: that a restarted executor +/// rejoins a cluster that has already noticed it was gone. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn restarted_executor_rejoins_and_serves_queries(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + run.cluster.kill_executor(0).expect("kill executor 0"); + run.cluster + .await_executor_count(1) + .await + .expect("scheduler must reap the killed executor before we restart it"); + + run.cluster + .restart_executor(0) + .await + .expect("restarted executor must re-register"); + + assert_eq!( + run.cluster.registered_executors().await.unwrap(), + 2, + "both executors must be registered after the restart" + ); + + let actual = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must serve queries after an executor restart"); + assert_eq!(actual, expected); +} + +/// Scenario G: every executor is killed mid-query. +/// +/// There is no executor left to recover onto, so the job cannot succeed. The only +/// requirement is that it *terminates* — a scheduler that waits forever for tasks +/// that can never be scheduled is a hang, and a hang is the bug this detects. The +/// assertion is deliberately weak: this is a hang detector, not a correctness test. +/// +/// Both cases reproduce #2029: the job does not terminate within the 120s +/// timeout. The scheduler waits rather than failing the query once no executor +/// can ever satisfy the remaining tasks. +/// +/// Ignored, not deleted or weakened. Un-ignore it as the regression test when +/// #2029 is fixed; see chaos-testing/README.md, Finding 3. Note that this one +/// costs 120s of wall clock when it runs, since detecting a hang means waiting +/// out the timeout. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +#[ignore = "reproduces #2029: killing every executor hangs the job instead of failing it"] +async fn killing_every_executor_terminates_the_job(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 300)"); + + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + let sql = sql.clone(); + async move { ctx.sql(&sql).await?.collect().await } + }); + + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_stage_running(&job_id, 1) + .await + .expect("stage 1 must start running"); + + run.cluster.kill_executor(0).expect("kill executor 0"); + run.cluster.kill_executor(1).expect("kill executor 1"); + + // The job must end, one way or another, rather than hanging forever. + let outcome = tokio::time::timeout(Duration::from_secs(120), query).await; + assert!( + outcome.is_ok(), + "job did not terminate within 120s after every executor was killed — this is a hang" + ); +}