From a54ca2a5180f3da0ec91cba72d24de2f5fb0047b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 08:32:14 -0600 Subject: [PATCH 1/9] feat: add ballista.job.client_side_cleanup config flag Adds a boolean config key (default true) controlling whether the client sends a best-effort CleanJobData RPC to the scheduler when it finishes consuming a job's results, so on-disk result data can be reclaimed immediately instead of waiting for the scheduler's timed cleanup. This task only adds the flag and typed getter; a later task wires up the actual cleanup call. --- ballista/core/src/config.rs | 24 ++++++++++++++++++++++++ docs/source/user-guide/configs.md | 1 + 2 files changed, 25 insertions(+) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fd3060297d..d2108ab81f 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -82,6 +82,11 @@ pub const BALLISTA_CLIENT_IO_RETRIES_TIMES: &str = "ballista.client.io_retries_t /// Wait time in milliseconds between IO retries in the Ballista client pub const BALLISTA_CLIENT_IO_RETRY_WAIT_TIME_MS: &str = "ballista.client.io_retry_wait_time_ms"; +/// Configuration key controlling whether the client sends a best-effort +/// `CleanJobData` RPC to the scheduler when it finishes consuming a job's +/// results, so on-disk result data is reclaimed immediately rather than +/// waiting for the scheduler's timed cleanup. +pub const BALLISTA_JOB_CLIENT_SIDE_CLEANUP: &str = "ballista.job.client_side_cleanup"; /// Enables adaptive query planning pub const BALLISTA_ADAPTIVE_PLANNER_ENABLED: &str = "ballista.planner.adaptive.enabled"; /// Configuration key for sort shuffle target batch size in rows. @@ -290,6 +295,14 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| "Wait time in milliseconds between IO retries in the Ballista client.".to_string(), DataType::UInt64, Some(3000.to_string())), + ConfigEntry::new(BALLISTA_JOB_CLIENT_SIDE_CLEANUP.to_string(), + "When enabled, the client sends a best-effort CleanJobData request \ + to the scheduler as soon as it finishes (or abandons) consuming a \ + job's results, reclaiming the job's on-disk result data on executors \ + immediately instead of waiting for the scheduler's timed cleanup. \ + The scheduler's timed cleanup remains as a safety net.".to_string(), + DataType::Boolean, + Some(true.to_string())), ConfigEntry::new(BALLISTA_COALESCE_ENABLED.to_string(), "Enables the AQE coalesce-shuffle-partitions rule. \ Disabled by default — opt in when fewer/larger \ @@ -678,6 +691,12 @@ impl BallistaConfig { self.get_bool_setting(BALLISTA_PROPAGATE_EMPTY_ENABLED) } + /// Whether the client eagerly reclaims a finished job's on-disk data via + /// the scheduler's `CleanJobData` RPC. See [`BALLISTA_JOB_CLIENT_SIDE_CLEANUP`]. + pub fn client_side_cleanup_enabled(&self) -> bool { + self.get_bool_setting(BALLISTA_JOB_CLIENT_SIDE_CLEANUP) + } + /// Returns compression codec that will be used during write stage of shuffle pub fn shuffle_compression_codec( &self, @@ -983,4 +1002,9 @@ mod tests { .expect("entry is registered"); assert_eq!(batch_size.doc_default(), Some("8192")); } + + #[test] + fn client_side_cleanup_default_enabled() { + assert!(BallistaConfig::default().client_side_cleanup_enabled()); + } } diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c40f1a582c..b1838a0327 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -118,6 +118,7 @@ standard DataFusion settings. | ballista.client.io_retry_wait_time_ms | UInt64 | 3000 | Wait time in milliseconds between IO retries in the Ballista client. | | ballista.client.pull | Boolean | false | Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients. | | ballista.client.use_tls | Boolean | false | Should connection between client, scheduler, and executors use TLS. | +| ballista.job.client_side_cleanup | Boolean | true | When enabled, the client sends a best-effort CleanJobData request to the scheduler as soon as it finishes (or abandons) consuming a job's results, reclaiming the job's on-disk result data on executors immediately instead of waiting for the scheduler's timed cleanup. The scheduler's timed cleanup remains as a safety net. | | ballista.job.name | Utf8 | (none) | Sets the job name that will appear in the web user interface for any submitted jobs | | ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern. Governs broadcast selection under both the static distributed planner and adaptive query planning (AQE). Set to 0 to disable promotion. | | ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern, used as a fallback when byte-size statistics are unavailable. Applies to adaptive query planning (AQE). Set to 0 to disable promotion via the row-count path. | From 055441b5592d3ba8e0b36364439e8129cc866bb2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 08:37:16 -0600 Subject: [PATCH 2/9] feat: add JobCleanupGuard and GuardedStream for client-side job cleanup --- .../src/execution_plans/distributed_query.rs | 131 +++++++++++++++++- 1 file changed, 130 insertions(+), 1 deletion(-) diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 20f3234cfe..cc3f75b369 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -52,6 +52,7 @@ use log::{debug, error, info}; use parking_lot::Mutex; use std::fmt::Debug; use std::marker::PhantomData; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; use url::Url; @@ -792,6 +793,56 @@ fn get_client_host_port( } } } + +/// Boxed one-shot cleanup action run when a [`JobCleanupGuard`] is dropped. +type JobCleanupFn = Box; + +/// Runs a best-effort cleanup action exactly once when dropped. +/// +/// Constructed with `Some(closure)` when client-side cleanup is enabled, or via +/// [`JobCleanupGuard::disabled`] (no closure) when it is turned off. Dropping the +/// guard takes the closure and invokes it; a disabled guard does nothing. +struct JobCleanupGuard { + cleanup: Option, +} + +impl JobCleanupGuard { + fn new(cleanup: Option) -> Self { + Self { cleanup } + } + + fn disabled() -> Self { + Self { cleanup: None } + } +} + +impl Drop for JobCleanupGuard { + fn drop(&mut self) { + if let Some(cleanup) = self.cleanup.take() { + cleanup(); + } + } +} + +/// Wraps a result stream so that when the consumer finishes or abandons it, the +/// owned [`JobCleanupGuard`] is dropped and fires its cleanup action. Polling is +/// delegated verbatim to the inner stream. +struct GuardedStream { + inner: Pin> + Send>>, + _guard: JobCleanupGuard, +} + +impl Stream for GuardedStream { + type Item = Result; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.get_mut().inner.as_mut().poll_next(cx) + } +} + #[allow(clippy::too_many_arguments)] async fn fetch_partition( location: PartitionLocation, @@ -849,7 +900,7 @@ mod test { use crate::JobId; use crate::config::BallistaConfig; use crate::execution_plans::distributed_query::{ - DistributedQueryExec, get_client_host_port, + DistributedQueryExec, GuardedStream, JobCleanupGuard, get_client_host_port, }; use crate::serde::protobuf::ExecutorMetadata; use crate::serde::protobuf::get_job_status_result::FlightProxy; @@ -930,4 +981,82 @@ mod test { assert_eq!(new_exec.job_id(), Some(JobId::new("job-123"))); } + + #[test] + fn job_cleanup_guard_fires_once_on_drop() { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let count = Arc::new(AtomicUsize::new(0)); + let c = count.clone(); + let guard = JobCleanupGuard::new(Some(Box::new(move || { + c.fetch_add(1, Ordering::SeqCst); + }))); + drop(guard); + assert_eq!(count.load(Ordering::SeqCst), 1); + } + + #[test] + fn job_cleanup_guard_disabled_does_not_fire() { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let count = Arc::new(AtomicUsize::new(0)); + let c = count.clone(); + // Simulate the production "disabled" branch: no closure is built. + let _c = c; // the counter is intentionally never wired to a closure + let guard = JobCleanupGuard::disabled(); + drop(guard); + assert_eq!(count.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn guarded_stream_fires_after_full_consumption() { + use futures::StreamExt; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let count = Arc::new(AtomicUsize::new(0)); + let c = count.clone(); + let guard = JobCleanupGuard::new(Some(Box::new(move || { + c.fetch_add(1, Ordering::SeqCst); + }))); + let inner = futures::stream::iter(vec![]); + let mut s = GuardedStream { + inner: Box::pin(inner), + _guard: guard, + }; + while s.next().await.is_some() {} + assert_eq!(count.load(Ordering::SeqCst), 0, "must not fire before drop"); + drop(s); + assert_eq!( + count.load(Ordering::SeqCst), + 1, + "fires on drop after consumption" + ); + } + + #[tokio::test] + async fn guarded_stream_fires_on_early_drop() { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let count = Arc::new(AtomicUsize::new(0)); + let c = count.clone(); + let guard = JobCleanupGuard::new(Some(Box::new(move || { + c.fetch_add(1, Ordering::SeqCst); + }))); + let inner = futures::stream::iter(vec![Ok( + datafusion::arrow::array::RecordBatch::new_empty(std::sync::Arc::new( + datafusion::arrow::datatypes::Schema::empty(), + )), + )]); + let s = GuardedStream { + inner: Box::pin(inner), + _guard: guard, + }; + // Drop without consuming a single item. + drop(s); + assert_eq!(count.load(Ordering::SeqCst), 1); + } } From 1fd7a20fc407eb5771ff13402b60d7f63abbd2f6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 08:42:20 -0600 Subject: [PATCH 3/9] test: strengthen disabled JobCleanupGuard assertion --- .../src/execution_plans/distributed_query.rs | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index cc3f75b369..d948afe489 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -1001,13 +1001,23 @@ mod test { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + // A disabled guard must hold no cleanup action... + let guard = JobCleanupGuard::disabled(); + assert!( + guard.cleanup.is_none(), + "disabled guard must hold no cleanup action" + ); + drop(guard); // ...and dropping it must not panic or fire anything. + + // Contrast: an enabled guard built the same way production gates it DOES fire, + // proving the counter wiring is real and the disabled case is a genuine no-op. let count = Arc::new(AtomicUsize::new(0)); let c = count.clone(); - // Simulate the production "disabled" branch: no closure is built. - let _c = c; // the counter is intentionally never wired to a closure - let guard = JobCleanupGuard::disabled(); - drop(guard); - assert_eq!(count.load(Ordering::SeqCst), 0); + let enabled = JobCleanupGuard::new(Some(Box::new(move || { + c.fetch_add(1, Ordering::SeqCst); + }))); + drop(enabled); + assert_eq!(count.load(Ordering::SeqCst), 1); } #[tokio::test] From 62cbb43f0695a0fc3593bb743438636daeeb0d84 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 08:48:05 -0600 Subject: [PATCH 4/9] feat: fire CleanJobData when client finishes consuming results Wire the JobCleanupGuard/GuardedStream from the previous commit into execute_query_pull and execute_query_push: when client-side cleanup is enabled, wrap the flattened result stream so dropping it fires a best-effort CleanJobData RPC for the finished job. --- .../src/execution_plans/distributed_query.rs | 60 +++++++++++++++++-- 1 file changed, 55 insertions(+), 5 deletions(-) diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index d948afe489..838465861a 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -21,9 +21,9 @@ use crate::config::BallistaConfig; use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt}; use crate::serde::protobuf::get_job_status_result::FlightProxy; use crate::serde::protobuf::{ - ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult, KeyValuePair, - PartitionLocation, execute_query_params::Query, execute_query_result, job_status, - scheduler_grpc_client::SchedulerGrpcClient, + CleanJobDataParams, ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult, + KeyValuePair, PartitionLocation, execute_query_params::Query, execute_query_result, + job_status, scheduler_grpc_client::SchedulerGrpcClient, }; use crate::serde::protobuf::{ExecutorMetadata, SuccessfulJob}; use crate::utils::{GrpcClientConfig, create_grpc_client_endpoint}; @@ -429,6 +429,9 @@ async fn execute_query_pull( let use_tls = session_config.ballista_use_tls(); let io_retries_times = grpc_config.io_retries_times; let io_retry_wait_time_ms = grpc_config.io_retry_wait_time_ms; + let client_side_cleanup = session_config + .ballista_config() + .client_side_cleanup_enabled(); // Capture query submission time for total_query_time_ms let query_start_time = std::time::Instant::now(); @@ -577,7 +580,29 @@ async fn execute_query_pull( futures::stream::once(f).try_flatten() }); - break Ok(futures::stream::iter(streams).flatten()); + let result_stream = futures::stream::iter(streams).flatten(); + let guard = if client_side_cleanup { + let handle = tokio::runtime::Handle::current(); + let mut cleanup_client = scheduler.clone(); + let cleanup_job_id = job_id.clone(); + JobCleanupGuard::new(Some(Box::new(move || { + handle.spawn(async move { + let params = CleanJobDataParams { + job_id: cleanup_job_id.into_inner(), + remove_stage_ids: vec![], + }; + if let Err(e) = cleanup_client.clean_job_data(params).await { + debug!("client-side job data cleanup RPC failed: {e:?}"); + } + }); + }))) + } else { + JobCleanupGuard::disabled() + }; + break Ok(GuardedStream { + inner: Box::pin(result_stream), + _guard: guard, + }); } }; } @@ -602,6 +627,9 @@ async fn execute_query_push( let use_tls = session_config.ballista_use_tls(); let io_retries_times = grpc_config.io_retries_times; let io_retry_wait_time_ms = grpc_config.io_retry_wait_time_ms; + let client_side_cleanup = session_config + .ballista_config() + .client_side_cleanup_enabled(); // Capture query submission time for total_query_time_ms let query_start_time = std::time::Instant::now(); @@ -743,7 +771,29 @@ async fn execute_query_push( futures::stream::once(f).try_flatten() }); - break Ok(futures::stream::iter(streams).flatten()); + let result_stream = futures::stream::iter(streams).flatten(); + let guard = if client_side_cleanup { + let handle = tokio::runtime::Handle::current(); + let mut cleanup_client = scheduler.clone(); + let cleanup_job_id = job_id.clone(); + JobCleanupGuard::new(Some(Box::new(move || { + handle.spawn(async move { + let params = CleanJobDataParams { + job_id: cleanup_job_id.into_inner(), + remove_stage_ids: vec![], + }; + if let Err(e) = cleanup_client.clean_job_data(params).await { + debug!("client-side job data cleanup RPC failed: {e:?}"); + } + }); + }))) + } else { + JobCleanupGuard::disabled() + }; + break Ok(GuardedStream { + inner: Box::pin(result_stream), + _guard: guard, + }); } }; } From 4f00116dca4bbb53236b77b21f7b42db121ec7ba Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 09:03:48 -0600 Subject: [PATCH 5/9] test: verify client-side job data cleanup on standalone cluster --- ballista/client/tests/client_side_cleanup.rs | 212 +++++++++++++++++++ 1 file changed, 212 insertions(+) create mode 100644 ballista/client/tests/client_side_cleanup.rs diff --git a/ballista/client/tests/client_side_cleanup.rs b/ballista/client/tests/client_side_cleanup.rs new file mode 100644 index 0000000000..ede64430d8 --- /dev/null +++ b/ballista/client/tests/client_side_cleanup.rs @@ -0,0 +1,212 @@ +// 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 integration tests for client-triggered job data cleanup. +//! +//! When the Ballista client finishes consuming a job's results, +//! `DistributedQueryExec` fires the scheduler's `CleanJobData` RPC +//! (fire-and-forget, from the dropped result stream), which deletes the +//! job's on-disk data on the executor immediately instead of waiting for the +//! scheduler's timed cleanup (default ~300s). This is gated by the +//! `ballista.job.client_side_cleanup` config flag (default enabled). +//! +//! The standalone executor used by [`ballista::prelude::SessionContextExt`] +//! does not expose its work dir through any public API: it always creates a +//! fresh `tempfile::TempDir` internally (see +//! `ballista-executor::standalone::new_standalone_executor_from_builder`). +//! To locate it from the outside, these tests snapshot the process-wide temp +//! directory before and after starting the standalone cluster and diff the +//! two snapshots to find the directory the executor just created. + +mod common; + +#[cfg(test)] +#[cfg(feature = "standalone")] +mod client_side_cleanup_tests { + use ballista::prelude::{SessionConfigExt, SessionContextExt}; + use ballista_core::config::BALLISTA_JOB_CLIENT_SIDE_CLEANUP; + use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; + use std::collections::HashSet; + use std::path::PathBuf; + use std::time::{Duration, Instant}; + + /// Starting a standalone cluster spawns exactly one new `tempfile::TempDir` + /// (the executor's work dir) under the process temp directory. Serialize + /// cluster startup across the tests in this file so that the "snapshot the + /// temp dir before/after starting the cluster" trick below can't race + /// against another test in this binary doing the same thing at the same + /// time. + static CLUSTER_STARTUP_LOCK: tokio::sync::Mutex<()> = + tokio::sync::Mutex::const_new(()); + + /// Returns the set of directories directly under the process temp + /// directory (`std::env::temp_dir()`). + fn temp_dir_entries() -> HashSet { + std::fs::read_dir(std::env::temp_dir()) + .map(|read_dir| { + read_dir + .filter_map(|entry| entry.ok()) + .map(|entry| entry.path()) + .filter(|path| path.is_dir()) + .collect() + }) + .unwrap_or_default() + } + + /// Number of entries directly under `dir` (0 if `dir` doesn't exist). + fn dir_entry_count(dir: &PathBuf) -> usize { + std::fs::read_dir(dir).map(|rd| rd.count()).unwrap_or(0) + } + + /// Starts a standalone cluster (optionally with a custom session config) + /// and returns it together with the set of directories that appeared + /// under the process temp dir while it was starting up. In the common + /// case that set contains exactly one entry: the executor's work dir. + async fn start_standalone_and_diff_temp_dir( + config: Option, + ) -> (SessionContext, Vec) { + let _lock = CLUSTER_STARTUP_LOCK.lock().await; + + let before = temp_dir_entries(); + + let ctx = match config { + Some(config) => { + let state = datafusion::execution::SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + SessionContext::standalone_with_state(state) + .await + .expect("standalone cluster should start") + } + None => SessionContext::standalone() + .await + .expect("standalone cluster should start"), + }; + + let after = temp_dir_entries(); + let candidates: Vec = after.difference(&before).cloned().collect(); + assert!( + !candidates.is_empty(), + "expected the standalone executor to create a new directory under {:?}, found none", + std::env::temp_dir() + ); + + (ctx, candidates) + } + + /// Registers the shared test parquet fixture used by the sibling + /// standalone tests in this crate. + async fn register_test_data(ctx: &SessionContext) { + ctx.register_parquet( + "test", + "testdata/alltypes_plain.parquet", + ParquetReadOptions::default(), + ) + .await + .unwrap(); + } + + /// Runs a query with a `GROUP BY`, forcing a shuffle stage so that the + /// executor writes job data (shuffle files) to its work dir. + async fn run_shuffling_query(ctx: &SessionContext) { + let df = ctx + .sql("SELECT bool_col, COUNT(*) as cnt FROM test GROUP BY bool_col") + .await + .unwrap(); + // Dropping the collected batches drops the underlying client result + // stream, which is what fires the (fire-and-forget) cleanup guard. + let _ = df.collect().await.unwrap(); + } + + /// Out of the candidate directories produced by starting the cluster, + /// picks the one that currently contains job data. This is also used + /// right after `collect()` returns (before any further `.await`) to + /// disambiguate which candidate is really the executor's work dir, since + /// on the single-threaded `#[tokio::test]` runtime the fire-and-forget + /// cleanup task spawned by the dropped result stream cannot have run yet + /// at that point: the executor's work dir must still hold the job's + /// shuffle output. + fn find_populated_dir(candidates: &[PathBuf]) -> PathBuf { + let populated: Vec<&PathBuf> = candidates + .iter() + .filter(|dir| dir_entry_count(dir) > 0) + .collect(); + assert_eq!( + populated.len(), + 1, + "expected exactly one candidate temp dir to contain job output \ + right after collect(), found {populated:?} (all candidates: {candidates:?})" + ); + populated[0].clone() + } + + /// With client-side cleanup enabled (the default), a job's on-disk data + /// on the executor should be removed shortly after the client finishes + /// consuming results -- long before the scheduler's timed cleanup + /// (default 300s). + #[tokio::test] + async fn job_data_removed_after_client_consumes_results() { + let (ctx, candidates) = start_standalone_and_diff_temp_dir(None).await; + register_test_data(&ctx).await; + + run_shuffling_query(&ctx).await; + + // No `.await` has happened since `collect()` returned above, so the + // cleanup task the dropped stream spawned cannot have executed yet: + // the job dir must still be present, which lets us reliably identify + // the executor's work dir among the candidates. + let work_dir = find_populated_dir(&candidates); + + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if dir_entry_count(&work_dir) == 0 { + break; + } + assert!( + Instant::now() < deadline, + "job dir under {work_dir:?} was not cleaned up within timeout" + ); + tokio::time::sleep(Duration::from_millis(200)).await; + } + } + + /// With `ballista.job.client_side_cleanup` disabled, the client must not + /// send the `CleanJobData` RPC, so the job's on-disk data should still be + /// present a short time after the client finishes consuming results. + #[tokio::test] + async fn job_data_kept_when_client_side_cleanup_disabled() { + let config = SessionConfig::new_with_ballista() + .set_bool(BALLISTA_JOB_CLIENT_SIDE_CLEANUP, false); + let (ctx, candidates) = start_standalone_and_diff_temp_dir(Some(config)).await; + register_test_data(&ctx).await; + + run_shuffling_query(&ctx).await; + + let work_dir = find_populated_dir(&candidates); + + // Keep this well short of the scheduler's ~300s timed cleanup: we're + // only proving the client didn't clean up on its own. + tokio::time::sleep(Duration::from_secs(2)).await; + + assert!( + dir_entry_count(&work_dir) > 0, + "job dir under {work_dir:?} was removed even though \ + client_side_cleanup was disabled" + ); + } +} From 5aade4c518c05b0595d60a230c2399ad2f00c5af Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 09:21:04 -0600 Subject: [PATCH 6/9] fix: locate executor work dir after query, not at startup TempDir::new() in ballista-executor::standalone is never bound, so it drops (and deletes the directory) at the end of that statement; the work dir only exists on disk once a query lazily recreates it via create_dir_all in ShuffleWriterExec. Diffing the shared system temp dir around cluster startup therefore always found zero candidates. Pin TMPDIR to a private per-test root before starting the cluster and scan that root after the query runs instead, avoiding both the startup-timing bug and any race with the shared system temp dir used by other standalone-cluster tests in this crate. --- ballista/client/tests/client_side_cleanup.rs | 234 +++++++++++-------- 1 file changed, 132 insertions(+), 102 deletions(-) diff --git a/ballista/client/tests/client_side_cleanup.rs b/ballista/client/tests/client_side_cleanup.rs index ede64430d8..b74994ae72 100644 --- a/ballista/client/tests/client_side_cleanup.rs +++ b/ballista/client/tests/client_side_cleanup.rs @@ -25,12 +25,25 @@ //! `ballista.job.client_side_cleanup` config flag (default enabled). //! //! The standalone executor used by [`ballista::prelude::SessionContextExt`] -//! does not expose its work dir through any public API: it always creates a -//! fresh `tempfile::TempDir` internally (see -//! `ballista-executor::standalone::new_standalone_executor_from_builder`). -//! To locate it from the outside, these tests snapshot the process-wide temp -//! directory before and after starting the standalone cluster and diff the -//! two snapshots to find the directory the executor just created. +//! does not expose its work dir through any public API, and the work dir is +//! not even created at cluster startup: `new_standalone_executor_from_builder` +//! (`ballista-executor::standalone`) does +//! `let work_dir = TempDir::new()?.path().to_str().unwrap().to_string();` -- +//! the `TempDir` guard is never bound to a variable, so it drops (and +//! deletes the directory it just created) at the end of that statement. Only +//! the *path string* survives. The directory tree is lazily recreated later +//! by `create_dir_all` in `ShuffleWriterExec` +//! (`ballista/core/src/execution_plans/shuffle_writer.rs`) the first time a +//! query actually writes shuffle output there. +//! +//! To locate that path from the outside without racing the shared system +//! temp directory (which every other standalone-cluster test in this crate +//! also uses), these tests pin `TMPDIR` to a private, per-test root before +//! starting the cluster. `std::env::temp_dir()` (used by both +//! `tempfile::TempDir::new()` and DataFusion's own disk manager) honors +//! `TMPDIR` on macOS/Linux, so the executor's work dir path ends up +//! somewhere under that root, and it's the only thing that will ever be +//! written there. mod common; @@ -39,120 +52,123 @@ mod common; mod client_side_cleanup_tests { use ballista::prelude::{SessionConfigExt, SessionContextExt}; use ballista_core::config::BALLISTA_JOB_CLIENT_SIDE_CLEANUP; + use datafusion::common::Result; + use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; - use std::collections::HashSet; - use std::path::PathBuf; + use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; + use tempfile::TempDir; - /// Starting a standalone cluster spawns exactly one new `tempfile::TempDir` - /// (the executor's work dir) under the process temp directory. Serialize - /// cluster startup across the tests in this file so that the "snapshot the - /// temp dir before/after starting the cluster" trick below can't race - /// against another test in this binary doing the same thing at the same - /// time. - static CLUSTER_STARTUP_LOCK: tokio::sync::Mutex<()> = - tokio::sync::Mutex::const_new(()); - - /// Returns the set of directories directly under the process temp - /// directory (`std::env::temp_dir()`). - fn temp_dir_entries() -> HashSet { - std::fs::read_dir(std::env::temp_dir()) - .map(|read_dir| { - read_dir - .filter_map(|entry| entry.ok()) - .map(|entry| entry.path()) - .filter(|path| path.is_dir()) - .collect() - }) - .unwrap_or_default() - } + /// Mutating `TMPDIR` is process-global, so the two tests in this file + /// serialize on this lock for as long as `TMPDIR` matters: from just + /// before it's set to just after it's restored in + /// [`start_standalone_under`]. + static CLUSTER_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); /// Number of entries directly under `dir` (0 if `dir` doesn't exist). - fn dir_entry_count(dir: &PathBuf) -> usize { + fn dir_entry_count(dir: &Path) -> usize { std::fs::read_dir(dir).map(|rd| rd.count()).unwrap_or(0) } - /// Starts a standalone cluster (optionally with a custom session config) - /// and returns it together with the set of directories that appeared - /// under the process temp dir while it was starting up. In the common - /// case that set contains exactly one entry: the executor's work dir. - async fn start_standalone_and_diff_temp_dir( - config: Option, - ) -> (SessionContext, Vec) { - let _lock = CLUSTER_STARTUP_LOCK.lock().await; + /// True if `path` is a directory that itself contains at least one + /// sub-directory. Used to pick the executor's work dir + /// (`{work_dir}/{job_id}/...`) out from under `root` once a job has + /// written to it. + fn is_populated_dir(path: &Path) -> bool { + path.is_dir() + && std::fs::read_dir(path) + .map(|mut entries| { + entries.any(|e| e.map(|e| e.path().is_dir()).unwrap_or(false)) + }) + .unwrap_or(false) + } + + /// Finds the single top-level directory under `root` that currently + /// holds a job subdirectory. Only meaningful *after* a query has run -- + /// see the module docs for why the work dir doesn't exist right after + /// cluster startup. + fn find_populated_work_dir(root: &Path) -> PathBuf { + let populated: Vec = std::fs::read_dir(root) + .expect("temp root should be readable") + .filter_map(|entry| entry.ok()) + .map(|entry| entry.path()) + .filter(|path| is_populated_dir(path)) + .collect(); + assert_eq!( + populated.len(), + 1, + "expected exactly one populated work dir under {root:?}, found {populated:?}" + ); + populated[0].clone() + } - let before = temp_dir_entries(); + /// Starts a standalone cluster with its temp-file root pinned to `root` + /// for the duration of startup, so the executor's work dir path (chosen + /// once, at startup) ends up under `root` where we can find it later. + /// `TMPDIR` is restored immediately after startup completes -- the + /// executor keeps using its already-resolved absolute work dir path + /// regardless of `TMPDIR`'s value from then on, so this doesn't need to + /// stay set for the life of the test. + async fn start_standalone_under( + root: &Path, + config: Option, + ) -> SessionContext { + let _lock = CLUSTER_LOCK.lock().await; + + let previous_tmpdir = std::env::var("TMPDIR").ok(); + // SAFETY: `env::set_var`/`remove_var` are unsafe because mutating + // the environment races with other threads reading it. `CLUSTER_LOCK` + // is held for this entire function, and this is the only place in + // this test binary that touches `TMPDIR`, so no such race exists. + unsafe { + std::env::set_var("TMPDIR", root); + } - let ctx = match config { + let result = match config { Some(config) => { - let state = datafusion::execution::SessionStateBuilder::new() + let state = SessionStateBuilder::new() .with_config(config) .with_default_features() .build(); - SessionContext::standalone_with_state(state) - .await - .expect("standalone cluster should start") + SessionContext::standalone_with_state(state).await } - None => SessionContext::standalone() - .await - .expect("standalone cluster should start"), + None => SessionContext::standalone().await, }; - let after = temp_dir_entries(); - let candidates: Vec = after.difference(&before).cloned().collect(); - assert!( - !candidates.is_empty(), - "expected the standalone executor to create a new directory under {:?}, found none", - std::env::temp_dir() - ); + // SAFETY: see above. + unsafe { + match &previous_tmpdir { + Some(value) => std::env::set_var("TMPDIR", value), + None => std::env::remove_var("TMPDIR"), + } + } - (ctx, candidates) + result.expect("standalone cluster should start") } /// Registers the shared test parquet fixture used by the sibling /// standalone tests in this crate. - async fn register_test_data(ctx: &SessionContext) { + async fn register_test_data(ctx: &SessionContext) -> Result<()> { ctx.register_parquet( "test", "testdata/alltypes_plain.parquet", ParquetReadOptions::default(), ) - .await - .unwrap(); + .await?; + Ok(()) } /// Runs a query with a `GROUP BY`, forcing a shuffle stage so that the - /// executor writes job data (shuffle files) to its work dir. - async fn run_shuffling_query(ctx: &SessionContext) { + /// executor writes job data (shuffle files) to its work dir, and + /// consumes all results. + async fn run_shuffling_query(ctx: &SessionContext) -> Result<()> { let df = ctx .sql("SELECT bool_col, COUNT(*) as cnt FROM test GROUP BY bool_col") - .await - .unwrap(); + .await?; // Dropping the collected batches drops the underlying client result // stream, which is what fires the (fire-and-forget) cleanup guard. - let _ = df.collect().await.unwrap(); - } - - /// Out of the candidate directories produced by starting the cluster, - /// picks the one that currently contains job data. This is also used - /// right after `collect()` returns (before any further `.await`) to - /// disambiguate which candidate is really the executor's work dir, since - /// on the single-threaded `#[tokio::test]` runtime the fire-and-forget - /// cleanup task spawned by the dropped result stream cannot have run yet - /// at that point: the executor's work dir must still hold the job's - /// shuffle output. - fn find_populated_dir(candidates: &[PathBuf]) -> PathBuf { - let populated: Vec<&PathBuf> = candidates - .iter() - .filter(|dir| dir_entry_count(dir) > 0) - .collect(); - assert_eq!( - populated.len(), - 1, - "expected exactly one candidate temp dir to contain job output \ - right after collect(), found {populated:?} (all candidates: {candidates:?})" - ); - populated[0].clone() + let _ = df.collect().await?; + Ok(()) } /// With client-side cleanup enabled (the default), a job's on-disk data @@ -160,17 +176,28 @@ mod client_side_cleanup_tests { /// consuming results -- long before the scheduler's timed cleanup /// (default 300s). #[tokio::test] - async fn job_data_removed_after_client_consumes_results() { - let (ctx, candidates) = start_standalone_and_diff_temp_dir(None).await; - register_test_data(&ctx).await; - - run_shuffling_query(&ctx).await; - - // No `.await` has happened since `collect()` returned above, so the - // cleanup task the dropped stream spawned cannot have executed yet: - // the job dir must still be present, which lets us reliably identify - // the executor's work dir among the candidates. - let work_dir = find_populated_dir(&candidates); + async fn job_data_removed_after_client_consumes_results() -> Result<()> { + let root = TempDir::new().expect("temp root should be created"); + let ctx = start_standalone_under(root.path(), None).await; + register_test_data(&ctx).await?; + + run_shuffling_query(&ctx).await?; + + // IMPORTANT / internal ordering detail: no `.await` has happened + // since `collect()` returned above. The cleanup RPC is fired from + // `JobCleanupGuard::drop` + // (`ballista/core/src/execution_plans/distributed_query.rs`), which + // synchronously calls `Handle::spawn` -- it enqueues a task to send + // the RPC, it does not send it inline. `#[tokio::test]` defaults to + // the single-threaded runtime, and the standalone executor's own + // background tasks (poll loop, flight service) run on that same + // runtime, so that enqueued task cannot have been polled even once + // yet at this point: the job's shuffle output is guaranteed to + // still be on disk here. This is correct today but is an internal + // scheduling detail, not a public guarantee -- it's only used to + // reliably locate the work dir below, not as the assertion that the + // feature works (that's the poll loop further down). + let work_dir = find_populated_work_dir(root.path()); let deadline = Instant::now() + Duration::from_secs(10); loop { @@ -183,21 +210,23 @@ mod client_side_cleanup_tests { ); tokio::time::sleep(Duration::from_millis(200)).await; } + Ok(()) } /// With `ballista.job.client_side_cleanup` disabled, the client must not /// send the `CleanJobData` RPC, so the job's on-disk data should still be /// present a short time after the client finishes consuming results. #[tokio::test] - async fn job_data_kept_when_client_side_cleanup_disabled() { + async fn job_data_kept_when_client_side_cleanup_disabled() -> Result<()> { + let root = TempDir::new().expect("temp root should be created"); let config = SessionConfig::new_with_ballista() .set_bool(BALLISTA_JOB_CLIENT_SIDE_CLEANUP, false); - let (ctx, candidates) = start_standalone_and_diff_temp_dir(Some(config)).await; - register_test_data(&ctx).await; + let ctx = start_standalone_under(root.path(), Some(config)).await; + register_test_data(&ctx).await?; - run_shuffling_query(&ctx).await; + run_shuffling_query(&ctx).await?; - let work_dir = find_populated_dir(&candidates); + let work_dir = find_populated_work_dir(root.path()); // Keep this well short of the scheduler's ~300s timed cleanup: we're // only proving the client didn't clean up on its own. @@ -208,5 +237,6 @@ mod client_side_cleanup_tests { "job dir under {work_dir:?} was removed even though \ client_side_cleanup was disabled" ); + Ok(()) } } From 2912ae145cd1c174766a4d99053b0f54f651ebbd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 09:27:07 -0600 Subject: [PATCH 7/9] test: restore TMPDIR via RAII guard for exception safety --- ballista/client/tests/client_side_cleanup.rs | 53 ++++++++++++++------ 1 file changed, 37 insertions(+), 16 deletions(-) diff --git a/ballista/client/tests/client_side_cleanup.rs b/ballista/client/tests/client_side_cleanup.rs index b74994ae72..988e2169de 100644 --- a/ballista/client/tests/client_side_cleanup.rs +++ b/ballista/client/tests/client_side_cleanup.rs @@ -102,27 +102,54 @@ mod client_side_cleanup_tests { populated[0].clone() } + /// RAII guard that restores the process `TMPDIR` (to whatever it was + /// before this test overrode it) when dropped -- including on unwind, if + /// `SessionContext::standalone[_with_state]()` were to panic instead of + /// returning an `Err`. A plain "set, await, restore" sequence would skip + /// the restore in that case, since unwinding jumps straight past it. + struct TmpdirGuard { + previous: Option, + } + + impl Drop for TmpdirGuard { + fn drop(&mut self) { + // SAFETY: single-threaded w.r.t. `TMPDIR` writes -- the two + // tests in this file are serialized by `CLUSTER_LOCK` (held for + // the guard's entire lifetime, see `start_standalone_under`), + // and this is the only code in this binary that touches + // `TMPDIR`. + unsafe { + match &self.previous { + Some(value) => std::env::set_var("TMPDIR", value), + None => std::env::remove_var("TMPDIR"), + } + } + } + } + /// Starts a standalone cluster with its temp-file root pinned to `root` /// for the duration of startup, so the executor's work dir path (chosen /// once, at startup) ends up under `root` where we can find it later. - /// `TMPDIR` is restored immediately after startup completes -- the - /// executor keeps using its already-resolved absolute work dir path - /// regardless of `TMPDIR`'s value from then on, so this doesn't need to - /// stay set for the life of the test. + /// `TMPDIR` is restored (via [`TmpdirGuard`]) immediately after startup + /// completes -- the executor keeps using its already-resolved absolute + /// work dir path regardless of `TMPDIR`'s value from then on, so this + /// doesn't need to stay set for the life of the test. async fn start_standalone_under( root: &Path, config: Option, ) -> SessionContext { let _lock = CLUSTER_LOCK.lock().await; - let previous_tmpdir = std::env::var("TMPDIR").ok(); - // SAFETY: `env::set_var`/`remove_var` are unsafe because mutating - // the environment races with other threads reading it. `CLUSTER_LOCK` - // is held for this entire function, and this is the only place in - // this test binary that touches `TMPDIR`, so no such race exists. + let previous = std::env::var_os("TMPDIR"); + // SAFETY: see `TmpdirGuard::drop`. unsafe { std::env::set_var("TMPDIR", root); } + // Constructed *after* the set above and *before* the panic-capable + // `.await` below, so it's guaranteed to run the restore on every + // exit path -- normal return, `Err`, or panic/unwind -- once it + // drops at the end of this scope. + let tmpdir_guard = TmpdirGuard { previous }; let result = match config { Some(config) => { @@ -135,13 +162,7 @@ mod client_side_cleanup_tests { None => SessionContext::standalone().await, }; - // SAFETY: see above. - unsafe { - match &previous_tmpdir { - Some(value) => std::env::set_var("TMPDIR", value), - None => std::env::remove_var("TMPDIR"), - } - } + drop(tmpdir_guard); result.expect("standalone cluster should start") } From ba0389f75eeb46adfa5332695c50947cc45509e1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 09:33:47 -0600 Subject: [PATCH 8/9] chore: address final-review minors (warn-level cleanup log, drop unused test module) --- ballista/client/tests/client_side_cleanup.rs | 2 -- ballista/core/src/execution_plans/distributed_query.rs | 6 +++--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/ballista/client/tests/client_side_cleanup.rs b/ballista/client/tests/client_side_cleanup.rs index 988e2169de..571b99b1e1 100644 --- a/ballista/client/tests/client_side_cleanup.rs +++ b/ballista/client/tests/client_side_cleanup.rs @@ -45,8 +45,6 @@ //! somewhere under that root, and it's the only thing that will ever be //! written there. -mod common; - #[cfg(test)] #[cfg(feature = "standalone")] mod client_side_cleanup_tests { diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 838465861a..54d2280e74 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -48,7 +48,7 @@ use datafusion_proto::logical_plan::{ }; use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec}; use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt}; -use log::{debug, error, info}; +use log::{debug, error, info, warn}; use parking_lot::Mutex; use std::fmt::Debug; use std::marker::PhantomData; @@ -592,7 +592,7 @@ async fn execute_query_pull( remove_stage_ids: vec![], }; if let Err(e) = cleanup_client.clean_job_data(params).await { - debug!("client-side job data cleanup RPC failed: {e:?}"); + warn!("client-side job data cleanup RPC failed: {e:?}"); } }); }))) @@ -783,7 +783,7 @@ async fn execute_query_push( remove_stage_ids: vec![], }; if let Err(e) = cleanup_client.clean_job_data(params).await { - debug!("client-side job data cleanup RPC failed: {e:?}"); + warn!("client-side job data cleanup RPC failed: {e:?}"); } }); }))) From d481ebf33e35fc5d06e5ce1ac5ba9023a8c1a73e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 10:04:06 -0600 Subject: [PATCH 9/9] refactor: extract shared guard_result_stream helper for pull/push paths --- .../src/execution_plans/distributed_query.rs | 106 ++++++++++-------- 1 file changed, 61 insertions(+), 45 deletions(-) diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 54d2280e74..9976c0a53f 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -18,7 +18,9 @@ use crate::JobId; use crate::client::BallistaClient; use crate::config::BallistaConfig; -use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt}; +use crate::extension::{ + BallistaConfigGrpcEndpoint, BallistaGrpcMetadataInterceptor, SessionConfigExt, +}; use crate::serde::protobuf::get_job_status_result::FlightProxy; use crate::serde::protobuf::{ CleanJobDataParams, ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult, @@ -55,6 +57,8 @@ use std::marker::PhantomData; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; +use tonic::service::interceptor::InterceptedService; +use tonic::transport::Channel; use url::Url; /// This operator sends a logical plan to a Ballista scheduler for execution and @@ -581,28 +585,12 @@ async fn execute_query_pull( }); let result_stream = futures::stream::iter(streams).flatten(); - let guard = if client_side_cleanup { - let handle = tokio::runtime::Handle::current(); - let mut cleanup_client = scheduler.clone(); - let cleanup_job_id = job_id.clone(); - JobCleanupGuard::new(Some(Box::new(move || { - handle.spawn(async move { - let params = CleanJobDataParams { - job_id: cleanup_job_id.into_inner(), - remove_stage_ids: vec![], - }; - if let Err(e) = cleanup_client.clean_job_data(params).await { - warn!("client-side job data cleanup RPC failed: {e:?}"); - } - }); - }))) - } else { - JobCleanupGuard::disabled() - }; - break Ok(GuardedStream { - inner: Box::pin(result_stream), - _guard: guard, - }); + break Ok(guard_result_stream( + result_stream, + client_side_cleanup, + &scheduler, + &job_id, + )); } }; } @@ -772,28 +760,12 @@ async fn execute_query_push( }); let result_stream = futures::stream::iter(streams).flatten(); - let guard = if client_side_cleanup { - let handle = tokio::runtime::Handle::current(); - let mut cleanup_client = scheduler.clone(); - let cleanup_job_id = job_id.clone(); - JobCleanupGuard::new(Some(Box::new(move || { - handle.spawn(async move { - let params = CleanJobDataParams { - job_id: cleanup_job_id.into_inner(), - remove_stage_ids: vec![], - }; - if let Err(e) = cleanup_client.clean_job_data(params).await { - warn!("client-side job data cleanup RPC failed: {e:?}"); - } - }); - }))) - } else { - JobCleanupGuard::disabled() - }; - break Ok(GuardedStream { - inner: Box::pin(result_stream), - _guard: guard, - }); + break Ok(guard_result_stream( + result_stream, + client_side_cleanup, + &scheduler, + &job_id, + )); } }; } @@ -893,6 +865,50 @@ impl Stream for GuardedStream { } } +/// Concrete scheduler-client type used on the query-execution path: a +/// [`SchedulerGrpcClient`] over a [`Channel`] with the Ballista metadata +/// interceptor applied. +type BallistaSchedulerClient = + SchedulerGrpcClient>; + +/// Wraps a job's result stream so that, when the client finishes consuming (or +/// drops) it, a best-effort `CleanJobData` RPC is fired at the scheduler to +/// reclaim the job's on-disk data immediately instead of waiting for the +/// scheduler's timed cleanup. When `client_side_cleanup` is false the stream is +/// returned with a no-op guard. Shared by the pull and push query paths. +fn guard_result_stream( + result_stream: S, + client_side_cleanup: bool, + scheduler: &BallistaSchedulerClient, + job_id: &JobId, +) -> GuardedStream +where + S: Stream> + Send + 'static, +{ + let guard = if client_side_cleanup { + let handle = tokio::runtime::Handle::current(); + let mut cleanup_client = scheduler.clone(); + let cleanup_job_id = job_id.clone(); + JobCleanupGuard::new(Some(Box::new(move || { + handle.spawn(async move { + let params = CleanJobDataParams { + job_id: cleanup_job_id.into_inner(), + remove_stage_ids: vec![], + }; + if let Err(e) = cleanup_client.clean_job_data(params).await { + warn!("client-side job data cleanup RPC failed: {e:?}"); + } + }); + }))) + } else { + JobCleanupGuard::disabled() + }; + GuardedStream { + inner: Box::pin(result_stream), + _guard: guard, + } +} + #[allow(clippy::too_many_arguments)] async fn fetch_partition( location: PartitionLocation,