feat(core): eagerly reclaim job data when the client finishes consuming results - #2170
Draft
andygrove wants to merge 9 commits into
Draft
feat(core): eagerly reclaim job data when the client finishes consuming results#2170andygrove wants to merge 9 commits into
andygrove wants to merge 9 commits into
Conversation
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.
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.
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.
github-actions
Bot
added
the
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
No linked issue.
Rationale for this change
When a job completes, the scheduler fires
JobFinishedas soon as the final stage's tasks finish on the executors — independent of whether the client has fetched the results. It then schedules the final-stage output for deletion only afterfinished_job_data_clean_up_interval_seconds(default 300s). As a result, a job's result files linger on executor disk for up to ~5 minutes after the client has actually finished consuming them, because nothing tells the scheduler the client is done.The scheduler already exposes a
CleanJobData(job_id)RPC that reclaims a job's on-disk data on executors, but nothing calls it. This PR wires the client to call it as soon as it finishes (or abandons) consuming results, so the data is reclaimed promptly. The scheduler's timed cleanup remains as a safety net, and scheduler-side job state is deliberately left on its own timer so jobs still appear in history/status.What changes are included in this PR?
DistributedQueryExecnow wraps its result stream in aGuardedStreamwhose drop-guard fires the existingCleanJobDataRPC (fire-and-forget) when the client finishes consuming results or drops the stream early (e.g.LIMIT, or a mid-stream error). This covers both the pull and push job-tracking paths.remove_stage_ids, i.e. "remove the whole job dir" — reclaiming the final-stage output. Scheduler job state is untouched.tokio::runtime::Handlein async context and spawns the RPC on drop, so dropping never blocks; RPC errors are best-effort (logged atwarn!).GuardedStreamdeclaresinnerbefore the guard, so the result-fetch substreams tear down before cleanup fires — no file is deleted out from under an in-flight fetch.ballista.job.client_side_cleanup(Boolean, defaulttrue) gates the behavior. When disabled, the stream is returned unchanged and only the scheduler's timer reclaims data.Are there any user-facing changes?
Yes:
ballista.job.client_side_cleanup(defaulttrue), documented in the generated config reference.ballista.job.client_side_cleanup=falseto preserve the old timing.No breaking public API changes (the additions are a new config key and a getter).