andygrove commented on code in PR #2425: URL: https://github.com/apache/datafusion-ballista/pull/2425#discussion_r3941388332
########## chaos-testing/src/rest.rs: ########## @@ -0,0 +1,159 @@ +// 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. + +//! Scheduler REST-API polling shared by both cluster backends. +//! +//! [`crate::cluster::TestCluster`] (local processes) and [`crate::k8s`] +//! (`kind` pods) both drive scenarios by polling the scheduler's REST API — +//! the same endpoints (`/api/executors`, `/api/jobs`, `/api/job/{id}/stages`) +//! and the same JSON shape, differing only in the base URL (loopback vs. the +//! port-forward). These free functions take that base URL so the polling logic +//! lives in one place; each backend wraps them in thin methods. + +use serde_json::Value; +use std::time::{Duration, Instant}; + +/// Per-request timeout. Short so a stalled connection (e.g. a k8s port-forward +/// that dropped) surfaces as a retryable error inside a polling loop rather than +/// hanging the whole wait; harmless for the loopback (process) backend. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); + +fn client() -> Result<reqwest::Client, String> { + reqwest::Client::builder() + .timeout(REQUEST_TIMEOUT) + .build() + .map_err(|e| e.to_string()) +} + +async fn get_json(url: String) -> Result<Value, String> { + client()? + .get(url) + .send() + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string()) +} + +/// How many executors the scheduler currently considers registered. +pub(crate) async fn registered_executors(rest_url: &str) -> Result<usize, String> { + let body = get_json(format!("{rest_url}/api/executors")).await?; + Ok(body.as_array().map(|a| a.len()).unwrap_or(0)) +} + +/// The ids of every executor the scheduler currently lists. Lets a scenario +/// prove a *new* executor (a fresh id) replaced a killed one after rescheduling. +/// Only the k8s backend reschedules automatically, so this is k8s-only. +#[cfg(feature = "k8s")] +pub(crate) async fn executor_ids(rest_url: &str) -> Result<Vec<String>, String> { + let body = get_json(format!("{rest_url}/api/executors")).await?; + Ok(body + .as_array() + .into_iter() + .flatten() + .filter_map(|e| e.get("id").and_then(|v| v.as_str()).map(String::from)) + .collect()) +} + +/// 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(crate) async fn running_job_id(rest_url: &str) -> Result<String, String> { + let deadline = Instant::now() + Duration::from_secs(30); + loop { + let body = get_json(format!("{rest_url}/api/jobs")).await?; + 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; + } +} + +/// The stage summary for `job_id`. +pub(crate) async fn stages(rest_url: &str, job_id: &str) -> Result<Value, String> { + get_json(format!("{rest_url}/api/job/{job_id}/stages")).await +} + +/// Block until any task in any stage is Running. +/// +/// Planner-agnostic sync point: the static and adaptive (AQE) planners number +/// and materialize stages differently, so rather than target a specific stage id +/// we wait until the job is genuinely executing a task somewhere. Used where the +/// scenario only needs a fault to land mid-flight. +pub(crate) async fn await_any_stage_running( + rest_url: &str, + job_id: &str, +) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = stages(rest_url, job_id).await?; Review Comment: `REQUEST_TIMEOUT` is documented as making a stalled connection "surface as a retryable error inside a polling loop rather than hanging the whole wait", but only `await_executor_count` actually treats it that way with its `if let Ok(count)`. Both `running_job_id` (L79) and `await_any_stage_running` (L109) do `get_json(...).await?`, so a single timed out or refused request ends the wait instead of retrying to the deadline. That was harmless while these two only ran over loopback, but scenario G is the first thing to drive them across the kubectl port-forward, and the harness has a whole supervisor that expects the forward to drop and come back (`PORT_FORWARD_RETRY_DELAY`). A restart landing inside either poll fails the test at `expect("job must appear")` with nothing useful to go on. Retrying on error until the deadline, and carrying the last error into the timeout message, would match what the comment already promises. ########## chaos-testing/src/k8s.rs: ########## @@ -259,29 +259,37 @@ impl K8sCluster { /// How many executors the scheduler currently considers registered. pub async fn registered_executors(&self) -> Result<usize, String> { - // A short timeout so a stalled port-forward surfaces as a retryable - // error in the polling loop rather than hanging the whole wait. - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(5)) - .build() - .map_err(|e| e.to_string())?; - let body: serde_json::Value = client - .get(format!("{}/api/executors", self.rest_url())) - .send() - .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)) - } - - /// Scale the executor Deployment. `0` is a total loss that stays lost (the - /// controller does not recreate the pods); scaling back up recovers. - /// - /// Not yet exercised by a scenario — this is the k8s primitive the executor - /// kill/loss scenarios (the #2029 follow-ups) will drive; the baseline test - /// only needs a healthy cluster. Kept here so the backend is complete. + crate::rest::registered_executors(&self.rest_url()).await + } + + /// The ids of every executor the scheduler currently lists. A killed + /// executor's pod is rescheduled with a fresh id, so a scenario can compare + /// this before and after a kill to prove the replacement is genuinely new. + pub async fn executor_ids(&self) -> Result<Vec<String>, String> { + crate::rest::executor_ids(&self.rest_url()).await + } + + /// The id of the single job the scheduler currently knows about. + pub async fn running_job_id(&self) -> Result<String, String> { + crate::rest::running_job_id(&self.rest_url()).await + } + + /// Block until any task in any stage is Running, so a kill lands mid-flight. + pub async fn await_any_stage_running(&self, job_id: &str) -> Result<(), String> { + crate::rest::await_any_stage_running(&self.rest_url(), job_id).await + } + + /// Block until the scheduler considers exactly `n` executors registered + /// (i.e. a killed executor has been reaped, or a rescheduled one re-joined). + pub async fn await_executor_count(&self, n: usize) -> Result<(), String> { Review Comment: I think this one is dead code, and a slightly worse copy of something already here. On the process backend the pair is a real distinction, because `TestCluster::await_executors` waits for `count >= n` while `await_executor_count` waits for `count == n`. But `K8sCluster::await_executors` (L199) already waits for `count == n`, on the same 120s deadline, and it also calls `dump_diagnostics()` on timeout. So this method is that same wait minus the diagnostics, and nothing calls it. Scenario F rolls its own `await_replacement` instead. Probably worth either dropping it or folding the two into one. ########## chaos-testing/src/rest.rs: ########## @@ -0,0 +1,159 @@ +// 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. + +//! Scheduler REST-API polling shared by both cluster backends. +//! +//! [`crate::cluster::TestCluster`] (local processes) and [`crate::k8s`] +//! (`kind` pods) both drive scenarios by polling the scheduler's REST API — +//! the same endpoints (`/api/executors`, `/api/jobs`, `/api/job/{id}/stages`) +//! and the same JSON shape, differing only in the base URL (loopback vs. the +//! port-forward). These free functions take that base URL so the polling logic +//! lives in one place; each backend wraps them in thin methods. + +use serde_json::Value; +use std::time::{Duration, Instant}; + +/// Per-request timeout. Short so a stalled connection (e.g. a k8s port-forward +/// that dropped) surfaces as a retryable error inside a polling loop rather than +/// hanging the whole wait; harmless for the loopback (process) backend. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); + +fn client() -> Result<reqwest::Client, String> { Review Comment: Nit: this builds a fresh `Client`, and so a fresh connection pool, on every request. `await_any_stage_running` polls every 50ms for up to 60s, so that is roughly 1200 new connections through the port-forward over one wait. Not a regression, since `reqwest::get` did the same thing before, but now that it is all in one place a `OnceLock<Client>` would fix it once and take a bit of pressure off the forward, which is also relevant to the retry comment above. ########## chaos-testing/src/k8s.rs: ########## @@ -294,6 +302,28 @@ impl K8sCluster { .map(|_| ()) } + /// Total, sustained executor loss *mid-query* (#2029): every executor dies at + /// once and none is rescheduled. Scales the Deployment to 0 first so the + /// controller will not recreate the pods, then force-deletes them + /// (`--grace-period=0 --force`) so they are SIGKILLed without the graceful + /// drain that [`Self::scale_executors`] alone allows — otherwise the + /// executors would finish the in-flight query and it would succeed. + pub async fn kill_all_executors_hard(&self) -> Result<(), String> { Review Comment: This races SIGTERM against SIGKILL, and I think it makes scenario G non-deterministic about which scheduler path it exercises. `kubectl scale --replicas=0` returns as soon as the API server accepts it. The ReplicaSet controller then marks the pods for deletion and kubelet sends SIGTERM, and all of that happens concurrently with the second kubectl round trip that force-deletes them. So the executors may or may not see SIGTERM before the SIGKILL lands, and which one wins will vary run to run. That matters because `start_executor_process` handles SIGTERM by setting `TERMINATING` and immediately sending a `Fenced` heartbeat to the scheduler (`ballista/executor/src/executor_process.rs`, the `tokio::select!` around L775 and the `notify_scheduler` block just after). The scheduler then reaps that executor through the `terminating && grace_period_expired` branch of `get_expired_executors` rather than the plain `expired` branch. Both end in a failed job, so the assertion passes either way, but some fraction of runs will be covering the graceful fence path instead of the abrupt-loss path #2029 is actually about. The `ha.rs` version uses a raw SIGKILL and has no such ambiguity, so k8s is the weaker of the two here, which is the opposite of the intent. If you want a clean abrupt total loss with no controller left to reschedule, `kubectl delete deploy ballista-executor --cascade=orphan` followed by a force-delete of the now orphaned pods gets there with no SIGTERM in the picture at all. ########## chaos-testing/tests/k8s.rs: ########## @@ -27,87 +27,278 @@ //! kind load docker-image ballista-chaos:test //! CHAOS_BACKEND=kind cargo test -p ballista-chaos --features k8s --test k8s -- --test-threads=1 //! ``` +//! +//! These are the scenarios that genuinely need a cluster — real pod lifecycle, +//! rescheduling, and the port-forward/flight-proxy path — rather than the +//! fault-injection scenarios in `ha.rs`, which are backend-agnostic and stay on +//! the fast process harness. Which planner each scenario runs under mirrors +//! `ha.rs`: both AQE settings only where the planner changes the code path. #![cfg(feature = "k8s")] +use std::time::{Duration, Instant}; + use ballista::prelude::{SessionConfigExt, SessionContextExt}; +use ballista_core::config::BALLISTA_ADAPTIVE_PLANNER_ENABLED; use chaos_testing::fixture::Fixture; -use chaos_testing::k8s::K8sCluster; +use chaos_testing::k8s::{K8sCluster, KillMode}; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; +use rstest::rstest; -/// The k8s scenarios need a running kind cluster with the chaos images loaded; +/// The k8s scenarios need a running kind cluster with the chaos image loaded; /// they are opt-in via `CHAOS_BACKEND=kind` so a plain `cargo test` skips them. fn kind_backend_selected() -> bool { if std::env::var("CHAOS_BACKEND").as_deref() == Ok("kind") { true } else { eprintln!( "skipping k8s scenario: set CHAOS_BACKEND=kind and provide a kind cluster \ - with the chaos images loaded (see the crate README runbook)" + with the chaos image loaded (see the crate README runbook)" ); false } } -/// The chaos-free baseline query, run on a fresh local DataFusion context. This -/// is the reference the cluster must reproduce exactly. -async fn local_baseline(fixture: &Fixture) -> String { - let ctx = SessionContext::new(); - for stmt in fixture.register_sql() { - ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); +/// One kind cluster plus its fixture and a connected client, wired for a single +/// scenario. The k8s counterpart of `ha.rs`'s `ChaosRun`: it centralises the +/// fixture write, the client connect, and the UDF-after-upgrade registration so +/// each scenario reads as just its fault and its assertions. +struct K8sRun { + cluster: K8sCluster, + fixture: Fixture, + ctx: SessionContext, +} + +impl K8sRun { + /// Deploy a cluster of `executors`, write the fixture into the shared mount, + /// and connect a client with AQE set to `aqe`. + async fn start(aqe: bool, executors: usize) -> Self { + let cluster = K8sCluster::start(executors) + .await + .expect("kind cluster must start"); + + // Written into the shared mount, so the scheduler and executor pods see it. + let fixture = Fixture::write(cluster.shared_dir()) + .await + .expect("fixture must be written to the shared mount"); + + 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`: `upgrade_for_ballista` rebuilds + // the state with `with_scalar_functions(...)`, which replaces rather than + // merges the scalar-function map and would drop a UDF registered before. + // (Same subtlety documented in `ha.rs`'s `ChaosRun`.) + 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 clone of the session context, for running a query concurrently with a + /// fault (the query is spawned while the main task kills executors). + fn clone_ctx(&self) -> SessionContext { + self.ctx.clone() + } + + /// Run a query on the cluster, returning the formatted result. + async fn sql(&self, sql: &str) -> Result<String, String> { + 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 over the same + /// fixture — the reference the cluster must reproduce exactly. + 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() } - let batches = ctx - .sql(Fixture::baseline_query()) - .await - .unwrap() - .collect() - .await - .unwrap(); - pretty_format_batches(&batches).unwrap().to_string() } /// Smoke test: a real query on a real kind cluster returns the same result as /// plain local DataFusion. Exercises the whole path — client → scheduler → pods /// → shuffle → result — with the fixture shared through the `hostPath` mount. +/// +/// Single planner: the wiring this smoke-tests (mount, port-forward, flight +/// proxy, shuffle) does not vary by planner, so there is nothing to gain from +/// running it under both. #[tokio::test] async fn baseline_matches_local_datafusion_on_k8s() { if !kind_backend_selected() { return; } - let cluster = K8sCluster::start(2).await.expect("kind cluster must start"); + let run = K8sRun::start(false, 2).await; + let expected = run.local_baseline().await; + let actual = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must serve the baseline query"); + + assert_eq!( + actual, expected, + "cluster result must match plain local DataFusion" + ); +} + +/// Scenario G (#2029) on k8s: every executor is lost mid-query. +/// +/// The kill force-deletes every executor pod while holding the Deployment at +/// zero replicas (`kill_all_executors_hard`). Scaling to zero alone is not +/// enough: it is a graceful SIGTERM, so the executors drain the in-flight query +/// to completion and it *succeeds*. Force-deleting the pods SIGKILLs them so the +/// work is genuinely lost, and holding the Deployment at zero means the +/// controller will not reschedule replacements. With no executor left, the job +/// cannot succeed: once the last one is reaped the scheduler waits the bounded +/// no-executors grace period and then fails the job rather than hanging forever. +/// +/// Both AQE settings: #2029 had an AQE-on-only second hang path (a task-launch +/// failure that removed the last executor without arming the grace timer), so +/// the planner genuinely changes the code path here. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn killing_every_executor_terminates_the_job_on_k8s(#[case] aqe: bool) { + if !kind_backend_selected() { + return; + } + + let run = K8sRun::start(aqe, 2).await; + // A long per-batch delay keeps the query in flight while the hard kill runs + // its two kubectl round-trips (scale to zero, then force-delete), so the + // pods die mid-query rather than after it would have finished. + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 2000)"); + + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + async move { ctx.sql(&sql).await?.collect().await } + }); - // Written into the shared mount, so the scheduler and executor pods see it. - let fixture = Fixture::write(cluster.shared_dir()) + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_any_stage_running(&job_id) .await - .expect("fixture must be written to the shared mount"); + .expect("the job must start running a task before we remove its executors"); - let expected = local_baseline(&fixture).await; + // Total loss that stays lost and is not drained: force-kill every pod and + // hold the Deployment at zero so no replacement is scheduled. + run.cluster + .kill_all_executors_hard() + .await + .expect("force-kill every executor"); - let config = SessionConfig::new_with_ballista(); - let state = SessionStateBuilder::new() - .with_config(config) - .with_default_features() - .build(); - let ctx = SessionContext::remote_with_state(&cluster.scheduler_url(), state) + let result = tokio::time::timeout(Duration::from_secs(120), query) .await - .expect("client must connect to the scheduler"); + .expect("job must terminate, not hang, after every executor is lost") Review Comment: Scenario G never dumps diagnostics when it fails. `await_replacement` in scenario F calls `cluster.dump_diagnostics()` before it panics, which is right, but G's three failure points (`running_job_id`, `await_any_stage_running`, and this 120s timeout) all just panic, and the namespace is deleted on drop immediately afterwards. A CI failure there leaves nothing behind to look at. Worth a dump on at least this timeout, since it is the assertion the whole PR exists for. ########## chaos-testing/tests/k8s.rs: ########## @@ -27,87 +27,278 @@ //! kind load docker-image ballista-chaos:test //! CHAOS_BACKEND=kind cargo test -p ballista-chaos --features k8s --test k8s -- --test-threads=1 //! ``` +//! +//! These are the scenarios that genuinely need a cluster — real pod lifecycle, +//! rescheduling, and the port-forward/flight-proxy path — rather than the +//! fault-injection scenarios in `ha.rs`, which are backend-agnostic and stay on +//! the fast process harness. Which planner each scenario runs under mirrors +//! `ha.rs`: both AQE settings only where the planner changes the code path. #![cfg(feature = "k8s")] +use std::time::{Duration, Instant}; + use ballista::prelude::{SessionConfigExt, SessionContextExt}; +use ballista_core::config::BALLISTA_ADAPTIVE_PLANNER_ENABLED; use chaos_testing::fixture::Fixture; -use chaos_testing::k8s::K8sCluster; +use chaos_testing::k8s::{K8sCluster, KillMode}; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; +use rstest::rstest; -/// The k8s scenarios need a running kind cluster with the chaos images loaded; +/// The k8s scenarios need a running kind cluster with the chaos image loaded; /// they are opt-in via `CHAOS_BACKEND=kind` so a plain `cargo test` skips them. fn kind_backend_selected() -> bool { if std::env::var("CHAOS_BACKEND").as_deref() == Ok("kind") { true } else { eprintln!( "skipping k8s scenario: set CHAOS_BACKEND=kind and provide a kind cluster \ - with the chaos images loaded (see the crate README runbook)" + with the chaos image loaded (see the crate README runbook)" ); false } } -/// The chaos-free baseline query, run on a fresh local DataFusion context. This -/// is the reference the cluster must reproduce exactly. -async fn local_baseline(fixture: &Fixture) -> String { - let ctx = SessionContext::new(); - for stmt in fixture.register_sql() { - ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); +/// One kind cluster plus its fixture and a connected client, wired for a single +/// scenario. The k8s counterpart of `ha.rs`'s `ChaosRun`: it centralises the +/// fixture write, the client connect, and the UDF-after-upgrade registration so +/// each scenario reads as just its fault and its assertions. +struct K8sRun { + cluster: K8sCluster, + fixture: Fixture, + ctx: SessionContext, +} + +impl K8sRun { + /// Deploy a cluster of `executors`, write the fixture into the shared mount, + /// and connect a client with AQE set to `aqe`. + async fn start(aqe: bool, executors: usize) -> Self { + let cluster = K8sCluster::start(executors) + .await + .expect("kind cluster must start"); + + // Written into the shared mount, so the scheduler and executor pods see it. + let fixture = Fixture::write(cluster.shared_dir()) + .await + .expect("fixture must be written to the shared mount"); + + 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`: `upgrade_for_ballista` rebuilds + // the state with `with_scalar_functions(...)`, which replaces rather than + // merges the scalar-function map and would drop a UDF registered before. + // (Same subtlety documented in `ha.rs`'s `ChaosRun`.) + 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 clone of the session context, for running a query concurrently with a + /// fault (the query is spawned while the main task kills executors). + fn clone_ctx(&self) -> SessionContext { + self.ctx.clone() + } + + /// Run a query on the cluster, returning the formatted result. + async fn sql(&self, sql: &str) -> Result<String, String> { + 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 over the same + /// fixture — the reference the cluster must reproduce exactly. + 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() } - let batches = ctx - .sql(Fixture::baseline_query()) - .await - .unwrap() - .collect() - .await - .unwrap(); - pretty_format_batches(&batches).unwrap().to_string() } /// Smoke test: a real query on a real kind cluster returns the same result as /// plain local DataFusion. Exercises the whole path — client → scheduler → pods /// → shuffle → result — with the fixture shared through the `hostPath` mount. +/// +/// Single planner: the wiring this smoke-tests (mount, port-forward, flight +/// proxy, shuffle) does not vary by planner, so there is nothing to gain from +/// running it under both. #[tokio::test] async fn baseline_matches_local_datafusion_on_k8s() { if !kind_backend_selected() { return; } - let cluster = K8sCluster::start(2).await.expect("kind cluster must start"); + let run = K8sRun::start(false, 2).await; + let expected = run.local_baseline().await; + let actual = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must serve the baseline query"); + + assert_eq!( + actual, expected, + "cluster result must match plain local DataFusion" + ); +} + +/// Scenario G (#2029) on k8s: every executor is lost mid-query. +/// +/// The kill force-deletes every executor pod while holding the Deployment at +/// zero replicas (`kill_all_executors_hard`). Scaling to zero alone is not +/// enough: it is a graceful SIGTERM, so the executors drain the in-flight query +/// to completion and it *succeeds*. Force-deleting the pods SIGKILLs them so the +/// work is genuinely lost, and holding the Deployment at zero means the +/// controller will not reschedule replacements. With no executor left, the job +/// cannot succeed: once the last one is reaped the scheduler waits the bounded +/// no-executors grace period and then fails the job rather than hanging forever. +/// +/// Both AQE settings: #2029 had an AQE-on-only second hang path (a task-launch +/// failure that removed the last executor without arming the grace timer), so +/// the planner genuinely changes the code path here. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn killing_every_executor_terminates_the_job_on_k8s(#[case] aqe: bool) { + if !kind_backend_selected() { + return; + } + + let run = K8sRun::start(aqe, 2).await; + // A long per-batch delay keeps the query in flight while the hard kill runs + // its two kubectl round-trips (scale to zero, then force-delete), so the + // pods die mid-query rather than after it would have finished. + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 2000)"); + + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + async move { ctx.sql(&sql).await?.collect().await } + }); - // Written into the shared mount, so the scheduler and executor pods see it. - let fixture = Fixture::write(cluster.shared_dir()) + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_any_stage_running(&job_id) .await - .expect("fixture must be written to the shared mount"); + .expect("the job must start running a task before we remove its executors"); - let expected = local_baseline(&fixture).await; + // Total loss that stays lost and is not drained: force-kill every pod and + // hold the Deployment at zero so no replacement is scheduled. + run.cluster + .kill_all_executors_hard() + .await + .expect("force-kill every executor"); - let config = SessionConfig::new_with_ballista(); - let state = SessionStateBuilder::new() - .with_config(config) - .with_default_features() - .build(); - let ctx = SessionContext::remote_with_state(&cluster.scheduler_url(), state) + let result = tokio::time::timeout(Duration::from_secs(120), query) .await - .expect("client must connect to the scheduler"); + .expect("job must terminate, not hang, after every executor is lost") + .expect("query task should not panic"); + let err = result.expect_err("query must fail once every executor is lost"); + let msg = err.to_string().to_lowercase(); + assert!( + msg.contains("executor"), + "failure should name the executor loss, got: {err}" + ); +} - for stmt in fixture.register_sql() { - ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); +/// Scenario F on k8s: an executor pod is killed and the cluster reabsorbs its +/// replacement. +/// +/// Unlike the process harness — where the test spawns a fresh executor itself — +/// deleting a pod lets the Deployment reschedule a replacement automatically, +/// with a brand-new executor id. That is the k8s-unique behaviour this asserts: +/// after a forced pod delete, the scheduler settles back to two executors, one +/// of which is genuinely new (an id not present before), and the cluster still +/// serves queries. +/// +/// Single planner (aqe off): rescheduling and re-registration are +/// planner-independent, and the post-restart query correctness is already +/// covered by the baseline scenario. +#[tokio::test] +async fn restarted_executor_rejoins_and_serves_queries_on_k8s() { + if !kind_backend_selected() { + return; } - let batches = ctx - .sql(Fixture::baseline_query()) + + let run = K8sRun::start(false, 2).await; + let expected = run.local_baseline().await; + + let before = run + .cluster + .executor_ids() .await - .unwrap() - .collect() + .expect("must list executors before the kill"); + assert_eq!(before.len(), 2, "expected two executors to start"); + + run.cluster + .kill_one_executor(KillMode::Forced) .await - .unwrap(); - let actual = pretty_format_batches(&batches).unwrap().to_string(); + .expect("force-delete one executor pod"); - assert_eq!( - actual, expected, - "cluster result must match plain local DataFusion" + // Wait for steady state: exactly two executors again, one of which is new + // (a fresh id). A plain count==2 wait is not enough — the killed executor's + // heartbeat lingers, so the count passes transiently through 2 (ghost + + // survivor) and 3 (ghost + survivor + replacement) before settling. + let after = await_replacement(&run.cluster, &before).await; + assert_eq!(after.len(), 2, "cluster must settle back to two executors"); + assert!( + after.iter().any(|id| !before.contains(id)), + "a rescheduled executor with a new id must have joined; before={before:?} after={after:?}" ); + + let actual = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must serve queries after an executor is rescheduled"); + assert_eq!(actual, expected); +} + +/// Poll until the scheduler lists exactly two executors and at least one is not +/// in `before` — i.e. the killed executor has been reaped and its rescheduled +/// replacement (new id) has registered. +async fn await_replacement(cluster: &K8sCluster, before: &[String]) -> Vec<String> { + let deadline = Instant::now() + Duration::from_secs(120); + loop { + if let Ok(ids) = cluster.executor_ids().await + && ids.len() == 2 Review Comment: Nit: the expected count of 2 is hardcoded here while the caller already knows it from `before.len()`. Passing it in would let this helper serve a scenario with a different executor count later. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
