Skip to content

[Query] Fix scheduler stop races and query attempt context leaks - #19224

Open
JackieTien97 wants to merge 1 commit into
masterfrom
ty/query-attempt-cleanup
Open

JackieTien97 wants to merge 1 commit into
masterfrom
ty/query-attempt-cleanup

Conversation

@JackieTien97

Copy link
Copy Markdown
Contributor

Description

Fix two cleanup failures after partial query fragment dispatch.

Planners previously started a scheduler before returning it to QueryExecution. A terminal-state listener could consume the stop flag while the scheduler reference was still null or referred to the previous attempt, leaving the current scheduler unstopped. Separately, DataNodeQueryContext removal depended on the planned FI count reaching zero, so an FI that was never dispatched could retain the context indefinitely.

Example

Query Q plans three FIs on DN1: FI1 starts, FI2 fails with status 719, and FI3 is never dispatched. Previously, termination could be lost and the shared context would continue waiting for FI3. With this change, closing (Q, 0) cancels the actual registered FIs and releases their shared context once its users finish cleanup. A retry uses (Q, 1), so delayed cancellation and resource-release callbacks from the first attempt cannot remove the retry's context.

Scheduler and attempt lifecycle

  • Create schedulers without starting them, publish them through a per-attempt handoff, and retain stop requests that arrive before publication.
  • Make cluster scheduler stop idempotent, abort further dispatch, and wake topology dispatch waits. Preserve the normal completion grace period and terminate immediately when stopping during dispatch or after a dispatch failure.
  • Adapt LOAD scheduler cleanup so stopping before start also releases constructor-owned memory exactly once.
  • Register shared DataNode state by (queryId, attemptId). Keep it across gaps between sequential FI arrivals, and close a failed/cancelled attempt without waiting for never-dispatched FIs. Planning failures and execution cleanup release registrations idempotently.

Late requests and compatibility

After its actual users release their resources, a closed attempt retains only a lightweight admission record, without LAST caches or other shared query data. Admission deadlines are separate from execution deadlines: the admission window follows the internal DataNode RPC timeout (60 seconds when the socket timeout is unlimited), capped by the query deadline. Transport retries reuse the original FI deadline; stopping dispatch freezes the latest issued cutoff. Expired requests remain rejected after the corresponding record is removed, while an already admitted long-running query keeps its execution timeout. Deadline comparisons use coordinator-provided wall-clock timestamps.

Attempt metadata is carried in optional TSendFragmentInstanceReq and TCancelQueryReq fields; the binary FI body is unchanged. Legacy requests retain their fallback behavior, and older receivers can ignore the new fields. Complete attempt isolation requires both endpoints to support the metadata.

Validation

  • JDK 17 targeted reactor clean test: 79 tests passed across 14 classes, with zero failures, errors, or skips. Coverage includes scheduler publication/stop races, partial dispatch, retry isolation, late requests, bounded admission records, RPC retries, legacy wire compatibility, LOAD cleanup, LAST cache behavior, and execution memory budgets.
  • mvn test-compile -DskipTests: passed for the full reactor.
  • mvn test-compile -P with-zh-locale -DskipTests: passed for the full reactor.
  • Formatting and Checkstyle passed in the reactor builds; English/Chinese message keys and the new message's placeholders match.

Multi-node cluster and rolling-upgrade integration tests have not been run.

This PR has:

  • been self-reviewed.
  • added comments explaining lifecycle ownership and concurrency decisions.
  • added or updated unit tests for the changed code paths.
Key changed/added classes

QueryExecution, ClusterScheduler, FragmentInstanceDispatcherImpl, SimpleQueryTerminator, QueryAttemptId, QueryAttemptRegistry, FragmentInstanceManager, and FragmentInstanceContext.

Publish schedulers before starting them and preserve stop requests across creation. Isolate DataNode contexts by query attempt, close partially dispatched attempts, and bound late-dispatch fences with admission deadlines.

Add regression coverage for concurrent lifecycle transitions, retry isolation, partial dispatch cleanup, RPC compatibility, and admission expiry.
@sonarqubecloud

Copy link
Copy Markdown

Quality Gate Failed Quality Gate failed

Failed conditions
B Reliability Rating on New Code (required ≥ A)

See analysis details on SonarQube Cloud

Catch issues before they fail your Quality Gate with our IDE extension SonarQube for IDE

@JackieTien97 JackieTien97 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Four correctness issues remain in the scheduler stop and admission paths; please address the inline comments before merging.

Validation on 0030ac2: a targeted JDK 17 clean reactor test run passed all 38 tests across 9 classes. Four additional focused unit probes reproduced the behaviors described below. Multi-node cluster and rolling-upgrade integration tests were not run.

Comment on lines +152 to +153
if (stopped) {
return;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Complete the query state when terminal cleanup stops dispatch

stopped does not imply that the query state machine is terminal: Coordinator.cleanUpStaleQueries() calls cleanupQueryExecution() -> stopAndCleanup(timeout) directly, without calling cancel() or transitioning the query to FAILED. If that cleanup runs while dispatch is blocked, this branch discards the dispatch result when it eventually returns, leaving the query in DISPATCHING. The state tracker has also been aborted, so the subsequent QueryExecution.getStatus() waits indefinitely and retains the RPC thread. A latch-controlled test with a READ query reproduced this: stopAndCleanup during dispatch, let dispatch return, then getStatus never completes. Terminal cleanup needs to complete the query state or otherwise wake its status waiters, while preserving the separate stop semantics used for retries.

Comment on lines +812 to +816
if (deadline == 0) {
deadline =
Math.min(
queryDeadline, LongMath.saturatedAdd(System.currentTimeMillis(), admissionWindowMs));
instance.setAdmissionDeadline(deadline);

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve recovery after the first RPC exhausts its socket timeout

The admission window equals getDnConnectionTimeoutInMS(), but this branch only initializes the deadline once. When the first sendFragmentInstance call fails after waiting for its socket timeout, dispatchRemote() reconnects and resends the FI with an already expired admission deadline. The receiver then returns QUERY_TIMEOUT (720), which ClusterScheduler.needRetry() does not retry, even if the query still has ample execution budget and the first request never reached the receiver. This was reproduced with a 20 ms socket timeout and a one-hour query timeout. Preserve the immutable cutoff for an old request, but distinguish admission expiry from execution timeout and allow a fresh query attempt, or otherwise budget admission for the supported transport retry.

Comment on lines +90 to +94
long currentTime = now();
if (admissionDeadline <= currentTime
|| (existing != null
&& (existing.closed || existing.executionDeadline <= currentTime))) {
throw rejected(key, admissionDeadline <= currentTime);

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Avoid comparing coordinator deadlines directly with receiver wall time

queryDeadline and admissionDeadline are produced using the coordinator's wall clock, whereas now() reads the receiving DataNode's clock. With the default 60-second admission window, a receiver whose clock is 61 seconds ahead rejects a freshly issued FI immediately, even when its query timeout is one hour and there is no network delay. A clock-injected test reproduced this; the same request using the legacy relative-timeout path is admitted. This introduces a cluster clock-synchronization requirement that the repository's clock assumptions explicitly do not provide (THREAT_MODEL.md, section 5). Deadline comparisons need a common time basis or an explicit clock-offset-safe admission protocol.

}

private long now() {
return observedTime.accumulateAndGet(clock.getAsLong(), Math::max);

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Do not apply the historical maximum clock to unrelated new queries

observedTime is shared by every attempt and never decreases. After a wall-clock rollback, even a coordinator and receiver in the same JVM generate fresh request deadlines from the adjusted clock while registration continues comparing against the old maximum. For example, after a 120-second rollback, a new query's default 60-second admission deadline is already below observedTime, so new queries are rejected for approximately a minute despite having a one-hour execution budget. This was reproduced by calling removeExpired(), moving the injected clock backwards, and registering a completely new queryId. Keep the protection against admitting old requests scoped to the relevant attempt so that it does not reject fresh query IDs after clock correction.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant