From af73e70e7fb232974f4447e2187aa32469337080 Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 19:09:44 +0000 Subject: [PATCH 01/15] Add PPL asynchronous query lifecycle core Signed-off-by: Peng Huo --- .../sql/common/setting/Settings.java | 12 + .../sql/executor/AsyncQueryExecution.java | 34 + .../setting/OpenSearchSettings.java | 66 ++ .../setting/OpenSearchSettingsTest.java | 46 + plugin/build.gradle | 1 + .../sql/plugin/PPLQueryErrorHandler.java | 67 ++ .../asyncquery/PPLAsyncQueryJob.java | 430 ++++++++ .../asyncquery/PPLAsyncQueryJobId.java | 141 +++ .../asyncquery/PPLAsyncQueryService.java | 778 +++++++++++++++ .../asyncquery/PPLAsyncQueryUser.java | 86 ++ .../asyncquery/PPLAsyncQueryJobIdTest.java | 28 + .../asyncquery/PPLAsyncQueryServiceTest.java | 917 ++++++++++++++++++ .../asyncquery/PPLAsyncQueryUserTest.java | 85 ++ .../sql/ppl/DefaultAsyncQueryExecution.java | 51 + .../org/opensearch/sql/ppl/PPLService.java | 35 + .../sql/ppl/domain/PPLQueryRequest.java | 74 ++ .../ppl/DefaultAsyncQueryExecutionTest.java | 62 ++ .../opensearch/sql/ppl/PPLServiceTest.java | 22 + .../sql/ppl/domain/PPLQueryRequestTest.java | 64 ++ 19 files changed, 2999 insertions(+) create mode 100644 core/src/main/java/org/opensearch/sql/executor/AsyncQueryExecution.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/PPLQueryErrorHandler.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java create mode 100644 ppl/src/main/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecution.java create mode 100644 ppl/src/test/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecutionTest.java diff --git a/common/src/main/java/org/opensearch/sql/common/setting/Settings.java b/common/src/main/java/org/opensearch/sql/common/setting/Settings.java index 67643a80add..fcdf84d9e8f 100644 --- a/common/src/main/java/org/opensearch/sql/common/setting/Settings.java +++ b/common/src/main/java/org/opensearch/sql/common/setting/Settings.java @@ -38,6 +38,18 @@ public enum Key { PPL_JOIN_SUBSEARCH_MAXOUT("plugins.ppl.join.subsearch_maxout"), PPL_REST_ALLOWED_ENDPOINTS("plugins.ppl.rest.allowed_endpoints"), + /** Maximum number of asynchronous PPL queries running concurrently on one node. */ + PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES("plugins.ppl.async.node_concurrent_running_queries"), + + /** Maximum number of asynchronous PPL jobs retained on one node. */ + PPL_ASYNC_MAX_RETAINED_JOBS("plugins.ppl.async.max_retained_jobs"), + + /** Maximum accepted submit wait-for-completion timeout. */ + PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT("plugins.ppl.async.max_wait_for_completion_timeout"), + + /** Maximum accepted asynchronous PPL job lease. */ + PPL_ASYNC_MAX_KEEP_ALIVE("plugins.ppl.async.max_keep_alive"), + /** Enable Calcite as execution engine */ CALCITE_ENGINE_ENABLED("plugins.calcite.enabled"), CALCITE_FALLBACK_ALLOWED("plugins.calcite.fallback.allowed"), diff --git a/core/src/main/java/org/opensearch/sql/executor/AsyncQueryExecution.java b/core/src/main/java/org/opensearch/sql/executor/AsyncQueryExecution.java new file mode 100644 index 00000000000..c52c124a099 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/AsyncQueryExecution.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import java.util.Optional; +import java.util.concurrent.CompletionStage; +import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; + +/** + * Lifecycle-facing handle for one asynchronous query execution. + * + *

The execution module owns result production and execution-specific resources. The lifecycle + * module owns this handle after submission and uses it to read the current result, observe terminal + * completion, and release those resources. + * + *

On successful completion, the authoritative final result must be visible through {@link + * #currentResult()} before {@link #completion()} completes normally. Implementations must make + * {@link #close()} idempotent and safe to call concurrently with {@link #currentResult()}. + */ +public interface AsyncQueryExecution extends AutoCloseable { + + /** Returns the complete result currently visible, or empty before a result is available. */ + Optional currentResult(); + + /** Completes normally on query success and exceptionally on query failure. */ + CompletionStage completion(); + + /** Releases execution-owned result resources. */ + @Override + void close(); +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java index ad6bdcb7c9c..97f4c390b74 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java @@ -36,6 +36,8 @@ /** Setting implementation on OpenSearch. */ @Log4j2 public class OpenSearchSettings extends Settings { + private static final TimeValue MAX_PPL_ASYNC_KEEP_ALIVE = TimeValue.timeValueHours(24); + /** Default settings. */ private final Map> defaultSettings; @@ -86,6 +88,42 @@ public class OpenSearchSettings extends Settings { Setting.Property.NodeScope, Setting.Property.Dynamic); + /** Maximum number of asynchronous PPL queries running concurrently on one node. */ + public static final Setting PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES_SETTING = + Setting.intSetting( + Key.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES.getKeyValue(), + 20, + 1, + Setting.Property.NodeScope, + Setting.Property.Dynamic); + + /** Maximum number of asynchronous PPL jobs retained on one node. */ + public static final Setting PPL_ASYNC_MAX_RETAINED_JOBS_SETTING = + Setting.intSetting( + Key.PPL_ASYNC_MAX_RETAINED_JOBS.getKeyValue(), + 100, + 1, + Setting.Property.NodeScope, + Setting.Property.Dynamic); + + /** Maximum accepted submit wait-for-completion timeout. */ + public static final Setting PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT_SETTING = + Setting.positiveTimeSetting( + Key.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT.getKeyValue(), + TimeValue.timeValueSeconds(60), + Setting.Property.NodeScope, + Setting.Property.Dynamic); + + /** Configurable asynchronous PPL job lease limit, capped at 24 hours. */ + public static final Setting PPL_ASYNC_MAX_KEEP_ALIVE_SETTING = + Setting.timeSetting( + Key.PPL_ASYNC_MAX_KEEP_ALIVE.getKeyValue(), + MAX_PPL_ASYNC_KEEP_ALIVE, + TimeValue.ZERO, + MAX_PPL_ASYNC_KEEP_ALIVE, + Setting.Property.NodeScope, + Setting.Property.Dynamic); + public static final Setting PPL_SYNTAX_LEGACY_PREFERRED_SETTING = Setting.boolSetting( Key.PPL_SYNTAX_LEGACY_PREFERRED.getKeyValue(), @@ -444,6 +482,30 @@ public OpenSearchSettings(ClusterSettings clusterSettings) { Key.PPL_QUERY_TIMEOUT, PPL_QUERY_TIMEOUT_SETTING, new Updater(Key.PPL_QUERY_TIMEOUT)); + register( + settingBuilder, + clusterSettings, + Key.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES, + PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES_SETTING, + new Updater(Key.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES)); + register( + settingBuilder, + clusterSettings, + Key.PPL_ASYNC_MAX_RETAINED_JOBS, + PPL_ASYNC_MAX_RETAINED_JOBS_SETTING, + new Updater(Key.PPL_ASYNC_MAX_RETAINED_JOBS)); + register( + settingBuilder, + clusterSettings, + Key.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT, + PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT_SETTING, + new Updater(Key.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT)); + register( + settingBuilder, + clusterSettings, + Key.PPL_ASYNC_MAX_KEEP_ALIVE, + PPL_ASYNC_MAX_KEEP_ALIVE_SETTING, + new Updater(Key.PPL_ASYNC_MAX_KEEP_ALIVE)); register( settingBuilder, clusterSettings, @@ -767,6 +829,10 @@ public static List> pluginSettings() { .add(DESERIALIZATION_MAX_BYTES_SETTING) .add(PPL_ENABLED_SETTING) .add(PPL_QUERY_TIMEOUT_SETTING) + .add(PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES_SETTING) + .add(PPL_ASYNC_MAX_RETAINED_JOBS_SETTING) + .add(PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT_SETTING) + .add(PPL_ASYNC_MAX_KEEP_ALIVE_SETTING) .add(PPL_SYNTAX_LEGACY_PREFERRED_SETTING) .add(CALCITE_ENGINE_ENABLED_SETTING) .add(CALCITE_FALLBACK_ALLOWED_SETTING) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/setting/OpenSearchSettingsTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/setting/OpenSearchSettingsTest.java index 63678051d3d..7c9e3cb3f9e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/setting/OpenSearchSettingsTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/setting/OpenSearchSettingsTest.java @@ -9,6 +9,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.AdditionalMatchers.not; import static org.mockito.AdditionalMatchers.or; @@ -28,6 +29,7 @@ import org.opensearch.cluster.ClusterName; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; +import org.opensearch.common.unit.TimeValue; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.utils.DeserializationFilterUtil; @@ -168,6 +170,50 @@ void deserializationStructuralLimitSettingsAreDynamicAndRegistered() { assertTrue(pluginSettings.contains(OpenSearchSettings.DESERIALIZATION_MAX_BYTES_SETTING)); } + @Test + void pplAsyncSettingsHaveExpectedDefaultsAndAreDynamic() { + org.opensearch.common.settings.Settings empty = org.opensearch.common.settings.Settings.EMPTY; + + assertEquals( + 20, OpenSearchSettings.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES_SETTING.get(empty)); + assertEquals(100, OpenSearchSettings.PPL_ASYNC_MAX_RETAINED_JOBS_SETTING.get(empty)); + assertEquals( + TimeValue.timeValueSeconds(60), + OpenSearchSettings.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT_SETTING.get(empty)); + assertEquals( + TimeValue.timeValueHours(24), + OpenSearchSettings.PPL_ASYNC_MAX_KEEP_ALIVE_SETTING.get(empty)); + + List> pluginSettings = OpenSearchSettings.pluginSettings(); + assertTrue( + pluginSettings.contains( + OpenSearchSettings.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES_SETTING)); + assertTrue(pluginSettings.contains(OpenSearchSettings.PPL_ASYNC_MAX_RETAINED_JOBS_SETTING)); + assertTrue( + pluginSettings.contains( + OpenSearchSettings.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT_SETTING)); + assertTrue(pluginSettings.contains(OpenSearchSettings.PPL_ASYNC_MAX_KEEP_ALIVE_SETTING)); + assertTrue(OpenSearchSettings.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES_SETTING.isDynamic()); + assertTrue(OpenSearchSettings.PPL_ASYNC_MAX_RETAINED_JOBS_SETTING.isDynamic()); + assertTrue(OpenSearchSettings.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT_SETTING.isDynamic()); + assertTrue(OpenSearchSettings.PPL_ASYNC_MAX_KEEP_ALIVE_SETTING.isDynamic()); + } + + @Test + void pplAsyncMaxKeepAliveIsCappedAt24Hours() { + String key = Settings.Key.PPL_ASYNC_MAX_KEEP_ALIVE.getKeyValue(); + + assertEquals( + TimeValue.timeValueHours(24), + OpenSearchSettings.PPL_ASYNC_MAX_KEEP_ALIVE_SETTING.get( + org.opensearch.common.settings.Settings.builder().put(key, "24h").build())); + assertThrows( + IllegalArgumentException.class, + () -> + OpenSearchSettings.PPL_ASYNC_MAX_KEEP_ALIVE_SETTING.get( + org.opensearch.common.settings.Settings.builder().put(key, "25h").build())); + } + @Test void getSparkExecutionEngineConfigSetting() { // Default is empty string diff --git a/plugin/build.gradle b/plugin/build.gradle index d0c825424be..64ac97f561a 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -168,6 +168,7 @@ dependencies { api project(':datasources') api project(':async-query') api project(':direct-query') + implementation "org.opensearch:common-utils:${opensearch_build}" testImplementation group: 'net.bytebuddy', name: 'byte-buddy-agent', version: '1.15.11' testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: "${hamcrest_version}" diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/PPLQueryErrorHandler.java b/plugin/src/main/java/org/opensearch/sql/plugin/PPLQueryErrorHandler.java new file mode 100644 index 00000000000..9d690dd3e67 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/PPLQueryErrorHandler.java @@ -0,0 +1,67 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.OpenSearchException; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.sql.common.antlr.SyntaxCheckException; +import org.opensearch.sql.common.error.ErrorReport; +import org.opensearch.sql.datasources.exceptions.DataSourceClientException; +import org.opensearch.sql.exception.QueryEngineException; +import org.opensearch.sql.legacy.metrics.MetricName; +import org.opensearch.sql.legacy.metrics.Metrics; + +/** Classifies PPL failures and records the corresponding customer or system error metric. */ +public final class PPLQueryErrorHandler { + private static final Logger LOG = LogManager.getLogger(PPLQueryErrorHandler.class); + + private PPLQueryErrorHandler() {} + + /** + * Records a PPL failure and returns the HTTP status associated with it. + * + * @param exception query failure + * @return client or system error status + */ + public static RestStatus recordFailure(Exception exception) { + int code = rawStatusCode(exception); + if (400 <= code && code < 500) { + increment(MetricName.PPL_FAILED_REQ_COUNT_CUS); + } else if (500 <= code && code < 600) { + increment(MetricName.PPL_FAILED_REQ_COUNT_SYS); + } else { + LOG.warn( + "Got an exception returning non-error status {}", RestStatus.fromCode(code), exception); + } + return RestStatus.fromCode(code); + } + + private static int rawStatusCode(Exception exception) { + if (exception instanceof ErrorReport errorReport) { + return rawStatusCode(errorReport.getCause()); + } + if (exception instanceof OpenSearchException openSearchException) { + return openSearchException.status().getStatus(); + } + return isClientError(exception) ? 400 : 500; + } + + private static boolean isClientError(Exception exception) { + return exception instanceof IllegalArgumentException + || exception instanceof IndexNotFoundException + || exception instanceof QueryEngineException + || exception instanceof SyntaxCheckException + || exception instanceof DataSourceClientException + || exception instanceof IllegalAccessException; + } + + private static void increment(MetricName metricName) { + Metrics.getInstance().getNumericalMetric(metricName).increment(); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java new file mode 100644 index 00000000000..771b1b5cfc0 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java @@ -0,0 +1,430 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import java.util.Objects; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.sql.executor.AsyncQueryExecution; +import org.opensearch.tasks.CancellableTask; + +/** + * Mutable state machine for one asynchronous PPL query. + * + *

Every lifecycle and lease transition is synchronized on this object. Methods mutate only + * job-owned state and return immutable transition values; they never call listeners, mutate the + * service map, update capacity counters, or cancel tasks while holding the lock. + * + *

After attachment, the job owns one {@link AsyncQueryExecution}. Reads borrow it through a + * {@link ResponseContext}; removal transitions detach it so the service can close it outside the + * lock. + * + * 

+ * Every job starts in RUNNING. It becomes retained when wait_for_completion_timeout expires.
+ *
+ * Current state       Event              Next state             Response
+ * RUNNING             success            REMOVED                final result without ID
+ * RUNNING             failure            REMOVED                failure without ID
+ * RUNNING             retain             RETAINED_RUNNING       running status with ID
+ * RETAINED_RUNNING    success            RETAINED_SUCCEEDED     none
+ * RETAINED_RUNNING    failure            RETAINED_FAILED        none
+ * RUNNING             abort/close        REMOVED                none
+ * RETAINED_*          delete/expire/abort/close REMOVED         none
+ *
+ * + *

GET lease renewal and execution attachment do not change the lifecycle state. Events received + * after {@code REMOVED} are ignored or reported as not found; a late execution handle is rejected + * so the service can close it. + */ +final class PPLAsyncQueryJob { + private final String id; + private final PPLAsyncQueryUser owner; + private final long startTimeMillis; + private JobTask task; + + private long keepAliveMillis; + private long expirationTimeMillis; + private State state; + private AsyncQueryExecution execution; + private PPLAsyncQueryService.Failure failure; + private long completionTimeMillis = -1L; + + /** + * Creates an unretained job in the {@link State#RUNNING} state. + * + * @param id opaque ID assigned by the owner node + * @param owner caller that is allowed to access the retained job + * @param startTimeMillis execution start time + * @param keepAliveMillis lease duration applied when the job is retained + * @param task independently cancellable task owned by the job + */ + PPLAsyncQueryJob( + String id, + PPLAsyncQueryUser owner, + long startTimeMillis, + long keepAliveMillis, + JobTask task) { + this.id = id; + this.owner = owner; + this.startTimeMillis = startTimeMillis; + this.keepAliveMillis = keepAliveMillis; + this.expirationTimeMillis = startTimeMillis + keepAliveMillis; + this.task = Objects.requireNonNull(task); + this.state = State.RUNNING; + } + + /** + * Returns the opaque ID assigned to this job. + * + * @return job ID used as the service registry key + */ + String id() { + return id; + } + + /** + * Returns the immutable identity captured when this job was created. + * + * @return job owner used by the service for GET and DELETE authorization + */ + PPLAsyncQueryUser owner() { + return owner; + } + + /** + * Retains a running job after its initial response wait expires. + * + * @param now time at which the job becomes visible to GET and DELETE + * @return transition that publishes the job ID, or {@code null} if the job already finished + */ + synchronized Transition retain(long now) { + if (state != State.RUNNING) { + return null; + } + state = State.RETAINED_RUNNING; + expirationTimeMillis = now + keepAliveMillis; + return Transition.retain(retainedResponse()); + } + + /** + * Transfers ownership of an execution handle to this job. + * + * @param execution handle that produces current and final query results + * @return {@code true} when attached; {@code false} when the caller must close the rejected + * handle + */ + synchronized boolean tryAttachExecution(AsyncQueryExecution execution) { + Objects.requireNonNull(execution); + if (!isExecuting() || this.execution != null) { + return false; + } + this.execution = execution; + return true; + } + + /** + * Records successful execution completion. + * + * @param now completion time + * @return direct-response or retained-completion transition, or {@code null} after removal + * @throws IllegalStateException if successful completion is reported before execution attachment + */ + synchronized Transition complete(long now) { + if (!isExecuting()) { + return null; + } + if (execution == null) { + throw new IllegalStateException( + "PPL asynchronous execution must be attached before successful completion"); + } + return finish(State.RETAINED_SUCCEEDED, now); + } + + /** + * Records failed execution completion. + * + * @param failure client-visible failure retained with the job + * @param now completion time + * @return direct-response or retained-completion transition, or {@code null} after removal + */ + synchronized Transition fail(PPLAsyncQueryService.Failure failure, long now) { + if (!isExecuting()) { + return null; + } + this.failure = failure; + return finish(State.RETAINED_FAILED, now); + } + + private Transition finish(State terminalState, long now) { + boolean returnDirect = state == State.RUNNING; + JobTask taskToClose = detachTask(); + state = terminalState; + completionTimeMillis = now; + if (returnDirect) { + ResponseContext response = directResponse(); + state = State.REMOVED; + return Transition.returnDirect(response, detachExecution(), taskToClose); + } + return Transition.finishRetained( + terminalState == State.RETAINED_FAILED ? detachExecution() : null, taskToClose); + } + + /** + * Returns the current retained response. + * + *

A supplied {@code requestedKeepAlive} starts a new lease from {@code now}. A request at or + * after the existing expiration time returns an expiration removal instead of data. + * + * @param now request time + * @param requestedKeepAlive replacement lease, or {@code null} to keep the current expiration + * @return current response or the removal required for an expired job + * @throws ResourceNotFoundException if the job was already removed + */ + synchronized GetResult get(long now, TimeValue requestedKeepAlive) { + ensurePresent(); + if (now >= expirationTimeMillis) { + return new GetResult.Expired(expireLocked("PPL asynchronous query expired")); + } + if (requestedKeepAlive != null) { + keepAliveMillis = requestedKeepAlive.millis(); + expirationTimeMillis = now + keepAliveMillis; + } + return new GetResult.Found(retainedResponse()); + } + + /** + * Cancels if still running and removes this job. + * + * @param now request time + * @return removal containing the response status and detached resources + * @throws ResourceNotFoundException if the job was already removed + */ + synchronized Removal delete(long now) { + ensurePresent(); + if (now >= expirationTimeMillis) { + return expireLocked("PPL asynchronous query expired"); + } + return remove( + isExecuting() ? PPLAsyncQueryService.Status.CANCELLED : responseStatus(), + "PPL asynchronous query cancelled by user"); + } + + /** + * Removes a retained job whose lease has expired. + * + * @param now expiration check time + * @return removal with detached resources, or {@code null} if no expiration is due + */ + synchronized Removal expire(long now) { + if (state == State.REMOVED || state == State.RUNNING || now < expirationTimeMillis) { + return null; + } + return expireLocked("PPL asynchronous query expired"); + } + + private Removal expireLocked(String reason) { + return remove(responseStatus(), reason, true); + } + + /** + * Removes a job whose submission could not be completed. + * + * @return removal with detached resources, or {@code null} if already removed + */ + synchronized Removal abort() { + return removeIfPresent("PPL asynchronous query startup failed"); + } + + /** + * Removes this job during service shutdown. + * + * @param reason cancellation reason used if execution is still running + * @return removal with detached resources, or {@code null} if already removed + */ + synchronized Removal close(String reason) { + return removeIfPresent(reason); + } + + private Removal removeIfPresent(String reason) { + if (state == State.REMOVED) { + return null; + } + return remove(responseStatus(), reason); + } + + private Removal remove(PPLAsyncQueryService.Status responseStatus, String reason) { + return remove(responseStatus, reason, false); + } + + private Removal remove( + PPLAsyncQueryService.Status responseStatus, String reason, boolean expired) { + Removal removal = new Removal(responseStatus, detachTask(), detachExecution(), reason, expired); + state = State.REMOVED; + return removal; + } + + private ResponseContext directResponse() { + return responseContext(null); + } + + private ResponseContext retainedResponse() { + return responseContext(id); + } + + private ResponseContext responseContext(String responseId) { + long tookMillis = + completionTimeMillis < 0 ? -1L : Math.max(0L, completionTimeMillis - startTimeMillis); + return new ResponseContext(responseId, responseStatus(), execution, failure, tookMillis); + } + + private boolean isExecuting() { + return state == State.RUNNING || state == State.RETAINED_RUNNING; + } + + private PPLAsyncQueryService.Status responseStatus() { + return switch (state) { + case RUNNING, RETAINED_RUNNING -> PPLAsyncQueryService.Status.RUNNING; + case RETAINED_SUCCEEDED -> PPLAsyncQueryService.Status.SUCCEEDED; + case RETAINED_FAILED -> PPLAsyncQueryService.Status.FAILED; + case REMOVED -> throw new IllegalStateException("PPL asynchronous query was removed"); + }; + } + + private AsyncQueryExecution detachExecution() { + AsyncQueryExecution detached = execution; + execution = null; + return detached; + } + + private JobTask detachTask() { + JobTask detached = task; + task = null; + return detached; + } + + private void ensurePresent() { + if (state == State.REMOVED) { + throw new ResourceNotFoundException("PPL asynchronous query not found"); + } + } + + /** + * Lifecycle data captured under the job lock for later response materialization. + * + * @param id job ID included in a retained response, or {@code null} for a direct POST response + * @param status public lifecycle status + * @param execution execution handle borrowed for result materialization + * @param failure failure returned for a failed job + * @param tookMillis elapsed execution time, or {@code -1} while running + */ + record ResponseContext( + String id, + PPLAsyncQueryService.Status status, + AsyncQueryExecution execution, + PPLAsyncQueryService.Failure failure, + long tookMillis) {} + + /** + * Cancellable task and the cleanup that releases its TaskManager registrations. + * + * @param task task used to cancel query execution + * @param release registration cleanup + */ + record JobTask(CancellableTask task, Runnable release) { + /** Releases the task and child-node registrations owned by this wrapper. */ + void close() { + release.run(); + } + } + + /** Internal lifecycle; unlike the response status, this includes retention and removal. */ + enum State { + RUNNING, + RETAINED_RUNNING, + RETAINED_SUCCEEDED, + RETAINED_FAILED, + REMOVED + } + + /** Whether a transition keeps the job in or removes it from the service registry. */ + enum Retention { + RETAIN, + REMOVE + } + + /** + * State-machine output consumed by {@link PPLAsyncQueryService}. + * + * @param response response to materialize and publish, or {@code null} when none is due + * @param retention registry action + * @param executionToClose execution handle detached by the transition + * @param taskToClose completed task registration detached by the transition + */ + record Transition( + ResponseContext response, + Retention retention, + AsyncQueryExecution executionToClose, + JobTask taskToClose) { + + private static Transition retain(ResponseContext response) { + return new Transition(response, Retention.RETAIN, null, null); + } + + private static Transition returnDirect( + ResponseContext response, AsyncQueryExecution executionToClose, JobTask taskToClose) { + return new Transition(response, Retention.REMOVE, executionToClose, taskToClose); + } + + private static Transition finishRetained( + AsyncQueryExecution executionToClose, JobTask taskToClose) { + return new Transition(null, Retention.RETAIN, executionToClose, taskToClose); + } + + /** Returns whether this transition releases a running-query capacity slot. */ + boolean releasesRunningSlot() { + return taskToClose != null; + } + } + + /** Result of an authorized GET attempt. */ + sealed interface GetResult { + /** + * GET result for a live retained job. + * + * @param response current response context + */ + record Found(ResponseContext response) implements GetResult {} + + /** + * GET result when the lease expired before the request. + * + * @param removal cleanup required for the expired job + */ + record Expired(Removal removal) implements GetResult {} + } + + /** + * Resources and accounting changes produced when a job leaves the registry. + * + * @param responseStatus status returned to DELETE when the job has not expired + * @param task running task to cancel + * @param execution execution handle to close + * @param reason cancellation or removal reason + * @param expired whether expiration caused the removal + */ + record Removal( + PPLAsyncQueryService.Status responseStatus, + JobTask task, + AsyncQueryExecution execution, + String reason, + boolean expired) { + + /** Returns whether this removal releases a running-query capacity slot. */ + boolean releasesRunningSlot() { + return task != null; + } + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java new file mode 100644 index 00000000000..48c0fb19dea --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java @@ -0,0 +1,141 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Base64; +import java.util.UUID; + +/** + * Opaque, URL-safe identifier for an asynchronous PPL job. + * + *

The encoded value contains a format version, the owner node ID used for request routing, and a + * random per-job context ID. The complete encoded value is the key in the owner node's job + * registry. It contains no query text, user identity, or result data, and clients must treat it as + * opaque. + * + *

Before Base64 URL encoding, the binary layout is: + * + * 

+ * int formatVersion
+ * int ownerNodeIdLength + UTF-8 ownerNodeId
+ * int contextIdLength   + UTF-8 contextId
+ *
+ * + *

The ID provides routing, not authorization. GET and DELETE still authorize the current caller + * against the owner stored in the job. + * + * @param ownerNodeId node that owns the in-memory job + * @param contextId random identifier for one job on the owner node + */ +record PPLAsyncQueryJobId(String ownerNodeId, String contextId) { + private static final int FORMAT_VERSION = 1; + private static final int MAX_ENCODED_LENGTH = 2_048; + private static final int MAX_OWNER_NODE_ID_BYTES = 1_024; + private static final int MAX_CONTEXT_ID_BYTES = 128; + + PPLAsyncQueryJobId { + if (ownerNodeId == null || ownerNodeId.isBlank()) { + throw new IllegalArgumentException("PPL asynchronous query owner node must not be empty"); + } + if (contextId == null || contextId.isBlank()) { + throw new IllegalArgumentException("PPL asynchronous query context must not be empty"); + } + } + + /** + * Creates a new job ID for the given owner node. + * + * @param ownerNodeId node that will own the job + * @return unencoded job ID with a random context ID + */ + static PPLAsyncQueryJobId create(String ownerNodeId) { + return new PPLAsyncQueryJobId(ownerNodeId, UUID.randomUUID().toString()); + } + + /** + * Serializes this ID using the versioned binary format and URL-safe Base64 without padding. + * + * @return opaque value returned by the asynchronous PPL API + */ + String encode() { + try (ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + DataOutputStream output = new DataOutputStream(bytes)) { + output.writeInt(FORMAT_VERSION); + writeString(output, ownerNodeId, MAX_OWNER_NODE_ID_BYTES); + writeString(output, contextId, MAX_CONTEXT_ID_BYTES); + output.flush(); + return Base64.getUrlEncoder().withoutPadding().encodeToString(bytes.toByteArray()); + } catch (IOException e) { + throw new IllegalStateException("Failed to encode PPL asynchronous query ID", e); + } + } + + /** + * Decodes and validates an opaque job ID. + * + *

Parsing rejects empty or oversized input, unsupported versions, invalid component lengths, + * truncated data, and trailing bytes. All malformed input is reported uniformly to avoid exposing + * details of the internal encoding. + * + * @param encoded opaque value supplied by the client + * @return decoded owner node and context IDs + * @throws IllegalArgumentException when the value is not a valid job ID + */ + static PPLAsyncQueryJobId parse(String encoded) { + try { + if (encoded == null || encoded.isBlank() || encoded.length() > MAX_ENCODED_LENGTH) { + throw new IllegalArgumentException("Invalid PPL asynchronous query ID length"); + } + byte[] bytes = Base64.getUrlDecoder().decode(encoded); + try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(bytes))) { + int version = input.readInt(); + if (version != FORMAT_VERSION) { + throw new IllegalArgumentException( + "Unsupported PPL asynchronous query ID version [" + version + "]"); + } + PPLAsyncQueryJobId id = + new PPLAsyncQueryJobId( + readString(input, MAX_OWNER_NODE_ID_BYTES), + readString(input, MAX_CONTEXT_ID_BYTES)); + if (input.available() != 0) { + throw new IllegalArgumentException( + "Unexpected trailing bytes in PPL asynchronous query ID"); + } + return id; + } + } catch (Exception e) { + throw new IllegalArgumentException("Invalid PPL asynchronous query ID", e); + } + } + + private static void writeString(DataOutputStream output, String value, int maxLength) + throws IOException { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + if (bytes.length > maxLength) { + throw new IllegalArgumentException("PPL asynchronous query ID component is too long"); + } + output.writeInt(bytes.length); + output.write(bytes); + } + + private static String readString(DataInputStream input, int maxLength) throws IOException { + int length = input.readInt(); + if (length < 0 || length > maxLength) { + throw new IllegalArgumentException("Invalid PPL asynchronous query ID component length"); + } + byte[] bytes = input.readNBytes(length); + if (bytes.length != length) { + throw new IllegalArgumentException("Truncated PPL asynchronous query ID"); + } + return new String(bytes, StandardCharsets.UTF_8); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java new file mode 100644 index 00000000000..871ce62280d --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java @@ -0,0 +1,778 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import java.io.IOException; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.function.Function; +import java.util.function.IntSupplier; +import java.util.function.LongSupplier; +import java.util.function.Supplier; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.OpenSearchStatusException; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.lease.Releasable; +import org.opensearch.common.lifecycle.AbstractLifecycleComponent; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.executor.AsyncQueryExecution; +import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; +import org.opensearch.sql.executor.ExecutionEngine.Schema; +import org.opensearch.sql.plugin.PPLQueryErrorHandler; +import org.opensearch.sql.plugin.transport.PPLQueryAction; +import org.opensearch.sql.plugin.transport.PPLQueryTask; +import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.GetResult; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.JobTask; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Removal; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.ResponseContext; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Retention; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition; +import org.opensearch.sql.ppl.domain.PPLQueryRequest; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskManager; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.ThreadPool; + +/** + * Owns the asynchronous PPL jobs assigned to the local node. + * + *

{@link PPLAsyncQueryJob} owns the mutable state of one job. This service owns the job registry + * and performs the work requested by each state transition: capacity accounting, task management, + * result materialization, listener notification, and execution cleanup. Those side effects happen + * after the job lock is released. + * + *

POST races query completion against {@code wait_for_completion_timeout}. Completion wins by + * returning the final result directly and removing the job. The timeout wins by retaining the job + * and returning its opaque ID. GET and DELETE are then routed to this owner node. + * + *

This class is thread-safe. The registry is concurrent, capacity counters are guarded by {@code + * admissionLock}, and each job synchronizes its own lifecycle transitions. + */ +public final class PPLAsyncQueryService extends AbstractLifecycleComponent { + private static final Logger LOG = LogManager.getLogger(PPLAsyncQueryService.class); + + static final TimeValue DEFAULT_WAIT_FOR_COMPLETION = + TimeValue.parseTimeValue( + PPLQueryRequest.DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT, + PPLQueryRequest.WAIT_FOR_COMPLETION_TIMEOUT_FIELD); + static final TimeValue DEFAULT_KEEP_ALIVE = + TimeValue.parseTimeValue( + PPLQueryRequest.DEFAULT_KEEP_ALIVE, PPLQueryRequest.KEEP_ALIVE_FIELD); + private static final TimeValue REAPER_INTERVAL = TimeValue.timeValueMinutes(1); + private static final TimeoutHandle NO_TIMEOUT = () -> {}; + + /** Lifecycle state exposed in asynchronous PPL responses. */ + public enum Status { + /** Query execution is still running. */ + RUNNING, + + /** Query execution completed successfully. */ + SUCCEEDED, + + /** Query execution failed. */ + FAILED, + + /** Query execution was cancelled by DELETE. */ + CANCELLED + } + + @FunctionalInterface + interface TimeoutHandle { + void cancel(); + } + + @FunctionalInterface + interface TimeoutScheduler { + TimeoutHandle schedule(TimeValue delay, Runnable task); + } + + /** + * Immutable point-in-time response view of a job. + * + *

The formatter converts this internal model to the public JSON response. Query data is copied + * from the execution after releasing the job lock, so formatting never observes mutable job + * state. + * + * @param id opaque job ID, or {@code null} for a terminal response returned directly by POST + * @param status lifecycle state captured with the result + * @param response current query result, or {@code null} before a result is available + * @param failure client-visible failure for {@link Status#FAILED}, otherwise {@code null} + * @param tookMillis elapsed execution time, available for a completed job + */ + public record JobSnapshot( + String id, Status status, QueryResponse response, Failure failure, long tookMillis) {} + + /** Response model returned after DELETE removes a retained job. */ + record DeleteResult(String id, Status status) {} + + /** + * Client-visible failure retained by a job and returned only after owner authorization. + * + * @param type exception type + * @param reason client-facing failure reason + */ + public record Failure(String type, String reason) { + private static Failure from(Exception exception) { + String type = + exception.getClass().getSimpleName().isBlank() + ? exception.getClass().getName() + : exception.getClass().getSimpleName(); + String reason = + exception.getMessage() == null || exception.getMessage().isBlank() + ? "query execution failed" + : exception.getMessage(); + return new Failure(type, reason); + } + } + + private final Supplier ownerNodeIdSupplier; + private final LongSupplier currentTimeMillis; + private final TimeoutScheduler timeoutScheduler; + private final IntSupplier maxRunningQueries; + private final IntSupplier maxRetainedJobs; + private final Supplier maxWaitForCompletion; + private final Supplier maxKeepAlive; + private final ThreadPool threadPool; + private final ConcurrentMap jobs = new ConcurrentHashMap<>(); + private final Object admissionLock = new Object(); + + private int runningQueries; + private int retainedJobs; + private volatile boolean acceptingNewJobs = true; + private volatile Scheduler.Cancellable reaper; + private volatile TaskManager taskManager; + + /** + * Creates the owner-node lifecycle service. + * + * @param ownerNodeIdSupplier supplies the current local node ID + * @param threadPool schedules retention deadlines and expiration reaping + * @param settings supplies asynchronous query capacity and duration limits + */ + public PPLAsyncQueryService( + Supplier ownerNodeIdSupplier, ThreadPool threadPool, Settings settings) { + this( + ownerNodeIdSupplier, + System::currentTimeMillis, + (delay, task) -> { + Scheduler.ScheduledCancellable cancellable = + threadPool.schedule(task, delay, ThreadPool.Names.GENERIC); + return cancellable::cancel; + }, + () -> + (Integer) + settings.getSettingValue(Settings.Key.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES), + () -> (Integer) settings.getSettingValue(Settings.Key.PPL_ASYNC_MAX_RETAINED_JOBS), + () -> + (TimeValue) + settings.getSettingValue(Settings.Key.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT), + () -> (TimeValue) settings.getSettingValue(Settings.Key.PPL_ASYNC_MAX_KEEP_ALIVE), + threadPool); + } + + PPLAsyncQueryService( + String ownerNodeId, + LongSupplier currentTimeMillis, + TimeoutScheduler timeoutScheduler, + IntSupplier maxRunningQueries, + IntSupplier maxRetainedJobs, + Supplier maxWaitForCompletion, + Supplier maxKeepAlive) { + this( + () -> ownerNodeId, + currentTimeMillis, + timeoutScheduler, + maxRunningQueries, + maxRetainedJobs, + maxWaitForCompletion, + maxKeepAlive, + null); + } + + private PPLAsyncQueryService( + Supplier ownerNodeIdSupplier, + LongSupplier currentTimeMillis, + TimeoutScheduler timeoutScheduler, + IntSupplier maxRunningQueries, + IntSupplier maxRetainedJobs, + Supplier maxWaitForCompletion, + Supplier maxKeepAlive, + ThreadPool threadPool) { + this.ownerNodeIdSupplier = Objects.requireNonNull(ownerNodeIdSupplier); + this.currentTimeMillis = Objects.requireNonNull(currentTimeMillis); + this.timeoutScheduler = Objects.requireNonNull(timeoutScheduler); + this.maxRunningQueries = Objects.requireNonNull(maxRunningQueries); + this.maxRetainedJobs = Objects.requireNonNull(maxRetainedJobs); + this.maxWaitForCompletion = Objects.requireNonNull(maxWaitForCompletion); + this.maxKeepAlive = Objects.requireNonNull(maxKeepAlive); + this.threadPool = threadPool; + } + + /** + * Starts an asynchronous PPL query and produces its POST response. + * + *

The service registers and owns a job task, schedules its retention deadline, starts + * execution, and attaches the returned execution handle to the same job. If execution finishes + * before the deadline, {@code responseListener} receives the final result without a job ID. + * Otherwise, the job becomes retained and the listener receives its current result with an opaque + * ID. + * + * @param owner submit caller retained with the job for later authorization + * @param requestedKeepAlive requested job lease + * @param requestedWaitForCompletion maximum time to wait for a direct result + * @param request transport request used to register the job task + * @param requestTask task associated with the POST request + * @param executionStarter starts execution using the job-owned cancellable task + * @param responseListener receives either the direct result or retained job ID + */ + public void start( + PPLAsyncQueryUser owner, + String requestedKeepAlive, + String requestedWaitForCompletion, + TransportPPLQueryRequest request, + PPLQueryTask requestTask, + Function executionStarter, + ActionListener responseListener) { + TimeValue keepAlive = + TimeValue.parseTimeValue(requestedKeepAlive, PPLQueryRequest.KEEP_ALIVE_FIELD); + TimeValue waitForCompletion = + TimeValue.parseTimeValue( + requestedWaitForCompletion, PPLQueryRequest.WAIT_FOR_COMPLETION_TIMEOUT_FIELD); + validateKeepAlive(keepAlive); + validateWaitForCompletion(waitForCompletion); + JobTask jobTask = registerJobTask(request, requestTask); + start(owner, keepAlive, waitForCompletion, jobTask, executionStarter, responseListener); + } + + /** + * Starts a job using an already registered task. + * + *

This package-private entry point keeps task registration separate for tests while preserving + * the same production lifecycle: create the job, establish its retention deadline, start + * execution, and transfer ownership of the execution handle to the job. + * + * @param owner submit caller retained with the job for later authorization + * @param keepAlive validated job lease + * @param waitForCompletion validated direct-result wait + * @param task job-owned task and registration cleanup + * @param executionStarter starts execution using the job task + * @param responseListener receives the POST response + */ + void start( + PPLAsyncQueryUser owner, + TimeValue keepAlive, + TimeValue waitForCompletion, + JobTask task, + Function executionStarter, + ActionListener responseListener) { + PPLAsyncQueryJob job; + try { + Objects.requireNonNull(task); + Objects.requireNonNull(executionStarter); + Objects.requireNonNull(responseListener); + job = createJob(owner, keepAlive, task); + } catch (RuntimeException | Error e) { + closeTask(task); + throw e; + } + + TimeoutHandle retentionDeadline; + try { + retentionDeadline = scheduleRetention(job, waitForCompletion, responseListener); + } catch (RuntimeException | Error e) { + applyRemoval(job, job.abort()); + throw e; + } + startExecution(job, task.task(), executionStarter, responseListener, retentionDeadline); + } + + private PPLAsyncQueryJob createJob(PPLAsyncQueryUser owner, TimeValue keepAlive, JobTask task) { + Objects.requireNonNull(owner); + reserveCapacity(); + try { + String ownerNodeId = + Objects.requireNonNull(ownerNodeIdSupplier.get(), "Local node ID is not initialized"); + long now = currentTimeMillis.getAsLong(); + while (true) { + PPLAsyncQueryJobId jobId = PPLAsyncQueryJobId.create(ownerNodeId); + String encodedId = jobId.encode(); + PPLAsyncQueryJob job = + new PPLAsyncQueryJob(encodedId, owner, now, keepAlive.millis(), task); + if (jobs.putIfAbsent(encodedId, job) == null) { + return job; + } + } + } catch (RuntimeException | Error e) { + releaseCapacity(); + throw e; + } + } + + private TimeoutHandle scheduleRetention( + PPLAsyncQueryJob job, + TimeValue waitForCompletion, + ActionListener responseListener) { + if (waitForCompletion.millis() == 0) { + applyTransition(job, job.retain(currentTimeMillis.getAsLong()), responseListener); + return NO_TIMEOUT; + } + + return timeoutScheduler.schedule( + waitForCompletion, + () -> applyTransition(job, job.retain(currentTimeMillis.getAsLong()), responseListener)); + } + + private void startExecution( + PPLAsyncQueryJob job, + CancellableTask task, + Function executionStarter, + ActionListener responseListener, + TimeoutHandle retentionDeadline) { + try { + AsyncQueryExecution execution = Objects.requireNonNull(executionStarter.apply(task)); + attachExecution(job, execution, responseListener, retentionDeadline); + } catch (RuntimeException e) { + retentionDeadline.cancel(); + fail(job, e, responseListener); + } + } + + private void attachExecution( + PPLAsyncQueryJob job, + AsyncQueryExecution execution, + ActionListener responseListener, + TimeoutHandle retentionDeadline) { + if (!job.tryAttachExecution(execution)) { + retentionDeadline.cancel(); + closeExecution(execution); + return; + } + execution + .completion() + .whenComplete( + (ignored, failure) -> { + retentionDeadline.cancel(); + if (failure == null) { + complete(job, responseListener); + } else { + fail(job, asException(failure), responseListener); + } + }); + } + + private void complete(PPLAsyncQueryJob job, ActionListener responseListener) { + applyTransition(job, job.complete(currentTimeMillis.getAsLong()), responseListener); + } + + private void fail( + PPLAsyncQueryJob job, Exception failure, ActionListener responseListener) { + Objects.requireNonNull(failure); + Transition transition = job.fail(Failure.from(failure), currentTimeMillis.getAsLong()); + if (transition == null) { + return; + } + if (transition.response() == null) { + PPLQueryErrorHandler.recordFailure(failure); + applyTransition(job, transition, responseListener); + return; + } + applyTransition( + job, + transition, + ActionListener.wrap( + ignored -> responseListener.onFailure(failure), responseListener::onFailure)); + } + + /** + * Returns the current snapshot of a retained job. + * + *

The service authorizes the caller against the immutable job owner before requesting a + * synchronized lease transition. Result materialization happens afterward and therefore cannot + * block lifecycle transitions. + * + * @param id opaque job ID owned by this node + * @param caller caller to compare with the stored job owner + * @param requestedKeepAlive new lease duration, or {@code null} to leave the lease unchanged + * @return immutable current job snapshot + */ + JobSnapshot get(String id, PPLAsyncQueryUser caller, TimeValue requestedKeepAlive) { + if (requestedKeepAlive != null) { + validateKeepAlive(requestedKeepAlive); + } + PPLAsyncQueryJob job = findLocal(id); + job.owner().authorize(caller); + GetResult result = job.get(currentTimeMillis.getAsLong(), requestedKeepAlive); + if (result instanceof GetResult.Expired expired) { + applyRemoval(job, expired.removal()); + throw notFound(); + } + return materialize(((GetResult.Found) result).response()); + } + + /** + * Cancels and removes a retained job. + * + * @param id opaque job ID owned by this node + * @param caller caller to compare with the stored job owner + * @return the status observed when the job was removed + */ + DeleteResult delete(String id, PPLAsyncQueryUser caller) { + PPLAsyncQueryJob job = findLocal(id); + job.owner().authorize(caller); + Removal removal = job.delete(currentTimeMillis.getAsLong()); + applyRemoval(job, removal); + if (removal.expired()) { + throw notFound(); + } + return new DeleteResult(id, removal.responseStatus()); + } + + void reapExpired() { + long now = currentTimeMillis.getAsLong(); + jobs.forEach((id, job) -> applyRemoval(job, job.expire(now))); + } + + int runningQueryCount() { + synchronized (admissionLock) { + return runningQueries; + } + } + + int retainedJobCount() { + synchronized (admissionLock) { + return retainedJobs; + } + } + + /** + * Attaches the node task manager after transport actions have been initialized. + * + * @param taskManager task manager used to register and cancel retained query tasks + */ + public void attachTaskManager(TaskManager taskManager) { + this.taskManager = Objects.requireNonNull(taskManager); + } + + private void reserveCapacity() { + synchronized (admissionLock) { + if (!acceptingNewJobs) { + throw new OpenSearchStatusException( + "PPL asynchronous query service is stopping", RestStatus.SERVICE_UNAVAILABLE); + } + if (runningQueries >= maxRunningQueries.getAsInt() + || retainedJobs >= maxRetainedJobs.getAsInt()) { + throw new OpenSearchStatusException( + "PPL asynchronous query capacity is exhausted", RestStatus.TOO_MANY_REQUESTS); + } + runningQueries++; + retainedJobs++; + } + } + + private void releaseRunning() { + synchronized (admissionLock) { + if (runningQueries > 0) { + runningQueries--; + } + } + } + + private void releaseRetained() { + synchronized (admissionLock) { + if (retainedJobs > 0) { + retainedJobs--; + } + } + } + + private void releaseCapacity() { + synchronized (admissionLock) { + if (runningQueries > 0) { + runningQueries--; + } + if (retainedJobs > 0) { + retainedJobs--; + } + } + } + + /** + * Applies side effects selected by a {@link PPLAsyncQueryJob} transition. + * + *

The job decides its state change while holding the job lock, then returns a value describing + * the required side effects. Map mutation, capacity accounting, and listener callbacks happen + * here after the lock has been released. + */ + private void applyTransition( + PPLAsyncQueryJob job, Transition transition, ActionListener responseListener) { + if (transition == null) { + return; + } + if (transition.releasesRunningSlot()) { + releaseRunning(); + } + if (transition.retention() == Retention.REMOVE && jobs.remove(job.id(), job)) { + releaseRetained(); + } + + JobSnapshot snapshot = null; + RuntimeException materializationFailure = null; + try { + if (transition.response() != null) { + snapshot = materialize(transition.response()); + } + } catch (RuntimeException e) { + materializationFailure = e; + } finally { + closeExecution(transition.executionToClose()); + closeTask(transition.taskToClose()); + } + + if (transition.response() != null) { + if (materializationFailure == null) { + responseListener.onResponse(snapshot); + } else { + if (transition.retention() == Retention.RETAIN) { + applyRemoval(job, job.abort()); + } + responseListener.onFailure(materializationFailure); + } + } + } + + private void applyRemoval(PPLAsyncQueryJob job, Removal removal) { + if (removal == null) { + return; + } + if (jobs.remove(job.id(), job)) { + if (removal.releasesRunningSlot()) { + releaseRunning(); + } + releaseRetained(); + } + cancel(removal.task(), removal.reason()); + closeExecution(removal.execution()); + } + + private JobSnapshot materialize(ResponseContext context) { + QueryResponse response = null; + if (context.status() == Status.SUCCEEDED || context.status() == Status.RUNNING) { + response = + context.execution() == null + ? null + : context + .execution() + .currentResult() + .map(PPLAsyncQueryService::snapshotResponse) + .orElse(null); + } + if (context.status() == Status.SUCCEEDED && response == null) { + throw new IllegalStateException( + "Successful PPL asynchronous execution completed without a final result"); + } + return new JobSnapshot( + context.id(), context.status(), response, context.failure(), context.tookMillis()); + } + + private static void closeExecution(AsyncQueryExecution execution) { + if (execution == null) { + return; + } + try { + execution.close(); + } catch (RuntimeException e) { + LOG.warn( + "Failed to close PPL asynchronous query execution ({})", e.getClass().getSimpleName()); + } + } + + private static void closeTask(JobTask task) { + if (task != null) { + task.close(); + } + } + + private PPLAsyncQueryJob findLocal(String encodedId) { + PPLAsyncQueryJob job = jobs.get(encodedId); + if (job == null) { + throw notFound(); + } + return job; + } + + void validateKeepAlive(TimeValue keepAlive) { + TimeValue maximum = maxKeepAlive.get(); + if (keepAlive == null || keepAlive.millis() <= 0 || keepAlive.millis() > maximum.millis()) { + throw new IllegalArgumentException( + "[keep_alive] must be greater than 0 and no more than " + maximum); + } + } + + void validateWaitForCompletion(TimeValue waitForCompletion) { + TimeValue maximum = maxWaitForCompletion.get(); + if (waitForCompletion == null + || waitForCompletion.millis() < 0 + || waitForCompletion.millis() > maximum.millis()) { + throw new IllegalArgumentException( + "[wait_for_completion_timeout] must be between 0 and " + maximum); + } + } + + /** + * Copies an execution-owned response before exposing it through a job snapshot. + * + *

This prevents later execution updates from changing a response already handed to a caller. + */ + private static QueryResponse snapshotResponse(QueryResponse response) { + Schema schema = new Schema(List.copyOf(response.getSchema().getColumns())); + QueryResponse copy = + new QueryResponse(schema, List.copyOf(response.getResults()), response.getCursor()); + copy.setWarnings(List.copyOf(response.getWarnings())); + return copy; + } + + private void cancel(JobTask task, String reason) { + if (task == null) { + return; + } + CancellableTask cancellableTask = task.task(); + if (cancellableTask == null || cancellableTask.isCancelled()) { + task.close(); + return; + } + try { + TaskManager currentTaskManager = taskManager; + if (currentTaskManager == null) { + cancellableTask.cancel(reason); + task.close(); + } else { + currentTaskManager.cancelTaskAndDescendants( + cancellableTask, + reason, + false, + ActionListener.wrap( + ignored -> task.close(), + failure -> { + task.close(); + LOG.warn( + "Failed to cancel descendants of PPL asynchronous query task ({})", + failure.getClass().getSimpleName()); + })); + } + } catch (RuntimeException e) { + task.close(); + LOG.warn("Failed to cancel PPL asynchronous query task ({})", e.getClass().getSimpleName()); + } + } + + /** + * Registers the independently cancellable task owned by an asynchronous job. + * + *

The POST request task is assigned as parent during registration so cancellation can reach + * startup work. The returned {@link JobTask} owns both task unregistration and child-node + * registration cleanup. + * + * @param request request used by {@link TaskManager} to create the job task + * @param requestTask task associated with the POST request + * @return job-owned task and its registration cleanup + */ + private JobTask registerJobTask(TransportPPLQueryRequest request, PPLQueryTask requestTask) { + TaskManager currentTaskManager = + Objects.requireNonNull( + taskManager, "PPL asynchronous query task manager is not initialized"); + Objects.requireNonNull(requestTask, "PPL asynchronous query request task is not initialized"); + DiscoveryNode localNode = + Objects.requireNonNull(currentTaskManager.localNode(), "Local node is not initialized"); + + // This task is registered directly rather than through TransportAction.execute(), so reproduce + // the two pieces of OpenSearch child-task bookkeeping that TransportAction normally performs. + // The child-node registration lets parent cancellation send a ban to this node; parentTaskId + // lets that ban find and cancel the retained task. + Releasable childNodeRegistration = + currentTaskManager.registerChildNode(requestTask.getId(), localNode); + TaskId originalParent = request.getParentTask(); + boolean registered = false; + try { + request.setParentTask(localNode.getId(), requestTask.getId()); + Task task = currentTaskManager.register("transport", PPLQueryAction.NAME, request); + if (!(task instanceof PPLQueryTask pplQueryTask)) { + currentTaskManager.unregister(task); + throw new IllegalStateException("Failed to create PPL asynchronous query task"); + } + registered = true; + return new JobTask( + pplQueryTask, + () -> { + try { + currentTaskManager.unregister(pplQueryTask); + } finally { + childNodeRegistration.close(); + } + }); + } finally { + request.setParentTask(originalParent); + if (!registered) { + childNodeRegistration.close(); + } + } + } + + private static Exception asException(Throwable failure) { + Throwable cause = + failure instanceof java.util.concurrent.CompletionException && failure.getCause() != null + ? failure.getCause() + : failure; + return cause instanceof Exception exception ? exception : new RuntimeException(cause); + } + + private static ResourceNotFoundException notFound() { + return new ResourceNotFoundException("PPL asynchronous query not found"); + } + + /** Starts accepting submissions and schedules periodic retained-job expiration. */ + @Override + protected void doStart() { + acceptingNewJobs = true; + if (threadPool != null) { + reaper = + threadPool.scheduleWithFixedDelay( + this::reapExpired, REAPER_INTERVAL, ThreadPool.Names.GENERIC); + } + } + + /** Stops accepting submissions and cancels the periodic expiration task. */ + @Override + protected void doStop() { + acceptingNewJobs = false; + Scheduler.Cancellable scheduledReaper = reaper; + if (scheduledReaper != null) { + scheduledReaper.cancel(); + reaper = null; + } + } + + /** + * Removes all remaining jobs and releases their task and execution resources. + * + * @throws IOException if lifecycle shutdown fails + */ + @Override + protected void doClose() throws IOException { + acceptingNewJobs = false; + jobs.forEach( + (id, job) -> applyRemoval(job, job.close("PPL asynchronous query service is closing"))); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java new file mode 100644 index 00000000000..e060962a3ff --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java @@ -0,0 +1,86 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import java.util.List; +import java.util.Objects; +import org.opensearch.OpenSearchSecurityException; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.commons.ConfigConstants; +import org.opensearch.commons.authuser.User; +import org.opensearch.core.rest.RestStatus; + +/** + * Immutable owner identity used to authorize retained asynchronous query state. + * + * @param name authenticated principal, or {@code null} when no identity was supplied + * @param requestedTenant requested security tenant + * @param backendRoles backend roles captured when the job starts + */ +public record PPLAsyncQueryUser(String name, String requestedTenant, List backendRoles) { + + /** + * Creates an immutable asynchronous query identity. + * + * @param name authenticated principal, or {@code null} when no identity was supplied + * @param requestedTenant requested security tenant + * @param backendRoles backend roles captured when the job starts + * @throws IllegalArgumentException if {@code name} is blank + */ + public PPLAsyncQueryUser { + backendRoles = backendRoles == null ? List.of() : List.copyOf(backendRoles); + if (name != null && name.isBlank()) { + throw new IllegalArgumentException("PPL asynchronous query user must not be blank"); + } + } + + /** + * Captures the current caller from the OpenSearch thread context. + * + * @param threadContext current request thread context + * @return immutable caller identity + * @throws OpenSearchSecurityException if the security identity cannot be parsed + */ + public static PPLAsyncQueryUser current(ThreadContext threadContext) { + try { + Object serialized = + threadContext.getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT); + if (serialized == null) { + return new PPLAsyncQueryUser(null, null, List.of()); + } + User user = + serialized instanceof User currentUser + ? currentUser + : serialized instanceof String value ? User.parse(value) : null; + if (user == null) { + throw forbidden(); + } + return new PPLAsyncQueryUser( + user.getName(), user.getRequestedTenant(), user.getBackendRoles()); + } catch (RuntimeException e) { + throw forbidden(); + } + } + + /** + * Verifies that a caller may access asynchronous query state owned by this identity. + * + * @param caller identity of the caller requesting access + * @throws OpenSearchSecurityException if the caller does not match the owner identity + */ + void authorize(PPLAsyncQueryUser caller) { + if (!Objects.equals(name, caller.name) + || !Objects.equals(requestedTenant, caller.requestedTenant) + || !caller.backendRoles.containsAll(backendRoles)) { + throw forbidden(); + } + } + + private static OpenSearchSecurityException forbidden() { + return new OpenSearchSecurityException( + "Not authorized to access PPL asynchronous query", RestStatus.FORBIDDEN); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java new file mode 100644 index 00000000000..bca7f1fa50a --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java @@ -0,0 +1,28 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import org.junit.Test; + +public class PPLAsyncQueryJobIdTest { + + @Test + public void roundTripsOwnerAndContext() { + PPLAsyncQueryJobId id = PPLAsyncQueryJobId.create("node-a"); + + assertEquals(id, PPLAsyncQueryJobId.parse(id.encode())); + } + + @Test + public void rejectsMalformedIds() { + assertThrows(IllegalArgumentException.class, () -> PPLAsyncQueryJobId.parse("")); + assertThrows(IllegalArgumentException.class, () -> PPLAsyncQueryJobId.parse("not-an-id")); + assertThrows(IllegalArgumentException.class, () -> new PPLAsyncQueryJobId("", "context-id")); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java new file mode 100644 index 00000000000..7b2381f88ff --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java @@ -0,0 +1,917 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; +import org.mockito.InOrder; +import org.opensearch.OpenSearchSecurityException; +import org.opensearch.OpenSearchStatusException; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.lease.Releasable; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.sql.data.model.ExprValueUtils; +import org.opensearch.sql.data.type.ExprCoreType; +import org.opensearch.sql.executor.AsyncQueryExecution; +import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; +import org.opensearch.sql.executor.ExecutionEngine.Schema; +import org.opensearch.sql.executor.ExecutionEngine.Schema.Column; +import org.opensearch.sql.legacy.metrics.BasicCounter; +import org.opensearch.sql.legacy.metrics.MetricName; +import org.opensearch.sql.legacy.metrics.Metrics; +import org.opensearch.sql.legacy.metrics.NumericMetric; +import org.opensearch.sql.plugin.transport.PPLQueryAction; +import org.opensearch.sql.plugin.transport.PPLQueryTask; +import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.TaskManager; + +public class PPLAsyncQueryServiceTest { + private static final PPLAsyncQueryUser OWNER = new PPLAsyncQueryUser(null, null, List.of()); + + private final AtomicLong now = new AtomicLong(1_000); + private final AtomicReference timeoutTask = new AtomicReference<>(); + private final AtomicBoolean timeoutCancelled = new AtomicBoolean(); + private final PPLAsyncQueryService service = service(20, 100); + + @Test + public void fastSuccessReturnsDirectResultWithoutRetainingJob() { + AtomicReference result = new AtomicReference<>(); + AtomicInteger responses = new AtomicInteger(); + TrackingExecution execution = new TrackingExecution(response(2)); + startQuery( + service, + null, + TimeValue.timeValueSeconds(5), + execution, + listener( + snapshot -> { + result.set(snapshot); + responses.incrementAndGet(); + })); + now.addAndGet(25); + execution.complete(); + + assertEquals(1, responses.get()); + assertNull(result.get().id()); + assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, result.get().status()); + assertEquals(2, result.get().response().getResults().size()); + assertEquals(25, result.get().tookMillis()); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + assertEquals(1, execution.reads.get()); + assertEquals(1, execution.closes.get()); + assertTrue(timeoutCancelled.get()); + + timeoutTask.get().run(); + assertEquals(1, responses.get()); + } + + @Test + public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() { + AtomicReference retainedResponse = new AtomicReference<>(); + TrackingExecution execution = new TrackingExecution(null); + startQuery( + service, null, TimeValue.timeValueSeconds(5), execution, listener(retainedResponse::set)); + + timeoutTask.get().run(); + + String id = retainedResponse.get().id(); + assertEquals(PPLAsyncQueryService.Status.RUNNING, retainedResponse.get().status()); + assertNull(retainedResponse.get().response()); + assertEquals(1, service.runningQueryCount()); + assertEquals(1, service.retainedJobCount()); + + now.addAndGet(25); + execution.succeed(response(2)); + PPLAsyncQueryService.JobSnapshot completed = service.get(id, OWNER, null); + + assertEquals(id, completed.id()); + assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, completed.status()); + assertEquals(2, completed.response().getResults().size()); + assertEquals(0, service.runningQueryCount()); + assertEquals(1, service.retainedJobCount()); + assertEquals(0, execution.closes.get()); + } + + @Test + public void retentionWaitPreventsExpiryAndLeaseStartsWhenIdIsReturned() { + AtomicReference retainedResponse = new AtomicReference<>(); + startQuery( + service, + null, + TimeValue.timeValueSeconds(1), + TimeValue.timeValueSeconds(5), + new TrackingExecution(null), + listener(retainedResponse::set)); + + now.addAndGet(TimeValue.timeValueSeconds(2).millis()); + service.reapExpired(); + + assertNull(retainedResponse.get()); + assertEquals(1, service.runningQueryCount()); + assertEquals(1, service.retainedJobCount()); + + timeoutTask.get().run(); + assertEquals(PPLAsyncQueryService.Status.RUNNING, retainedResponse.get().status()); + + now.addAndGet(TimeValue.timeValueSeconds(1).millis() + 1); + service.reapExpired(); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void fastFailureReturnsDirectFailureWithoutId() { + AtomicReference failure = new AtomicReference<>(); + TrackingExecution execution = new TrackingExecution(response(1)); + startQuery( + service, + null, + TimeValue.timeValueSeconds(5), + execution, + ActionListener.wrap( + ignored -> { + throw new AssertionError("Expected direct query failure"); + }, + failure::set)); + + execution.fail(new IllegalStateException("boom")); + + assertEquals("boom", failure.get().getMessage()); + assertEquals(0, execution.reads.get()); + assertEquals(1, execution.closes.get()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void getWithoutKeepAliveDoesNotRenewLease() { + CancellableTask task = mock(CancellableTask.class); + when(task.isCancelled()).thenReturn(false); + TrackingExecution execution = new TrackingExecution(null); + String id = startRetainedQuery(service, task, execution); + + now.addAndGet(TimeValue.timeValueMinutes(4).millis()); + assertEquals(PPLAsyncQueryService.Status.RUNNING, service.get(id, OWNER, null).status()); + + now.addAndGet(TimeValue.timeValueMinutes(1).millis() + 1); + assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); + verify(task).cancel("PPL asynchronous query expired"); + assertEquals(1, execution.closes.get()); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void getWithKeepAliveRenewsLease() { + CancellableTask task = mock(CancellableTask.class); + when(task.isCancelled()).thenReturn(false); + TrackingExecution execution = new TrackingExecution(null); + String id = startRetainedQuery(service, task, execution); + + now.addAndGet(TimeValue.timeValueMinutes(4).millis()); + assertEquals( + PPLAsyncQueryService.Status.RUNNING, + service.get(id, OWNER, TimeValue.timeValueMinutes(5)).status()); + + now.addAndGet(TimeValue.timeValueMinutes(4).millis()); + assertEquals(PPLAsyncQueryService.Status.RUNNING, service.get(id, OWNER, null).status()); + + now.addAndGet(TimeValue.timeValueMinutes(1).millis() + 1); + assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); + verify(task).cancel("PPL asynchronous query expired"); + } + + @Test + public void missingLocalJobReturnsNotFound() { + String remoteId = PPLAsyncQueryJobId.create("node-b").encode(); + + assertThrows(ResourceNotFoundException.class, () -> service.get(remoteId, OWNER, null)); + assertThrows(ResourceNotFoundException.class, () -> service.delete(remoteId, OWNER)); + } + + @Test + public void deleteCancelsRunningJobAndReleasesState() { + CancellableTask task = mock(CancellableTask.class); + when(task.isCancelled()).thenReturn(false); + TrackingExecution execution = new TrackingExecution(null); + String id = startRetainedQuery(service, task, execution); + + PPLAsyncQueryService.DeleteResult result = service.delete(id, OWNER); + + assertEquals(id, result.id()); + assertEquals(PPLAsyncQueryService.Status.CANCELLED, result.status()); + verify(task).cancel("PPL asynchronous query cancelled by user"); + assertEquals(1, execution.closes.get()); + assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void deleteReturnsExistingTerminalStatus() { + CancellableTask task = mock(CancellableTask.class); + TrackingExecution execution = new TrackingExecution(response(1)); + String id = startRetainedQuery(service, task, execution); + execution.complete(); + + PPLAsyncQueryService.DeleteResult result = service.delete(id, OWNER); + + assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, result.status()); + verify(task, never()).cancel(org.mockito.ArgumentMatchers.anyString()); + assertEquals(1, execution.closes.get()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void cancellationUsesTaskManagerWhenAttached() { + TaskManager taskManager = mock(TaskManager.class); + CancellableTask task = mock(CancellableTask.class); + when(task.isCancelled()).thenReturn(false); + service.attachTaskManager(taskManager); + String id = startRetainedQuery(service, task, new TrackingExecution(null)); + + service.delete(id, OWNER); + + verify(taskManager) + .cancelTaskAndDescendants( + org.mockito.ArgumentMatchers.eq(task), + org.mockito.ArgumentMatchers.eq("PPL asynchronous query cancelled by user"), + org.mockito.ArgumentMatchers.eq(false), + org.mockito.ArgumentMatchers.any()); + } + + @Test + public void startRegistersTaskAndCompletionReleasesIt() { + TaskManager taskManager = mock(TaskManager.class); + PPLQueryTask task = mock(PPLQueryTask.class); + TransportPPLQueryRequest request = + new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl"); + RequestTaskRegistration registration = registerRequestTask(taskManager, request, task); + service.attachTaskManager(taskManager); + + TrackingExecution execution = new TrackingExecution(null); + service.start( + OWNER, + "5m", + "0s", + request, + registration.requestTask(), + ignored -> execution, + listener(snapshot -> assertEquals(PPLAsyncQueryService.Status.RUNNING, snapshot.status()))); + execution.succeed(response(1)); + + InOrder registrationOrder = inOrder(taskManager); + registrationOrder + .verify(taskManager) + .registerChildNode(registration.requestTask().getId(), registration.localNode()); + registrationOrder.verify(taskManager).register("transport", PPLQueryAction.NAME, request); + assertEquals(TaskId.EMPTY_TASK_ID, request.getParentTask()); + verify(taskManager).unregister(task); + verify(registration.childNodeRegistration()).close(); + assertEquals(0, service.runningQueryCount()); + assertEquals(1, service.retainedJobCount()); + } + + @Test + public void executionStartFailureCompletesJobAndReleasesTask() { + TaskManager taskManager = mock(TaskManager.class); + PPLQueryTask task = mock(PPLQueryTask.class); + TransportPPLQueryRequest request = + new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl"); + RequestTaskRegistration registration = registerRequestTask(taskManager, request, task); + service.attachTaskManager(taskManager); + AtomicReference failure = new AtomicReference<>(); + + service.start( + OWNER, + "5m", + "5s", + request, + registration.requestTask(), + ignored -> { + throw new IllegalStateException("execution did not start"); + }, + ActionListener.wrap( + ignored -> { + throw new AssertionError("Expected execution startup failure"); + }, + failure::set)); + + assertEquals("execution did not start", failure.get().getMessage()); + verify(taskManager, times(1)).unregister(task); + verify(registration.childNodeRegistration()).close(); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void deleteCancelsAndReleasesServiceOwnedTask() { + TaskManager taskManager = mock(TaskManager.class); + PPLQueryTask task = mock(PPLQueryTask.class); + when(task.isCancelled()).thenReturn(false); + TransportPPLQueryRequest request = + new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl"); + RequestTaskRegistration registration = registerRequestTask(taskManager, request, task); + doAnswer( + invocation -> { + ActionListener listener = invocation.getArgument(3); + listener.onResponse(null); + return null; + }) + .when(taskManager) + .cancelTaskAndDescendants( + org.mockito.ArgumentMatchers.eq(task), + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.eq(false), + org.mockito.ArgumentMatchers.any()); + service.attachTaskManager(taskManager); + AtomicReference id = new AtomicReference<>(); + + service.start( + OWNER, + "5m", + "0s", + request, + registration.requestTask(), + ignored -> new TrackingExecution(null), + listener(snapshot -> id.set(snapshot.id()))); + + service.delete(id.get(), OWNER); + + verify(taskManager) + .cancelTaskAndDescendants( + org.mockito.ArgumentMatchers.eq(task), + org.mockito.ArgumentMatchers.eq("PPL asynchronous query cancelled by user"), + org.mockito.ArgumentMatchers.eq(false), + org.mockito.ArgumentMatchers.any()); + verify(taskManager, times(1)).unregister(task); + verify(registration.childNodeRegistration()).close(); + } + + @Test + public void startupFailureAfterJobCreationReleasesServiceOwnedTask() { + PPLAsyncQueryService abortingService = + new PPLAsyncQueryService( + "node-a", + now::get, + (delay, task) -> { + throw new IllegalStateException("scheduler unavailable"); + }, + () -> 20, + () -> 100, + () -> TimeValue.timeValueSeconds(60), + () -> TimeValue.timeValueHours(24)); + TaskManager taskManager = mock(TaskManager.class); + PPLQueryTask task = mock(PPLQueryTask.class); + when(task.isCancelled()).thenReturn(true); + TransportPPLQueryRequest request = + new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl"); + RequestTaskRegistration registration = registerRequestTask(taskManager, request, task); + abortingService.attachTaskManager(taskManager); + + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> + abortingService.start( + OWNER, + "5m", + "5s", + request, + registration.requestTask(), + ignored -> new TrackingExecution(null), + listener(snapshot -> {}))); + + assertEquals("scheduler unavailable", failure.getMessage()); + verify(taskManager, times(1)).unregister(task); + verify(registration.childNodeRegistration()).close(); + assertEquals(0, abortingService.runningQueryCount()); + assertEquals(0, abortingService.retainedJobCount()); + } + + @Test + public void childTrackingFailurePreventsRetainedTaskRegistration() { + TaskManager taskManager = mock(TaskManager.class); + PPLQueryTask requestTask = mock(PPLQueryTask.class); + DiscoveryNode localNode = mock(DiscoveryNode.class); + TransportPPLQueryRequest request = + new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl"); + when(requestTask.getId()).thenReturn(42L); + when(taskManager.localNode()).thenReturn(localNode); + when(taskManager.registerChildNode(42L, localNode)) + .thenThrow(new IllegalStateException("channel closed")); + service.attachTaskManager(taskManager); + + assertThrows( + IllegalStateException.class, + () -> + service.start( + OWNER, + "5m", + "5s", + request, + requestTask, + ignored -> new TrackingExecution(null), + listener(snapshot -> {}))); + + verify(taskManager, never()).register("transport", PPLQueryAction.NAME, request); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void rejectsUnauthorizedCallerWithoutRenewingOrDeleting() { + PPLAsyncQueryUser securedOwner = new PPLAsyncQueryUser("alice", "tenant", List.of("role-a")); + PPLAsyncQueryUser otherUser = new PPLAsyncQueryUser("bob", "tenant", List.of("role-a")); + AtomicReference id = new AtomicReference<>(); + service.start( + securedOwner, + PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + TimeValue.ZERO, + jobTask(null), + ignored -> new TrackingExecution(null), + listener(snapshot -> id.set(snapshot.id()))); + + assertThrows(OpenSearchSecurityException.class, () -> service.get(id.get(), otherUser, null)); + assertThrows(OpenSearchSecurityException.class, () -> service.delete(id.get(), otherUser)); + assertEquals( + PPLAsyncQueryService.Status.RUNNING, service.get(id.get(), securedOwner, null).status()); + } + + @Test + public void enforcesRunningAndRetainedCapacity() { + PPLAsyncQueryService limited = service(1, 1); + limited.start( + OWNER, + PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + TimeValue.ZERO, + jobTask(null), + ignored -> new TrackingExecution(null), + listener(snapshot -> {})); + + OpenSearchStatusException exception = + assertThrows( + OpenSearchStatusException.class, + () -> + limited.start( + OWNER, + PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + TimeValue.ZERO, + jobTask(null), + ignored -> new TrackingExecution(null), + listener(snapshot -> {}))); + + assertEquals(429, exception.status().getStatus()); + } + + @Test + public void createFailureReleasesReservedCapacity() { + PPLAsyncQueryService missingOwnerNode = + new PPLAsyncQueryService( + (String) null, + now::get, + (delay, task) -> () -> {}, + () -> 1, + () -> 1, + () -> TimeValue.timeValueSeconds(60), + () -> TimeValue.timeValueHours(24)); + + assertThrows( + NullPointerException.class, + () -> + missingOwnerNode.start( + OWNER, + PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + TimeValue.ZERO, + jobTask(null), + ignored -> new TrackingExecution(null), + listener(snapshot -> {}))); + + assertEquals(0, missingOwnerNode.runningQueryCount()); + assertEquals(0, missingOwnerNode.retainedJobCount()); + } + + @Test + public void finalSnapshotDefensivelyCopiesRows() { + AtomicReference result = new AtomicReference<>(); + List rows = new ArrayList<>(); + rows.add(ExprValueUtils.stringValue("first")); + QueryResponse response = + new QueryResponse( + new Schema(List.of(new Column("state", null, ExprCoreType.STRING))), rows, null); + + TrackingExecution execution = new TrackingExecution(response); + startQuery(service, null, TimeValue.timeValueSeconds(5), execution, listener(result::set)); + execution.complete(); + rows.add(ExprValueUtils.stringValue("second")); + + assertEquals(1, result.get().response().getResults().size()); + } + + @Test + public void completedExecutionCanBeAttachedBeforeCompletionIsObserved() { + AtomicReference result = new AtomicReference<>(); + TrackingExecution execution = new TrackingExecution(null); + + execution.succeed(response(2)); + startQuery(service, null, TimeValue.timeValueSeconds(5), execution, listener(result::set)); + + assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, result.get().status()); + assertEquals(2, result.get().response().getResults().size()); + assertEquals(1, execution.closes.get()); + } + + @Test + public void runningGetMaterializesCurrentResultOutsideJob() { + TrackingExecution execution = new TrackingExecution(response(1)); + String id = startRetainedQuery(service, null, execution); + + PPLAsyncQueryService.JobSnapshot first = service.get(id, OWNER, null); + execution.setCurrent(response(3)); + PPLAsyncQueryService.JobSnapshot second = service.get(id, OWNER, null); + + assertEquals(PPLAsyncQueryService.Status.RUNNING, first.status()); + assertEquals(1, first.response().getResults().size()); + assertEquals(3, second.response().getResults().size()); + } + + @Test + public void failedRetainedJobReturnsNoProvisionalRowsAndClosesExecution() { + NumericMetric failures = + new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_SYS.getName(), new BasicCounter()); + Metrics.getInstance().registerMetric(failures); + try { + TrackingExecution execution = new TrackingExecution(response(1)); + String id = startRetainedQuery(service, null, execution); + + execution.fail(new IllegalStateException("boom")); + PPLAsyncQueryService.JobSnapshot failed = service.get(id, OWNER, null); + + assertEquals(PPLAsyncQueryService.Status.FAILED, failed.status()); + assertEquals("boom", failed.failure().reason()); + assertNull(failed.response()); + assertEquals(0, execution.reads.get()); + assertEquals(1, execution.closes.get()); + } finally { + Metrics.getInstance().unregisterMetric(failures.getName()); + } + } + + @Test + public void failedRetainedJobRecordsFailureMetric() { + NumericMetric failures = + new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_CUS.getName(), new BasicCounter()); + Metrics.getInstance().registerMetric(failures); + try { + TrackingExecution execution = new TrackingExecution(null); + startRetainedQuery(service, null, execution); + + execution.fail(new IllegalArgumentException("invalid query")); + + assertEquals(Long.valueOf(1), failures.getValue()); + } finally { + Metrics.getInstance().unregisterMetric(failures.getName()); + } + } + + @Test + public void deleteBeforeExecutionAttachmentClosesLateHandle() { + TrackingExecution execution = new TrackingExecution(response(1)); + AtomicReference id = new AtomicReference<>(); + + service.start( + OWNER, + PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + TimeValue.ZERO, + jobTask(null), + ignored -> { + service.delete(id.get(), OWNER); + return execution; + }, + listener(snapshot -> id.set(snapshot.id()))); + + assertEquals(1, execution.closes.get()); + assertThrows(ResourceNotFoundException.class, () -> service.get(id.get(), OWNER, null)); + } + + @Test + public void concurrentGetDoesNotBlockDeleteOnResultMaterialization() throws Exception { + BlockingExecution execution = new BlockingExecution(response(1)); + String id = startRetainedQuery(service, null, execution); + + CompletableFuture get = + CompletableFuture.supplyAsync(() -> service.get(id, OWNER, null)); + assertTrue(execution.readStarted.await(5, TimeUnit.SECONDS)); + CompletableFuture delete = + CompletableFuture.supplyAsync(() -> service.delete(id, OWNER)); + + try { + assertEquals(PPLAsyncQueryService.Status.CANCELLED, delete.get(5, TimeUnit.SECONDS).status()); + } finally { + execution.allowRead.countDown(); + } + + assertEquals(PPLAsyncQueryService.Status.RUNNING, get.get(5, TimeUnit.SECONDS).status()); + assertEquals(1, execution.closes.get()); + assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); + } + + @Test + public void shutdownClosesRetainedExecutionExactlyOnce() throws Exception { + TrackingExecution execution = new TrackingExecution(response(1)); + startRetainedQuery(service, null, execution); + + service.close(); + service.close(); + + assertEquals(1, execution.closes.get()); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void successfulCompletionRequiresFinalResultToBeVisible() { + AtomicReference failure = new AtomicReference<>(); + TrackingExecution execution = new TrackingExecution(null); + startQuery( + service, + null, + TimeValue.timeValueSeconds(5), + execution, + ActionListener.wrap(snapshot -> {}, failure::set)); + + execution.complete(); + + assertTrue(failure.get() instanceof IllegalStateException); + assertEquals(1, execution.closes.get()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void retentionResponseMaterializationFailureAbortsUndeliverableJob() { + CancellableTask task = mock(CancellableTask.class); + when(task.isCancelled()).thenReturn(false); + AtomicReference failure = new AtomicReference<>(); + ThrowingExecution execution = new ThrowingExecution(); + startQuery( + service, + task, + TimeValue.timeValueSeconds(5), + execution, + ActionListener.wrap(snapshot -> {}, failure::set)); + + timeoutTask.get().run(); + + assertTrue(failure.get() instanceof IllegalStateException); + verify(task).cancel("PPL asynchronous query startup failed"); + assertEquals(1, execution.closes.get()); + assertEquals(0, service.runningQueryCount()); + assertEquals(0, service.retainedJobCount()); + } + + @Test + public void validatesDurationBounds() { + assertThrows(IllegalArgumentException.class, () -> service.validateKeepAlive(TimeValue.ZERO)); + assertThrows( + IllegalArgumentException.class, + () -> service.validateKeepAlive(TimeValue.timeValueHours(25))); + assertThrows( + IllegalArgumentException.class, + () -> service.validateWaitForCompletion(TimeValue.timeValueSeconds(61))); + + service.validateWaitForCompletion(TimeValue.ZERO); + service.validateWaitForCompletion(TimeValue.timeValueSeconds(60)); + service.validateKeepAlive(TimeValue.timeValueHours(24)); + } + + private PPLAsyncQueryService service(int maxRunning, int maxRetained) { + return new PPLAsyncQueryService( + "node-a", + now::get, + (delay, task) -> { + timeoutTask.set(task); + timeoutCancelled.set(false); + return () -> timeoutCancelled.set(true); + }, + () -> maxRunning, + () -> maxRetained, + () -> TimeValue.timeValueSeconds(60), + () -> TimeValue.timeValueHours(24)); + } + + private String startRetainedQuery( + PPLAsyncQueryService targetService, CancellableTask task, AsyncQueryExecution execution) { + AtomicReference id = new AtomicReference<>(); + startQuery( + targetService, + task, + TimeValue.ZERO, + execution, + listener(snapshot -> id.set(snapshot.id()))); + return id.get(); + } + + private void startQuery( + PPLAsyncQueryService targetService, + CancellableTask task, + TimeValue waitForCompletion, + AsyncQueryExecution execution, + ActionListener responseListener) { + startQuery( + targetService, + task, + PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + waitForCompletion, + execution, + responseListener); + } + + private void startQuery( + PPLAsyncQueryService targetService, + CancellableTask task, + TimeValue keepAlive, + TimeValue waitForCompletion, + AsyncQueryExecution execution, + ActionListener responseListener) { + targetService.start( + OWNER, keepAlive, waitForCompletion, jobTask(task), ignored -> execution, responseListener); + } + + private static PPLAsyncQueryJob.JobTask jobTask(CancellableTask task) { + return new PPLAsyncQueryJob.JobTask(task, () -> {}); + } + + private static RequestTaskRegistration registerRequestTask( + TaskManager taskManager, TransportPPLQueryRequest request, PPLQueryTask retainedTask) { + PPLQueryTask requestTask = mock(PPLQueryTask.class); + DiscoveryNode localNode = mock(DiscoveryNode.class); + Releasable childNodeRegistration = mock(Releasable.class); + when(requestTask.getId()).thenReturn(42L); + when(localNode.getId()).thenReturn("node-a"); + when(taskManager.localNode()).thenReturn(localNode); + when(taskManager.registerChildNode(42L, localNode)).thenReturn(childNodeRegistration); + when(taskManager.register("transport", PPLQueryAction.NAME, request)) + .thenAnswer( + invocation -> { + assertEquals(new TaskId("node-a", 42L), request.getParentTask()); + return retainedTask; + }); + return new RequestTaskRegistration(requestTask, localNode, childNodeRegistration); + } + + private record RequestTaskRegistration( + PPLQueryTask requestTask, DiscoveryNode localNode, Releasable childNodeRegistration) {} + + private static ActionListener listener( + java.util.function.Consumer consumer) { + return ActionListener.wrap( + snapshot -> consumer.accept(snapshot), + failure -> { + throw new AssertionError(failure); + }); + } + + private static QueryResponse response(int rowCount) { + Schema schema = new Schema(List.of(new Column("state", null, ExprCoreType.STRING))); + return new QueryResponse( + schema, + java.util.stream.IntStream.range(0, rowCount) + .mapToObj(i -> ExprValueUtils.stringValue("state-" + i)) + .toList(), + null); + } + + private static final class TrackingExecution implements AsyncQueryExecution { + private final AtomicReference current; + private final CompletableFuture completion = new CompletableFuture<>(); + private final AtomicBoolean closed = new AtomicBoolean(); + private final AtomicInteger reads = new AtomicInteger(); + private final AtomicInteger closes = new AtomicInteger(); + + private TrackingExecution(QueryResponse current) { + this.current = new AtomicReference<>(current); + } + + private void setCurrent(QueryResponse response) { + current.set(response); + } + + private void succeed(QueryResponse response) { + current.set(response); + completion.complete(null); + } + + private void complete() { + completion.complete(null); + } + + private void fail(Exception failure) { + completion.completeExceptionally(failure); + } + + @Override + public Optional currentResult() { + reads.incrementAndGet(); + return Optional.ofNullable(current.get()); + } + + @Override + public CompletionStage completion() { + return completion; + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + closes.incrementAndGet(); + } + } + } + + private static final class BlockingExecution implements AsyncQueryExecution { + private final QueryResponse response; + private final CountDownLatch readStarted = new CountDownLatch(1); + private final CountDownLatch allowRead = new CountDownLatch(1); + private final AtomicBoolean closed = new AtomicBoolean(); + private final AtomicInteger closes = new AtomicInteger(); + + private BlockingExecution(QueryResponse response) { + this.response = response; + } + + @Override + public Optional currentResult() { + readStarted.countDown(); + try { + assertTrue(allowRead.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + return Optional.of(response); + } + + @Override + public CompletionStage completion() { + return new CompletableFuture<>(); + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + closes.incrementAndGet(); + } + } + } + + private static final class ThrowingExecution implements AsyncQueryExecution { + private final AtomicBoolean closed = new AtomicBoolean(); + private final AtomicInteger closes = new AtomicInteger(); + + @Override + public Optional currentResult() { + throw new IllegalStateException("materialization failed"); + } + + @Override + public CompletionStage completion() { + return new CompletableFuture<>(); + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + closes.incrementAndGet(); + } + } + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java new file mode 100644 index 00000000000..6c9a0061849 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java @@ -0,0 +1,85 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; + +import java.util.List; +import org.junit.Test; +import org.opensearch.OpenSearchSecurityException; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.commons.ConfigConstants; +import org.opensearch.commons.authuser.User; + +public class PPLAsyncQueryUserTest { + + @Test + public void capturesSecurityIdentity() { + ThreadContext context = new ThreadContext(Settings.EMPTY); + context.putTransient( + ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, + "alice|backend-a|ppl-role|tenant-a"); + + PPLAsyncQueryUser identity = PPLAsyncQueryUser.current(context); + + assertEquals("alice", identity.name()); + assertEquals("tenant-a", identity.requestedTenant()); + assertEquals(List.of("backend-a"), identity.backendRoles()); + } + + @Test + public void acceptsUserObjectAndRejectsUnknownSecurityContext() { + ThreadContext objectContext = new ThreadContext(Settings.EMPTY); + objectContext.putTransient( + ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, + new User("alice", List.of("backend-a"), List.of("ppl-role"), null, "tenant-a")); + assertEquals("alice", PPLAsyncQueryUser.current(objectContext).name()); + + ThreadContext invalidContext = new ThreadContext(Settings.EMPTY); + invalidContext.putTransient( + ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, new Object()); + assertThrows( + OpenSearchSecurityException.class, () -> PPLAsyncQueryUser.current(invalidContext)); + } + + @Test + public void missingIdentityRepresentsAnUnsecuredCaller() { + ThreadContext context = new ThreadContext(Settings.EMPTY); + + assertNull(PPLAsyncQueryUser.current(context).name()); + } + + @Test + public void requiresSamePrincipalTenantAndOriginalBackendRoles() { + PPLAsyncQueryUser owner = new PPLAsyncQueryUser("alice", "tenant-a", List.of("role-a")); + + owner.authorize( + new PPLAsyncQueryUser("alice", "tenant-a", List.of("role-a", "newly-added-role"))); + + assertThrows( + OpenSearchSecurityException.class, + () -> owner.authorize(new PPLAsyncQueryUser("bob", "tenant-a", List.of("role-a")))); + assertThrows( + OpenSearchSecurityException.class, + () -> owner.authorize(new PPLAsyncQueryUser("alice", "tenant-b", List.of("role-a")))); + assertThrows( + OpenSearchSecurityException.class, + () -> owner.authorize(new PPLAsyncQueryUser("alice", "tenant-a", List.of()))); + } + + @Test + public void unsecuredModeRequiresAnUnsecuredCaller() { + PPLAsyncQueryUser unsecured = new PPLAsyncQueryUser(null, null, List.of()); + unsecured.authorize(new PPLAsyncQueryUser(null, null, List.of())); + + assertThrows( + OpenSearchSecurityException.class, + () -> unsecured.authorize(new PPLAsyncQueryUser("alice", null, List.of("role-a")))); + } +} diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecution.java b/ppl/src/main/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecution.java new file mode 100644 index 00000000000..d0c1e5b3921 --- /dev/null +++ b/ppl/src/main/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecution.java @@ -0,0 +1,51 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.AsyncQueryExecution; +import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; + +/** + * Final-result implementation of the asynchronous execution contract. + * + *

This adapter keeps the existing callback-based query execution unchanged. Until partial-result + * producers are added, {@link #currentResult()} is empty while execution is running and exposes the + * final response immediately before successful completion is published. + */ +final class DefaultAsyncQueryExecution + implements AsyncQueryExecution, ResponseListener { + private final CompletableFuture completion = new CompletableFuture<>(); + private volatile QueryResponse finalResult; + + @Override + public void onResponse(QueryResponse response) { + finalResult = Objects.requireNonNull(response); + completion.complete(null); + } + + @Override + public void onFailure(Exception failure) { + completion.completeExceptionally(Objects.requireNonNull(failure)); + } + + @Override + public Optional currentResult() { + return Optional.ofNullable(finalResult); + } + + @Override + public CompletionStage completion() { + return completion; + } + + @Override + public void close() {} +} diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index e2572b5f0da..b8ac519fbbe 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -18,6 +18,7 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; import org.opensearch.sql.executor.AnalyzeResponse; +import org.opensearch.sql.executor.AsyncQueryExecution; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import org.opensearch.sql.executor.QueryManager; import org.opensearch.sql.executor.QueryType; @@ -88,6 +89,40 @@ public void execute( } } + /** + * Starts query execution and immediately returns its lifecycle-facing handle. + * + *

The existing callback execution remains internal to the PPL execution module. This + * final-only implementation exposes no current result until the callback publishes the + * authoritative response. + * + * @param request PPL query request + * @param anonymizedQuerySink receives anonymized query text for metrics + * @return lifecycle-facing asynchronous execution handle + */ + public AsyncQueryExecution executeAsync( + PPLQueryRequest request, Consumer anonymizedQuerySink) { + DefaultAsyncQueryExecution execution = new DefaultAsyncQueryExecution(); + execute( + request, + execution, + new ResponseListener<>() { + @Override + public void onResponse(ExplainResponse response) { + execution.onFailure( + new IllegalStateException( + "Asynchronous query execution received an explain response")); + } + + @Override + public void onFailure(Exception e) { + execution.onFailure(e); + } + }, + anonymizedQuerySink); + return execution; + } + /** * Explain the query in {@link PPLQueryRequest} using {@link ResponseListener} to get and format * explain response. diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java index 67fff2db5e5..0849b086b36 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java @@ -37,6 +37,19 @@ public class PPLQueryRequest { private static final String START_TIME_FIELD = "start_time"; private static final String END_TIME_FIELD = "end_time"; private static final String TIME_FIELD_FIELD = "time_field"; + + /** JSON field selecting the asynchronous job lease. */ + public static final String KEEP_ALIVE_FIELD = "keep_alive"; + + /** JSON field selecting how long submit waits for direct completion. */ + public static final String WAIT_FOR_COMPLETION_TIMEOUT_FIELD = "wait_for_completion_timeout"; + + /** Default asynchronous job lease. */ + public static final String DEFAULT_KEEP_ALIVE = "5m"; + + /** Default submit wait-for-completion timeout. */ + public static final String DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT = "5s"; + private static final int MAX_HIGHLIGHT_FIELDS = 100; private static final int MAX_TAG_ENTRIES = 10; @@ -48,6 +61,11 @@ public class PPLQueryRequest { @Getter private String format = ""; @Getter private String explainMode; + @Setter + @Getter + @Accessors(fluent = true) + private boolean formatExplicitlySpecified = false; + @Setter @Getter @Accessors(fluent = true) @@ -141,6 +159,62 @@ public boolean isExplainRequest() { return path.endsWith("/_explain"); } + /** + * Returns whether the request contains an asynchronous lifecycle field. + * + * @return {@code true} when keep-alive or wait-for-completion was explicitly requested + */ + public boolean isAsyncQueryRequest() { + return jsonContent != null + && (jsonContent.has(WAIT_FOR_COMPLETION_TIMEOUT_FIELD) + || jsonContent.has(KEEP_ALIVE_FIELD)); + } + + /** + * Returns the requested asynchronous job lease. + * + * @return requested lease or {@link #DEFAULT_KEEP_ALIVE} + */ + public String getKeepAlive() { + if (jsonContent == null || !jsonContent.has(KEEP_ALIVE_FIELD)) { + return DEFAULT_KEEP_ALIVE; + } + return stringLifecycleField(KEEP_ALIVE_FIELD); + } + + /** + * Returns how long asynchronous submit waits for direct completion. + * + * @return requested timeout or {@link #DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT} + */ + public String getWaitForCompletionTimeout() { + if (jsonContent == null || !jsonContent.has(WAIT_FOR_COMPLETION_TIMEOUT_FIELD)) { + return DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT; + } + return stringLifecycleField(WAIT_FOR_COMPLETION_TIMEOUT_FIELD); + } + + private String stringLifecycleField(String field) { + Object value = jsonContent.get(field); + if (!(value instanceof String text)) { + throw new IllegalArgumentException("[" + field + "] must be a string"); + } + return text; + } + + /** + * Returns whether this request mode and response format support asynchronous execution. + * + * @return {@code true} for a normal query using the default JSON response format + */ + public boolean supportsAsyncExecution() { + return !isExplainRequest() + && !profile + && !analyze + && !pplQuery.trim().toLowerCase(Locale.ROOT).startsWith("explain") + && !formatExplicitlySpecified; + } + /** Decide on the formatter by the requested format. */ public Format format() { Optional optionalFormat = Format.of(format); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecutionTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecutionTest.java new file mode 100644 index 00000000000..f4a93f7fc05 --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/DefaultAsyncQueryExecutionTest.java @@ -0,0 +1,62 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.Test; +import org.opensearch.sql.data.model.ExprValueUtils; +import org.opensearch.sql.data.type.ExprCoreType; +import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; +import org.opensearch.sql.executor.ExecutionEngine.Schema; +import org.opensearch.sql.executor.ExecutionEngine.Schema.Column; + +public class DefaultAsyncQueryExecutionTest { + + @Test + public void finalResultIsVisibleBeforeSuccessfulCompletionNotification() { + DefaultAsyncQueryExecution execution = new DefaultAsyncQueryExecution(); + AtomicBoolean visibleFromCompletion = new AtomicBoolean(); + execution + .completion() + .whenComplete( + (ignored, failure) -> + visibleFromCompletion.set( + failure == null + && execution.currentResult().orElseThrow().getResults().size() == 1)); + + execution.onResponse(response("final")); + + assertTrue(execution.completion().toCompletableFuture().isDone()); + assertFalse(execution.completion().toCompletableFuture().isCompletedExceptionally()); + assertTrue(visibleFromCompletion.get()); + } + + @Test + public void failureCompletesExceptionallyWithoutPublishingRows() { + DefaultAsyncQueryExecution execution = new DefaultAsyncQueryExecution(); + + execution.onFailure(new IllegalStateException("boom")); + + CompletionException failure = + assertThrows( + CompletionException.class, () -> execution.completion().toCompletableFuture().join()); + assertTrue(failure.getCause() instanceof IllegalStateException); + assertTrue(execution.currentResult().isEmpty()); + } + + private static QueryResponse response(String value) { + return new QueryResponse( + new Schema(List.of(new Column("state", null, ExprCoreType.STRING))), + List.of(ExprValueUtils.stringValue(value)), + null); + } +} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index fea5e3d4030..8a83dc6923b 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -20,6 +20,7 @@ import org.mockito.junit.MockitoJUnitRunner; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.executor.AsyncQueryExecution; import org.opensearch.sql.executor.DefaultQueryManager; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; @@ -115,6 +116,27 @@ public void testExecuteShouldPass() { getExplainListener(false)); } + @Test + public void testExecuteAsyncReturnsCompletedFinalResult() { + QueryResponse response = new QueryResponse(schema, Collections.emptyList(), Cursor.None); + doAnswer( + invocation -> { + ResponseListener listener = invocation.getArgument(4); + listener.onResponse(response); + return null; + }) + .when(queryService) + .execute(any(), any(), any(), anyBoolean(), any()); + + AsyncQueryExecution execution = + pplService.executeAsync( + new PPLQueryRequest("search source=t a=1", null, QUERY), + PPLService.NO_ANONYMIZED_QUERY_SINK); + + execution.completion().toCompletableFuture().join(); + Assert.assertSame(response, execution.currentResult().orElseThrow()); + } + @Test public void testExecuteCsvFormatShouldPass() { doAnswer( diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java index d4c45e2ece4..25efa4be1fc 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java @@ -6,6 +6,7 @@ package org.opensearch.sql.ppl.domain; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; @@ -63,6 +64,69 @@ public void testUnsupportedFormat() { request.format(); } + @Test + public void asyncFieldsUseDefaultsOrRequestedValues() { + PPLQueryRequest defaults = + new PPLQueryRequest( + "source=t", + new JSONObject().put("wait_for_completion_timeout", "1s"), + "/_plugins/_ppl"); + assertTrue(defaults.isAsyncQueryRequest()); + assertEquals("5m", defaults.getKeepAlive()); + assertEquals("1s", defaults.getWaitForCompletionTimeout()); + + PPLQueryRequest requested = + new PPLQueryRequest("source=t", new JSONObject().put("keep_alive", "2m"), "/_plugins/_ppl"); + assertTrue(requested.isAsyncQueryRequest()); + assertEquals("2m", requested.getKeepAlive()); + assertEquals("5s", requested.getWaitForCompletionTimeout()); + } + + @Test + public void asyncLifecycleFieldsMustBeStrings() { + PPLQueryRequest numericKeepAlive = + new PPLQueryRequest("source=t", new JSONObject().put("keep_alive", 300), "/_plugins/_ppl"); + PPLQueryRequest numericWait = + new PPLQueryRequest( + "source=t", new JSONObject().put("wait_for_completion_timeout", 1), "/_plugins/_ppl"); + + IllegalArgumentException keepAliveFailure = + assertThrows(IllegalArgumentException.class, numericKeepAlive::getKeepAlive); + assertEquals("[keep_alive] must be a string", keepAliveFailure.getMessage()); + IllegalArgumentException waitFailure = + assertThrows(IllegalArgumentException.class, numericWait::getWaitForCompletionTimeout); + assertEquals("[wait_for_completion_timeout] must be a string", waitFailure.getMessage()); + } + + @Test + public void defaultFormatQuerySupportsAsyncExecution() { + assertTrue( + new PPLQueryRequest("source=t", null, "/_plugins/_ppl", "jdbc").supportsAsyncExecution()); + } + + @Test + public void specialModesAndFormatsDoNotSupportAsyncExecution() { + assertFalse( + new PPLQueryRequest("source=t", null, "/_plugins/_ppl/_explain").supportsAsyncExecution()); + assertFalse( + new PPLQueryRequest("explain source=t", null, "/_plugins/_ppl").supportsAsyncExecution()); + PPLQueryRequest explicitCsv = new PPLQueryRequest("source=t", null, "/_plugins/_ppl", "csv"); + explicitCsv.formatExplicitlySpecified(true); + assertFalse(explicitCsv.supportsAsyncExecution()); + + PPLQueryRequest explicitJdbc = new PPLQueryRequest("source=t", null, "/_plugins/_ppl", "jdbc"); + explicitJdbc.formatExplicitlySpecified(true); + assertFalse(explicitJdbc.supportsAsyncExecution()); + + PPLQueryRequest profile = new PPLQueryRequest("source=t", null, "/_plugins/_ppl"); + profile.profile(true); + assertFalse(profile.supportsAsyncExecution()); + + PPLQueryRequest analyze = new PPLQueryRequest("source=t", null, "/_plugins/_ppl"); + analyze.analyze(true); + assertFalse(analyze.supportsAsyncExecution()); + } + @Test public void testGetFetchSizeReturnsValueFromJson() { JSONObject json = new JSONObject("{\"query\": \"source=t\", \"fetch_size\": 100}"); From 419de0133f64aba08f391f4f9d4e0142f7f55107 Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 20:08:35 +0000 Subject: [PATCH 02/15] Keep async request parsing in API layer Signed-off-by: Peng Huo --- .../asyncquery/PPLAsyncQueryService.java | 21 +----- .../asyncquery/PPLAsyncQueryServiceTest.java | 39 +++++----- .../sql/ppl/domain/PPLQueryRequest.java | 74 ------------------- .../sql/ppl/domain/PPLQueryRequestTest.java | 64 ---------------- 4 files changed, 21 insertions(+), 177 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java index 871ce62280d..381f42eaf5d 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java @@ -39,7 +39,6 @@ import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.ResponseContext; import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Retention; import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition; -import org.opensearch.sql.ppl.domain.PPLQueryRequest; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; @@ -64,13 +63,6 @@ public final class PPLAsyncQueryService extends AbstractLifecycleComponent { private static final Logger LOG = LogManager.getLogger(PPLAsyncQueryService.class); - static final TimeValue DEFAULT_WAIT_FOR_COMPLETION = - TimeValue.parseTimeValue( - PPLQueryRequest.DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT, - PPLQueryRequest.WAIT_FOR_COMPLETION_TIMEOUT_FIELD); - static final TimeValue DEFAULT_KEEP_ALIVE = - TimeValue.parseTimeValue( - PPLQueryRequest.DEFAULT_KEEP_ALIVE, PPLQueryRequest.KEEP_ALIVE_FIELD); private static final TimeValue REAPER_INTERVAL = TimeValue.timeValueMinutes(1); private static final TimeoutHandle NO_TIMEOUT = () -> {}; @@ -231,8 +223,8 @@ private PPLAsyncQueryService( * ID. * * @param owner submit caller retained with the job for later authorization - * @param requestedKeepAlive requested job lease - * @param requestedWaitForCompletion maximum time to wait for a direct result + * @param keepAlive requested job lease + * @param waitForCompletion maximum time to wait for a direct result * @param request transport request used to register the job task * @param requestTask task associated with the POST request * @param executionStarter starts execution using the job-owned cancellable task @@ -240,17 +232,12 @@ private PPLAsyncQueryService( */ public void start( PPLAsyncQueryUser owner, - String requestedKeepAlive, - String requestedWaitForCompletion, + TimeValue keepAlive, + TimeValue waitForCompletion, TransportPPLQueryRequest request, PPLQueryTask requestTask, Function executionStarter, ActionListener responseListener) { - TimeValue keepAlive = - TimeValue.parseTimeValue(requestedKeepAlive, PPLQueryRequest.KEEP_ALIVE_FIELD); - TimeValue waitForCompletion = - TimeValue.parseTimeValue( - requestedWaitForCompletion, PPLQueryRequest.WAIT_FOR_COMPLETION_TIMEOUT_FIELD); validateKeepAlive(keepAlive); validateWaitForCompletion(waitForCompletion); JobTask jobTask = registerJobTask(request, requestTask); diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java index 7b2381f88ff..74a399a58e7 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java @@ -56,6 +56,7 @@ public class PPLAsyncQueryServiceTest { private static final PPLAsyncQueryUser OWNER = new PPLAsyncQueryUser(null, null, List.of()); + private static final TimeValue KEEP_ALIVE = TimeValue.timeValueMinutes(5); private final AtomicLong now = new AtomicLong(1_000); private final AtomicReference timeoutTask = new AtomicReference<>(); @@ -281,8 +282,8 @@ public void startRegistersTaskAndCompletionReleasesIt() { TrackingExecution execution = new TrackingExecution(null); service.start( OWNER, - "5m", - "0s", + KEEP_ALIVE, + TimeValue.ZERO, request, registration.requestTask(), ignored -> execution, @@ -313,8 +314,8 @@ public void executionStartFailureCompletesJobAndReleasesTask() { service.start( OWNER, - "5m", - "5s", + KEEP_ALIVE, + TimeValue.timeValueSeconds(5), request, registration.requestTask(), ignored -> { @@ -358,8 +359,8 @@ public void deleteCancelsAndReleasesServiceOwnedTask() { service.start( OWNER, - "5m", - "0s", + KEEP_ALIVE, + TimeValue.ZERO, request, registration.requestTask(), ignored -> new TrackingExecution(null), @@ -404,8 +405,8 @@ public void startupFailureAfterJobCreationReleasesServiceOwnedTask() { () -> abortingService.start( OWNER, - "5m", - "5s", + KEEP_ALIVE, + TimeValue.timeValueSeconds(5), request, registration.requestTask(), ignored -> new TrackingExecution(null), @@ -436,8 +437,8 @@ public void childTrackingFailurePreventsRetainedTaskRegistration() { () -> service.start( OWNER, - "5m", - "5s", + KEEP_ALIVE, + TimeValue.timeValueSeconds(5), request, requestTask, ignored -> new TrackingExecution(null), @@ -455,7 +456,7 @@ public void rejectsUnauthorizedCallerWithoutRenewingOrDeleting() { AtomicReference id = new AtomicReference<>(); service.start( securedOwner, - PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + KEEP_ALIVE, TimeValue.ZERO, jobTask(null), ignored -> new TrackingExecution(null), @@ -472,7 +473,7 @@ public void enforcesRunningAndRetainedCapacity() { PPLAsyncQueryService limited = service(1, 1); limited.start( OWNER, - PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + KEEP_ALIVE, TimeValue.ZERO, jobTask(null), ignored -> new TrackingExecution(null), @@ -484,7 +485,7 @@ public void enforcesRunningAndRetainedCapacity() { () -> limited.start( OWNER, - PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + KEEP_ALIVE, TimeValue.ZERO, jobTask(null), ignored -> new TrackingExecution(null), @@ -510,7 +511,7 @@ public void createFailureReleasesReservedCapacity() { () -> missingOwnerNode.start( OWNER, - PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + KEEP_ALIVE, TimeValue.ZERO, jobTask(null), ignored -> new TrackingExecution(null), @@ -610,7 +611,7 @@ public void deleteBeforeExecutionAttachmentClosesLateHandle() { service.start( OWNER, - PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, + KEEP_ALIVE, TimeValue.ZERO, jobTask(null), ignored -> { @@ -746,13 +747,7 @@ private void startQuery( TimeValue waitForCompletion, AsyncQueryExecution execution, ActionListener responseListener) { - startQuery( - targetService, - task, - PPLAsyncQueryService.DEFAULT_KEEP_ALIVE, - waitForCompletion, - execution, - responseListener); + startQuery(targetService, task, KEEP_ALIVE, waitForCompletion, execution, responseListener); } private void startQuery( diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java index 0849b086b36..67fff2db5e5 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java @@ -37,19 +37,6 @@ public class PPLQueryRequest { private static final String START_TIME_FIELD = "start_time"; private static final String END_TIME_FIELD = "end_time"; private static final String TIME_FIELD_FIELD = "time_field"; - - /** JSON field selecting the asynchronous job lease. */ - public static final String KEEP_ALIVE_FIELD = "keep_alive"; - - /** JSON field selecting how long submit waits for direct completion. */ - public static final String WAIT_FOR_COMPLETION_TIMEOUT_FIELD = "wait_for_completion_timeout"; - - /** Default asynchronous job lease. */ - public static final String DEFAULT_KEEP_ALIVE = "5m"; - - /** Default submit wait-for-completion timeout. */ - public static final String DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT = "5s"; - private static final int MAX_HIGHLIGHT_FIELDS = 100; private static final int MAX_TAG_ENTRIES = 10; @@ -61,11 +48,6 @@ public class PPLQueryRequest { @Getter private String format = ""; @Getter private String explainMode; - @Setter - @Getter - @Accessors(fluent = true) - private boolean formatExplicitlySpecified = false; - @Setter @Getter @Accessors(fluent = true) @@ -159,62 +141,6 @@ public boolean isExplainRequest() { return path.endsWith("/_explain"); } - /** - * Returns whether the request contains an asynchronous lifecycle field. - * - * @return {@code true} when keep-alive or wait-for-completion was explicitly requested - */ - public boolean isAsyncQueryRequest() { - return jsonContent != null - && (jsonContent.has(WAIT_FOR_COMPLETION_TIMEOUT_FIELD) - || jsonContent.has(KEEP_ALIVE_FIELD)); - } - - /** - * Returns the requested asynchronous job lease. - * - * @return requested lease or {@link #DEFAULT_KEEP_ALIVE} - */ - public String getKeepAlive() { - if (jsonContent == null || !jsonContent.has(KEEP_ALIVE_FIELD)) { - return DEFAULT_KEEP_ALIVE; - } - return stringLifecycleField(KEEP_ALIVE_FIELD); - } - - /** - * Returns how long asynchronous submit waits for direct completion. - * - * @return requested timeout or {@link #DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT} - */ - public String getWaitForCompletionTimeout() { - if (jsonContent == null || !jsonContent.has(WAIT_FOR_COMPLETION_TIMEOUT_FIELD)) { - return DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT; - } - return stringLifecycleField(WAIT_FOR_COMPLETION_TIMEOUT_FIELD); - } - - private String stringLifecycleField(String field) { - Object value = jsonContent.get(field); - if (!(value instanceof String text)) { - throw new IllegalArgumentException("[" + field + "] must be a string"); - } - return text; - } - - /** - * Returns whether this request mode and response format support asynchronous execution. - * - * @return {@code true} for a normal query using the default JSON response format - */ - public boolean supportsAsyncExecution() { - return !isExplainRequest() - && !profile - && !analyze - && !pplQuery.trim().toLowerCase(Locale.ROOT).startsWith("explain") - && !formatExplicitlySpecified; - } - /** Decide on the formatter by the requested format. */ public Format format() { Optional optionalFormat = Format.of(format); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java index 25efa4be1fc..d4c45e2ece4 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/domain/PPLQueryRequestTest.java @@ -6,7 +6,6 @@ package org.opensearch.sql.ppl.domain; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; @@ -64,69 +63,6 @@ public void testUnsupportedFormat() { request.format(); } - @Test - public void asyncFieldsUseDefaultsOrRequestedValues() { - PPLQueryRequest defaults = - new PPLQueryRequest( - "source=t", - new JSONObject().put("wait_for_completion_timeout", "1s"), - "/_plugins/_ppl"); - assertTrue(defaults.isAsyncQueryRequest()); - assertEquals("5m", defaults.getKeepAlive()); - assertEquals("1s", defaults.getWaitForCompletionTimeout()); - - PPLQueryRequest requested = - new PPLQueryRequest("source=t", new JSONObject().put("keep_alive", "2m"), "/_plugins/_ppl"); - assertTrue(requested.isAsyncQueryRequest()); - assertEquals("2m", requested.getKeepAlive()); - assertEquals("5s", requested.getWaitForCompletionTimeout()); - } - - @Test - public void asyncLifecycleFieldsMustBeStrings() { - PPLQueryRequest numericKeepAlive = - new PPLQueryRequest("source=t", new JSONObject().put("keep_alive", 300), "/_plugins/_ppl"); - PPLQueryRequest numericWait = - new PPLQueryRequest( - "source=t", new JSONObject().put("wait_for_completion_timeout", 1), "/_plugins/_ppl"); - - IllegalArgumentException keepAliveFailure = - assertThrows(IllegalArgumentException.class, numericKeepAlive::getKeepAlive); - assertEquals("[keep_alive] must be a string", keepAliveFailure.getMessage()); - IllegalArgumentException waitFailure = - assertThrows(IllegalArgumentException.class, numericWait::getWaitForCompletionTimeout); - assertEquals("[wait_for_completion_timeout] must be a string", waitFailure.getMessage()); - } - - @Test - public void defaultFormatQuerySupportsAsyncExecution() { - assertTrue( - new PPLQueryRequest("source=t", null, "/_plugins/_ppl", "jdbc").supportsAsyncExecution()); - } - - @Test - public void specialModesAndFormatsDoNotSupportAsyncExecution() { - assertFalse( - new PPLQueryRequest("source=t", null, "/_plugins/_ppl/_explain").supportsAsyncExecution()); - assertFalse( - new PPLQueryRequest("explain source=t", null, "/_plugins/_ppl").supportsAsyncExecution()); - PPLQueryRequest explicitCsv = new PPLQueryRequest("source=t", null, "/_plugins/_ppl", "csv"); - explicitCsv.formatExplicitlySpecified(true); - assertFalse(explicitCsv.supportsAsyncExecution()); - - PPLQueryRequest explicitJdbc = new PPLQueryRequest("source=t", null, "/_plugins/_ppl", "jdbc"); - explicitJdbc.formatExplicitlySpecified(true); - assertFalse(explicitJdbc.supportsAsyncExecution()); - - PPLQueryRequest profile = new PPLQueryRequest("source=t", null, "/_plugins/_ppl"); - profile.profile(true); - assertFalse(profile.supportsAsyncExecution()); - - PPLQueryRequest analyze = new PPLQueryRequest("source=t", null, "/_plugins/_ppl"); - analyze.analyze(true); - assertFalse(analyze.supportsAsyncExecution()); - } - @Test public void testGetFetchSizeReturnsValueFromJson() { JSONObject json = new JSONObject("{\"query\": \"source=t\", \"fetch_size\": 100}"); From c2c5674f676e7d5fa7897ec3b37cfed869fc3bdf Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 20:20:29 +0000 Subject: [PATCH 03/15] Clarify async query user tests Signed-off-by: Peng Huo --- .../asyncquery/PPLAsyncQueryUserTest.java | 44 ++++++++++--------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java index 6c9a0061849..ac8fc32f41b 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java @@ -6,7 +6,6 @@ package org.opensearch.sql.plugin.transport.asyncquery; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import java.util.List; @@ -21,38 +20,35 @@ public class PPLAsyncQueryUserTest { @Test public void capturesSecurityIdentity() { - ThreadContext context = new ThreadContext(Settings.EMPTY); - context.putTransient( - ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, - "alice|backend-a|ppl-role|tenant-a"); + ThreadContext context = contextWith("alice|backend-a|ppl-role|tenant-a"); + + assertEquals( + new PPLAsyncQueryUser("alice", "tenant-a", List.of("backend-a")), + PPLAsyncQueryUser.current(context)); + } - PPLAsyncQueryUser identity = PPLAsyncQueryUser.current(context); + @Test + public void acceptsUserObject() { + ThreadContext context = + contextWith(new User("alice", List.of("backend-a"), List.of("ppl-role"), null, "tenant-a")); - assertEquals("alice", identity.name()); - assertEquals("tenant-a", identity.requestedTenant()); - assertEquals(List.of("backend-a"), identity.backendRoles()); + assertEquals( + new PPLAsyncQueryUser("alice", "tenant-a", List.of("backend-a")), + PPLAsyncQueryUser.current(context)); } @Test - public void acceptsUserObjectAndRejectsUnknownSecurityContext() { - ThreadContext objectContext = new ThreadContext(Settings.EMPTY); - objectContext.putTransient( - ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, - new User("alice", List.of("backend-a"), List.of("ppl-role"), null, "tenant-a")); - assertEquals("alice", PPLAsyncQueryUser.current(objectContext).name()); - - ThreadContext invalidContext = new ThreadContext(Settings.EMPTY); - invalidContext.putTransient( - ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, new Object()); + public void rejectsUnknownSecurityContext() { assertThrows( - OpenSearchSecurityException.class, () -> PPLAsyncQueryUser.current(invalidContext)); + OpenSearchSecurityException.class, + () -> PPLAsyncQueryUser.current(contextWith(new Object()))); } @Test public void missingIdentityRepresentsAnUnsecuredCaller() { ThreadContext context = new ThreadContext(Settings.EMPTY); - assertNull(PPLAsyncQueryUser.current(context).name()); + assertEquals(new PPLAsyncQueryUser(null, null, List.of()), PPLAsyncQueryUser.current(context)); } @Test @@ -82,4 +78,10 @@ public void unsecuredModeRequiresAnUnsecuredCaller() { OpenSearchSecurityException.class, () -> unsecured.authorize(new PPLAsyncQueryUser("alice", null, List.of("role-a")))); } + + private static ThreadContext contextWith(Object userInfo) { + ThreadContext context = new ThreadContext(Settings.EMPTY); + context.putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, userInfo); + return context; + } } From 7b73c9f795ea01ee2e6e984235e959d3fb419617 Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 20:24:33 +0000 Subject: [PATCH 04/15] Name the unsecured async query identity Signed-off-by: Peng Huo --- .../plugin/transport/asyncquery/PPLAsyncQueryUser.java | 10 +++++++++- .../transport/asyncquery/PPLAsyncQueryServiceTest.java | 2 +- .../transport/asyncquery/PPLAsyncQueryUserTest.java | 9 +++++---- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java index e060962a3ff..c8dcc7c8b35 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java @@ -22,6 +22,14 @@ */ public record PPLAsyncQueryUser(String name, String requestedTenant, List backendRoles) { + /** + * Identity used when OpenSearch Security does not provide caller information. + * + *

This identity does not bypass job ownership checks. A job owned by {@code UNSECURED} can be + * accessed only by a caller represented by the same identity. + */ + public static final PPLAsyncQueryUser UNSECURED = new PPLAsyncQueryUser(null, null, List.of()); + /** * Creates an immutable asynchronous query identity. * @@ -49,7 +57,7 @@ public static PPLAsyncQueryUser current(ThreadContext threadContext) { Object serialized = threadContext.getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT); if (serialized == null) { - return new PPLAsyncQueryUser(null, null, List.of()); + return UNSECURED; } User user = serialized instanceof User currentUser diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java index 74a399a58e7..13c893c20b3 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java @@ -55,7 +55,7 @@ import org.opensearch.tasks.TaskManager; public class PPLAsyncQueryServiceTest { - private static final PPLAsyncQueryUser OWNER = new PPLAsyncQueryUser(null, null, List.of()); + private static final PPLAsyncQueryUser OWNER = PPLAsyncQueryUser.UNSECURED; private static final TimeValue KEEP_ALIVE = TimeValue.timeValueMinutes(5); private final AtomicLong now = new AtomicLong(1_000); diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java index ac8fc32f41b..fb028eea779 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java @@ -48,7 +48,7 @@ public void rejectsUnknownSecurityContext() { public void missingIdentityRepresentsAnUnsecuredCaller() { ThreadContext context = new ThreadContext(Settings.EMPTY); - assertEquals(new PPLAsyncQueryUser(null, null, List.of()), PPLAsyncQueryUser.current(context)); + assertEquals(PPLAsyncQueryUser.UNSECURED, PPLAsyncQueryUser.current(context)); } @Test @@ -71,12 +71,13 @@ public void requiresSamePrincipalTenantAndOriginalBackendRoles() { @Test public void unsecuredModeRequiresAnUnsecuredCaller() { - PPLAsyncQueryUser unsecured = new PPLAsyncQueryUser(null, null, List.of()); - unsecured.authorize(new PPLAsyncQueryUser(null, null, List.of())); + PPLAsyncQueryUser.UNSECURED.authorize(PPLAsyncQueryUser.UNSECURED); assertThrows( OpenSearchSecurityException.class, - () -> unsecured.authorize(new PPLAsyncQueryUser("alice", null, List.of("role-a")))); + () -> + PPLAsyncQueryUser.UNSECURED.authorize( + new PPLAsyncQueryUser("alice", null, List.of("role-a")))); } private static ThreadContext contextWith(Object userInfo) { From eed8cb22c4d95bb14deb33510b8e4d507b867a2e Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 20:34:58 +0000 Subject: [PATCH 05/15] Split async job state machine tests Signed-off-by: Peng Huo --- .../asyncquery/PPLAsyncQueryJobTest.java | 256 ++++++++++++++++++ .../asyncquery/PPLAsyncQueryServiceTest.java | 82 +++--- 2 files changed, 294 insertions(+), 44 deletions(-) create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java new file mode 100644 index 00000000000..97bc6306f99 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java @@ -0,0 +1,256 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport.asyncquery; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +import org.junit.Test; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.sql.executor.AsyncQueryExecution; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.GetResult; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.JobTask; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Removal; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.ResponseContext; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Retention; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryService.Failure; +import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryService.Status; +import org.opensearch.tasks.CancellableTask; + +public class PPLAsyncQueryJobTest { + private static final String ID = "job-id"; + private static final long START_TIME = 1_000; + private static final long RETAINED_TIME = 2_000; + private static final long COMPLETION_TIME = 2_025; + private static final long KEEP_ALIVE_MILLIS = TimeValue.timeValueMinutes(5).millis(); + + @Test + public void retainPublishesRunningResponseWithId() { + JobFixture fixture = newJob(); + assertTrue(fixture.job().tryAttachExecution(fixture.execution())); + + Transition transition = fixture.job().retain(RETAINED_TIME); + + assertEquals( + new Transition( + new ResponseContext(ID, Status.RUNNING, fixture.execution(), null, -1), + Retention.RETAIN, + null, + null), + transition); + assertNull(fixture.job().retain(RETAINED_TIME + 1)); + } + + @Test + public void completionBeforeRetentionReturnsDirectResponseAndRemovesJob() { + JobFixture fixture = newJob(); + assertTrue(fixture.job().tryAttachExecution(fixture.execution())); + + Transition transition = fixture.job().complete(COMPLETION_TIME); + + assertEquals( + new Transition( + new ResponseContext( + null, Status.SUCCEEDED, fixture.execution(), null, COMPLETION_TIME - START_TIME), + Retention.REMOVE, + fixture.execution(), + fixture.task()), + transition); + assertTrue(transition.releasesRunningSlot()); + assertNull(fixture.job().complete(COMPLETION_TIME + 1)); + assertThrows(ResourceNotFoundException.class, () -> fixture.job().get(COMPLETION_TIME, null)); + } + + @Test + public void failureBeforeRetentionReturnsDirectResponseAndRemovesJob() { + JobFixture fixture = newJob(); + Failure failure = new Failure("IllegalArgumentException", "invalid query"); + assertTrue(fixture.job().tryAttachExecution(fixture.execution())); + + Transition transition = fixture.job().fail(failure, COMPLETION_TIME); + + assertEquals( + new Transition( + new ResponseContext( + null, Status.FAILED, fixture.execution(), failure, COMPLETION_TIME - START_TIME), + Retention.REMOVE, + fixture.execution(), + fixture.task()), + transition); + assertTrue(transition.releasesRunningSlot()); + } + + @Test + public void retainedSuccessRemainsReadableUntilDeleted() { + JobFixture fixture = retainedJob(); + + Transition transition = fixture.job().complete(COMPLETION_TIME); + + assertEquals(new Transition(null, Retention.RETAIN, null, fixture.task()), transition); + assertEquals( + new GetResult.Found( + new ResponseContext( + ID, Status.SUCCEEDED, fixture.execution(), null, COMPLETION_TIME - START_TIME)), + fixture.job().get(COMPLETION_TIME, null)); + assertEquals( + new Removal( + Status.SUCCEEDED, + null, + fixture.execution(), + "PPL asynchronous query cancelled by user", + false), + fixture.job().delete(COMPLETION_TIME)); + } + + @Test + public void retainedFailureDetachesExecutionAndRemainsReadable() { + JobFixture fixture = retainedJob(); + Failure failure = new Failure("IllegalStateException", "query failed"); + + Transition transition = fixture.job().fail(failure, COMPLETION_TIME); + + assertEquals( + new Transition(null, Retention.RETAIN, fixture.execution(), fixture.task()), transition); + assertEquals( + new GetResult.Found( + new ResponseContext(ID, Status.FAILED, null, failure, COMPLETION_TIME - START_TIME)), + fixture.job().get(COMPLETION_TIME, null)); + } + + @Test + public void getWithoutKeepAlivePreservesExpiration() { + JobFixture fixture = retainedJob(); + long expirationTime = RETAINED_TIME + KEEP_ALIVE_MILLIS; + + assertTrue(fixture.job().get(expirationTime - 1, null) instanceof GetResult.Found); + GetResult result = fixture.job().get(expirationTime, null); + + assertEquals( + new GetResult.Expired( + new Removal( + Status.RUNNING, + fixture.task(), + fixture.execution(), + "PPL asynchronous query expired", + true)), + result); + } + + @Test + public void getWithKeepAliveReplacesExpiration() { + JobFixture fixture = retainedJob(); + long renewalTime = RETAINED_TIME + 500; + TimeValue requestedKeepAlive = TimeValue.timeValueSeconds(1); + + assertTrue(fixture.job().get(renewalTime, requestedKeepAlive) instanceof GetResult.Found); + assertTrue(fixture.job().get(renewalTime + 999, null) instanceof GetResult.Found); + assertTrue( + fixture.job().get(renewalTime + requestedKeepAlive.millis(), null) + instanceof GetResult.Expired); + } + + @Test + public void expirationStartsWhenJobIsRetained() { + JobFixture fixture = newJob(); + assertTrue(fixture.job().tryAttachExecution(fixture.execution())); + long retainedTime = START_TIME + KEEP_ALIVE_MILLIS + 1; + + assertNull(fixture.job().expire(retainedTime - 1)); + fixture.job().retain(retainedTime); + assertNull(fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS - 1)); + assertEquals( + new Removal( + Status.RUNNING, + fixture.task(), + fixture.execution(), + "PPL asynchronous query expired", + true), + fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS)); + } + + @Test + public void deleteRunningJobReturnsCancellationAndDetachesResources() { + JobFixture fixture = retainedJob(); + + Removal removal = fixture.job().delete(RETAINED_TIME + 1); + + assertEquals( + new Removal( + Status.CANCELLED, + fixture.task(), + fixture.execution(), + "PPL asynchronous query cancelled by user", + false), + removal); + assertTrue(removal.releasesRunningSlot()); + assertThrows(ResourceNotFoundException.class, () -> fixture.job().delete(RETAINED_TIME + 2)); + } + + @Test + public void abortDetachesResourcesOnlyOnce() { + JobFixture fixture = newJob(); + assertTrue(fixture.job().tryAttachExecution(fixture.execution())); + + assertEquals( + new Removal( + Status.RUNNING, + fixture.task(), + fixture.execution(), + "PPL asynchronous query startup failed", + false), + fixture.job().abort()); + assertNull(fixture.job().abort()); + } + + @Test + public void closeDetachesResourcesOnlyOnce() { + JobFixture fixture = retainedJob(); + + assertEquals( + new Removal(Status.RUNNING, fixture.task(), fixture.execution(), "service closing", false), + fixture.job().close("service closing")); + assertNull(fixture.job().close("service closing")); + } + + @Test + public void lateExecutionAttachmentAfterRemovalIsRejected() { + JobFixture fixture = newJob(); + fixture.job().abort(); + + assertFalse(fixture.job().tryAttachExecution(fixture.execution())); + } + + @Test + public void successfulCompletionRequiresAttachedExecution() { + JobFixture fixture = newJob(); + + assertThrows(IllegalStateException.class, () -> fixture.job().complete(COMPLETION_TIME)); + } + + private static JobFixture retainedJob() { + JobFixture fixture = newJob(); + assertTrue(fixture.job().tryAttachExecution(fixture.execution())); + fixture.job().retain(RETAINED_TIME); + return fixture; + } + + private static JobFixture newJob() { + JobTask task = new JobTask(mock(CancellableTask.class), () -> {}); + AsyncQueryExecution execution = mock(AsyncQueryExecution.class); + return new JobFixture( + new PPLAsyncQueryJob(ID, PPLAsyncQueryUser.UNSECURED, START_TIME, KEEP_ALIVE_MILLIS, task), + task, + execution); + } + + private record JobFixture(PPLAsyncQueryJob job, JobTask task, AsyncQueryExecution execution) {} +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java index 13c893c20b3..05872c706a1 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java @@ -82,10 +82,10 @@ public void fastSuccessReturnsDirectResultWithoutRetainingJob() { execution.complete(); assertEquals(1, responses.get()); - assertNull(result.get().id()); - assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, result.get().status()); - assertEquals(2, result.get().response().getResults().size()); - assertEquals(25, result.get().tookMillis()); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + null, PPLAsyncQueryService.Status.SUCCEEDED, response(2), null, 25), + result.get()); assertEquals(0, service.runningQueryCount()); assertEquals(0, service.retainedJobCount()); assertEquals(1, execution.reads.get()); @@ -106,8 +106,10 @@ public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() { timeoutTask.get().run(); String id = retainedResponse.get().id(); - assertEquals(PPLAsyncQueryService.Status.RUNNING, retainedResponse.get().status()); - assertNull(retainedResponse.get().response()); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + id, PPLAsyncQueryService.Status.RUNNING, null, null, -1), + retainedResponse.get()); assertEquals(1, service.runningQueryCount()); assertEquals(1, service.retainedJobCount()); @@ -115,9 +117,10 @@ public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() { execution.succeed(response(2)); PPLAsyncQueryService.JobSnapshot completed = service.get(id, OWNER, null); - assertEquals(id, completed.id()); - assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, completed.status()); - assertEquals(2, completed.response().getResults().size()); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + id, PPLAsyncQueryService.Status.SUCCEEDED, response(2), null, 25), + completed); assertEquals(0, service.runningQueryCount()); assertEquals(1, service.retainedJobCount()); assertEquals(0, execution.closes.get()); @@ -174,16 +177,14 @@ public void fastFailureReturnsDirectFailureWithoutId() { } @Test - public void getWithoutKeepAliveDoesNotRenewLease() { + public void expiredGetCancelsAndRemovesJob() { CancellableTask task = mock(CancellableTask.class); when(task.isCancelled()).thenReturn(false); TrackingExecution execution = new TrackingExecution(null); String id = startRetainedQuery(service, task, execution); - now.addAndGet(TimeValue.timeValueMinutes(4).millis()); - assertEquals(PPLAsyncQueryService.Status.RUNNING, service.get(id, OWNER, null).status()); + now.addAndGet(KEEP_ALIVE.millis()); - now.addAndGet(TimeValue.timeValueMinutes(1).millis() + 1); assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); verify(task).cancel("PPL asynchronous query expired"); assertEquals(1, execution.closes.get()); @@ -191,26 +192,6 @@ public void getWithoutKeepAliveDoesNotRenewLease() { assertEquals(0, service.retainedJobCount()); } - @Test - public void getWithKeepAliveRenewsLease() { - CancellableTask task = mock(CancellableTask.class); - when(task.isCancelled()).thenReturn(false); - TrackingExecution execution = new TrackingExecution(null); - String id = startRetainedQuery(service, task, execution); - - now.addAndGet(TimeValue.timeValueMinutes(4).millis()); - assertEquals( - PPLAsyncQueryService.Status.RUNNING, - service.get(id, OWNER, TimeValue.timeValueMinutes(5)).status()); - - now.addAndGet(TimeValue.timeValueMinutes(4).millis()); - assertEquals(PPLAsyncQueryService.Status.RUNNING, service.get(id, OWNER, null).status()); - - now.addAndGet(TimeValue.timeValueMinutes(1).millis() + 1); - assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); - verify(task).cancel("PPL asynchronous query expired"); - } - @Test public void missingLocalJobReturnsNotFound() { String remoteId = PPLAsyncQueryJobId.create("node-b").encode(); @@ -228,8 +209,8 @@ public void deleteCancelsRunningJobAndReleasesState() { PPLAsyncQueryService.DeleteResult result = service.delete(id, OWNER); - assertEquals(id, result.id()); - assertEquals(PPLAsyncQueryService.Status.CANCELLED, result.status()); + assertEquals( + new PPLAsyncQueryService.DeleteResult(id, PPLAsyncQueryService.Status.CANCELLED), result); verify(task).cancel("PPL asynchronous query cancelled by user"); assertEquals(1, execution.closes.get()); assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null)); @@ -246,7 +227,8 @@ public void deleteReturnsExistingTerminalStatus() { PPLAsyncQueryService.DeleteResult result = service.delete(id, OWNER); - assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, result.status()); + assertEquals( + new PPLAsyncQueryService.DeleteResult(id, PPLAsyncQueryService.Status.SUCCEEDED), result); verify(task, never()).cancel(org.mockito.ArgumentMatchers.anyString()); assertEquals(1, execution.closes.get()); assertEquals(0, service.retainedJobCount()); @@ -546,8 +528,10 @@ public void completedExecutionCanBeAttachedBeforeCompletionIsObserved() { execution.succeed(response(2)); startQuery(service, null, TimeValue.timeValueSeconds(5), execution, listener(result::set)); - assertEquals(PPLAsyncQueryService.Status.SUCCEEDED, result.get().status()); - assertEquals(2, result.get().response().getResults().size()); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + null, PPLAsyncQueryService.Status.SUCCEEDED, response(2), null, 0), + result.get()); assertEquals(1, execution.closes.get()); } @@ -560,9 +544,14 @@ public void runningGetMaterializesCurrentResultOutsideJob() { execution.setCurrent(response(3)); PPLAsyncQueryService.JobSnapshot second = service.get(id, OWNER, null); - assertEquals(PPLAsyncQueryService.Status.RUNNING, first.status()); - assertEquals(1, first.response().getResults().size()); - assertEquals(3, second.response().getResults().size()); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + id, PPLAsyncQueryService.Status.RUNNING, response(1), null, -1), + first); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + id, PPLAsyncQueryService.Status.RUNNING, response(3), null, -1), + second); } @Test @@ -577,9 +566,14 @@ public void failedRetainedJobReturnsNoProvisionalRowsAndClosesExecution() { execution.fail(new IllegalStateException("boom")); PPLAsyncQueryService.JobSnapshot failed = service.get(id, OWNER, null); - assertEquals(PPLAsyncQueryService.Status.FAILED, failed.status()); - assertEquals("boom", failed.failure().reason()); - assertNull(failed.response()); + assertEquals( + new PPLAsyncQueryService.JobSnapshot( + id, + PPLAsyncQueryService.Status.FAILED, + null, + new PPLAsyncQueryService.Failure("IllegalStateException", "boom"), + 0), + failed); assertEquals(0, execution.reads.get()); assertEquals(1, execution.closes.get()); } finally { From 712d6e5390117cec67498afe8128f0c3167fb805 Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 21:00:12 +0000 Subject: [PATCH 06/15] Refine async query lifecycle state model Signed-off-by: Peng Huo --- .../asyncquery/PPLAsyncQueryJob.java | 229 ++++++++++-------- .../asyncquery/PPLAsyncQueryService.java | 220 +++++++++++------ .../asyncquery/PPLAsyncQueryJobTest.java | 146 ++++++----- .../asyncquery/PPLAsyncQueryServiceTest.java | 43 ++-- 4 files changed, 369 insertions(+), 269 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java index 771b1b5cfc0..41308646411 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java @@ -6,6 +6,7 @@ package org.opensearch.sql.plugin.transport.asyncquery; import java.util.Objects; +import java.util.Optional; import org.opensearch.ResourceNotFoundException; import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.executor.AsyncQueryExecution; @@ -19,7 +20,7 @@ * service map, update capacity counters, or cancel tasks while holding the lock. * *

After attachment, the job owns one {@link AsyncQueryExecution}. Reads borrow it through a - * {@link ResponseContext}; removal transitions detach it so the service can close it outside the + * {@link SnapshotSource}; removal transitions detach it so the service can close it outside the * lock. * * 

@@ -106,7 +107,7 @@ synchronized Transition retain(long now) {
     }
     state = State.RETAINED_RUNNING;
     expirationTimeMillis = now + keepAliveMillis;
-    return Transition.retain(retainedResponse());
+    return new Transition.Retained(runningSnapshot());
   }
 
   /**
@@ -164,12 +165,14 @@ private Transition finish(State terminalState, long now) {
     state = terminalState;
     completionTimeMillis = now;
     if (returnDirect) {
-      ResponseContext response = directResponse();
+      SnapshotSource response = snapshot(Optional.empty());
       state = State.REMOVED;
-      return Transition.returnDirect(response, detachExecution(), taskToClose);
+      return new Transition.DirectResponse(
+          response, taskToClose, Optional.ofNullable(detachExecution()));
     }
-    return Transition.finishRetained(
-        terminalState == State.RETAINED_FAILED ? detachExecution() : null, taskToClose);
+    AsyncQueryExecution executionToClose =
+        terminalState == State.RETAINED_FAILED ? detachExecution() : null;
+    return new Transition.ExecutionFinished(taskToClose, Optional.ofNullable(executionToClose));
   }
 
   /**
@@ -192,7 +195,7 @@ synchronized GetResult get(long now, TimeValue requestedKeepAlive) {
       keepAliveMillis = requestedKeepAlive.millis();
       expirationTimeMillis = now + keepAliveMillis;
     }
-    return new GetResult.Found(retainedResponse());
+    return new GetResult.Found(retainedSnapshot());
   }
 
   /**
@@ -207,7 +210,7 @@ synchronized Removal delete(long now) {
     if (now >= expirationTimeMillis) {
       return expireLocked("PPL asynchronous query expired");
     }
-    return remove(
+    return delete(
         isExecuting() ? PPLAsyncQueryService.Status.CANCELLED : responseStatus(),
         "PPL asynchronous query cancelled by user");
   }
@@ -225,8 +228,10 @@ synchronized Removal expire(long now) {
     return expireLocked("PPL asynchronous query expired");
   }
 
-  private Removal expireLocked(String reason) {
-    return remove(responseStatus(), reason, true);
+  private Removal.Expired expireLocked(String reason) {
+    Removal.Expired removal = new Removal.Expired(detachResources(), reason);
+    state = State.REMOVED;
+    return removal;
   }
 
   /**
@@ -235,7 +240,7 @@ private Removal expireLocked(String reason) {
    * @return removal with detached resources, or {@code null} if already removed
    */
   synchronized Removal abort() {
-    return removeIfPresent("PPL asynchronous query startup failed");
+    return discardIfPresent("PPL asynchronous query startup failed");
   }
 
   /**
@@ -245,39 +250,47 @@ synchronized Removal abort() {
    * @return removal with detached resources, or {@code null} if already removed
    */
   synchronized Removal close(String reason) {
-    return removeIfPresent(reason);
+    return discardIfPresent(reason);
   }
 
-  private Removal removeIfPresent(String reason) {
+  private Removal discardIfPresent(String reason) {
     if (state == State.REMOVED) {
       return null;
     }
-    return remove(responseStatus(), reason);
-  }
-
-  private Removal remove(PPLAsyncQueryService.Status responseStatus, String reason) {
-    return remove(responseStatus, reason, false);
+    Removal.Discarded removal = new Removal.Discarded(detachResources(), reason);
+    state = State.REMOVED;
+    return removal;
   }
 
-  private Removal remove(
-      PPLAsyncQueryService.Status responseStatus, String reason, boolean expired) {
-    Removal removal = new Removal(responseStatus, detachTask(), detachExecution(), reason, expired);
+  private Removal delete(PPLAsyncQueryService.Status responseStatus, String reason) {
+    Removal.Deleted removal = new Removal.Deleted(responseStatus, detachResources(), reason);
     state = State.REMOVED;
     return removal;
   }
 
-  private ResponseContext directResponse() {
-    return responseContext(null);
+  private DetachedResources detachResources() {
+    return new DetachedResources(detachTask(), detachExecution());
+  }
+
+  private SnapshotSource retainedSnapshot() {
+    return snapshot(Optional.of(id));
   }
 
-  private ResponseContext retainedResponse() {
-    return responseContext(id);
+  private SnapshotSource.Running runningSnapshot() {
+    return new SnapshotSource.Running(id, Optional.ofNullable(execution));
   }
 
-  private ResponseContext responseContext(String responseId) {
+  private SnapshotSource snapshot(Optional responseId) {
     long tookMillis =
         completionTimeMillis < 0 ? -1L : Math.max(0L, completionTimeMillis - startTimeMillis);
-    return new ResponseContext(responseId, responseStatus(), execution, failure, tookMillis);
+    return switch (state) {
+      case RUNNING, RETAINED_RUNNING -> runningSnapshot();
+      case RETAINED_SUCCEEDED ->
+          new SnapshotSource.Succeeded(responseId, Objects.requireNonNull(execution), tookMillis);
+      case RETAINED_FAILED ->
+          new SnapshotSource.Failed(responseId, Objects.requireNonNull(failure), tookMillis);
+      case REMOVED -> throw new IllegalStateException("PPL asynchronous query was removed");
+    };
   }
 
   private boolean isExecuting() {
@@ -311,21 +324,36 @@ private void ensurePresent() {
     }
   }
 
-  /**
-   * Lifecycle data captured under the job lock for later response materialization.
-   *
-   * @param id job ID included in a retained response, or {@code null} for a direct POST response
-   * @param status public lifecycle status
-   * @param execution execution handle borrowed for result materialization
-   * @param failure failure returned for a failed job
-   * @param tookMillis elapsed execution time, or {@code -1} while running
-   */
-  record ResponseContext(
-      String id,
-      PPLAsyncQueryService.Status status,
-      AsyncQueryExecution execution,
-      PPLAsyncQueryService.Failure failure,
-      long tookMillis) {}
+  /** Lifecycle state captured under the job lock for later result materialization. */
+  sealed interface SnapshotSource {
+    /**
+     * Source for a retained running response.
+     *
+     * @param id retained job ID
+     * @param execution attached execution, or empty when retention wins before attachment
+     */
+    record Running(String id, Optional execution) implements SnapshotSource {}
+
+    /**
+     * Source for a successful final response.
+     *
+     * @param id retained job ID, or empty for a direct POST response
+     * @param execution completed execution that owns the final result
+     * @param tookMillis elapsed execution time
+     */
+    record Succeeded(Optional id, AsyncQueryExecution execution, long tookMillis)
+        implements SnapshotSource {}
+
+    /**
+     * Source for a failed final response.
+     *
+     * @param id retained job ID, or empty for a direct POST response
+     * @param failure client-visible failure
+     * @param tookMillis elapsed execution time
+     */
+    record Failed(Optional id, PPLAsyncQueryService.Failure failure, long tookMillis)
+        implements SnapshotSource {}
+  }
 
   /**
    * Cancellable task and the cleanup that releases its TaskManager registrations.
@@ -349,44 +377,34 @@ enum State {
     REMOVED
   }
 
-  /** Whether a transition keeps the job in or removes it from the service registry. */
-  enum Retention {
-    RETAIN,
-    REMOVE
-  }
-
-  /**
-   * State-machine output consumed by {@link PPLAsyncQueryService}.
-   *
-   * @param response response to materialize and publish, or {@code null} when none is due
-   * @param retention registry action
-   * @param executionToClose execution handle detached by the transition
-   * @param taskToClose completed task registration detached by the transition
-   */
-  record Transition(
-      ResponseContext response,
-      Retention retention,
-      AsyncQueryExecution executionToClose,
-      JobTask taskToClose) {
-
-    private static Transition retain(ResponseContext response) {
-      return new Transition(response, Retention.RETAIN, null, null);
-    }
-
-    private static Transition returnDirect(
-        ResponseContext response, AsyncQueryExecution executionToClose, JobTask taskToClose) {
-      return new Transition(response, Retention.REMOVE, executionToClose, taskToClose);
-    }
+  /** State-machine output consumed by {@link PPLAsyncQueryService}. */
+  sealed interface Transition {
+    /**
+     * The initial wait expired and the running job became retained.
+     *
+     * @param response response published with the retained job ID
+     */
+    record Retained(SnapshotSource.Running response) implements Transition {}
 
-    private static Transition finishRetained(
-        AsyncQueryExecution executionToClose, JobTask taskToClose) {
-      return new Transition(null, Retention.RETAIN, executionToClose, taskToClose);
-    }
+    /**
+     * Execution finished before retention and the initial request receives the final response.
+     *
+     * @param response successful or failed final response source
+     * @param task completed task registration to close
+     * @param executionToClose detached execution to close after materialization
+     */
+    record DirectResponse(
+        SnapshotSource response, JobTask task, Optional executionToClose)
+        implements Transition {}
 
-    /** Returns whether this transition releases a running-query capacity slot. */
-    boolean releasesRunningSlot() {
-      return taskToClose != null;
-    }
+    /**
+     * A retained execution finished without producing an immediate response.
+     *
+     * @param task completed task registration to close
+     * @param executionToClose failed execution to close; successful execution remains retained
+     */
+    record ExecutionFinished(JobTask task, Optional executionToClose)
+        implements Transition {}
   }
 
   /** Result of an authorized GET attempt. */
@@ -394,37 +412,58 @@ sealed interface GetResult {
     /**
      * GET result for a live retained job.
      *
-     * @param response current response context
+     * @param response current lifecycle source for result materialization
      */
-    record Found(ResponseContext response) implements GetResult {}
+    record Found(SnapshotSource response) implements GetResult {}
 
     /**
      * GET result when the lease expired before the request.
      *
      * @param removal cleanup required for the expired job
      */
-    record Expired(Removal removal) implements GetResult {}
+    record Expired(Removal.Expired removal) implements GetResult {}
   }
 
   /**
-   * Resources and accounting changes produced when a job leaves the registry.
+   * Task and execution handles detached when a job leaves the registry.
    *
-   * @param responseStatus status returned to DELETE when the job has not expired
-   * @param task running task to cancel
-   * @param execution execution handle to close
-   * @param reason cancellation or removal reason
-   * @param expired whether expiration caused the removal
+   * @param task running task to cancel, or {@code null} after execution already finished
+   * @param execution execution handle to close, or {@code null} before attachment
    */
-  record Removal(
-      PPLAsyncQueryService.Status responseStatus,
-      JobTask task,
-      AsyncQueryExecution execution,
-      String reason,
-      boolean expired) {
-
-    /** Returns whether this removal releases a running-query capacity slot. */
+  record DetachedResources(JobTask task, AsyncQueryExecution execution) {
+    /** Returns whether removing these resources releases a running-query capacity slot. */
     boolean releasesRunningSlot() {
       return task != null;
     }
   }
+
+  /** Reason-specific removal consumed by {@link PPLAsyncQueryService}. */
+  sealed interface Removal {
+    /**
+     * Removal requested by DELETE.
+     *
+     * @param responseStatus status returned to the caller
+     * @param resources detached task and execution handles
+     * @param reason task cancellation reason
+     */
+    record Deleted(
+        PPLAsyncQueryService.Status responseStatus, DetachedResources resources, String reason)
+        implements Removal {}
+
+    /**
+     * Removal caused by lease expiration.
+     *
+     * @param resources detached task and execution handles
+     * @param reason task cancellation reason
+     */
+    record Expired(DetachedResources resources, String reason) implements Removal {}
+
+    /**
+     * Internal removal caused by startup failure or service shutdown.
+     *
+     * @param resources detached task and execution handles
+     * @param reason task cancellation reason
+     */
+    record Discarded(DetachedResources resources, String reason) implements Removal {}
+  }
 }
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java
index 381f42eaf5d..e44ebbe1bac 100644
--- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java
+++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java
@@ -8,6 +8,7 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.function.Function;
@@ -36,8 +37,7 @@
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.GetResult;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.JobTask;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Removal;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.ResponseContext;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Retention;
+import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.SnapshotSource;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
@@ -92,20 +92,76 @@ interface TimeoutScheduler {
   }
 
   /**
-   * Immutable point-in-time response view of a job.
+   * Immutable point-in-time response produced after releasing the job lock.
    *
-   * 
The formatter converts this internal model to the public JSON response. Query data is copied
-   * from the execution after releasing the job lock, so formatting never observes mutable job
-   * state.
-   *
-   * @param id opaque job ID, or {@code null} for a terminal response returned directly by POST
-   * @param status lifecycle state captured with the result
-   * @param response current query result, or {@code null} before a result is available
-   * @param failure client-visible failure for {@link Status#FAILED}, otherwise {@code null}
-   * @param tookMillis elapsed execution time, available for a completed job
+   * 
Each implementation exposes only the data valid for its lifecycle state.
    */
-  public record JobSnapshot(
-      String id, Status status, QueryResponse response, Failure failure, long tookMillis) {}
+  public sealed interface JobSnapshot {
+    /**
+     * Returns the retained job ID.
+     *
+     * @return job ID, or empty for a final response returned directly by POST
+     */
+    Optional id();
+
+    /**
+     * Returns the public lifecycle state.
+     *
+     * @return response status
+     */
+    Status status();
+
+    /**
+     * Snapshot of a retained query that is still running.
+     *
+     * @param jobId retained job ID
+     * @param response current result, or empty before a result is available
+     */
+    record Running(String jobId, Optional response) implements JobSnapshot {
+      /** {@inheritDoc} */
+      @Override
+      public Optional id() {
+        return Optional.of(jobId);
+      }
+
+      /** {@inheritDoc} */
+      @Override
+      public Status status() {
+        return Status.RUNNING;
+      }
+    }
+
+    /**
+     * Snapshot of a successfully completed query.
+     *
+     * @param id retained job ID, or empty for a direct POST response
+     * @param response final query result
+     * @param tookMillis elapsed execution time
+     */
+    record Succeeded(Optional id, QueryResponse response, long tookMillis)
+        implements JobSnapshot {
+      /** {@inheritDoc} */
+      @Override
+      public Status status() {
+        return Status.SUCCEEDED;
+      }
+    }
+
+    /**
+     * Snapshot of a failed query.
+     *
+     * @param id retained job ID, or empty for a direct POST response
+     * @param failure client-visible failure
+     * @param tookMillis elapsed execution time
+     */
+    record Failed(Optional id, Failure failure, long tookMillis) implements JobSnapshot {
+      /** {@inheritDoc} */
+      @Override
+      public Status status() {
+        return Status.FAILED;
+      }
+    }
+  }
 
   /** Response model returned after DELETE removes a retained job. */
   record DeleteResult(String id, Status status) {}
@@ -371,16 +427,20 @@ private void fail(
     if (transition == null) {
       return;
     }
-    if (transition.response() == null) {
-      PPLQueryErrorHandler.recordFailure(failure);
-      applyTransition(job, transition, responseListener);
-      return;
+    switch (transition) {
+      case Transition.ExecutionFinished finished -> {
+        PPLQueryErrorHandler.recordFailure(failure);
+        applyTransition(job, finished, responseListener);
+      }
+      case Transition.DirectResponse direct ->
+          applyTransition(
+              job,
+              direct,
+              ActionListener.wrap(
+                  ignored -> responseListener.onFailure(failure), responseListener::onFailure));
+      case Transition.Retained retained ->
+          throw new IllegalStateException("Query failure cannot retain a job");
     }
-    applyTransition(
-        job,
-        transition,
-        ActionListener.wrap(
-            ignored -> responseListener.onFailure(failure), responseListener::onFailure));
   }
 
   /**
@@ -421,10 +481,12 @@ DeleteResult delete(String id, PPLAsyncQueryUser caller) {
     job.owner().authorize(caller);
     Removal removal = job.delete(currentTimeMillis.getAsLong());
     applyRemoval(job, removal);
-    if (removal.expired()) {
-      throw notFound();
-    }
-    return new DeleteResult(id, removal.responseStatus());
+    return switch (removal) {
+      case Removal.Deleted deleted -> new DeleteResult(id, deleted.responseStatus());
+      case Removal.Expired expired -> throw notFound();
+      case Removal.Discarded discarded ->
+          throw new IllegalStateException("DELETE cannot discard a job");
+    };
   }
 
   void reapExpired() {
@@ -508,34 +570,33 @@ private void applyTransition(
     if (transition == null) {
       return;
     }
-    if (transition.releasesRunningSlot()) {
-      releaseRunning();
-    }
-    if (transition.retention() == Retention.REMOVE && jobs.remove(job.id(), job)) {
-      releaseRetained();
-    }
-
-    JobSnapshot snapshot = null;
-    RuntimeException materializationFailure = null;
-    try {
-      if (transition.response() != null) {
-        snapshot = materialize(transition.response());
-      }
-    } catch (RuntimeException e) {
-      materializationFailure = e;
-    } finally {
-      closeExecution(transition.executionToClose());
-      closeTask(transition.taskToClose());
-    }
-
-    if (transition.response() != null) {
-      if (materializationFailure == null) {
-        responseListener.onResponse(snapshot);
-      } else {
-        if (transition.retention() == Retention.RETAIN) {
+    switch (transition) {
+      case Transition.Retained retained -> {
+        try {
+          responseListener.onResponse(materialize(retained.response()));
+        } catch (RuntimeException e) {
           applyRemoval(job, job.abort());
+          responseListener.onFailure(e);
         }
-        responseListener.onFailure(materializationFailure);
+      }
+      case Transition.DirectResponse direct -> {
+        releaseRunning();
+        if (jobs.remove(job.id(), job)) {
+          releaseRetained();
+        }
+        try {
+          responseListener.onResponse(materialize(direct.response()));
+        } catch (RuntimeException e) {
+          responseListener.onFailure(e);
+        } finally {
+          direct.executionToClose().ifPresent(PPLAsyncQueryService::closeExecution);
+          closeTask(direct.task());
+        }
+      }
+      case Transition.ExecutionFinished finished -> {
+        releaseRunning();
+        finished.executionToClose().ifPresent(PPLAsyncQueryService::closeExecution);
+        closeTask(finished.task());
       }
     }
   }
@@ -544,34 +605,51 @@ private void applyRemoval(PPLAsyncQueryJob job, Removal removal) {
     if (removal == null) {
       return;
     }
+    switch (removal) {
+      case Removal.Deleted deleted -> applyRemoval(job, deleted.resources(), deleted.reason());
+      case Removal.Expired expired -> applyRemoval(job, expired.resources(), expired.reason());
+      case Removal.Discarded discarded ->
+          applyRemoval(job, discarded.resources(), discarded.reason());
+    }
+  }
+
+  private void applyRemoval(
+      PPLAsyncQueryJob job, PPLAsyncQueryJob.DetachedResources resources, String reason) {
     if (jobs.remove(job.id(), job)) {
-      if (removal.releasesRunningSlot()) {
+      if (resources.releasesRunningSlot()) {
         releaseRunning();
       }
       releaseRetained();
     }
-    cancel(removal.task(), removal.reason());
-    closeExecution(removal.execution());
+    cancel(resources.task(), reason);
+    closeExecution(resources.execution());
   }
 
-  private JobSnapshot materialize(ResponseContext context) {
-    QueryResponse response = null;
-    if (context.status() == Status.SUCCEEDED || context.status() == Status.RUNNING) {
-      response =
-          context.execution() == null
-              ? null
-              : context
+  private JobSnapshot materialize(SnapshotSource source) {
+    return switch (source) {
+      case SnapshotSource.Running running ->
+          new JobSnapshot.Running(
+              running.id(),
+              running
+                  .execution()
+                  .flatMap(AsyncQueryExecution::currentResult)
+                  .map(PPLAsyncQueryService::snapshotResponse));
+      case SnapshotSource.Succeeded succeeded ->
+          new JobSnapshot.Succeeded(
+              succeeded.id(),
+              succeeded
                   .execution()
                   .currentResult()
                   .map(PPLAsyncQueryService::snapshotResponse)
-                  .orElse(null);
-    }
-    if (context.status() == Status.SUCCEEDED && response == null) {
-      throw new IllegalStateException(
-          "Successful PPL asynchronous execution completed without a final result");
-    }
-    return new JobSnapshot(
-        context.id(), context.status(), response, context.failure(), context.tookMillis());
+                  .orElseThrow(
+                      () ->
+                          new IllegalStateException(
+                              "Successful PPL asynchronous execution completed without a final"
+                                  + " result")),
+              succeeded.tookMillis());
+      case SnapshotSource.Failed failed ->
+          new JobSnapshot.Failed(failed.id(), failed.failure(), failed.tookMillis());
+    };
   }
 
   private static void closeExecution(AsyncQueryExecution execution) {
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
index 97bc6306f99..70ba31c8739 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
@@ -8,10 +8,12 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 
+import java.util.Optional;
 import org.junit.Test;
 import org.opensearch.ResourceNotFoundException;
 import org.opensearch.common.unit.TimeValue;
@@ -19,8 +21,7 @@
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.GetResult;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.JobTask;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Removal;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.ResponseContext;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Retention;
+import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.SnapshotSource;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryService.Failure;
 import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryService.Status;
@@ -41,11 +42,7 @@ public void retainPublishesRunningResponseWithId() {
     Transition transition = fixture.job().retain(RETAINED_TIME);
 
     assertEquals(
-        new Transition(
-            new ResponseContext(ID, Status.RUNNING, fixture.execution(), null, -1),
-            Retention.RETAIN,
-            null,
-            null),
+        new Transition.Retained(new SnapshotSource.Running(ID, Optional.of(fixture.execution()))),
         transition);
     assertNull(fixture.job().retain(RETAINED_TIME + 1));
   }
@@ -57,15 +54,14 @@ public void completionBeforeRetentionReturnsDirectResponseAndRemovesJob() {
 
     Transition transition = fixture.job().complete(COMPLETION_TIME);
 
+    assertTrue(transition instanceof Transition.DirectResponse);
+    Transition.DirectResponse direct = (Transition.DirectResponse) transition;
     assertEquals(
-        new Transition(
-            new ResponseContext(
-                null, Status.SUCCEEDED, fixture.execution(), null, COMPLETION_TIME - START_TIME),
-            Retention.REMOVE,
-            fixture.execution(),
-            fixture.task()),
-        transition);
-    assertTrue(transition.releasesRunningSlot());
+        new SnapshotSource.Succeeded(
+            Optional.empty(), fixture.execution(), COMPLETION_TIME - START_TIME),
+        direct.response());
+    assertSame(fixture.task(), direct.task());
+    assertEquals(Optional.of(fixture.execution()), direct.executionToClose());
     assertNull(fixture.job().complete(COMPLETION_TIME + 1));
     assertThrows(ResourceNotFoundException.class, () -> fixture.job().get(COMPLETION_TIME, null));
   }
@@ -78,15 +74,13 @@ public void failureBeforeRetentionReturnsDirectResponseAndRemovesJob() {
 
     Transition transition = fixture.job().fail(failure, COMPLETION_TIME);
 
+    assertTrue(transition instanceof Transition.DirectResponse);
+    Transition.DirectResponse direct = (Transition.DirectResponse) transition;
     assertEquals(
-        new Transition(
-            new ResponseContext(
-                null, Status.FAILED, fixture.execution(), failure, COMPLETION_TIME - START_TIME),
-            Retention.REMOVE,
-            fixture.execution(),
-            fixture.task()),
-        transition);
-    assertTrue(transition.releasesRunningSlot());
+        new SnapshotSource.Failed(Optional.empty(), failure, COMPLETION_TIME - START_TIME),
+        direct.response());
+    assertSame(fixture.task(), direct.task());
+    assertEquals(Optional.of(fixture.execution()), direct.executionToClose());
   }
 
   @Test
@@ -95,20 +89,22 @@ public void retainedSuccessRemainsReadableUntilDeleted() {
 
     Transition transition = fixture.job().complete(COMPLETION_TIME);
 
-    assertEquals(new Transition(null, Retention.RETAIN, null, fixture.task()), transition);
-    assertEquals(
-        new GetResult.Found(
-            new ResponseContext(
-                ID, Status.SUCCEEDED, fixture.execution(), null, COMPLETION_TIME - START_TIME)),
-        fixture.job().get(COMPLETION_TIME, null));
+    assertEquals(new Transition.ExecutionFinished(fixture.task(), Optional.empty()), transition);
+
+    GetResult result = fixture.job().get(COMPLETION_TIME, null);
+    assertTrue(result instanceof GetResult.Found);
     assertEquals(
-        new Removal(
-            Status.SUCCEEDED,
-            null,
-            fixture.execution(),
-            "PPL asynchronous query cancelled by user",
-            false),
-        fixture.job().delete(COMPLETION_TIME));
+        new SnapshotSource.Succeeded(
+            Optional.of(ID), fixture.execution(), COMPLETION_TIME - START_TIME),
+        ((GetResult.Found) result).response());
+
+    Removal removal = fixture.job().delete(COMPLETION_TIME);
+    assertTrue(removal instanceof Removal.Deleted);
+    Removal.Deleted deleted = (Removal.Deleted) removal;
+    assertEquals(Status.SUCCEEDED, deleted.responseStatus());
+    assertNull(deleted.resources().task());
+    assertSame(fixture.execution(), deleted.resources().execution());
+    assertEquals("PPL asynchronous query cancelled by user", deleted.reason());
   }
 
   @Test
@@ -119,11 +115,14 @@ public void retainedFailureDetachesExecutionAndRemainsReadable() {
     Transition transition = fixture.job().fail(failure, COMPLETION_TIME);
 
     assertEquals(
-        new Transition(null, Retention.RETAIN, fixture.execution(), fixture.task()), transition);
+        new Transition.ExecutionFinished(fixture.task(), Optional.of(fixture.execution())),
+        transition);
+
+    GetResult result = fixture.job().get(COMPLETION_TIME, null);
+    assertTrue(result instanceof GetResult.Found);
     assertEquals(
-        new GetResult.Found(
-            new ResponseContext(ID, Status.FAILED, null, failure, COMPLETION_TIME - START_TIME)),
-        fixture.job().get(COMPLETION_TIME, null));
+        new SnapshotSource.Failed(Optional.of(ID), failure, COMPLETION_TIME - START_TIME),
+        ((GetResult.Found) result).response());
   }
 
   @Test
@@ -134,15 +133,11 @@ public void getWithoutKeepAlivePreservesExpiration() {
     assertTrue(fixture.job().get(expirationTime - 1, null) instanceof GetResult.Found);
     GetResult result = fixture.job().get(expirationTime, null);
 
-    assertEquals(
-        new GetResult.Expired(
-            new Removal(
-                Status.RUNNING,
-                fixture.task(),
-                fixture.execution(),
-                "PPL asynchronous query expired",
-                true)),
-        result);
+    assertTrue(result instanceof GetResult.Expired);
+    Removal.Expired removal = ((GetResult.Expired) result).removal();
+    assertSame(fixture.task(), removal.resources().task());
+    assertSame(fixture.execution(), removal.resources().execution());
+    assertEquals("PPL asynchronous query expired", removal.reason());
   }
 
   @Test
@@ -167,14 +162,12 @@ public void expirationStartsWhenJobIsRetained() {
     assertNull(fixture.job().expire(retainedTime - 1));
     fixture.job().retain(retainedTime);
     assertNull(fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS - 1));
-    assertEquals(
-        new Removal(
-            Status.RUNNING,
-            fixture.task(),
-            fixture.execution(),
-            "PPL asynchronous query expired",
-            true),
-        fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS));
+    Removal removal = fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS);
+    assertTrue(removal instanceof Removal.Expired);
+    Removal.Expired expired = (Removal.Expired) removal;
+    assertSame(fixture.task(), expired.resources().task());
+    assertSame(fixture.execution(), expired.resources().execution());
+    assertEquals("PPL asynchronous query expired", expired.reason());
   }
 
   @Test
@@ -183,15 +176,13 @@ public void deleteRunningJobReturnsCancellationAndDetachesResources() {
 
     Removal removal = fixture.job().delete(RETAINED_TIME + 1);
 
-    assertEquals(
-        new Removal(
-            Status.CANCELLED,
-            fixture.task(),
-            fixture.execution(),
-            "PPL asynchronous query cancelled by user",
-            false),
-        removal);
-    assertTrue(removal.releasesRunningSlot());
+    assertTrue(removal instanceof Removal.Deleted);
+    Removal.Deleted deleted = (Removal.Deleted) removal;
+    assertEquals(Status.CANCELLED, deleted.responseStatus());
+    assertSame(fixture.task(), deleted.resources().task());
+    assertSame(fixture.execution(), deleted.resources().execution());
+    assertEquals("PPL asynchronous query cancelled by user", deleted.reason());
+    assertTrue(deleted.resources().releasesRunningSlot());
     assertThrows(ResourceNotFoundException.class, () -> fixture.job().delete(RETAINED_TIME + 2));
   }
 
@@ -200,14 +191,12 @@ public void abortDetachesResourcesOnlyOnce() {
     JobFixture fixture = newJob();
     assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
 
-    assertEquals(
-        new Removal(
-            Status.RUNNING,
-            fixture.task(),
-            fixture.execution(),
-            "PPL asynchronous query startup failed",
-            false),
-        fixture.job().abort());
+    Removal removal = fixture.job().abort();
+    assertTrue(removal instanceof Removal.Discarded);
+    Removal.Discarded discarded = (Removal.Discarded) removal;
+    assertSame(fixture.task(), discarded.resources().task());
+    assertSame(fixture.execution(), discarded.resources().execution());
+    assertEquals("PPL asynchronous query startup failed", discarded.reason());
     assertNull(fixture.job().abort());
   }
 
@@ -215,9 +204,12 @@ public void abortDetachesResourcesOnlyOnce() {
   public void closeDetachesResourcesOnlyOnce() {
     JobFixture fixture = retainedJob();
 
-    assertEquals(
-        new Removal(Status.RUNNING, fixture.task(), fixture.execution(), "service closing", false),
-        fixture.job().close("service closing"));
+    Removal removal = fixture.job().close("service closing");
+    assertTrue(removal instanceof Removal.Discarded);
+    Removal.Discarded discarded = (Removal.Discarded) removal;
+    assertSame(fixture.task(), discarded.resources().task());
+    assertSame(fixture.execution(), discarded.resources().execution());
+    assertEquals("service closing", discarded.reason());
     assertNull(fixture.job().close("service closing"));
   }
 
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
index 05872c706a1..6ee65ef7bf8 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
@@ -83,8 +83,7 @@ public void fastSuccessReturnsDirectResultWithoutRetainingJob() {
 
     assertEquals(1, responses.get());
     assertEquals(
-        new PPLAsyncQueryService.JobSnapshot(
-            null, PPLAsyncQueryService.Status.SUCCEEDED, response(2), null, 25),
+        new PPLAsyncQueryService.JobSnapshot.Succeeded(Optional.empty(), response(2), 25),
         result.get());
     assertEquals(0, service.runningQueryCount());
     assertEquals(0, service.retainedJobCount());
@@ -105,11 +104,9 @@ public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() {
 
     timeoutTask.get().run();
 
-    String id = retainedResponse.get().id();
+    String id = retainedResponse.get().id().orElseThrow();
     assertEquals(
-        new PPLAsyncQueryService.JobSnapshot(
-            id, PPLAsyncQueryService.Status.RUNNING, null, null, -1),
-        retainedResponse.get());
+        new PPLAsyncQueryService.JobSnapshot.Running(id, Optional.empty()), retainedResponse.get());
     assertEquals(1, service.runningQueryCount());
     assertEquals(1, service.retainedJobCount());
 
@@ -118,8 +115,7 @@ public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() {
     PPLAsyncQueryService.JobSnapshot completed = service.get(id, OWNER, null);
 
     assertEquals(
-        new PPLAsyncQueryService.JobSnapshot(
-            id, PPLAsyncQueryService.Status.SUCCEEDED, response(2), null, 25),
+        new PPLAsyncQueryService.JobSnapshot.Succeeded(Optional.of(id), response(2), 25),
         completed);
     assertEquals(0, service.runningQueryCount());
     assertEquals(1, service.retainedJobCount());
@@ -346,7 +342,7 @@ public void deleteCancelsAndReleasesServiceOwnedTask() {
         request,
         registration.requestTask(),
         ignored -> new TrackingExecution(null),
-        listener(snapshot -> id.set(snapshot.id())));
+        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
 
     service.delete(id.get(), OWNER);
 
@@ -442,7 +438,7 @@ public void rejectsUnauthorizedCallerWithoutRenewingOrDeleting() {
         TimeValue.ZERO,
         jobTask(null),
         ignored -> new TrackingExecution(null),
-        listener(snapshot -> id.set(snapshot.id())));
+        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
 
     assertThrows(OpenSearchSecurityException.class, () -> service.get(id.get(), otherUser, null));
     assertThrows(OpenSearchSecurityException.class, () -> service.delete(id.get(), otherUser));
@@ -517,7 +513,10 @@ public void finalSnapshotDefensivelyCopiesRows() {
     execution.complete();
     rows.add(ExprValueUtils.stringValue("second"));
 
-    assertEquals(1, result.get().response().getResults().size());
+    assertTrue(result.get() instanceof PPLAsyncQueryService.JobSnapshot.Succeeded);
+    PPLAsyncQueryService.JobSnapshot.Succeeded succeeded =
+        (PPLAsyncQueryService.JobSnapshot.Succeeded) result.get();
+    assertEquals(1, succeeded.response().getResults().size());
   }
 
   @Test
@@ -529,8 +528,7 @@ public void completedExecutionCanBeAttachedBeforeCompletionIsObserved() {
     startQuery(service, null, TimeValue.timeValueSeconds(5), execution, listener(result::set));
 
     assertEquals(
-        new PPLAsyncQueryService.JobSnapshot(
-            null, PPLAsyncQueryService.Status.SUCCEEDED, response(2), null, 0),
+        new PPLAsyncQueryService.JobSnapshot.Succeeded(Optional.empty(), response(2), 0),
         result.get());
     assertEquals(1, execution.closes.get());
   }
@@ -544,14 +542,9 @@ public void runningGetMaterializesCurrentResultOutsideJob() {
     execution.setCurrent(response(3));
     PPLAsyncQueryService.JobSnapshot second = service.get(id, OWNER, null);
 
+    assertEquals(new PPLAsyncQueryService.JobSnapshot.Running(id, Optional.of(response(1))), first);
     assertEquals(
-        new PPLAsyncQueryService.JobSnapshot(
-            id, PPLAsyncQueryService.Status.RUNNING, response(1), null, -1),
-        first);
-    assertEquals(
-        new PPLAsyncQueryService.JobSnapshot(
-            id, PPLAsyncQueryService.Status.RUNNING, response(3), null, -1),
-        second);
+        new PPLAsyncQueryService.JobSnapshot.Running(id, Optional.of(response(3))), second);
   }
 
   @Test
@@ -567,10 +560,8 @@ public void failedRetainedJobReturnsNoProvisionalRowsAndClosesExecution() {
       PPLAsyncQueryService.JobSnapshot failed = service.get(id, OWNER, null);
 
       assertEquals(
-          new PPLAsyncQueryService.JobSnapshot(
-              id,
-              PPLAsyncQueryService.Status.FAILED,
-              null,
+          new PPLAsyncQueryService.JobSnapshot.Failed(
+              Optional.of(id),
               new PPLAsyncQueryService.Failure("IllegalStateException", "boom"),
               0),
           failed);
@@ -612,7 +603,7 @@ public void deleteBeforeExecutionAttachmentClosesLateHandle() {
           service.delete(id.get(), OWNER);
           return execution;
         },
-        listener(snapshot -> id.set(snapshot.id())));
+        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
 
     assertEquals(1, execution.closes.get());
     assertThrows(ResourceNotFoundException.class, () -> service.get(id.get(), OWNER, null));
@@ -731,7 +722,7 @@ private String startRetainedQuery(
         task,
         TimeValue.ZERO,
         execution,
-        listener(snapshot -> id.set(snapshot.id())));
+        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
     return id.get();
   }
 

From 557108fa59f1d57572ce42fc006485ff32cbeed3 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Thu, 24 Sep 2026 21:10:09 +0000
Subject: [PATCH 07/15] Clarify async job state assertions

Signed-off-by: Peng Huo 
---
 .../asyncquery/PPLAsyncQueryJobTest.java         | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
index 70ba31c8739..6ec88cbb0e0 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
@@ -41,6 +41,7 @@ public void retainPublishesRunningResponseWithId() {
 
     Transition transition = fixture.job().retain(RETAINED_TIME);
 
+    // Retention publishes the job ID once; later retention events are ignored.
     assertEquals(
         new Transition.Retained(new SnapshotSource.Running(ID, Optional.of(fixture.execution()))),
         transition);
@@ -54,6 +55,7 @@ public void completionBeforeRetentionReturnsDirectResponseAndRemovesJob() {
 
     Transition transition = fixture.job().complete(COMPLETION_TIME);
 
+    // A direct response removes the job: duplicate completion is ignored and GET returns not found.
     assertTrue(transition instanceof Transition.DirectResponse);
     Transition.DirectResponse direct = (Transition.DirectResponse) transition;
     assertEquals(
@@ -74,6 +76,7 @@ public void failureBeforeRetentionReturnsDirectResponseAndRemovesJob() {
 
     Transition transition = fixture.job().fail(failure, COMPLETION_TIME);
 
+    // A failure before retention returns without an ID and detaches all execution resources.
     assertTrue(transition instanceof Transition.DirectResponse);
     Transition.DirectResponse direct = (Transition.DirectResponse) transition;
     assertEquals(
@@ -89,6 +92,7 @@ public void retainedSuccessRemainsReadableUntilDeleted() {
 
     Transition transition = fixture.job().complete(COMPLETION_TIME);
 
+    // Retained success closes the task but keeps the execution readable until DELETE.
     assertEquals(new Transition.ExecutionFinished(fixture.task(), Optional.empty()), transition);
 
     GetResult result = fixture.job().get(COMPLETION_TIME, null);
@@ -114,6 +118,7 @@ public void retainedFailureDetachesExecutionAndRemainsReadable() {
 
     Transition transition = fixture.job().fail(failure, COMPLETION_TIME);
 
+    // Retained failure closes its execution and keeps only the failure snapshot readable.
     assertEquals(
         new Transition.ExecutionFinished(fixture.task(), Optional.of(fixture.execution())),
         transition);
@@ -133,6 +138,7 @@ public void getWithoutKeepAlivePreservesExpiration() {
     assertTrue(fixture.job().get(expirationTime - 1, null) instanceof GetResult.Found);
     GetResult result = fixture.job().get(expirationTime, null);
 
+    // Without renewal, the lease remains valid before but not at its expiration boundary.
     assertTrue(result instanceof GetResult.Expired);
     Removal.Expired removal = ((GetResult.Expired) result).removal();
     assertSame(fixture.task(), removal.resources().task());
@@ -146,6 +152,7 @@ public void getWithKeepAliveReplacesExpiration() {
     long renewalTime = RETAINED_TIME + 500;
     TimeValue requestedKeepAlive = TimeValue.timeValueSeconds(1);
 
+    // Renewal replaces the lease and the job expires exactly one requested keep-alive later.
     assertTrue(fixture.job().get(renewalTime, requestedKeepAlive) instanceof GetResult.Found);
     assertTrue(fixture.job().get(renewalTime + 999, null) instanceof GetResult.Found);
     assertTrue(
@@ -159,6 +166,7 @@ public void expirationStartsWhenJobIsRetained() {
     assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
     long retainedTime = START_TIME + KEEP_ALIVE_MILLIS + 1;
 
+    // Keep-alive starts when the ID is retained, not when the unretained job is created.
     assertNull(fixture.job().expire(retainedTime - 1));
     fixture.job().retain(retainedTime);
     assertNull(fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS - 1));
@@ -176,6 +184,8 @@ public void deleteRunningJobReturnsCancellationAndDetachesResources() {
 
     Removal removal = fixture.job().delete(RETAINED_TIME + 1);
 
+    // DELETE cancels a running job, transfers its resources, and makes later DELETE return not
+    // found.
     assertTrue(removal instanceof Removal.Deleted);
     Removal.Deleted deleted = (Removal.Deleted) removal;
     assertEquals(Status.CANCELLED, deleted.responseStatus());
@@ -192,6 +202,8 @@ public void abortDetachesResourcesOnlyOnce() {
     assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
 
     Removal removal = fixture.job().abort();
+
+    // Abort transfers resource ownership once; repeated abort calls have no transition.
     assertTrue(removal instanceof Removal.Discarded);
     Removal.Discarded discarded = (Removal.Discarded) removal;
     assertSame(fixture.task(), discarded.resources().task());
@@ -205,6 +217,8 @@ public void closeDetachesResourcesOnlyOnce() {
     JobFixture fixture = retainedJob();
 
     Removal removal = fixture.job().close("service closing");
+
+    // Close transfers resource ownership once; repeated close calls have no transition.
     assertTrue(removal instanceof Removal.Discarded);
     Removal.Discarded discarded = (Removal.Discarded) removal;
     assertSame(fixture.task(), discarded.resources().task());
@@ -218,6 +232,7 @@ public void lateExecutionAttachmentAfterRemovalIsRejected() {
     JobFixture fixture = newJob();
     fixture.job().abort();
 
+    // Removal wins the race, so a late execution handle remains owned by the caller.
     assertFalse(fixture.job().tryAttachExecution(fixture.execution()));
   }
 
@@ -225,6 +240,7 @@ public void lateExecutionAttachmentAfterRemovalIsRejected() {
   public void successfulCompletionRequiresAttachedExecution() {
     JobFixture fixture = newJob();
 
+    // Success cannot be published until an execution handle can provide the final result.
     assertThrows(IllegalStateException.class, () -> fixture.job().complete(COMPLETION_TIME));
   }
 

From 9050a7f02a5b535bf8b1e4f0c7a436e4acc23f29 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Thu, 24 Sep 2026 21:12:58 +0000
Subject: [PATCH 08/15] Verify direct failure removes async job

Signed-off-by: Peng Huo 
---
 .../sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
index 6ec88cbb0e0..0f3ed8cf1b4 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
@@ -84,6 +84,10 @@ public void failureBeforeRetentionReturnsDirectResponseAndRemovesJob() {
         direct.response());
     assertSame(fixture.task(), direct.task());
     assertEquals(Optional.of(fixture.execution()), direct.executionToClose());
+
+    // A direct failure removes the job: later completion is ignored and GET returns not found.
+    assertNull(fixture.job().complete(COMPLETION_TIME + 1));
+    assertThrows(ResourceNotFoundException.class, () -> fixture.job().get(COMPLETION_TIME, null));
   }
 
   @Test

From 83eaff43def5753d5477f75f3a4d9e9b5981ec17 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Thu, 24 Sep 2026 21:32:51 +0000
Subject: [PATCH 09/15] Clarify async job lease boundary test

Signed-off-by: Peng Huo 
---
 .../asyncquery/PPLAsyncQueryJobTest.java      | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
index 0f3ed8cf1b4..2c57d37fa49 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
@@ -155,13 +155,18 @@ public void getWithKeepAliveReplacesExpiration() {
     JobFixture fixture = retainedJob();
     long renewalTime = RETAINED_TIME + 500;
     TimeValue requestedKeepAlive = TimeValue.timeValueSeconds(1);
-
-    // Renewal replaces the lease and the job expires exactly one requested keep-alive later.
-    assertTrue(fixture.job().get(renewalTime, requestedKeepAlive) instanceof GetResult.Found);
-    assertTrue(fixture.job().get(renewalTime + 999, null) instanceof GetResult.Found);
-    assertTrue(
-        fixture.job().get(renewalTime + requestedKeepAlive.millis(), null)
-            instanceof GetResult.Expired);
+    long renewedExpiration = renewalTime + requestedKeepAlive.millis();
+
+    // At 2500ms, GET replaces the original lease with a one-second lease ending at 3500ms.
+    GetResult renewed = fixture.job().get(renewalTime, requestedKeepAlive);
+    // At 3499ms, the renewed lease has not expired.
+    GetResult beforeExpiration = fixture.job().get(renewedExpiration - 1, null);
+    // At 3500ms, the renewed lease expires.
+    GetResult atExpiration = fixture.job().get(renewedExpiration, null);
+
+    assertTrue(renewed instanceof GetResult.Found);
+    assertTrue(beforeExpiration instanceof GetResult.Found);
+    assertTrue(atExpiration instanceof GetResult.Expired);
   }
 
   @Test

From b2beea0ede2ced1bee41da100c734ec5a5b651d6 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Thu, 24 Sep 2026 21:54:30 +0000
Subject: [PATCH 10/15] Use fluent scenarios for async query service tests

Signed-off-by: Peng Huo 
---
 .../asyncquery/PPLAsyncQueryServiceTest.java  | 764 ++++++++++--------
 1 file changed, 407 insertions(+), 357 deletions(-)

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
index 6ee65ef7bf8..3beb9556d44 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
@@ -28,6 +28,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 import org.junit.Test;
 import org.mockito.InOrder;
 import org.opensearch.OpenSearchSecurityException;
@@ -65,127 +66,90 @@ public class PPLAsyncQueryServiceTest {
 
   @Test
   public void fastSuccessReturnsDirectResultWithoutRetainingJob() {
-    AtomicReference result = new AtomicReference<>();
-    AtomicInteger responses = new AtomicInteger();
-    TrackingExecution execution = new TrackingExecution(response(2));
-    startQuery(
-        service,
-        null,
-        TimeValue.timeValueSeconds(5),
-        execution,
-        listener(
-            snapshot -> {
-              result.set(snapshot);
-              responses.incrementAndGet();
-            }));
-    now.addAndGet(25);
-    execution.complete();
-
-    assertEquals(1, responses.get());
-    assertEquals(
-        new PPLAsyncQueryService.JobSnapshot.Succeeded(Optional.empty(), response(2), 25),
-        result.get());
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
-    assertEquals(1, execution.reads.get());
-    assertEquals(1, execution.closes.get());
-    assertTrue(timeoutCancelled.get());
+    AsyncQueryScenario scenario =
+        scenario()
+            .withCurrentResult(response(2))
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start();
+
+    scenario.completeAfterMillis(25);
 
-    timeoutTask.get().run();
-    assertEquals(1, responses.get());
+    scenario
+        .assertDirectSuccess(response(2), 25)
+        .assertNotRetained()
+        .assertExecutionReadOnceAndClosed()
+        .assertTimeoutCancelled();
+
+    scenario.fireTimeout().assertResponseCount(1);
   }
 
   @Test
   public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() {
-    AtomicReference retainedResponse = new AtomicReference<>();
-    TrackingExecution execution = new TrackingExecution(null);
-    startQuery(
-        service, null, TimeValue.timeValueSeconds(5), execution, listener(retainedResponse::set));
-
-    timeoutTask.get().run();
-
-    String id = retainedResponse.get().id().orElseThrow();
-    assertEquals(
-        new PPLAsyncQueryService.JobSnapshot.Running(id, Optional.empty()), retainedResponse.get());
-    assertEquals(1, service.runningQueryCount());
-    assertEquals(1, service.retainedJobCount());
-
-    now.addAndGet(25);
-    execution.succeed(response(2));
-    PPLAsyncQueryService.JobSnapshot completed = service.get(id, OWNER, null);
-
-    assertEquals(
-        new PPLAsyncQueryService.JobSnapshot.Succeeded(Optional.of(id), response(2), 25),
-        completed);
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(1, service.retainedJobCount());
-    assertEquals(0, execution.closes.get());
+    AsyncQueryScenario scenario =
+        scenario()
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start()
+            .fireTimeout()
+            .assertRetainedRunning(Optional.empty())
+            .assertCapacity(1, 1);
+
+    scenario
+        .succeedAfterMillis(25, response(2))
+        .get()
+        .assertGetResponse(
+            new PPLAsyncQueryService.JobSnapshot.Succeeded(
+                Optional.of(scenario.id()), response(2), 25))
+        .assertCapacity(0, 1)
+        .assertExecutionReadCount(2)
+        .assertExecutionCloseCount(0);
   }
 
   @Test
   public void retentionWaitPreventsExpiryAndLeaseStartsWhenIdIsReturned() {
-    AtomicReference retainedResponse = new AtomicReference<>();
-    startQuery(
-        service,
-        null,
-        TimeValue.timeValueSeconds(1),
-        TimeValue.timeValueSeconds(5),
-        new TrackingExecution(null),
-        listener(retainedResponse::set));
-
-    now.addAndGet(TimeValue.timeValueSeconds(2).millis());
-    service.reapExpired();
+    AsyncQueryScenario scenario =
+        scenario()
+            .withKeepAlive(TimeValue.timeValueSeconds(1))
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start();
 
-    assertNull(retainedResponse.get());
-    assertEquals(1, service.runningQueryCount());
-    assertEquals(1, service.retainedJobCount());
+    scenario
+        .advanceMillis(TimeValue.timeValueSeconds(2).millis())
+        .reapExpired()
+        .assertNoInitialResponse()
+        .assertCapacity(1, 1);
 
-    timeoutTask.get().run();
-    assertEquals(PPLAsyncQueryService.Status.RUNNING, retainedResponse.get().status());
+    scenario.fireTimeout().assertRetainedRunning(Optional.empty());
 
-    now.addAndGet(TimeValue.timeValueSeconds(1).millis() + 1);
-    service.reapExpired();
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
+    scenario
+        .advanceMillis(TimeValue.timeValueSeconds(1).millis() + 1)
+        .reapExpired()
+        .assertCapacity(0, 0);
   }
 
   @Test
   public void fastFailureReturnsDirectFailureWithoutId() {
-    AtomicReference failure = new AtomicReference<>();
-    TrackingExecution execution = new TrackingExecution(response(1));
-    startQuery(
-        service,
-        null,
-        TimeValue.timeValueSeconds(5),
-        execution,
-        ActionListener.wrap(
-            ignored -> {
-              throw new AssertionError("Expected direct query failure");
-            },
-            failure::set));
-
-    execution.fail(new IllegalStateException("boom"));
-
-    assertEquals("boom", failure.get().getMessage());
-    assertEquals(0, execution.reads.get());
-    assertEquals(1, execution.closes.get());
-    assertEquals(0, service.retainedJobCount());
+    scenario()
+        .withCurrentResult(response(1))
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .start()
+        .fail(new IllegalStateException("boom"))
+        .assertFailure(IllegalStateException.class, "boom")
+        .assertExecutionReadCount(0)
+        .assertExecutionCloseCount(1)
+        .assertCapacity(0, 0);
   }
 
   @Test
   public void expiredGetCancelsAndRemovesJob() {
     CancellableTask task = mock(CancellableTask.class);
     when(task.isCancelled()).thenReturn(false);
-    TrackingExecution execution = new TrackingExecution(null);
-    String id = startRetainedQuery(service, task, execution);
+    AsyncQueryScenario scenario = scenario().withTask(task).start();
 
-    now.addAndGet(KEEP_ALIVE.millis());
+    scenario.advanceMillis(KEEP_ALIVE.millis());
 
-    assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null));
+    assertThrows(ResourceNotFoundException.class, scenario::get);
     verify(task).cancel("PPL asynchronous query expired");
-    assertEquals(1, execution.closes.get());
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
+    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
   }
 
   @Test
@@ -200,34 +164,23 @@ public void missingLocalJobReturnsNotFound() {
   public void deleteCancelsRunningJobAndReleasesState() {
     CancellableTask task = mock(CancellableTask.class);
     when(task.isCancelled()).thenReturn(false);
-    TrackingExecution execution = new TrackingExecution(null);
-    String id = startRetainedQuery(service, task, execution);
-
-    PPLAsyncQueryService.DeleteResult result = service.delete(id, OWNER);
+    AsyncQueryScenario scenario = scenario().withTask(task).start().delete();
 
-    assertEquals(
-        new PPLAsyncQueryService.DeleteResult(id, PPLAsyncQueryService.Status.CANCELLED), result);
+    scenario.assertDeleteStatus(PPLAsyncQueryService.Status.CANCELLED);
     verify(task).cancel("PPL asynchronous query cancelled by user");
-    assertEquals(1, execution.closes.get());
-    assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null));
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
+    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
+    assertThrows(ResourceNotFoundException.class, scenario::get);
   }
 
   @Test
   public void deleteReturnsExistingTerminalStatus() {
     CancellableTask task = mock(CancellableTask.class);
-    TrackingExecution execution = new TrackingExecution(response(1));
-    String id = startRetainedQuery(service, task, execution);
-    execution.complete();
+    AsyncQueryScenario scenario =
+        scenario().withTask(task).withCurrentResult(response(1)).start().complete().delete();
 
-    PPLAsyncQueryService.DeleteResult result = service.delete(id, OWNER);
-
-    assertEquals(
-        new PPLAsyncQueryService.DeleteResult(id, PPLAsyncQueryService.Status.SUCCEEDED), result);
+    scenario.assertDeleteStatus(PPLAsyncQueryService.Status.SUCCEEDED);
     verify(task, never()).cancel(org.mockito.ArgumentMatchers.anyString());
-    assertEquals(1, execution.closes.get());
-    assertEquals(0, service.retainedJobCount());
+    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
   }
 
   @Test
@@ -236,9 +189,7 @@ public void cancellationUsesTaskManagerWhenAttached() {
     CancellableTask task = mock(CancellableTask.class);
     when(task.isCancelled()).thenReturn(false);
     service.attachTaskManager(taskManager);
-    String id = startRetainedQuery(service, task, new TrackingExecution(null));
-
-    service.delete(id, OWNER);
+    scenario().withTask(task).start().delete();
 
     verify(taskManager)
         .cancelTaskAndDescendants(
@@ -257,16 +208,11 @@ public void startRegistersTaskAndCompletionReleasesIt() {
     RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
     service.attachTaskManager(taskManager);
 
-    TrackingExecution execution = new TrackingExecution(null);
-    service.start(
-        OWNER,
-        KEEP_ALIVE,
-        TimeValue.ZERO,
-        request,
-        registration.requestTask(),
-        ignored -> execution,
-        listener(snapshot -> assertEquals(PPLAsyncQueryService.Status.RUNNING, snapshot.status())));
-    execution.succeed(response(1));
+    AsyncQueryScenario scenario =
+        scenario()
+            .start(request, registration.requestTask())
+            .assertRetainedRunning(Optional.empty())
+            .succeed(response(1));
 
     InOrder registrationOrder = inOrder(taskManager);
     registrationOrder
@@ -276,8 +222,7 @@ public void startRegistersTaskAndCompletionReleasesIt() {
     assertEquals(TaskId.EMPTY_TASK_ID, request.getParentTask());
     verify(taskManager).unregister(task);
     verify(registration.childNodeRegistration()).close();
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(1, service.retainedJobCount());
+    scenario.assertCapacity(0, 1);
   }
 
   @Test
@@ -288,28 +233,21 @@ public void executionStartFailureCompletesJobAndReleasesTask() {
         new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
     RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
     service.attachTaskManager(taskManager);
-    AtomicReference failure = new AtomicReference<>();
-
-    service.start(
-        OWNER,
-        KEEP_ALIVE,
-        TimeValue.timeValueSeconds(5),
-        request,
-        registration.requestTask(),
-        ignored -> {
-          throw new IllegalStateException("execution did not start");
-        },
-        ActionListener.wrap(
-            ignored -> {
-              throw new AssertionError("Expected execution startup failure");
-            },
-            failure::set));
 
-    assertEquals("execution did not start", failure.get().getMessage());
+    AsyncQueryScenario scenario =
+        scenario()
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .withExecutionStarter(
+                ignored -> {
+                  throw new IllegalStateException("execution did not start");
+                })
+            .start(request, registration.requestTask());
+
+    scenario
+        .assertFailure(IllegalStateException.class, "execution did not start")
+        .assertCapacity(0, 0);
     verify(taskManager, times(1)).unregister(task);
     verify(registration.childNodeRegistration()).close();
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
   }
 
   @Test
@@ -333,18 +271,8 @@ public void deleteCancelsAndReleasesServiceOwnedTask() {
             org.mockito.ArgumentMatchers.eq(false),
             org.mockito.ArgumentMatchers.any());
     service.attachTaskManager(taskManager);
-    AtomicReference id = new AtomicReference<>();
-
-    service.start(
-        OWNER,
-        KEEP_ALIVE,
-        TimeValue.ZERO,
-        request,
-        registration.requestTask(),
-        ignored -> new TrackingExecution(null),
-        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
 
-    service.delete(id.get(), OWNER);
+    scenario().start(request, registration.requestTask()).delete();
 
     verify(taskManager)
         .cancelTaskAndDescendants(
@@ -376,25 +304,19 @@ public void startupFailureAfterJobCreationReleasesServiceOwnedTask() {
         new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
     RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
     abortingService.attachTaskManager(taskManager);
+    AsyncQueryScenario scenario =
+        scenario()
+            .withService(abortingService)
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5));
 
     IllegalStateException failure =
         assertThrows(
-            IllegalStateException.class,
-            () ->
-                abortingService.start(
-                    OWNER,
-                    KEEP_ALIVE,
-                    TimeValue.timeValueSeconds(5),
-                    request,
-                    registration.requestTask(),
-                    ignored -> new TrackingExecution(null),
-                    listener(snapshot -> {})));
+            IllegalStateException.class, () -> scenario.start(request, registration.requestTask()));
 
     assertEquals("scheduler unavailable", failure.getMessage());
     verify(taskManager, times(1)).unregister(task);
     verify(registration.childNodeRegistration()).close();
-    assertEquals(0, abortingService.runningQueryCount());
-    assertEquals(0, abortingService.retainedJobCount());
+    scenario.assertCapacity(0, 0);
   }
 
   @Test
@@ -409,65 +331,36 @@ public void childTrackingFailurePreventsRetainedTaskRegistration() {
     when(taskManager.registerChildNode(42L, localNode))
         .thenThrow(new IllegalStateException("channel closed"));
     service.attachTaskManager(taskManager);
+    AsyncQueryScenario scenario = scenario().withWaitForCompletion(TimeValue.timeValueSeconds(5));
 
-    assertThrows(
-        IllegalStateException.class,
-        () ->
-            service.start(
-                OWNER,
-                KEEP_ALIVE,
-                TimeValue.timeValueSeconds(5),
-                request,
-                requestTask,
-                ignored -> new TrackingExecution(null),
-                listener(snapshot -> {})));
+    assertThrows(IllegalStateException.class, () -> scenario.start(request, requestTask));
 
     verify(taskManager, never()).register("transport", PPLQueryAction.NAME, request);
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
+    scenario.assertCapacity(0, 0);
   }
 
   @Test
   public void rejectsUnauthorizedCallerWithoutRenewingOrDeleting() {
     PPLAsyncQueryUser securedOwner = new PPLAsyncQueryUser("alice", "tenant", List.of("role-a"));
     PPLAsyncQueryUser otherUser = new PPLAsyncQueryUser("bob", "tenant", List.of("role-a"));
-    AtomicReference id = new AtomicReference<>();
-    service.start(
-        securedOwner,
-        KEEP_ALIVE,
-        TimeValue.ZERO,
-        jobTask(null),
-        ignored -> new TrackingExecution(null),
-        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
-
-    assertThrows(OpenSearchSecurityException.class, () -> service.get(id.get(), otherUser, null));
-    assertThrows(OpenSearchSecurityException.class, () -> service.delete(id.get(), otherUser));
+    AsyncQueryScenario scenario = scenario().withOwner(securedOwner).start();
+
+    assertThrows(
+        OpenSearchSecurityException.class, () -> service.get(scenario.id(), otherUser, null));
+    assertThrows(OpenSearchSecurityException.class, () -> service.delete(scenario.id(), otherUser));
     assertEquals(
-        PPLAsyncQueryService.Status.RUNNING, service.get(id.get(), securedOwner, null).status());
+        PPLAsyncQueryService.Status.RUNNING,
+        service.get(scenario.id(), securedOwner, null).status());
   }
 
   @Test
   public void enforcesRunningAndRetainedCapacity() {
     PPLAsyncQueryService limited = service(1, 1);
-    limited.start(
-        OWNER,
-        KEEP_ALIVE,
-        TimeValue.ZERO,
-        jobTask(null),
-        ignored -> new TrackingExecution(null),
-        listener(snapshot -> {}));
+    scenario().withService(limited).start();
 
     OpenSearchStatusException exception =
         assertThrows(
-            OpenSearchStatusException.class,
-            () ->
-                limited.start(
-                    OWNER,
-                    KEEP_ALIVE,
-                    TimeValue.ZERO,
-                    jobTask(null),
-                    ignored -> new TrackingExecution(null),
-                    listener(snapshot -> {})));
+            OpenSearchStatusException.class, () -> scenario().withService(limited).start());
 
     assertEquals(429, exception.status().getStatus());
   }
@@ -483,68 +376,54 @@ public void createFailureReleasesReservedCapacity() {
             () -> 1,
             () -> TimeValue.timeValueSeconds(60),
             () -> TimeValue.timeValueHours(24));
+    AsyncQueryScenario scenario = scenario().withService(missingOwnerNode);
 
-    assertThrows(
-        NullPointerException.class,
-        () ->
-            missingOwnerNode.start(
-                OWNER,
-                KEEP_ALIVE,
-                TimeValue.ZERO,
-                jobTask(null),
-                ignored -> new TrackingExecution(null),
-                listener(snapshot -> {})));
+    assertThrows(NullPointerException.class, scenario::start);
 
-    assertEquals(0, missingOwnerNode.runningQueryCount());
-    assertEquals(0, missingOwnerNode.retainedJobCount());
+    scenario.assertCapacity(0, 0);
   }
 
   @Test
   public void finalSnapshotDefensivelyCopiesRows() {
-    AtomicReference result = new AtomicReference<>();
     List rows = new ArrayList<>();
     rows.add(ExprValueUtils.stringValue("first"));
     QueryResponse response =
         new QueryResponse(
             new Schema(List.of(new Column("state", null, ExprCoreType.STRING))), rows, null);
 
-    TrackingExecution execution = new TrackingExecution(response);
-    startQuery(service, null, TimeValue.timeValueSeconds(5), execution, listener(result::set));
-    execution.complete();
+    AsyncQueryScenario scenario =
+        scenario()
+            .withCurrentResult(response)
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start()
+            .complete();
     rows.add(ExprValueUtils.stringValue("second"));
 
-    assertTrue(result.get() instanceof PPLAsyncQueryService.JobSnapshot.Succeeded);
-    PPLAsyncQueryService.JobSnapshot.Succeeded succeeded =
-        (PPLAsyncQueryService.JobSnapshot.Succeeded) result.get();
-    assertEquals(1, succeeded.response().getResults().size());
+    scenario.assertSucceededRowCount(1);
   }
 
   @Test
   public void completedExecutionCanBeAttachedBeforeCompletionIsObserved() {
-    AtomicReference result = new AtomicReference<>();
-    TrackingExecution execution = new TrackingExecution(null);
-
-    execution.succeed(response(2));
-    startQuery(service, null, TimeValue.timeValueSeconds(5), execution, listener(result::set));
-
-    assertEquals(
-        new PPLAsyncQueryService.JobSnapshot.Succeeded(Optional.empty(), response(2), 0),
-        result.get());
-    assertEquals(1, execution.closes.get());
+    scenario()
+        .succeed(response(2))
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .start()
+        .assertDirectSuccess(response(2), 0)
+        .assertExecutionReadOnceAndClosed();
   }
 
   @Test
   public void runningGetMaterializesCurrentResultOutsideJob() {
-    TrackingExecution execution = new TrackingExecution(response(1));
-    String id = startRetainedQuery(service, null, execution);
+    AsyncQueryScenario scenario = scenario().withCurrentResult(response(1)).start().get();
 
-    PPLAsyncQueryService.JobSnapshot first = service.get(id, OWNER, null);
-    execution.setCurrent(response(3));
-    PPLAsyncQueryService.JobSnapshot second = service.get(id, OWNER, null);
+    scenario.assertGetResponse(
+        new PPLAsyncQueryService.JobSnapshot.Running(scenario.id(), Optional.of(response(1))));
 
-    assertEquals(new PPLAsyncQueryService.JobSnapshot.Running(id, Optional.of(response(1))), first);
-    assertEquals(
-        new PPLAsyncQueryService.JobSnapshot.Running(id, Optional.of(response(3))), second);
+    scenario
+        .withCurrentResult(response(3))
+        .get()
+        .assertGetResponse(
+            new PPLAsyncQueryService.JobSnapshot.Running(scenario.id(), Optional.of(response(3))));
   }
 
   @Test
@@ -553,20 +432,21 @@ public void failedRetainedJobReturnsNoProvisionalRowsAndClosesExecution() {
         new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_SYS.getName(), new BasicCounter());
     Metrics.getInstance().registerMetric(failures);
     try {
-      TrackingExecution execution = new TrackingExecution(response(1));
-      String id = startRetainedQuery(service, null, execution);
-
-      execution.fail(new IllegalStateException("boom"));
-      PPLAsyncQueryService.JobSnapshot failed = service.get(id, OWNER, null);
-
-      assertEquals(
-          new PPLAsyncQueryService.JobSnapshot.Failed(
-              Optional.of(id),
-              new PPLAsyncQueryService.Failure("IllegalStateException", "boom"),
-              0),
-          failed);
-      assertEquals(0, execution.reads.get());
-      assertEquals(1, execution.closes.get());
+      AsyncQueryScenario scenario =
+          scenario()
+              .withCurrentResult(response(1))
+              .start()
+              .fail(new IllegalStateException("boom"))
+              .get();
+
+      scenario
+          .assertGetResponse(
+              new PPLAsyncQueryService.JobSnapshot.Failed(
+                  Optional.of(scenario.id()),
+                  new PPLAsyncQueryService.Failure("IllegalStateException", "boom"),
+                  0))
+          .assertExecutionReadCount(0)
+          .assertExecutionCloseCount(1);
     } finally {
       Metrics.getInstance().unregisterMetric(failures.getName());
     }
@@ -578,10 +458,7 @@ public void failedRetainedJobRecordsFailureMetric() {
         new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_CUS.getName(), new BasicCounter());
     Metrics.getInstance().registerMetric(failures);
     try {
-      TrackingExecution execution = new TrackingExecution(null);
-      startRetainedQuery(service, null, execution);
-
-      execution.fail(new IllegalArgumentException("invalid query"));
+      scenario().start().fail(new IllegalArgumentException("invalid query"));
 
       assertEquals(Long.valueOf(1), failures.getValue());
     } finally {
@@ -592,33 +469,29 @@ public void failedRetainedJobRecordsFailureMetric() {
   @Test
   public void deleteBeforeExecutionAttachmentClosesLateHandle() {
     TrackingExecution execution = new TrackingExecution(response(1));
-    AtomicReference id = new AtomicReference<>();
-
-    service.start(
-        OWNER,
-        KEEP_ALIVE,
-        TimeValue.ZERO,
-        jobTask(null),
-        ignored -> {
-          service.delete(id.get(), OWNER);
-          return execution;
-        },
-        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
+    AsyncQueryScenario scenario = scenario().withExecution(execution);
+    scenario
+        .withExecutionStarter(
+            ignored -> {
+              scenario.delete();
+              return execution;
+            })
+        .start();
 
     assertEquals(1, execution.closes.get());
-    assertThrows(ResourceNotFoundException.class, () -> service.get(id.get(), OWNER, null));
+    assertThrows(ResourceNotFoundException.class, scenario::get);
   }
 
   @Test
   public void concurrentGetDoesNotBlockDeleteOnResultMaterialization() throws Exception {
     BlockingExecution execution = new BlockingExecution(response(1));
-    String id = startRetainedQuery(service, null, execution);
+    AsyncQueryScenario scenario = scenario().withExecution(execution).start();
 
     CompletableFuture get =
-        CompletableFuture.supplyAsync(() -> service.get(id, OWNER, null));
+        CompletableFuture.supplyAsync(() -> service.get(scenario.id(), OWNER, null));
     assertTrue(execution.readStarted.await(5, TimeUnit.SECONDS));
     CompletableFuture delete =
-        CompletableFuture.supplyAsync(() -> service.delete(id, OWNER));
+        CompletableFuture.supplyAsync(() -> service.delete(scenario.id(), OWNER));
 
     try {
       assertEquals(PPLAsyncQueryService.Status.CANCELLED, delete.get(5, TimeUnit.SECONDS).status());
@@ -628,60 +501,46 @@ public void concurrentGetDoesNotBlockDeleteOnResultMaterialization() throws Exce
 
     assertEquals(PPLAsyncQueryService.Status.RUNNING, get.get(5, TimeUnit.SECONDS).status());
     assertEquals(1, execution.closes.get());
-    assertThrows(ResourceNotFoundException.class, () -> service.get(id, OWNER, null));
+    assertThrows(ResourceNotFoundException.class, scenario::get);
   }
 
   @Test
   public void shutdownClosesRetainedExecutionExactlyOnce() throws Exception {
-    TrackingExecution execution = new TrackingExecution(response(1));
-    startRetainedQuery(service, null, execution);
+    AsyncQueryScenario scenario = scenario().withCurrentResult(response(1)).start();
 
     service.close();
     service.close();
 
-    assertEquals(1, execution.closes.get());
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
+    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
   }
 
   @Test
   public void successfulCompletionRequiresFinalResultToBeVisible() {
-    AtomicReference failure = new AtomicReference<>();
-    TrackingExecution execution = new TrackingExecution(null);
-    startQuery(
-        service,
-        null,
-        TimeValue.timeValueSeconds(5),
-        execution,
-        ActionListener.wrap(snapshot -> {}, failure::set));
-
-    execution.complete();
-
-    assertTrue(failure.get() instanceof IllegalStateException);
-    assertEquals(1, execution.closes.get());
-    assertEquals(0, service.retainedJobCount());
+    scenario()
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .start()
+        .complete()
+        .assertFailure(IllegalStateException.class)
+        .assertExecutionReadOnceAndClosed()
+        .assertCapacity(0, 0);
   }
 
   @Test
   public void retentionResponseMaterializationFailureAbortsUndeliverableJob() {
     CancellableTask task = mock(CancellableTask.class);
     when(task.isCancelled()).thenReturn(false);
-    AtomicReference failure = new AtomicReference<>();
     ThrowingExecution execution = new ThrowingExecution();
-    startQuery(
-        service,
-        task,
-        TimeValue.timeValueSeconds(5),
-        execution,
-        ActionListener.wrap(snapshot -> {}, failure::set));
-
-    timeoutTask.get().run();
-
-    assertTrue(failure.get() instanceof IllegalStateException);
+    AsyncQueryScenario scenario =
+        scenario()
+            .withTask(task)
+            .withExecution(execution)
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start()
+            .fireTimeout();
+
+    scenario.assertFailure(IllegalStateException.class).assertCapacity(0, 0);
     verify(task).cancel("PPL asynchronous query startup failed");
     assertEquals(1, execution.closes.get());
-    assertEquals(0, service.runningQueryCount());
-    assertEquals(0, service.retainedJobCount());
   }
 
   @Test
@@ -714,36 +573,8 @@ private PPLAsyncQueryService service(int maxRunning, int maxRetained) {
         () -> TimeValue.timeValueHours(24));
   }
 
-  private String startRetainedQuery(
-      PPLAsyncQueryService targetService, CancellableTask task, AsyncQueryExecution execution) {
-    AtomicReference id = new AtomicReference<>();
-    startQuery(
-        targetService,
-        task,
-        TimeValue.ZERO,
-        execution,
-        listener(snapshot -> id.set(snapshot.id().orElseThrow())));
-    return id.get();
-  }
-
-  private void startQuery(
-      PPLAsyncQueryService targetService,
-      CancellableTask task,
-      TimeValue waitForCompletion,
-      AsyncQueryExecution execution,
-      ActionListener responseListener) {
-    startQuery(targetService, task, KEEP_ALIVE, waitForCompletion, execution, responseListener);
-  }
-
-  private void startQuery(
-      PPLAsyncQueryService targetService,
-      CancellableTask task,
-      TimeValue keepAlive,
-      TimeValue waitForCompletion,
-      AsyncQueryExecution execution,
-      ActionListener responseListener) {
-    targetService.start(
-        OWNER, keepAlive, waitForCompletion, jobTask(task), ignored -> execution, responseListener);
+  private AsyncQueryScenario scenario() {
+    return new AsyncQueryScenario();
   }
 
   private static PPLAsyncQueryJob.JobTask jobTask(CancellableTask task) {
@@ -771,15 +602,6 @@ private static RequestTaskRegistration registerRequestTask(
   private record RequestTaskRegistration(
       PPLQueryTask requestTask, DiscoveryNode localNode, Releasable childNodeRegistration) {}
 
-  private static ActionListener listener(
-      java.util.function.Consumer consumer) {
-    return ActionListener.wrap(
-        snapshot -> consumer.accept(snapshot),
-        failure -> {
-          throw new AssertionError(failure);
-        });
-  }
-
   private static QueryResponse response(int rowCount) {
     Schema schema = new Schema(List.of(new Column("state", null, ExprCoreType.STRING)));
     return new QueryResponse(
@@ -790,6 +612,234 @@ private static QueryResponse response(int rowCount) {
         null);
   }
 
+  private final class AsyncQueryScenario {
+    private PPLAsyncQueryService targetService = service;
+    private PPLAsyncQueryUser owner = OWNER;
+    private TimeValue keepAlive = KEEP_ALIVE;
+    private TimeValue waitForCompletion = TimeValue.ZERO;
+    private CancellableTask task;
+    private TrackingExecution trackingExecution = new TrackingExecution(null);
+    private AsyncQueryExecution execution = trackingExecution;
+    private Function executionStarter = ignored -> execution;
+    private final AtomicReference initialResponse =
+        new AtomicReference<>();
+    private final AtomicReference failure = new AtomicReference<>();
+    private final AtomicInteger responseCount = new AtomicInteger();
+    private PPLAsyncQueryService.JobSnapshot getResponse;
+    private PPLAsyncQueryService.DeleteResult deleteResponse;
+
+    private AsyncQueryScenario withService(PPLAsyncQueryService service) {
+      targetService = service;
+      return this;
+    }
+
+    private AsyncQueryScenario withOwner(PPLAsyncQueryUser owner) {
+      this.owner = owner;
+      return this;
+    }
+
+    private AsyncQueryScenario withKeepAlive(TimeValue keepAlive) {
+      this.keepAlive = keepAlive;
+      return this;
+    }
+
+    private AsyncQueryScenario withWaitForCompletion(TimeValue waitForCompletion) {
+      this.waitForCompletion = waitForCompletion;
+      return this;
+    }
+
+    private AsyncQueryScenario withTask(CancellableTask task) {
+      this.task = task;
+      return this;
+    }
+
+    private AsyncQueryScenario withCurrentResult(QueryResponse response) {
+      trackingExecution.setCurrent(response);
+      return this;
+    }
+
+    private AsyncQueryScenario withExecution(AsyncQueryExecution execution) {
+      this.execution = execution;
+      trackingExecution = execution instanceof TrackingExecution tracking ? tracking : null;
+      executionStarter = ignored -> this.execution;
+      return this;
+    }
+
+    private AsyncQueryScenario withExecutionStarter(
+        Function executionStarter) {
+      this.executionStarter = executionStarter;
+      return this;
+    }
+
+    private AsyncQueryScenario start() {
+      targetService.start(
+          owner, keepAlive, waitForCompletion, jobTask(task), executionStarter, responseListener());
+      return this;
+    }
+
+    private AsyncQueryScenario start(TransportPPLQueryRequest request, PPLQueryTask requestTask) {
+      targetService.start(
+          owner,
+          keepAlive,
+          waitForCompletion,
+          request,
+          requestTask,
+          executionStarter,
+          responseListener());
+      return this;
+    }
+
+    private ActionListener responseListener() {
+      return ActionListener.wrap(
+          snapshot -> {
+            initialResponse.set(snapshot);
+            responseCount.incrementAndGet();
+          },
+          failure::set);
+    }
+
+    private AsyncQueryScenario completeAfterMillis(long elapsedMillis) {
+      now.addAndGet(elapsedMillis);
+      trackingExecution.complete();
+      return this;
+    }
+
+    private AsyncQueryScenario succeedAfterMillis(long elapsedMillis, QueryResponse response) {
+      now.addAndGet(elapsedMillis);
+      trackingExecution.succeed(response);
+      return this;
+    }
+
+    private AsyncQueryScenario complete() {
+      trackingExecution.complete();
+      return this;
+    }
+
+    private AsyncQueryScenario succeed(QueryResponse response) {
+      trackingExecution.succeed(response);
+      return this;
+    }
+
+    private AsyncQueryScenario fail(Exception failure) {
+      trackingExecution.fail(failure);
+      return this;
+    }
+
+    private AsyncQueryScenario advanceMillis(long elapsedMillis) {
+      now.addAndGet(elapsedMillis);
+      return this;
+    }
+
+    private AsyncQueryScenario reapExpired() {
+      targetService.reapExpired();
+      return this;
+    }
+
+    private AsyncQueryScenario assertDirectSuccess(
+        QueryResponse expectedResponse, long expectedTookMillis) {
+      assertEquals(
+          new PPLAsyncQueryService.JobSnapshot.Succeeded(
+              Optional.empty(), expectedResponse, expectedTookMillis),
+          initialResponse.get());
+      return this;
+    }
+
+    private AsyncQueryScenario assertSucceededRowCount(int expected) {
+      assertTrue(initialResponse.get() instanceof PPLAsyncQueryService.JobSnapshot.Succeeded);
+      PPLAsyncQueryService.JobSnapshot.Succeeded succeeded =
+          (PPLAsyncQueryService.JobSnapshot.Succeeded) initialResponse.get();
+      assertEquals(expected, succeeded.response().getResults().size());
+      return this;
+    }
+
+    private AsyncQueryScenario assertRetainedRunning(Optional expectedResponse) {
+      assertEquals(
+          new PPLAsyncQueryService.JobSnapshot.Running(id(), expectedResponse),
+          initialResponse.get());
+      return this;
+    }
+
+    private AsyncQueryScenario assertGetResponse(
+        PPLAsyncQueryService.JobSnapshot expectedResponse) {
+      assertEquals(expectedResponse, getResponse);
+      return this;
+    }
+
+    private AsyncQueryScenario assertNotRetained() {
+      return assertCapacity(0, 0);
+    }
+
+    private AsyncQueryScenario assertCapacity(int running, int retained) {
+      assertEquals(running, targetService.runningQueryCount());
+      assertEquals(retained, targetService.retainedJobCount());
+      return this;
+    }
+
+    private AsyncQueryScenario assertExecutionReadOnceAndClosed() {
+      return assertExecutionReadCount(1).assertExecutionCloseCount(1);
+    }
+
+    private AsyncQueryScenario assertExecutionReadCount(int expected) {
+      assertEquals(expected, trackingExecution.reads.get());
+      return this;
+    }
+
+    private AsyncQueryScenario assertExecutionCloseCount(int expected) {
+      assertEquals(expected, trackingExecution.closes.get());
+      return this;
+    }
+
+    private AsyncQueryScenario assertTimeoutCancelled() {
+      assertTrue(timeoutCancelled.get());
+      return this;
+    }
+
+    private AsyncQueryScenario assertNoInitialResponse() {
+      assertNull(initialResponse.get());
+      return this;
+    }
+
+    private AsyncQueryScenario assertFailure(Class type, String message) {
+      assertTrue(type.isInstance(failure.get()));
+      assertEquals(message, failure.get().getMessage());
+      return this;
+    }
+
+    private AsyncQueryScenario assertFailure(Class type) {
+      assertTrue(type.isInstance(failure.get()));
+      return this;
+    }
+
+    private AsyncQueryScenario fireTimeout() {
+      timeoutTask.get().run();
+      return this;
+    }
+
+    private AsyncQueryScenario assertResponseCount(int expected) {
+      assertEquals(expected, responseCount.get());
+      return this;
+    }
+
+    private AsyncQueryScenario get() {
+      getResponse = targetService.get(id(), owner, null);
+      return this;
+    }
+
+    private AsyncQueryScenario delete() {
+      deleteResponse = targetService.delete(id(), owner);
+      return this;
+    }
+
+    private AsyncQueryScenario assertDeleteStatus(PPLAsyncQueryService.Status status) {
+      assertEquals(new PPLAsyncQueryService.DeleteResult(id(), status), deleteResponse);
+      return this;
+    }
+
+    private String id() {
+      return initialResponse.get().id().orElseThrow();
+    }
+  }
+
   private static final class TrackingExecution implements AsyncQueryExecution {
     private final AtomicReference current;
     private final CompletableFuture completion = new CompletableFuture<>();

From 178b674ac3a1fc3b3292d1e0934fe135b1ee3203 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Thu, 24 Sep 2026 22:35:12 +0000
Subject: [PATCH 11/15] Verify expired async jobs return not found

Signed-off-by: Peng Huo 
---
 .../transport/asyncquery/PPLAsyncQueryServiceTest.java       | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
index 3beb9556d44..0eea7b63db3 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
@@ -147,7 +147,10 @@ public void expiredGetCancelsAndRemovesJob() {
 
     scenario.advanceMillis(KEEP_ALIVE.millis());
 
-    assertThrows(ResourceNotFoundException.class, scenario::get);
+    ResourceNotFoundException error = assertThrows(ResourceNotFoundException.class, scenario::get);
+
+    // Expiration removes the job, so the client sees the same response as any missing job.
+    assertEquals("PPL asynchronous query not found", error.getMessage());
     verify(task).cancel("PPL asynchronous query expired");
     scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
   }

From 7abe947c8215b77dca3a00790258d960261b74d8 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Thu, 24 Sep 2026 22:38:14 +0000
Subject: [PATCH 12/15] Clarify deletion of completed async jobs

Signed-off-by: Peng Huo 
---
 .../plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java    | 1 +
 1 file changed, 1 insertion(+)

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
index 0eea7b63db3..584503bab4c 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
@@ -182,6 +182,7 @@ public void deleteReturnsExistingTerminalStatus() {
         scenario().withTask(task).withCurrentResult(response(1)).start().complete().delete();
 
     scenario.assertDeleteStatus(PPLAsyncQueryService.Status.SUCCEEDED);
+    // DELETE removes the retained result without cancelling an already completed task.
     verify(task, never()).cancel(org.mockito.ArgumentMatchers.anyString());
     scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
   }

From 197f446d0a1f6d8f275e60bdd4196f6a64321f2a Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Fri, 25 Sep 2026 16:11:40 +0000
Subject: [PATCH 13/15] refactor(async-query): consolidate lifecycle into
 QueryJob active object

Replaces PPLAsyncQueryService (~840 lines) and PPLAsyncQueryJob with a
single active-object design:

- QueryJob owns the state machine, retention deadline, keep-alive expiry
  timer, task registration and cancellation, execution attachment and
  completion cleanup. Public API: create(...), get(keepAlive),
  cancel(reason), getJobId, getOwner. Snapshot / Status / Failure are
  the only external types.
- QueryJobRegistry is a thin ConcurrentMap wrapper with add / get /
  remove / close. close() calls discard() on every job and rejects
  further adds.
- PPLAsyncQueryJobId -> QueryJobId (now public).
- PPLAsyncQueryUser  -> QueryJobOwner (authorize() now public).
- Renamed tests keep prior coverage; PPLAsyncQueryServiceTest and
  PPLAsyncQueryJobTest are removed together with the service.

AdmissionControl (running/retained capacity limits) and request
validation are deferred to a follow-up PR, matching the current PR
description's boundary.

Signed-off-by: Peng Huo 
---
 .../asyncquery/PPLAsyncQueryJob.java          | 469 ---------
 .../asyncquery/PPLAsyncQueryService.java      | 843 ---------------
 .../plugin/transport/asyncquery/QueryJob.java | 983 ++++++++++++++++++
 ...PLAsyncQueryJobId.java => QueryJobId.java} |  21 +-
 ...AsyncQueryUser.java => QueryJobOwner.java} |  13 +-
 .../asyncquery/QueryJobRegistry.java          | 103 ++
 .../asyncquery/PPLAsyncQueryJobIdTest.java    |  28 -
 .../asyncquery/PPLAsyncQueryJobTest.java      | 273 -----
 .../asyncquery/PPLAsyncQueryServiceTest.java  | 951 -----------------
 .../transport/asyncquery/QueryJobIdTest.java  |  28 +
 ...ryUserTest.java => QueryJobOwnerTest.java} |  31 +-
 11 files changed, 1147 insertions(+), 2596 deletions(-)
 delete mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java
 delete mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java
 create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJob.java
 rename plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/{PPLAsyncQueryJobId.java => QueryJobId.java} (91%)
 rename plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/{PPLAsyncQueryUser.java => QueryJobOwner.java} (86%)
 create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobRegistry.java
 delete mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java
 delete mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
 delete mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
 create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobIdTest.java
 rename plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/{PPLAsyncQueryUserTest.java => QueryJobOwnerTest.java} (60%)

diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java
deleted file mode 100644
index 41308646411..00000000000
--- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJob.java
+++ /dev/null
@@ -1,469 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.sql.plugin.transport.asyncquery;
-
-import java.util.Objects;
-import java.util.Optional;
-import org.opensearch.ResourceNotFoundException;
-import org.opensearch.common.unit.TimeValue;
-import org.opensearch.sql.executor.AsyncQueryExecution;
-import org.opensearch.tasks.CancellableTask;
-
-/**
- * Mutable state machine for one asynchronous PPL query.
- *
- * 
Every lifecycle and lease transition is synchronized on this object. Methods mutate only
- * job-owned state and return immutable transition values; they never call listeners, mutate the
- * service map, update capacity counters, or cancel tasks while holding the lock.
- *
- * 
After attachment, the job owns one {@link AsyncQueryExecution}. Reads borrow it through a
- * {@link SnapshotSource}; removal transitions detach it so the service can close it outside the
- * lock.
- *
- * 
- * Every job starts in RUNNING. It becomes retained when wait_for_completion_timeout expires.
- *
- * Current state       Event              Next state             Response
- * RUNNING             success            REMOVED                final result without ID
- * RUNNING             failure            REMOVED                failure without ID
- * RUNNING             retain             RETAINED_RUNNING       running status with ID
- * RETAINED_RUNNING    success            RETAINED_SUCCEEDED     none
- * RETAINED_RUNNING    failure            RETAINED_FAILED        none
- * RUNNING             abort/close        REMOVED                none
- * RETAINED_*          delete/expire/abort/close REMOVED         none
- * 

- *
- * 
GET lease renewal and execution attachment do not change the lifecycle state. Events received
- * after {@code REMOVED} are ignored or reported as not found; a late execution handle is rejected
- * so the service can close it.
- */
-final class PPLAsyncQueryJob {
-  private final String id;
-  private final PPLAsyncQueryUser owner;
-  private final long startTimeMillis;
-  private JobTask task;
-
-  private long keepAliveMillis;
-  private long expirationTimeMillis;
-  private State state;
-  private AsyncQueryExecution execution;
-  private PPLAsyncQueryService.Failure failure;
-  private long completionTimeMillis = -1L;
-
-  /**
-   * Creates an unretained job in the {@link State#RUNNING} state.
-   *
-   * @param id opaque ID assigned by the owner node
-   * @param owner caller that is allowed to access the retained job
-   * @param startTimeMillis execution start time
-   * @param keepAliveMillis lease duration applied when the job is retained
-   * @param task independently cancellable task owned by the job
-   */
-  PPLAsyncQueryJob(
-      String id,
-      PPLAsyncQueryUser owner,
-      long startTimeMillis,
-      long keepAliveMillis,
-      JobTask task) {
-    this.id = id;
-    this.owner = owner;
-    this.startTimeMillis = startTimeMillis;
-    this.keepAliveMillis = keepAliveMillis;
-    this.expirationTimeMillis = startTimeMillis + keepAliveMillis;
-    this.task = Objects.requireNonNull(task);
-    this.state = State.RUNNING;
-  }
-
-  /**
-   * Returns the opaque ID assigned to this job.
-   *
-   * @return job ID used as the service registry key
-   */
-  String id() {
-    return id;
-  }
-
-  /**
-   * Returns the immutable identity captured when this job was created.
-   *
-   * @return job owner used by the service for GET and DELETE authorization
-   */
-  PPLAsyncQueryUser owner() {
-    return owner;
-  }
-
-  /**
-   * Retains a running job after its initial response wait expires.
-   *
-   * @param now time at which the job becomes visible to GET and DELETE
-   * @return transition that publishes the job ID, or {@code null} if the job already finished
-   */
-  synchronized Transition retain(long now) {
-    if (state != State.RUNNING) {
-      return null;
-    }
-    state = State.RETAINED_RUNNING;
-    expirationTimeMillis = now + keepAliveMillis;
-    return new Transition.Retained(runningSnapshot());
-  }
-
-  /**
-   * Transfers ownership of an execution handle to this job.
-   *
-   * @param execution handle that produces current and final query results
-   * @return {@code true} when attached; {@code false} when the caller must close the rejected
-   *     handle
-   */
-  synchronized boolean tryAttachExecution(AsyncQueryExecution execution) {
-    Objects.requireNonNull(execution);
-    if (!isExecuting() || this.execution != null) {
-      return false;
-    }
-    this.execution = execution;
-    return true;
-  }
-
-  /**
-   * Records successful execution completion.
-   *
-   * @param now completion time
-   * @return direct-response or retained-completion transition, or {@code null} after removal
-   * @throws IllegalStateException if successful completion is reported before execution attachment
-   */
-  synchronized Transition complete(long now) {
-    if (!isExecuting()) {
-      return null;
-    }
-    if (execution == null) {
-      throw new IllegalStateException(
-          "PPL asynchronous execution must be attached before successful completion");
-    }
-    return finish(State.RETAINED_SUCCEEDED, now);
-  }
-
-  /**
-   * Records failed execution completion.
-   *
-   * @param failure client-visible failure retained with the job
-   * @param now completion time
-   * @return direct-response or retained-completion transition, or {@code null} after removal
-   */
-  synchronized Transition fail(PPLAsyncQueryService.Failure failure, long now) {
-    if (!isExecuting()) {
-      return null;
-    }
-    this.failure = failure;
-    return finish(State.RETAINED_FAILED, now);
-  }
-
-  private Transition finish(State terminalState, long now) {
-    boolean returnDirect = state == State.RUNNING;
-    JobTask taskToClose = detachTask();
-    state = terminalState;
-    completionTimeMillis = now;
-    if (returnDirect) {
-      SnapshotSource response = snapshot(Optional.empty());
-      state = State.REMOVED;
-      return new Transition.DirectResponse(
-          response, taskToClose, Optional.ofNullable(detachExecution()));
-    }
-    AsyncQueryExecution executionToClose =
-        terminalState == State.RETAINED_FAILED ? detachExecution() : null;
-    return new Transition.ExecutionFinished(taskToClose, Optional.ofNullable(executionToClose));
-  }
-
-  /**
-   * Returns the current retained response.
-   *
-   * 
A supplied {@code requestedKeepAlive} starts a new lease from {@code now}. A request at or
-   * after the existing expiration time returns an expiration removal instead of data.
-   *
-   * @param now request time
-   * @param requestedKeepAlive replacement lease, or {@code null} to keep the current expiration
-   * @return current response or the removal required for an expired job
-   * @throws ResourceNotFoundException if the job was already removed
-   */
-  synchronized GetResult get(long now, TimeValue requestedKeepAlive) {
-    ensurePresent();
-    if (now >= expirationTimeMillis) {
-      return new GetResult.Expired(expireLocked("PPL asynchronous query expired"));
-    }
-    if (requestedKeepAlive != null) {
-      keepAliveMillis = requestedKeepAlive.millis();
-      expirationTimeMillis = now + keepAliveMillis;
-    }
-    return new GetResult.Found(retainedSnapshot());
-  }
-
-  /**
-   * Cancels if still running and removes this job.
-   *
-   * @param now request time
-   * @return removal containing the response status and detached resources
-   * @throws ResourceNotFoundException if the job was already removed
-   */
-  synchronized Removal delete(long now) {
-    ensurePresent();
-    if (now >= expirationTimeMillis) {
-      return expireLocked("PPL asynchronous query expired");
-    }
-    return delete(
-        isExecuting() ? PPLAsyncQueryService.Status.CANCELLED : responseStatus(),
-        "PPL asynchronous query cancelled by user");
-  }
-
-  /**
-   * Removes a retained job whose lease has expired.
-   *
-   * @param now expiration check time
-   * @return removal with detached resources, or {@code null} if no expiration is due
-   */
-  synchronized Removal expire(long now) {
-    if (state == State.REMOVED || state == State.RUNNING || now < expirationTimeMillis) {
-      return null;
-    }
-    return expireLocked("PPL asynchronous query expired");
-  }
-
-  private Removal.Expired expireLocked(String reason) {
-    Removal.Expired removal = new Removal.Expired(detachResources(), reason);
-    state = State.REMOVED;
-    return removal;
-  }
-
-  /**
-   * Removes a job whose submission could not be completed.
-   *
-   * @return removal with detached resources, or {@code null} if already removed
-   */
-  synchronized Removal abort() {
-    return discardIfPresent("PPL asynchronous query startup failed");
-  }
-
-  /**
-   * Removes this job during service shutdown.
-   *
-   * @param reason cancellation reason used if execution is still running
-   * @return removal with detached resources, or {@code null} if already removed
-   */
-  synchronized Removal close(String reason) {
-    return discardIfPresent(reason);
-  }
-
-  private Removal discardIfPresent(String reason) {
-    if (state == State.REMOVED) {
-      return null;
-    }
-    Removal.Discarded removal = new Removal.Discarded(detachResources(), reason);
-    state = State.REMOVED;
-    return removal;
-  }
-
-  private Removal delete(PPLAsyncQueryService.Status responseStatus, String reason) {
-    Removal.Deleted removal = new Removal.Deleted(responseStatus, detachResources(), reason);
-    state = State.REMOVED;
-    return removal;
-  }
-
-  private DetachedResources detachResources() {
-    return new DetachedResources(detachTask(), detachExecution());
-  }
-
-  private SnapshotSource retainedSnapshot() {
-    return snapshot(Optional.of(id));
-  }
-
-  private SnapshotSource.Running runningSnapshot() {
-    return new SnapshotSource.Running(id, Optional.ofNullable(execution));
-  }
-
-  private SnapshotSource snapshot(Optional responseId) {
-    long tookMillis =
-        completionTimeMillis < 0 ? -1L : Math.max(0L, completionTimeMillis - startTimeMillis);
-    return switch (state) {
-      case RUNNING, RETAINED_RUNNING -> runningSnapshot();
-      case RETAINED_SUCCEEDED ->
-          new SnapshotSource.Succeeded(responseId, Objects.requireNonNull(execution), tookMillis);
-      case RETAINED_FAILED ->
-          new SnapshotSource.Failed(responseId, Objects.requireNonNull(failure), tookMillis);
-      case REMOVED -> throw new IllegalStateException("PPL asynchronous query was removed");
-    };
-  }
-
-  private boolean isExecuting() {
-    return state == State.RUNNING || state == State.RETAINED_RUNNING;
-  }
-
-  private PPLAsyncQueryService.Status responseStatus() {
-    return switch (state) {
-      case RUNNING, RETAINED_RUNNING -> PPLAsyncQueryService.Status.RUNNING;
-      case RETAINED_SUCCEEDED -> PPLAsyncQueryService.Status.SUCCEEDED;
-      case RETAINED_FAILED -> PPLAsyncQueryService.Status.FAILED;
-      case REMOVED -> throw new IllegalStateException("PPL asynchronous query was removed");
-    };
-  }
-
-  private AsyncQueryExecution detachExecution() {
-    AsyncQueryExecution detached = execution;
-    execution = null;
-    return detached;
-  }
-
-  private JobTask detachTask() {
-    JobTask detached = task;
-    task = null;
-    return detached;
-  }
-
-  private void ensurePresent() {
-    if (state == State.REMOVED) {
-      throw new ResourceNotFoundException("PPL asynchronous query not found");
-    }
-  }
-
-  /** Lifecycle state captured under the job lock for later result materialization. */
-  sealed interface SnapshotSource {
-    /**
-     * Source for a retained running response.
-     *
-     * @param id retained job ID
-     * @param execution attached execution, or empty when retention wins before attachment
-     */
-    record Running(String id, Optional execution) implements SnapshotSource {}
-
-    /**
-     * Source for a successful final response.
-     *
-     * @param id retained job ID, or empty for a direct POST response
-     * @param execution completed execution that owns the final result
-     * @param tookMillis elapsed execution time
-     */
-    record Succeeded(Optional id, AsyncQueryExecution execution, long tookMillis)
-        implements SnapshotSource {}
-
-    /**
-     * Source for a failed final response.
-     *
-     * @param id retained job ID, or empty for a direct POST response
-     * @param failure client-visible failure
-     * @param tookMillis elapsed execution time
-     */
-    record Failed(Optional id, PPLAsyncQueryService.Failure failure, long tookMillis)
-        implements SnapshotSource {}
-  }
-
-  /**
-   * Cancellable task and the cleanup that releases its TaskManager registrations.
-   *
-   * @param task task used to cancel query execution
-   * @param release registration cleanup
-   */
-  record JobTask(CancellableTask task, Runnable release) {
-    /** Releases the task and child-node registrations owned by this wrapper. */
-    void close() {
-      release.run();
-    }
-  }
-
-  /** Internal lifecycle; unlike the response status, this includes retention and removal. */
-  enum State {
-    RUNNING,
-    RETAINED_RUNNING,
-    RETAINED_SUCCEEDED,
-    RETAINED_FAILED,
-    REMOVED
-  }
-
-  /** State-machine output consumed by {@link PPLAsyncQueryService}. */
-  sealed interface Transition {
-    /**
-     * The initial wait expired and the running job became retained.
-     *
-     * @param response response published with the retained job ID
-     */
-    record Retained(SnapshotSource.Running response) implements Transition {}
-
-    /**
-     * Execution finished before retention and the initial request receives the final response.
-     *
-     * @param response successful or failed final response source
-     * @param task completed task registration to close
-     * @param executionToClose detached execution to close after materialization
-     */
-    record DirectResponse(
-        SnapshotSource response, JobTask task, Optional executionToClose)
-        implements Transition {}
-
-    /**
-     * A retained execution finished without producing an immediate response.
-     *
-     * @param task completed task registration to close
-     * @param executionToClose failed execution to close; successful execution remains retained
-     */
-    record ExecutionFinished(JobTask task, Optional executionToClose)
-        implements Transition {}
-  }
-
-  /** Result of an authorized GET attempt. */
-  sealed interface GetResult {
-    /**
-     * GET result for a live retained job.
-     *
-     * @param response current lifecycle source for result materialization
-     */
-    record Found(SnapshotSource response) implements GetResult {}
-
-    /**
-     * GET result when the lease expired before the request.
-     *
-     * @param removal cleanup required for the expired job
-     */
-    record Expired(Removal.Expired removal) implements GetResult {}
-  }
-
-  /**
-   * Task and execution handles detached when a job leaves the registry.
-   *
-   * @param task running task to cancel, or {@code null} after execution already finished
-   * @param execution execution handle to close, or {@code null} before attachment
-   */
-  record DetachedResources(JobTask task, AsyncQueryExecution execution) {
-    /** Returns whether removing these resources releases a running-query capacity slot. */
-    boolean releasesRunningSlot() {
-      return task != null;
-    }
-  }
-
-  /** Reason-specific removal consumed by {@link PPLAsyncQueryService}. */
-  sealed interface Removal {
-    /**
-     * Removal requested by DELETE.
-     *
-     * @param responseStatus status returned to the caller
-     * @param resources detached task and execution handles
-     * @param reason task cancellation reason
-     */
-    record Deleted(
-        PPLAsyncQueryService.Status responseStatus, DetachedResources resources, String reason)
-        implements Removal {}
-
-    /**
-     * Removal caused by lease expiration.
-     *
-     * @param resources detached task and execution handles
-     * @param reason task cancellation reason
-     */
-    record Expired(DetachedResources resources, String reason) implements Removal {}
-
-    /**
-     * Internal removal caused by startup failure or service shutdown.
-     *
-     * @param resources detached task and execution handles
-     * @param reason task cancellation reason
-     */
-    record Discarded(DetachedResources resources, String reason) implements Removal {}
-  }
-}
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java
deleted file mode 100644
index e44ebbe1bac..00000000000
--- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryService.java
+++ /dev/null
@@ -1,843 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.sql.plugin.transport.asyncquery;
-
-import java.io.IOException;
-import java.util.List;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.function.Function;
-import java.util.function.IntSupplier;
-import java.util.function.LongSupplier;
-import java.util.function.Supplier;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-import org.opensearch.OpenSearchStatusException;
-import org.opensearch.ResourceNotFoundException;
-import org.opensearch.cluster.node.DiscoveryNode;
-import org.opensearch.common.lease.Releasable;
-import org.opensearch.common.lifecycle.AbstractLifecycleComponent;
-import org.opensearch.common.unit.TimeValue;
-import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.rest.RestStatus;
-import org.opensearch.core.tasks.TaskId;
-import org.opensearch.sql.common.setting.Settings;
-import org.opensearch.sql.executor.AsyncQueryExecution;
-import org.opensearch.sql.executor.ExecutionEngine.QueryResponse;
-import org.opensearch.sql.executor.ExecutionEngine.Schema;
-import org.opensearch.sql.plugin.PPLQueryErrorHandler;
-import org.opensearch.sql.plugin.transport.PPLQueryAction;
-import org.opensearch.sql.plugin.transport.PPLQueryTask;
-import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.GetResult;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.JobTask;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Removal;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.SnapshotSource;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition;
-import org.opensearch.tasks.CancellableTask;
-import org.opensearch.tasks.Task;
-import org.opensearch.tasks.TaskManager;
-import org.opensearch.threadpool.Scheduler;
-import org.opensearch.threadpool.ThreadPool;
-
-/**
- * Owns the asynchronous PPL jobs assigned to the local node.
- *
- * 
{@link PPLAsyncQueryJob} owns the mutable state of one job. This service owns the job registry
- * and performs the work requested by each state transition: capacity accounting, task management,
- * result materialization, listener notification, and execution cleanup. Those side effects happen
- * after the job lock is released.
- *
- * 
POST races query completion against {@code wait_for_completion_timeout}. Completion wins by
- * returning the final result directly and removing the job. The timeout wins by retaining the job
- * and returning its opaque ID. GET and DELETE are then routed to this owner node.
- *
- * 
This class is thread-safe. The registry is concurrent, capacity counters are guarded by {@code
- * admissionLock}, and each job synchronizes its own lifecycle transitions.
- */
-public final class PPLAsyncQueryService extends AbstractLifecycleComponent {
-  private static final Logger LOG = LogManager.getLogger(PPLAsyncQueryService.class);
-
-  private static final TimeValue REAPER_INTERVAL = TimeValue.timeValueMinutes(1);
-  private static final TimeoutHandle NO_TIMEOUT = () -> {};
-
-  /** Lifecycle state exposed in asynchronous PPL responses. */
-  public enum Status {
-    /** Query execution is still running. */
-    RUNNING,
-
-    /** Query execution completed successfully. */
-    SUCCEEDED,
-
-    /** Query execution failed. */
-    FAILED,
-
-    /** Query execution was cancelled by DELETE. */
-    CANCELLED
-  }
-
-  @FunctionalInterface
-  interface TimeoutHandle {
-    void cancel();
-  }
-
-  @FunctionalInterface
-  interface TimeoutScheduler {
-    TimeoutHandle schedule(TimeValue delay, Runnable task);
-  }
-
-  /**
-   * Immutable point-in-time response produced after releasing the job lock.
-   *
-   * 
Each implementation exposes only the data valid for its lifecycle state.
-   */
-  public sealed interface JobSnapshot {
-    /**
-     * Returns the retained job ID.
-     *
-     * @return job ID, or empty for a final response returned directly by POST
-     */
-    Optional id();
-
-    /**
-     * Returns the public lifecycle state.
-     *
-     * @return response status
-     */
-    Status status();
-
-    /**
-     * Snapshot of a retained query that is still running.
-     *
-     * @param jobId retained job ID
-     * @param response current result, or empty before a result is available
-     */
-    record Running(String jobId, Optional response) implements JobSnapshot {
-      /** {@inheritDoc} */
-      @Override
-      public Optional id() {
-        return Optional.of(jobId);
-      }
-
-      /** {@inheritDoc} */
-      @Override
-      public Status status() {
-        return Status.RUNNING;
-      }
-    }
-
-    /**
-     * Snapshot of a successfully completed query.
-     *
-     * @param id retained job ID, or empty for a direct POST response
-     * @param response final query result
-     * @param tookMillis elapsed execution time
-     */
-    record Succeeded(Optional id, QueryResponse response, long tookMillis)
-        implements JobSnapshot {
-      /** {@inheritDoc} */
-      @Override
-      public Status status() {
-        return Status.SUCCEEDED;
-      }
-    }
-
-    /**
-     * Snapshot of a failed query.
-     *
-     * @param id retained job ID, or empty for a direct POST response
-     * @param failure client-visible failure
-     * @param tookMillis elapsed execution time
-     */
-    record Failed(Optional id, Failure failure, long tookMillis) implements JobSnapshot {
-      /** {@inheritDoc} */
-      @Override
-      public Status status() {
-        return Status.FAILED;
-      }
-    }
-  }
-
-  /** Response model returned after DELETE removes a retained job. */
-  record DeleteResult(String id, Status status) {}
-
-  /**
-   * Client-visible failure retained by a job and returned only after owner authorization.
-   *
-   * @param type exception type
-   * @param reason client-facing failure reason
-   */
-  public record Failure(String type, String reason) {
-    private static Failure from(Exception exception) {
-      String type =
-          exception.getClass().getSimpleName().isBlank()
-              ? exception.getClass().getName()
-              : exception.getClass().getSimpleName();
-      String reason =
-          exception.getMessage() == null || exception.getMessage().isBlank()
-              ? "query execution failed"
-              : exception.getMessage();
-      return new Failure(type, reason);
-    }
-  }
-
-  private final Supplier ownerNodeIdSupplier;
-  private final LongSupplier currentTimeMillis;
-  private final TimeoutScheduler timeoutScheduler;
-  private final IntSupplier maxRunningQueries;
-  private final IntSupplier maxRetainedJobs;
-  private final Supplier maxWaitForCompletion;
-  private final Supplier maxKeepAlive;
-  private final ThreadPool threadPool;
-  private final ConcurrentMap jobs = new ConcurrentHashMap<>();
-  private final Object admissionLock = new Object();
-
-  private int runningQueries;
-  private int retainedJobs;
-  private volatile boolean acceptingNewJobs = true;
-  private volatile Scheduler.Cancellable reaper;
-  private volatile TaskManager taskManager;
-
-  /**
-   * Creates the owner-node lifecycle service.
-   *
-   * @param ownerNodeIdSupplier supplies the current local node ID
-   * @param threadPool schedules retention deadlines and expiration reaping
-   * @param settings supplies asynchronous query capacity and duration limits
-   */
-  public PPLAsyncQueryService(
-      Supplier ownerNodeIdSupplier, ThreadPool threadPool, Settings settings) {
-    this(
-        ownerNodeIdSupplier,
-        System::currentTimeMillis,
-        (delay, task) -> {
-          Scheduler.ScheduledCancellable cancellable =
-              threadPool.schedule(task, delay, ThreadPool.Names.GENERIC);
-          return cancellable::cancel;
-        },
-        () ->
-            (Integer)
-                settings.getSettingValue(Settings.Key.PPL_ASYNC_NODE_CONCURRENT_RUNNING_QUERIES),
-        () -> (Integer) settings.getSettingValue(Settings.Key.PPL_ASYNC_MAX_RETAINED_JOBS),
-        () ->
-            (TimeValue)
-                settings.getSettingValue(Settings.Key.PPL_ASYNC_MAX_WAIT_FOR_COMPLETION_TIMEOUT),
-        () -> (TimeValue) settings.getSettingValue(Settings.Key.PPL_ASYNC_MAX_KEEP_ALIVE),
-        threadPool);
-  }
-
-  PPLAsyncQueryService(
-      String ownerNodeId,
-      LongSupplier currentTimeMillis,
-      TimeoutScheduler timeoutScheduler,
-      IntSupplier maxRunningQueries,
-      IntSupplier maxRetainedJobs,
-      Supplier maxWaitForCompletion,
-      Supplier maxKeepAlive) {
-    this(
-        () -> ownerNodeId,
-        currentTimeMillis,
-        timeoutScheduler,
-        maxRunningQueries,
-        maxRetainedJobs,
-        maxWaitForCompletion,
-        maxKeepAlive,
-        null);
-  }
-
-  private PPLAsyncQueryService(
-      Supplier ownerNodeIdSupplier,
-      LongSupplier currentTimeMillis,
-      TimeoutScheduler timeoutScheduler,
-      IntSupplier maxRunningQueries,
-      IntSupplier maxRetainedJobs,
-      Supplier maxWaitForCompletion,
-      Supplier maxKeepAlive,
-      ThreadPool threadPool) {
-    this.ownerNodeIdSupplier = Objects.requireNonNull(ownerNodeIdSupplier);
-    this.currentTimeMillis = Objects.requireNonNull(currentTimeMillis);
-    this.timeoutScheduler = Objects.requireNonNull(timeoutScheduler);
-    this.maxRunningQueries = Objects.requireNonNull(maxRunningQueries);
-    this.maxRetainedJobs = Objects.requireNonNull(maxRetainedJobs);
-    this.maxWaitForCompletion = Objects.requireNonNull(maxWaitForCompletion);
-    this.maxKeepAlive = Objects.requireNonNull(maxKeepAlive);
-    this.threadPool = threadPool;
-  }
-
-  /**
-   * Starts an asynchronous PPL query and produces its POST response.
-   *
-   * 
The service registers and owns a job task, schedules its retention deadline, starts
-   * execution, and attaches the returned execution handle to the same job. If execution finishes
-   * before the deadline, {@code responseListener} receives the final result without a job ID.
-   * Otherwise, the job becomes retained and the listener receives its current result with an opaque
-   * ID.
-   *
-   * @param owner submit caller retained with the job for later authorization
-   * @param keepAlive requested job lease
-   * @param waitForCompletion maximum time to wait for a direct result
-   * @param request transport request used to register the job task
-   * @param requestTask task associated with the POST request
-   * @param executionStarter starts execution using the job-owned cancellable task
-   * @param responseListener receives either the direct result or retained job ID
-   */
-  public void start(
-      PPLAsyncQueryUser owner,
-      TimeValue keepAlive,
-      TimeValue waitForCompletion,
-      TransportPPLQueryRequest request,
-      PPLQueryTask requestTask,
-      Function executionStarter,
-      ActionListener responseListener) {
-    validateKeepAlive(keepAlive);
-    validateWaitForCompletion(waitForCompletion);
-    JobTask jobTask = registerJobTask(request, requestTask);
-    start(owner, keepAlive, waitForCompletion, jobTask, executionStarter, responseListener);
-  }
-
-  /**
-   * Starts a job using an already registered task.
-   *
-   * 
This package-private entry point keeps task registration separate for tests while preserving
-   * the same production lifecycle: create the job, establish its retention deadline, start
-   * execution, and transfer ownership of the execution handle to the job.
-   *
-   * @param owner submit caller retained with the job for later authorization
-   * @param keepAlive validated job lease
-   * @param waitForCompletion validated direct-result wait
-   * @param task job-owned task and registration cleanup
-   * @param executionStarter starts execution using the job task
-   * @param responseListener receives the POST response
-   */
-  void start(
-      PPLAsyncQueryUser owner,
-      TimeValue keepAlive,
-      TimeValue waitForCompletion,
-      JobTask task,
-      Function executionStarter,
-      ActionListener responseListener) {
-    PPLAsyncQueryJob job;
-    try {
-      Objects.requireNonNull(task);
-      Objects.requireNonNull(executionStarter);
-      Objects.requireNonNull(responseListener);
-      job = createJob(owner, keepAlive, task);
-    } catch (RuntimeException | Error e) {
-      closeTask(task);
-      throw e;
-    }
-
-    TimeoutHandle retentionDeadline;
-    try {
-      retentionDeadline = scheduleRetention(job, waitForCompletion, responseListener);
-    } catch (RuntimeException | Error e) {
-      applyRemoval(job, job.abort());
-      throw e;
-    }
-    startExecution(job, task.task(), executionStarter, responseListener, retentionDeadline);
-  }
-
-  private PPLAsyncQueryJob createJob(PPLAsyncQueryUser owner, TimeValue keepAlive, JobTask task) {
-    Objects.requireNonNull(owner);
-    reserveCapacity();
-    try {
-      String ownerNodeId =
-          Objects.requireNonNull(ownerNodeIdSupplier.get(), "Local node ID is not initialized");
-      long now = currentTimeMillis.getAsLong();
-      while (true) {
-        PPLAsyncQueryJobId jobId = PPLAsyncQueryJobId.create(ownerNodeId);
-        String encodedId = jobId.encode();
-        PPLAsyncQueryJob job =
-            new PPLAsyncQueryJob(encodedId, owner, now, keepAlive.millis(), task);
-        if (jobs.putIfAbsent(encodedId, job) == null) {
-          return job;
-        }
-      }
-    } catch (RuntimeException | Error e) {
-      releaseCapacity();
-      throw e;
-    }
-  }
-
-  private TimeoutHandle scheduleRetention(
-      PPLAsyncQueryJob job,
-      TimeValue waitForCompletion,
-      ActionListener responseListener) {
-    if (waitForCompletion.millis() == 0) {
-      applyTransition(job, job.retain(currentTimeMillis.getAsLong()), responseListener);
-      return NO_TIMEOUT;
-    }
-
-    return timeoutScheduler.schedule(
-        waitForCompletion,
-        () -> applyTransition(job, job.retain(currentTimeMillis.getAsLong()), responseListener));
-  }
-
-  private void startExecution(
-      PPLAsyncQueryJob job,
-      CancellableTask task,
-      Function executionStarter,
-      ActionListener responseListener,
-      TimeoutHandle retentionDeadline) {
-    try {
-      AsyncQueryExecution execution = Objects.requireNonNull(executionStarter.apply(task));
-      attachExecution(job, execution, responseListener, retentionDeadline);
-    } catch (RuntimeException e) {
-      retentionDeadline.cancel();
-      fail(job, e, responseListener);
-    }
-  }
-
-  private void attachExecution(
-      PPLAsyncQueryJob job,
-      AsyncQueryExecution execution,
-      ActionListener responseListener,
-      TimeoutHandle retentionDeadline) {
-    if (!job.tryAttachExecution(execution)) {
-      retentionDeadline.cancel();
-      closeExecution(execution);
-      return;
-    }
-    execution
-        .completion()
-        .whenComplete(
-            (ignored, failure) -> {
-              retentionDeadline.cancel();
-              if (failure == null) {
-                complete(job, responseListener);
-              } else {
-                fail(job, asException(failure), responseListener);
-              }
-            });
-  }
-
-  private void complete(PPLAsyncQueryJob job, ActionListener responseListener) {
-    applyTransition(job, job.complete(currentTimeMillis.getAsLong()), responseListener);
-  }
-
-  private void fail(
-      PPLAsyncQueryJob job, Exception failure, ActionListener responseListener) {
-    Objects.requireNonNull(failure);
-    Transition transition = job.fail(Failure.from(failure), currentTimeMillis.getAsLong());
-    if (transition == null) {
-      return;
-    }
-    switch (transition) {
-      case Transition.ExecutionFinished finished -> {
-        PPLQueryErrorHandler.recordFailure(failure);
-        applyTransition(job, finished, responseListener);
-      }
-      case Transition.DirectResponse direct ->
-          applyTransition(
-              job,
-              direct,
-              ActionListener.wrap(
-                  ignored -> responseListener.onFailure(failure), responseListener::onFailure));
-      case Transition.Retained retained ->
-          throw new IllegalStateException("Query failure cannot retain a job");
-    }
-  }
-
-  /**
-   * Returns the current snapshot of a retained job.
-   *
-   * 
The service authorizes the caller against the immutable job owner before requesting a
-   * synchronized lease transition. Result materialization happens afterward and therefore cannot
-   * block lifecycle transitions.
-   *
-   * @param id opaque job ID owned by this node
-   * @param caller caller to compare with the stored job owner
-   * @param requestedKeepAlive new lease duration, or {@code null} to leave the lease unchanged
-   * @return immutable current job snapshot
-   */
-  JobSnapshot get(String id, PPLAsyncQueryUser caller, TimeValue requestedKeepAlive) {
-    if (requestedKeepAlive != null) {
-      validateKeepAlive(requestedKeepAlive);
-    }
-    PPLAsyncQueryJob job = findLocal(id);
-    job.owner().authorize(caller);
-    GetResult result = job.get(currentTimeMillis.getAsLong(), requestedKeepAlive);
-    if (result instanceof GetResult.Expired expired) {
-      applyRemoval(job, expired.removal());
-      throw notFound();
-    }
-    return materialize(((GetResult.Found) result).response());
-  }
-
-  /**
-   * Cancels and removes a retained job.
-   *
-   * @param id opaque job ID owned by this node
-   * @param caller caller to compare with the stored job owner
-   * @return the status observed when the job was removed
-   */
-  DeleteResult delete(String id, PPLAsyncQueryUser caller) {
-    PPLAsyncQueryJob job = findLocal(id);
-    job.owner().authorize(caller);
-    Removal removal = job.delete(currentTimeMillis.getAsLong());
-    applyRemoval(job, removal);
-    return switch (removal) {
-      case Removal.Deleted deleted -> new DeleteResult(id, deleted.responseStatus());
-      case Removal.Expired expired -> throw notFound();
-      case Removal.Discarded discarded ->
-          throw new IllegalStateException("DELETE cannot discard a job");
-    };
-  }
-
-  void reapExpired() {
-    long now = currentTimeMillis.getAsLong();
-    jobs.forEach((id, job) -> applyRemoval(job, job.expire(now)));
-  }
-
-  int runningQueryCount() {
-    synchronized (admissionLock) {
-      return runningQueries;
-    }
-  }
-
-  int retainedJobCount() {
-    synchronized (admissionLock) {
-      return retainedJobs;
-    }
-  }
-
-  /**
-   * Attaches the node task manager after transport actions have been initialized.
-   *
-   * @param taskManager task manager used to register and cancel retained query tasks
-   */
-  public void attachTaskManager(TaskManager taskManager) {
-    this.taskManager = Objects.requireNonNull(taskManager);
-  }
-
-  private void reserveCapacity() {
-    synchronized (admissionLock) {
-      if (!acceptingNewJobs) {
-        throw new OpenSearchStatusException(
-            "PPL asynchronous query service is stopping", RestStatus.SERVICE_UNAVAILABLE);
-      }
-      if (runningQueries >= maxRunningQueries.getAsInt()
-          || retainedJobs >= maxRetainedJobs.getAsInt()) {
-        throw new OpenSearchStatusException(
-            "PPL asynchronous query capacity is exhausted", RestStatus.TOO_MANY_REQUESTS);
-      }
-      runningQueries++;
-      retainedJobs++;
-    }
-  }
-
-  private void releaseRunning() {
-    synchronized (admissionLock) {
-      if (runningQueries > 0) {
-        runningQueries--;
-      }
-    }
-  }
-
-  private void releaseRetained() {
-    synchronized (admissionLock) {
-      if (retainedJobs > 0) {
-        retainedJobs--;
-      }
-    }
-  }
-
-  private void releaseCapacity() {
-    synchronized (admissionLock) {
-      if (runningQueries > 0) {
-        runningQueries--;
-      }
-      if (retainedJobs > 0) {
-        retainedJobs--;
-      }
-    }
-  }
-
-  /**
-   * Applies side effects selected by a {@link PPLAsyncQueryJob} transition.
-   *
-   * 
The job decides its state change while holding the job lock, then returns a value describing
-   * the required side effects. Map mutation, capacity accounting, and listener callbacks happen
-   * here after the lock has been released.
-   */
-  private void applyTransition(
-      PPLAsyncQueryJob job, Transition transition, ActionListener responseListener) {
-    if (transition == null) {
-      return;
-    }
-    switch (transition) {
-      case Transition.Retained retained -> {
-        try {
-          responseListener.onResponse(materialize(retained.response()));
-        } catch (RuntimeException e) {
-          applyRemoval(job, job.abort());
-          responseListener.onFailure(e);
-        }
-      }
-      case Transition.DirectResponse direct -> {
-        releaseRunning();
-        if (jobs.remove(job.id(), job)) {
-          releaseRetained();
-        }
-        try {
-          responseListener.onResponse(materialize(direct.response()));
-        } catch (RuntimeException e) {
-          responseListener.onFailure(e);
-        } finally {
-          direct.executionToClose().ifPresent(PPLAsyncQueryService::closeExecution);
-          closeTask(direct.task());
-        }
-      }
-      case Transition.ExecutionFinished finished -> {
-        releaseRunning();
-        finished.executionToClose().ifPresent(PPLAsyncQueryService::closeExecution);
-        closeTask(finished.task());
-      }
-    }
-  }
-
-  private void applyRemoval(PPLAsyncQueryJob job, Removal removal) {
-    if (removal == null) {
-      return;
-    }
-    switch (removal) {
-      case Removal.Deleted deleted -> applyRemoval(job, deleted.resources(), deleted.reason());
-      case Removal.Expired expired -> applyRemoval(job, expired.resources(), expired.reason());
-      case Removal.Discarded discarded ->
-          applyRemoval(job, discarded.resources(), discarded.reason());
-    }
-  }
-
-  private void applyRemoval(
-      PPLAsyncQueryJob job, PPLAsyncQueryJob.DetachedResources resources, String reason) {
-    if (jobs.remove(job.id(), job)) {
-      if (resources.releasesRunningSlot()) {
-        releaseRunning();
-      }
-      releaseRetained();
-    }
-    cancel(resources.task(), reason);
-    closeExecution(resources.execution());
-  }
-
-  private JobSnapshot materialize(SnapshotSource source) {
-    return switch (source) {
-      case SnapshotSource.Running running ->
-          new JobSnapshot.Running(
-              running.id(),
-              running
-                  .execution()
-                  .flatMap(AsyncQueryExecution::currentResult)
-                  .map(PPLAsyncQueryService::snapshotResponse));
-      case SnapshotSource.Succeeded succeeded ->
-          new JobSnapshot.Succeeded(
-              succeeded.id(),
-              succeeded
-                  .execution()
-                  .currentResult()
-                  .map(PPLAsyncQueryService::snapshotResponse)
-                  .orElseThrow(
-                      () ->
-                          new IllegalStateException(
-                              "Successful PPL asynchronous execution completed without a final"
-                                  + " result")),
-              succeeded.tookMillis());
-      case SnapshotSource.Failed failed ->
-          new JobSnapshot.Failed(failed.id(), failed.failure(), failed.tookMillis());
-    };
-  }
-
-  private static void closeExecution(AsyncQueryExecution execution) {
-    if (execution == null) {
-      return;
-    }
-    try {
-      execution.close();
-    } catch (RuntimeException e) {
-      LOG.warn(
-          "Failed to close PPL asynchronous query execution ({})", e.getClass().getSimpleName());
-    }
-  }
-
-  private static void closeTask(JobTask task) {
-    if (task != null) {
-      task.close();
-    }
-  }
-
-  private PPLAsyncQueryJob findLocal(String encodedId) {
-    PPLAsyncQueryJob job = jobs.get(encodedId);
-    if (job == null) {
-      throw notFound();
-    }
-    return job;
-  }
-
-  void validateKeepAlive(TimeValue keepAlive) {
-    TimeValue maximum = maxKeepAlive.get();
-    if (keepAlive == null || keepAlive.millis() <= 0 || keepAlive.millis() > maximum.millis()) {
-      throw new IllegalArgumentException(
-          "[keep_alive] must be greater than 0 and no more than " + maximum);
-    }
-  }
-
-  void validateWaitForCompletion(TimeValue waitForCompletion) {
-    TimeValue maximum = maxWaitForCompletion.get();
-    if (waitForCompletion == null
-        || waitForCompletion.millis() < 0
-        || waitForCompletion.millis() > maximum.millis()) {
-      throw new IllegalArgumentException(
-          "[wait_for_completion_timeout] must be between 0 and " + maximum);
-    }
-  }
-
-  /**
-   * Copies an execution-owned response before exposing it through a job snapshot.
-   *
-   * 
This prevents later execution updates from changing a response already handed to a caller.
-   */
-  private static QueryResponse snapshotResponse(QueryResponse response) {
-    Schema schema = new Schema(List.copyOf(response.getSchema().getColumns()));
-    QueryResponse copy =
-        new QueryResponse(schema, List.copyOf(response.getResults()), response.getCursor());
-    copy.setWarnings(List.copyOf(response.getWarnings()));
-    return copy;
-  }
-
-  private void cancel(JobTask task, String reason) {
-    if (task == null) {
-      return;
-    }
-    CancellableTask cancellableTask = task.task();
-    if (cancellableTask == null || cancellableTask.isCancelled()) {
-      task.close();
-      return;
-    }
-    try {
-      TaskManager currentTaskManager = taskManager;
-      if (currentTaskManager == null) {
-        cancellableTask.cancel(reason);
-        task.close();
-      } else {
-        currentTaskManager.cancelTaskAndDescendants(
-            cancellableTask,
-            reason,
-            false,
-            ActionListener.wrap(
-                ignored -> task.close(),
-                failure -> {
-                  task.close();
-                  LOG.warn(
-                      "Failed to cancel descendants of PPL asynchronous query task ({})",
-                      failure.getClass().getSimpleName());
-                }));
-      }
-    } catch (RuntimeException e) {
-      task.close();
-      LOG.warn("Failed to cancel PPL asynchronous query task ({})", e.getClass().getSimpleName());
-    }
-  }
-
-  /**
-   * Registers the independently cancellable task owned by an asynchronous job.
-   *
-   * 
The POST request task is assigned as parent during registration so cancellation can reach
-   * startup work. The returned {@link JobTask} owns both task unregistration and child-node
-   * registration cleanup.
-   *
-   * @param request request used by {@link TaskManager} to create the job task
-   * @param requestTask task associated with the POST request
-   * @return job-owned task and its registration cleanup
-   */
-  private JobTask registerJobTask(TransportPPLQueryRequest request, PPLQueryTask requestTask) {
-    TaskManager currentTaskManager =
-        Objects.requireNonNull(
-            taskManager, "PPL asynchronous query task manager is not initialized");
-    Objects.requireNonNull(requestTask, "PPL asynchronous query request task is not initialized");
-    DiscoveryNode localNode =
-        Objects.requireNonNull(currentTaskManager.localNode(), "Local node is not initialized");
-
-    // This task is registered directly rather than through TransportAction.execute(), so reproduce
-    // the two pieces of OpenSearch child-task bookkeeping that TransportAction normally performs.
-    // The child-node registration lets parent cancellation send a ban to this node; parentTaskId
-    // lets that ban find and cancel the retained task.
-    Releasable childNodeRegistration =
-        currentTaskManager.registerChildNode(requestTask.getId(), localNode);
-    TaskId originalParent = request.getParentTask();
-    boolean registered = false;
-    try {
-      request.setParentTask(localNode.getId(), requestTask.getId());
-      Task task = currentTaskManager.register("transport", PPLQueryAction.NAME, request);
-      if (!(task instanceof PPLQueryTask pplQueryTask)) {
-        currentTaskManager.unregister(task);
-        throw new IllegalStateException("Failed to create PPL asynchronous query task");
-      }
-      registered = true;
-      return new JobTask(
-          pplQueryTask,
-          () -> {
-            try {
-              currentTaskManager.unregister(pplQueryTask);
-            } finally {
-              childNodeRegistration.close();
-            }
-          });
-    } finally {
-      request.setParentTask(originalParent);
-      if (!registered) {
-        childNodeRegistration.close();
-      }
-    }
-  }
-
-  private static Exception asException(Throwable failure) {
-    Throwable cause =
-        failure instanceof java.util.concurrent.CompletionException && failure.getCause() != null
-            ? failure.getCause()
-            : failure;
-    return cause instanceof Exception exception ? exception : new RuntimeException(cause);
-  }
-
-  private static ResourceNotFoundException notFound() {
-    return new ResourceNotFoundException("PPL asynchronous query not found");
-  }
-
-  /** Starts accepting submissions and schedules periodic retained-job expiration. */
-  @Override
-  protected void doStart() {
-    acceptingNewJobs = true;
-    if (threadPool != null) {
-      reaper =
-          threadPool.scheduleWithFixedDelay(
-              this::reapExpired, REAPER_INTERVAL, ThreadPool.Names.GENERIC);
-    }
-  }
-
-  /** Stops accepting submissions and cancels the periodic expiration task. */
-  @Override
-  protected void doStop() {
-    acceptingNewJobs = false;
-    Scheduler.Cancellable scheduledReaper = reaper;
-    if (scheduledReaper != null) {
-      scheduledReaper.cancel();
-      reaper = null;
-    }
-  }
-
-  /**
-   * Removes all remaining jobs and releases their task and execution resources.
-   *
-   * @throws IOException if lifecycle shutdown fails
-   */
-  @Override
-  protected void doClose() throws IOException {
-    acceptingNewJobs = false;
-    jobs.forEach(
-        (id, job) -> applyRemoval(job, job.close("PPL asynchronous query service is closing")));
-  }
-}
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJob.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJob.java
new file mode 100644
index 00000000000..f105047ebe0
--- /dev/null
+++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJob.java
@@ -0,0 +1,983 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.plugin.transport.asyncquery;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.CompletionException;
+import java.util.function.Function;
+import java.util.function.LongSupplier;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.ResourceNotFoundException;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.common.lease.Releasable;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.tasks.TaskId;
+import org.opensearch.sql.executor.AsyncQueryExecution;
+import org.opensearch.sql.executor.ExecutionEngine.QueryResponse;
+import org.opensearch.sql.executor.ExecutionEngine.Schema;
+import org.opensearch.sql.plugin.PPLQueryErrorHandler;
+import org.opensearch.sql.plugin.transport.PPLQueryAction;
+import org.opensearch.sql.plugin.transport.PPLQueryTask;
+import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest;
+import org.opensearch.tasks.CancellableTask;
+import org.opensearch.tasks.Task;
+import org.opensearch.tasks.TaskManager;
+import org.opensearch.threadpool.Scheduler.ScheduledCancellable;
+import org.opensearch.threadpool.ThreadPool;
+
+/**
+ * Active object representing one asynchronous PPL query on the owner node.
+ *
+ * 
Each job owns its execution handle, cancellable task, retention deadline, and lease timer. It
+ * publishes itself to a {@link QueryJobRegistry} at creation and removes itself when it becomes
+ * unreachable through DELETE, expiry, direct completion, or shutdown. All state transitions are
+ * synchronized on the job; every observable side effect (registry mutation, listener callback, task
+ * cancellation, execution close) happens after the lock has been released.
+ *
+ * 
The state machine tracks two axes: retention and completion.
+ *
+ * 
+ * Current state       Event               Next state             Response
+ * RUNNING             direct success      REMOVED                final snapshot without ID
+ * RUNNING             direct failure      REMOVED                failure without ID
+ * RUNNING             retain deadline     RETAINED_RUNNING       running snapshot with ID
+ * RETAINED_RUNNING    success             RETAINED_SUCCEEDED     none (final result kept)
+ * RETAINED_RUNNING    failure             RETAINED_FAILED        none (failure kept)
+ * *                   DELETE              REMOVED                status observed at DELETE
+ * RETAINED_*          expire / shutdown   REMOVED                none
+ * 

+ */
+public final class QueryJob {
+  private static final Logger LOG = LogManager.getLogger(QueryJob.class);
+
+  private static final Cancellable NO_TIMEOUT = () -> {};
+
+  /** Lifecycle state exposed in asynchronous PPL responses. */
+  public enum Status {
+    /** Query execution is still running. */
+    RUNNING,
+
+    /** Query execution completed successfully. */
+    SUCCEEDED,
+
+    /** Query execution failed. */
+    FAILED,
+
+    /** Query execution was cancelled by DELETE. */
+    CANCELLED
+  }
+
+  /**
+   * Immutable point-in-time response for a job.
+   *
+   * 
Each variant exposes only the data valid for its lifecycle state.
+   */
+  public sealed interface Snapshot {
+    /**
+     * Returns the retained job ID.
+     *
+     * @return job ID, or empty for a final response returned directly by POST
+     */
+    Optional id();
+
+    /**
+     * Returns the public lifecycle state.
+     *
+     * @return response status
+     */
+    Status status();
+
+    /**
+     * Snapshot of a query that is still running.
+     *
+     * @param jobId retained job ID
+     * @param response current result, or empty before a result is available
+     */
+    record Running(QueryJobId jobId, Optional response) implements Snapshot {
+      @Override
+      public Optional id() {
+        return Optional.of(jobId);
+      }
+
+      @Override
+      public Status status() {
+        return Status.RUNNING;
+      }
+    }
+
+    /**
+     * Snapshot of a successfully completed query.
+     *
+     * @param id retained job ID, or empty for a direct POST response
+     * @param response final query result
+     * @param tookMillis elapsed execution time
+     */
+    record Succeeded(Optional id, QueryResponse response, long tookMillis)
+        implements Snapshot {
+      @Override
+      public Status status() {
+        return Status.SUCCEEDED;
+      }
+    }
+
+    /**
+     * Snapshot of a failed query.
+     *
+     * @param id retained job ID, or empty for a direct POST response
+     * @param failure client-visible failure
+     * @param tookMillis elapsed execution time
+     */
+    record Failed(Optional id, Failure failure, long tookMillis) implements Snapshot {
+      @Override
+      public Status status() {
+        return Status.FAILED;
+      }
+    }
+  }
+
+  /**
+   * Client-visible failure retained with a failed job.
+   *
+   * @param type exception type
+   * @param reason client-facing failure reason
+   */
+  public record Failure(String type, String reason) {
+    /**
+     * Derives a client-visible failure from an exception.
+     *
+     * @param exception exception raised by execution
+     * @return failure with the exception's simple type and message
+     */
+    public static Failure from(Exception exception) {
+      String type =
+          exception.getClass().getSimpleName().isBlank()
+              ? exception.getClass().getName()
+              : exception.getClass().getSimpleName();
+      String reason =
+          exception.getMessage() == null || exception.getMessage().isBlank()
+              ? "query execution failed"
+              : exception.getMessage();
+      return new Failure(type, reason);
+    }
+  }
+
+  /** Cancels a previously scheduled timer. Multiple cancellations are safe. */
+  @FunctionalInterface
+  public interface Cancellable {
+    /** Cancels the timer if it has not already fired. */
+    void cancel();
+  }
+
+  /** Schedules a one-shot runnable. Production wraps {@link ThreadPool}; tests use a fake. */
+  @FunctionalInterface
+  public interface Scheduler {
+    /**
+     * Schedules {@code task} to run once after {@code delayMillis} milliseconds.
+     *
+     * @param delayMillis delay before firing
+     * @param task action to run
+     * @return handle used to cancel the scheduled task
+     */
+    Cancellable schedule(long delayMillis, Runnable task);
+  }
+
+  /** Owned cancellable task and its registration cleanup. Package-private for tests. */
+  record JobTask(CancellableTask task, Runnable release) {
+    /** Releases task-manager registrations owned by this wrapper. Safe to call once. */
+    void close() {
+      release.run();
+    }
+  }
+
+  /** Internal lifecycle; unlike {@link Status}, this includes retention and removal. */
+  private enum State {
+    RUNNING,
+    RETAINED_RUNNING,
+    RETAINED_SUCCEEDED,
+    RETAINED_FAILED,
+    REMOVED
+  }
+
+  private final QueryJobId id;
+  private final QueryJobRegistry registry;
+  private final LongSupplier clock;
+  private final Scheduler scheduler;
+  private final TaskManager taskManager;
+  private final QueryJobOwner owner;
+  private final long startTimeMillis;
+  private final ActionListener responseListener;
+
+  private State state = State.RUNNING;
+  private long keepAliveMillis;
+  private long expirationTimeMillis;
+  private long completionTimeMillis = -1L;
+  private JobTask task;
+  private AsyncQueryExecution execution;
+  private Failure failure;
+  private Cancellable retentionTimer = NO_TIMEOUT;
+  private Cancellable expiryTimer = NO_TIMEOUT;
+
+  private QueryJob(
+      QueryJobId id,
+      QueryJobRegistry registry,
+      LongSupplier clock,
+      Scheduler scheduler,
+      TaskManager taskManager,
+      QueryJobOwner owner,
+      long keepAliveMillis,
+      JobTask task,
+      ActionListener responseListener) {
+    this.id = id;
+    this.registry = registry;
+    this.clock = clock;
+    this.scheduler = scheduler;
+    this.taskManager = taskManager;
+    this.owner = owner;
+    this.startTimeMillis = clock.getAsLong();
+    this.keepAliveMillis = keepAliveMillis;
+    this.expirationTimeMillis = startTimeMillis + keepAliveMillis;
+    this.task = task;
+    this.responseListener = responseListener;
+  }
+
+  /**
+   * Starts an asynchronous PPL query in production wiring.
+   *
+   * 
{@link TaskManager} registers the job's cancellable task as a child of {@code parentTask} so
+   * cancellation of the POST request cascades. The job schedules its own retention deadline on
+   * {@code threadPool} and calls {@code executionStarter} inline. If execution finishes before the
+   * deadline, {@code responseListener} receives the final result without a job ID. Otherwise the
+   * job becomes retained and the listener receives a running snapshot with the ID.
+   *
+   * @param registry registry that will hold the job while it is discoverable
+   * @param threadPool schedules the retention deadline and keep-alive expiry
+   * @param taskManager registers and cancels the job's own task
+   * @param parentTask POST request task; the job task becomes its child
+   * @param request transport request used to register the job task
+   * @param owner caller retained with the job for later authorization
+   * @param keepAlive lease duration applied when the job is retained
+   * @param waitForCompletion maximum time to wait for a direct result
+   * @param executionStarter starts execution using the job-owned cancellable task
+   * @param responseListener receives either the direct result or the retained running snapshot
+   * @return newly created job, already published to {@code registry}
+   * @throws NullPointerException if any required argument is null
+   */
+  public static QueryJob create(
+      QueryJobRegistry registry,
+      ThreadPool threadPool,
+      TaskManager taskManager,
+      PPLQueryTask parentTask,
+      TransportPPLQueryRequest request,
+      QueryJobOwner owner,
+      TimeValue keepAlive,
+      TimeValue waitForCompletion,
+      Function executionStarter,
+      ActionListener responseListener) {
+    Objects.requireNonNull(threadPool);
+    Objects.requireNonNull(taskManager);
+    JobTask jobTask = registerJobTask(taskManager, parentTask, request);
+    try {
+      return create(
+          taskManager.localNode().getId(),
+          registry,
+          System::currentTimeMillis,
+          wrapScheduler(threadPool),
+          taskManager,
+          jobTask,
+          owner,
+          keepAlive,
+          waitForCompletion,
+          executionStarter,
+          responseListener);
+    } catch (RuntimeException | Error e) {
+      jobTask.close();
+      throw e;
+    }
+  }
+
+  /**
+   * Test entry point that accepts injectable clock, scheduler, and pre-registered task.
+   *
+   * @param localNodeId local node ID used to route the ID back on future requests
+   * @param registry registry that will hold the job while it is discoverable
+   * @param clock time source
+   * @param scheduler timer used for retention deadline and keep-alive expiry
+   * @param taskManager cancels the job task with descendant cascade; may be {@code null} when the
+   *     job task should be cancelled directly
+   * @param task pre-registered cancellable task owned by the job
+   * @param owner caller retained with the job for later authorization
+   * @param keepAlive lease duration applied when the job is retained
+   * @param waitForCompletion maximum time to wait for a direct result
+   * @param executionStarter starts execution using the job-owned cancellable task
+   * @param responseListener receives either the direct result or the retained running snapshot
+   * @return newly created job, already published to {@code registry}
+   */
+  static QueryJob create(
+      String localNodeId,
+      QueryJobRegistry registry,
+      LongSupplier clock,
+      Scheduler scheduler,
+      TaskManager taskManager,
+      JobTask task,
+      QueryJobOwner owner,
+      TimeValue keepAlive,
+      TimeValue waitForCompletion,
+      Function executionStarter,
+      ActionListener responseListener) {
+    Objects.requireNonNull(localNodeId);
+    Objects.requireNonNull(registry);
+    Objects.requireNonNull(clock);
+    Objects.requireNonNull(scheduler);
+    Objects.requireNonNull(task);
+    Objects.requireNonNull(owner);
+    Objects.requireNonNull(keepAlive);
+    Objects.requireNonNull(waitForCompletion);
+    Objects.requireNonNull(executionStarter);
+    Objects.requireNonNull(responseListener);
+
+    QueryJob job =
+        publish(
+            localNodeId,
+            registry,
+            clock,
+            scheduler,
+            taskManager,
+            owner,
+            keepAlive.millis(),
+            task,
+            responseListener);
+    try {
+      job.scheduleRetention(waitForCompletion);
+    } catch (RuntimeException | Error e) {
+      job.discard("PPL asynchronous query startup failed");
+      throw e;
+    }
+    job.startExecution(executionStarter);
+    return job;
+  }
+
+  private static QueryJob publish(
+      String localNodeId,
+      QueryJobRegistry registry,
+      LongSupplier clock,
+      Scheduler scheduler,
+      TaskManager taskManager,
+      QueryJobOwner owner,
+      long keepAliveMillis,
+      JobTask task,
+      ActionListener responseListener) {
+    while (true) {
+      QueryJobId id = QueryJobId.create(localNodeId);
+      QueryJob job =
+          new QueryJob(
+              id,
+              registry,
+              clock,
+              scheduler,
+              taskManager,
+              owner,
+              keepAliveMillis,
+              task,
+              responseListener);
+      if (registry.add(job) == null) {
+        return job;
+      }
+    }
+  }
+
+  /**
+   * Returns the opaque ID assigned to this job.
+   *
+   * @return job ID used as the registry key
+   */
+  public QueryJobId getJobId() {
+    return id;
+  }
+
+  /**
+   * Returns the immutable owner identity captured when this job was created.
+   *
+   * @return job owner used for GET and DELETE authorization
+   */
+  public QueryJobOwner getOwner() {
+    return owner;
+  }
+
+  /**
+   * Returns the current snapshot of the job.
+   *
+   * 
A non-null {@code keepAlive} starts a new lease from the current time. A call at or after
+   * the current expiration removes the job and reports it as not found.
+   *
+   * @param keepAlive new lease duration, or {@code null} to leave the current lease unchanged
+   * @return immutable snapshot of the current state
+   * @throws ResourceNotFoundException if the job has been removed or has just expired
+   */
+  public Snapshot get(TimeValue keepAlive) {
+    long now = clock.getAsLong();
+    AsyncQueryExecution executionForResult;
+    State stateForResult;
+    long tookMillis;
+    Failure failureForResult;
+    boolean expired = false;
+    Cancellable retentionToCancel = NO_TIMEOUT;
+    Cancellable expiryToCancel = NO_TIMEOUT;
+    Cancellable rescheduledExpiry = NO_TIMEOUT;
+    JobTask taskToClose = null;
+    AsyncQueryExecution executionToClose = null;
+
+    synchronized (this) {
+      ensurePresent();
+      if (state != State.RUNNING && now >= expirationTimeMillis) {
+        expired = true;
+        retentionToCancel = retentionTimer;
+        expiryToCancel = expiryTimer;
+        retentionTimer = NO_TIMEOUT;
+        expiryTimer = NO_TIMEOUT;
+        taskToClose = detachTask();
+        executionToClose = detachExecution();
+        state = State.REMOVED;
+        executionForResult = null;
+        stateForResult = null;
+        tookMillis = 0L;
+        failureForResult = null;
+      } else {
+        if (keepAlive != null && state != State.RUNNING) {
+          keepAliveMillis = keepAlive.millis();
+          expirationTimeMillis = now + keepAliveMillis;
+          expiryToCancel = expiryTimer;
+          expiryTimer = NO_TIMEOUT;
+          rescheduledExpiry = scheduler.schedule(keepAliveMillis, this::expire);
+          expiryTimer = rescheduledExpiry;
+        }
+        executionForResult = execution;
+        stateForResult = state;
+        tookMillis =
+            completionTimeMillis < 0 ? -1L : Math.max(0L, completionTimeMillis - startTimeMillis);
+        failureForResult = failure;
+      }
+    }
+
+    if (expired) {
+      retentionToCancel.cancel();
+      expiryToCancel.cancel();
+      registry.remove(id, this);
+      cancelTaskAsync(taskToClose, "PPL asynchronous query expired");
+      closeExecution(executionToClose);
+      throw notFound();
+    }
+
+    expiryToCancel.cancel();
+    return buildSnapshot(stateForResult, executionForResult, failureForResult, tookMillis, true);
+  }
+
+  /**
+   * Cancels this job and returns the status observed at cancellation.
+   *
+   * @param reason human-readable reason recorded on the underlying task
+   * @return the {@link Status} observed when the job was removed
+   * @throws ResourceNotFoundException if the job has been removed or has just expired
+   */
+  public Status cancel(String reason) {
+    Objects.requireNonNull(reason);
+    Status observed;
+    Cancellable retentionToCancel;
+    Cancellable expiryToCancel;
+    JobTask taskToClose;
+    AsyncQueryExecution executionToClose;
+    String cancelReason;
+    boolean cancelRunning;
+
+    synchronized (this) {
+      ensurePresent();
+      if (state != State.RUNNING && clock.getAsLong() >= expirationTimeMillis) {
+        // Late DELETE races expiry; the job is gone.
+        retentionToCancel = retentionTimer;
+        expiryToCancel = expiryTimer;
+        retentionTimer = NO_TIMEOUT;
+        expiryTimer = NO_TIMEOUT;
+        taskToClose = detachTask();
+        executionToClose = detachExecution();
+        state = State.REMOVED;
+        cancelReason = "PPL asynchronous query expired";
+        cancelRunning = true;
+        observed = null;
+      } else {
+        observed =
+            switch (state) {
+              case RUNNING, RETAINED_RUNNING -> Status.CANCELLED;
+              case RETAINED_SUCCEEDED -> Status.SUCCEEDED;
+              case RETAINED_FAILED -> Status.FAILED;
+              case REMOVED -> throw new IllegalStateException("PPL asynchronous query was removed");
+            };
+        cancelRunning = (state == State.RUNNING || state == State.RETAINED_RUNNING);
+        retentionToCancel = retentionTimer;
+        expiryToCancel = expiryTimer;
+        retentionTimer = NO_TIMEOUT;
+        expiryTimer = NO_TIMEOUT;
+        taskToClose = detachTask();
+        executionToClose = detachExecution();
+        state = State.REMOVED;
+        cancelReason = reason;
+      }
+    }
+
+    retentionToCancel.cancel();
+    expiryToCancel.cancel();
+    registry.remove(id, this);
+    if (cancelRunning) {
+      cancelTaskAsync(taskToClose, cancelReason);
+    } else {
+      closeTask(taskToClose);
+    }
+    closeExecution(executionToClose);
+
+    if (observed == null) {
+      throw notFound();
+    }
+    return observed;
+  }
+
+  /**
+   * Removes the job as part of an internal cleanup (shutdown or startup abort). Idempotent.
+   *
+   * @param reason reason recorded on the underlying task
+   */
+  void discard(String reason) {
+    Objects.requireNonNull(reason);
+    Cancellable retentionToCancel;
+    Cancellable expiryToCancel;
+    JobTask taskToClose;
+    AsyncQueryExecution executionToClose;
+    boolean wasExecuting;
+
+    synchronized (this) {
+      if (state == State.REMOVED) {
+        return;
+      }
+      wasExecuting = (state == State.RUNNING || state == State.RETAINED_RUNNING);
+      retentionToCancel = retentionTimer;
+      expiryToCancel = expiryTimer;
+      retentionTimer = NO_TIMEOUT;
+      expiryTimer = NO_TIMEOUT;
+      taskToClose = detachTask();
+      executionToClose = detachExecution();
+      state = State.REMOVED;
+    }
+
+    retentionToCancel.cancel();
+    expiryToCancel.cancel();
+    registry.remove(id, this);
+    if (wasExecuting) {
+      cancelTaskAsync(taskToClose, reason);
+    } else {
+      closeTask(taskToClose);
+    }
+    closeExecution(executionToClose);
+  }
+
+  // ---------- Retention deadline ----------
+
+  private void scheduleRetention(TimeValue waitForCompletion) {
+    if (waitForCompletion.millis() == 0) {
+      onRetentionDeadline();
+      return;
+    }
+    Cancellable timer = scheduler.schedule(waitForCompletion.millis(), this::onRetentionDeadline);
+    synchronized (this) {
+      if (state == State.RUNNING) {
+        retentionTimer = timer;
+        return;
+      }
+    }
+    // Retention state already changed while the timer was being wired up. Cancel it.
+    timer.cancel();
+  }
+
+  private void onRetentionDeadline() {
+    Snapshot snapshotToDeliver;
+    Cancellable rescheduledExpiry;
+
+    synchronized (this) {
+      if (state != State.RUNNING) {
+        return;
+      }
+      state = State.RETAINED_RUNNING;
+      long now = clock.getAsLong();
+      expirationTimeMillis = now + keepAliveMillis;
+      snapshotToDeliver = new Snapshot.Running(id, currentResultCopy());
+      retentionTimer = NO_TIMEOUT;
+      rescheduledExpiry = scheduler.schedule(keepAliveMillis, this::expire);
+      expiryTimer = rescheduledExpiry;
+    }
+
+    try {
+      responseListener.onResponse(snapshotToDeliver);
+    } catch (RuntimeException e) {
+      // Cannot deliver retained ID; discard the job so no state leaks.
+      discard("PPL asynchronous query startup failed");
+      safeNotifyFailure(e);
+    }
+  }
+
+  // ---------- Execution attachment and completion ----------
+
+  private void startExecution(Function executionStarter) {
+    AsyncQueryExecution started;
+    try {
+      started = Objects.requireNonNull(executionStarter.apply(task().task()));
+    } catch (RuntimeException e) {
+      fail(e);
+      return;
+    }
+    if (!tryAttachExecution(started)) {
+      closeExecution(started);
+      return;
+    }
+    started
+        .completion()
+        .whenComplete(
+            (ignored, err) -> {
+              if (err == null) {
+                complete();
+              } else {
+                fail(asException(err));
+              }
+            });
+  }
+
+  private synchronized JobTask task() {
+    return task;
+  }
+
+  private synchronized boolean tryAttachExecution(AsyncQueryExecution incoming) {
+    if (state != State.RUNNING && state != State.RETAINED_RUNNING) {
+      return false;
+    }
+    if (execution != null) {
+      return false;
+    }
+    execution = incoming;
+    return true;
+  }
+
+  private void complete() {
+    long now = clock.getAsLong();
+    boolean directResponse;
+    JobTask taskToClose;
+    AsyncQueryExecution executionForSnapshot;
+    Cancellable retentionToCancel;
+
+    synchronized (this) {
+      if (state != State.RUNNING && state != State.RETAINED_RUNNING) {
+        return;
+      }
+      if (execution == null) {
+        throw new IllegalStateException(
+            "PPL asynchronous execution must be attached before successful completion");
+      }
+      completionTimeMillis = now;
+      directResponse = (state == State.RUNNING);
+      retentionToCancel = retentionTimer;
+      retentionTimer = NO_TIMEOUT;
+      taskToClose = detachTask();
+      if (directResponse) {
+        executionForSnapshot = detachExecution();
+        state = State.REMOVED;
+      } else {
+        executionForSnapshot = null;
+        state = State.RETAINED_SUCCEEDED;
+      }
+    }
+
+    retentionToCancel.cancel();
+    if (directResponse) {
+      registry.remove(id, this);
+      Snapshot snapshot;
+      try {
+        QueryResponse result =
+            executionForSnapshot
+                .currentResult()
+                .orElseThrow(
+                    () ->
+                        new IllegalStateException(
+                            "Successful PPL asynchronous execution completed without a final"
+                                + " result"));
+        snapshot =
+            new Snapshot.Succeeded(Optional.empty(), defensiveCopy(result), now - startTimeMillis);
+      } catch (RuntimeException e) {
+        safeNotifyFailure(e);
+        closeTask(taskToClose);
+        closeExecution(executionForSnapshot);
+        return;
+      }
+      try {
+        responseListener.onResponse(snapshot);
+      } catch (RuntimeException e) {
+        safeNotifyFailure(e);
+      } finally {
+        closeTask(taskToClose);
+        closeExecution(executionForSnapshot);
+      }
+    } else {
+      closeTask(taskToClose);
+    }
+  }
+
+  private void fail(Exception cause) {
+    Objects.requireNonNull(cause);
+    long now = clock.getAsLong();
+    boolean directResponse;
+    JobTask taskToClose;
+    AsyncQueryExecution executionToClose;
+    Cancellable retentionToCancel;
+    Failure captured = Failure.from(cause);
+
+    synchronized (this) {
+      if (state != State.RUNNING && state != State.RETAINED_RUNNING) {
+        return;
+      }
+      completionTimeMillis = now;
+      failure = captured;
+      directResponse = (state == State.RUNNING);
+      retentionToCancel = retentionTimer;
+      retentionTimer = NO_TIMEOUT;
+      taskToClose = detachTask();
+      executionToClose = detachExecution();
+      state = directResponse ? State.REMOVED : State.RETAINED_FAILED;
+    }
+
+    retentionToCancel.cancel();
+    if (directResponse) {
+      registry.remove(id, this);
+      try {
+        responseListener.onFailure(cause);
+      } catch (RuntimeException e) {
+        LOG.warn("PPL asynchronous listener rejected failure ({})", e.getClass().getSimpleName());
+      } finally {
+        closeTask(taskToClose);
+        closeExecution(executionToClose);
+      }
+    } else {
+      PPLQueryErrorHandler.recordFailure(cause);
+      closeTask(taskToClose);
+      closeExecution(executionToClose);
+    }
+  }
+
+  private void expire() {
+    long now = clock.getAsLong();
+    Cancellable retentionToCancel;
+    JobTask taskToClose;
+    AsyncQueryExecution executionToClose;
+
+    synchronized (this) {
+      if (state == State.REMOVED || state == State.RUNNING || now < expirationTimeMillis) {
+        return;
+      }
+      retentionToCancel = retentionTimer;
+      retentionTimer = NO_TIMEOUT;
+      expiryTimer = NO_TIMEOUT;
+      taskToClose = detachTask();
+      executionToClose = detachExecution();
+      state = State.REMOVED;
+    }
+
+    retentionToCancel.cancel();
+    registry.remove(id, this);
+    cancelTaskAsync(taskToClose, "PPL asynchronous query expired");
+    closeExecution(executionToClose);
+  }
+
+  // ---------- Helpers ----------
+
+  private void ensurePresent() {
+    if (state == State.REMOVED) {
+      throw notFound();
+    }
+  }
+
+  private JobTask detachTask() {
+    JobTask detached = task;
+    task = null;
+    return detached;
+  }
+
+  private AsyncQueryExecution detachExecution() {
+    AsyncQueryExecution detached = execution;
+    execution = null;
+    return detached;
+  }
+
+  private Optional currentResultCopy() {
+    if (execution == null) {
+      return Optional.empty();
+    }
+    return execution.currentResult().map(QueryJob::defensiveCopy);
+  }
+
+  private Snapshot buildSnapshot(
+      State from,
+      AsyncQueryExecution execForResult,
+      Failure failureForResult,
+      long tookMillis,
+      boolean withId) {
+    Optional snapshotId = withId ? Optional.of(id) : Optional.empty();
+    return switch (from) {
+      case RUNNING, RETAINED_RUNNING -> {
+        Optional current =
+            execForResult == null
+                ? Optional.empty()
+                : execForResult.currentResult().map(QueryJob::defensiveCopy);
+        yield new Snapshot.Running(id, current);
+      }
+      case RETAINED_SUCCEEDED -> {
+        QueryResponse result =
+            Objects.requireNonNull(execForResult, "attached execution required")
+                .currentResult()
+                .map(QueryJob::defensiveCopy)
+                .orElseThrow(
+                    () ->
+                        new IllegalStateException(
+                            "Successful PPL asynchronous execution completed without a final"
+                                + " result"));
+        yield new Snapshot.Succeeded(snapshotId, result, tookMillis);
+      }
+      case RETAINED_FAILED ->
+          new Snapshot.Failed(snapshotId, Objects.requireNonNull(failureForResult), tookMillis);
+      case REMOVED -> throw new IllegalStateException("PPL asynchronous query was removed");
+    };
+  }
+
+  private void safeNotifyFailure(Exception cause) {
+    try {
+      responseListener.onFailure(cause);
+    } catch (RuntimeException e) {
+      LOG.warn("PPL asynchronous listener rejected failure ({})", e.getClass().getSimpleName());
+    }
+  }
+
+  private static QueryResponse defensiveCopy(QueryResponse response) {
+    Schema schema = new Schema(List.copyOf(response.getSchema().getColumns()));
+    QueryResponse copy =
+        new QueryResponse(schema, List.copyOf(response.getResults()), response.getCursor());
+    copy.setWarnings(List.copyOf(response.getWarnings()));
+    return copy;
+  }
+
+  private static void closeTask(JobTask task) {
+    if (task != null) {
+      task.close();
+    }
+  }
+
+  private void cancelTaskAsync(JobTask jobTask, String reason) {
+    if (jobTask == null) {
+      return;
+    }
+    CancellableTask cancellable = jobTask.task();
+    if (cancellable == null || cancellable.isCancelled()) {
+      jobTask.close();
+      return;
+    }
+    try {
+      if (taskManager == null) {
+        cancellable.cancel(reason);
+        jobTask.close();
+      } else {
+        taskManager.cancelTaskAndDescendants(
+            cancellable,
+            reason,
+            false,
+            ActionListener.wrap(
+                ignored -> jobTask.close(),
+                error -> {
+                  jobTask.close();
+                  LOG.warn(
+                      "Failed to cancel descendants of PPL asynchronous query task ({})",
+                      error.getClass().getSimpleName());
+                }));
+      }
+    } catch (RuntimeException e) {
+      jobTask.close();
+      LOG.warn("Failed to cancel PPL asynchronous query task ({})", e.getClass().getSimpleName());
+    }
+  }
+
+  private static void closeExecution(AsyncQueryExecution execution) {
+    if (execution == null) {
+      return;
+    }
+    try {
+      execution.close();
+    } catch (RuntimeException e) {
+      LOG.warn(
+          "Failed to close PPL asynchronous query execution ({})", e.getClass().getSimpleName());
+    }
+  }
+
+  private static Exception asException(Throwable failure) {
+    Throwable cause =
+        failure instanceof CompletionException && failure.getCause() != null
+            ? failure.getCause()
+            : failure;
+    return cause instanceof Exception exception ? exception : new RuntimeException(cause);
+  }
+
+  private static ResourceNotFoundException notFound() {
+    return new ResourceNotFoundException("PPL asynchronous query not found");
+  }
+
+  private static Scheduler wrapScheduler(ThreadPool threadPool) {
+    return (delayMillis, task) -> {
+      ScheduledCancellable cancellable =
+          threadPool.schedule(
+              task, TimeValue.timeValueMillis(delayMillis), ThreadPool.Names.GENERIC);
+      return cancellable::cancel;
+    };
+  }
+
+  private static JobTask registerJobTask(
+      TaskManager taskManager, PPLQueryTask parentTask, TransportPPLQueryRequest request) {
+    Objects.requireNonNull(parentTask, "PPL asynchronous query request task is not initialized");
+    DiscoveryNode localNode =
+        Objects.requireNonNull(taskManager.localNode(), "Local node is not initialized");
+
+    // Registered directly rather than through TransportAction.execute(); reproduce the two child-
+    // task bookkeeping steps TransportAction would perform. childNode registration lets parent
+    // cancellation ban this node; parentTaskId lets the ban find and cancel the retained task.
+    Releasable childNodeRegistration = taskManager.registerChildNode(parentTask.getId(), localNode);
+    TaskId originalParent = request.getParentTask();
+    boolean registered = false;
+    try {
+      request.setParentTask(localNode.getId(), parentTask.getId());
+      Task task = taskManager.register("transport", PPLQueryAction.NAME, request);
+      if (!(task instanceof PPLQueryTask pplQueryTask)) {
+        taskManager.unregister(task);
+        throw new IllegalStateException("Failed to create PPL asynchronous query task");
+      }
+      registered = true;
+      return new JobTask(
+          pplQueryTask,
+          () -> {
+            try {
+              taskManager.unregister(pplQueryTask);
+            } finally {
+              childNodeRegistration.close();
+            }
+          });
+    } finally {
+      request.setParentTask(originalParent);
+      if (!registered) {
+        childNodeRegistration.close();
+      }
+    }
+  }
+}
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobId.java
similarity index 91%
rename from plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java
rename to plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobId.java
index 48c0fb19dea..05e4c97c7cd 100644
--- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobId.java
+++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobId.java
@@ -36,13 +36,18 @@
  * @param ownerNodeId node that owns the in-memory job
  * @param contextId random identifier for one job on the owner node
  */
-record PPLAsyncQueryJobId(String ownerNodeId, String contextId) {
+public record QueryJobId(String ownerNodeId, String contextId) {
   private static final int FORMAT_VERSION = 1;
   private static final int MAX_ENCODED_LENGTH = 2_048;
   private static final int MAX_OWNER_NODE_ID_BYTES = 1_024;
   private static final int MAX_CONTEXT_ID_BYTES = 128;
 
-  PPLAsyncQueryJobId {
+  /**
+   * Validates the component strings.
+   *
+   * @throws IllegalArgumentException when either component is null or blank
+   */
+  public QueryJobId {
     if (ownerNodeId == null || ownerNodeId.isBlank()) {
       throw new IllegalArgumentException("PPL asynchronous query owner node must not be empty");
     }
@@ -57,8 +62,8 @@ record PPLAsyncQueryJobId(String ownerNodeId, String contextId) {
    * @param ownerNodeId node that will own the job
    * @return unencoded job ID with a random context ID
    */
-  static PPLAsyncQueryJobId create(String ownerNodeId) {
-    return new PPLAsyncQueryJobId(ownerNodeId, UUID.randomUUID().toString());
+  public static QueryJobId create(String ownerNodeId) {
+    return new QueryJobId(ownerNodeId, UUID.randomUUID().toString());
   }
 
   /**
@@ -66,7 +71,7 @@ static PPLAsyncQueryJobId create(String ownerNodeId) {
    *
    * @return opaque value returned by the asynchronous PPL API
    */
-  String encode() {
+  public String encode() {
     try (ByteArrayOutputStream bytes = new ByteArrayOutputStream();
         DataOutputStream output = new DataOutputStream(bytes)) {
       output.writeInt(FORMAT_VERSION);
@@ -90,7 +95,7 @@ String encode() {
    * @return decoded owner node and context IDs
    * @throws IllegalArgumentException when the value is not a valid job ID
    */
-  static PPLAsyncQueryJobId parse(String encoded) {
+  public static QueryJobId parse(String encoded) {
     try {
       if (encoded == null || encoded.isBlank() || encoded.length() > MAX_ENCODED_LENGTH) {
         throw new IllegalArgumentException("Invalid PPL asynchronous query ID length");
@@ -102,8 +107,8 @@ static PPLAsyncQueryJobId parse(String encoded) {
           throw new IllegalArgumentException(
               "Unsupported PPL asynchronous query ID version [" + version + "]");
         }
-        PPLAsyncQueryJobId id =
-            new PPLAsyncQueryJobId(
+        QueryJobId id =
+            new QueryJobId(
                 readString(input, MAX_OWNER_NODE_ID_BYTES),
                 readString(input, MAX_CONTEXT_ID_BYTES));
         if (input.available() != 0) {
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobOwner.java
similarity index 86%
rename from plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java
rename to plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobOwner.java
index c8dcc7c8b35..243d5d6f602 100644
--- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUser.java
+++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobOwner.java
@@ -20,7 +20,7 @@
  * @param requestedTenant requested security tenant
  * @param backendRoles backend roles captured when the job starts
  */
-public record PPLAsyncQueryUser(String name, String requestedTenant, List backendRoles) {
+public record QueryJobOwner(String name, String requestedTenant, List backendRoles) {
 
   /**
    * Identity used when OpenSearch Security does not provide caller information.
@@ -28,7 +28,7 @@ public record PPLAsyncQueryUser(String name, String requestedTenant, ListThis identity does not bypass job ownership checks. A job owned by {@code UNSECURED} can be
    * accessed only by a caller represented by the same identity.
    */
-  public static final PPLAsyncQueryUser UNSECURED = new PPLAsyncQueryUser(null, null, List.of());
+  public static final QueryJobOwner UNSECURED = new QueryJobOwner(null, null, List.of());
 
   /**
    * Creates an immutable asynchronous query identity.
@@ -38,7 +38,7 @@ public record PPLAsyncQueryUser(String name, String requestedTenant, ListThe registry is a thin wrapper around a concurrent map. It exposes the four operations {@link
+ * QueryJob} needs to make itself discoverable, and no more: capacity accounting, listener
+ * notification, and task management remain inside {@link QueryJob}. All methods are safe for
+ * concurrent use.
+ *
+ * 
Only running jobs are registered. A job becomes visible when {@link QueryJob#create} publishes
+ * it and is removed on completion, DELETE, expiry, or shutdown. On {@link #close}, the registry
+ * stops accepting new jobs and asks each remaining job to release its resources.
+ */
+public final class QueryJobRegistry implements Closeable {
+  private final ConcurrentMap jobs = new ConcurrentHashMap<>();
+  private volatile boolean closed;
+
+  /**
+   * Publishes a job under its opaque ID.
+   *
+   * 
If a job with the same ID is already registered, the existing job is returned and the caller
+   * is expected to abort the incoming duplicate. In practice, {@link QueryJobId#create} draws its
+   * context ID from {@link java.util.UUID#randomUUID} and collisions are not observed.
+   *
+   * @param job job to publish; its ID is used as the registry key
+   * @return existing job if the ID was already registered, {@code null} on successful insertion
+   * @throws IllegalStateException when the registry has been closed
+   */
+  public QueryJob add(QueryJob job) {
+    Objects.requireNonNull(job);
+    if (closed) {
+      throw new IllegalStateException("PPL asynchronous query registry is closed");
+    }
+    return jobs.putIfAbsent(job.getJobId(), job);
+  }
+
+  /**
+   * Looks up a job by opaque ID.
+   *
+   * @param id opaque ID assigned by this node
+   * @return job when the ID resolves locally, or empty when the ID is unknown
+   */
+  public Optional get(QueryJobId id) {
+    Objects.requireNonNull(id);
+    return Optional.ofNullable(jobs.get(id));
+  }
+
+  /**
+   * Removes a job only when the given ID still resolves to the given job instance.
+   *
+   * @param id opaque ID assigned by this node
+   * @param job job that expects to still hold the registration
+   * @return {@code true} when the mapping was removed
+   */
+  public boolean remove(QueryJobId id, QueryJob job) {
+    Objects.requireNonNull(id);
+    Objects.requireNonNull(job);
+    return jobs.remove(id, job);
+  }
+
+  /**
+   * Snapshot of every job currently registered.
+   *
+   * 
Iteration order is unspecified. The returned collection is a defensive copy and does not
+   * reflect later registry mutations.
+   *
+   * @return jobs registered at the moment of the call
+   */
+  public Collection jobs() {
+    return List.copyOf(jobs.values());
+  }
+
+  /**
+   * Stops accepting new jobs and asks every registered job to release its resources.
+   *
+   * 
Each surviving job receives {@link QueryJob#discard} with a shutdown reason. The registry
+   * itself becomes idempotent: subsequent calls are no-ops and any future {@link #add} throws.
+   */
+  @Override
+  public void close() {
+    if (closed) {
+      return;
+    }
+    closed = true;
+    for (QueryJob job : jobs.values()) {
+      job.discard("PPL asynchronous query service is closing");
+    }
+  }
+}
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java
deleted file mode 100644
index bca7f1fa50a..00000000000
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobIdTest.java
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.sql.plugin.transport.asyncquery;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThrows;
-
-import org.junit.Test;
-
-public class PPLAsyncQueryJobIdTest {
-
-  @Test
-  public void roundTripsOwnerAndContext() {
-    PPLAsyncQueryJobId id = PPLAsyncQueryJobId.create("node-a");
-
-    assertEquals(id, PPLAsyncQueryJobId.parse(id.encode()));
-  }
-
-  @Test
-  public void rejectsMalformedIds() {
-    assertThrows(IllegalArgumentException.class, () -> PPLAsyncQueryJobId.parse(""));
-    assertThrows(IllegalArgumentException.class, () -> PPLAsyncQueryJobId.parse("not-an-id"));
-    assertThrows(IllegalArgumentException.class, () -> new PPLAsyncQueryJobId("", "context-id"));
-  }
-}
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
deleted file mode 100644
index 2c57d37fa49..00000000000
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryJobTest.java
+++ /dev/null
@@ -1,273 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.sql.plugin.transport.asyncquery;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.mock;
-
-import java.util.Optional;
-import org.junit.Test;
-import org.opensearch.ResourceNotFoundException;
-import org.opensearch.common.unit.TimeValue;
-import org.opensearch.sql.executor.AsyncQueryExecution;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.GetResult;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.JobTask;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Removal;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.SnapshotSource;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryJob.Transition;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryService.Failure;
-import org.opensearch.sql.plugin.transport.asyncquery.PPLAsyncQueryService.Status;
-import org.opensearch.tasks.CancellableTask;
-
-public class PPLAsyncQueryJobTest {
-  private static final String ID = "job-id";
-  private static final long START_TIME = 1_000;
-  private static final long RETAINED_TIME = 2_000;
-  private static final long COMPLETION_TIME = 2_025;
-  private static final long KEEP_ALIVE_MILLIS = TimeValue.timeValueMinutes(5).millis();
-
-  @Test
-  public void retainPublishesRunningResponseWithId() {
-    JobFixture fixture = newJob();
-    assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
-
-    Transition transition = fixture.job().retain(RETAINED_TIME);
-
-    // Retention publishes the job ID once; later retention events are ignored.
-    assertEquals(
-        new Transition.Retained(new SnapshotSource.Running(ID, Optional.of(fixture.execution()))),
-        transition);
-    assertNull(fixture.job().retain(RETAINED_TIME + 1));
-  }
-
-  @Test
-  public void completionBeforeRetentionReturnsDirectResponseAndRemovesJob() {
-    JobFixture fixture = newJob();
-    assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
-
-    Transition transition = fixture.job().complete(COMPLETION_TIME);
-
-    // A direct response removes the job: duplicate completion is ignored and GET returns not found.
-    assertTrue(transition instanceof Transition.DirectResponse);
-    Transition.DirectResponse direct = (Transition.DirectResponse) transition;
-    assertEquals(
-        new SnapshotSource.Succeeded(
-            Optional.empty(), fixture.execution(), COMPLETION_TIME - START_TIME),
-        direct.response());
-    assertSame(fixture.task(), direct.task());
-    assertEquals(Optional.of(fixture.execution()), direct.executionToClose());
-    assertNull(fixture.job().complete(COMPLETION_TIME + 1));
-    assertThrows(ResourceNotFoundException.class, () -> fixture.job().get(COMPLETION_TIME, null));
-  }
-
-  @Test
-  public void failureBeforeRetentionReturnsDirectResponseAndRemovesJob() {
-    JobFixture fixture = newJob();
-    Failure failure = new Failure("IllegalArgumentException", "invalid query");
-    assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
-
-    Transition transition = fixture.job().fail(failure, COMPLETION_TIME);
-
-    // A failure before retention returns without an ID and detaches all execution resources.
-    assertTrue(transition instanceof Transition.DirectResponse);
-    Transition.DirectResponse direct = (Transition.DirectResponse) transition;
-    assertEquals(
-        new SnapshotSource.Failed(Optional.empty(), failure, COMPLETION_TIME - START_TIME),
-        direct.response());
-    assertSame(fixture.task(), direct.task());
-    assertEquals(Optional.of(fixture.execution()), direct.executionToClose());
-
-    // A direct failure removes the job: later completion is ignored and GET returns not found.
-    assertNull(fixture.job().complete(COMPLETION_TIME + 1));
-    assertThrows(ResourceNotFoundException.class, () -> fixture.job().get(COMPLETION_TIME, null));
-  }
-
-  @Test
-  public void retainedSuccessRemainsReadableUntilDeleted() {
-    JobFixture fixture = retainedJob();
-
-    Transition transition = fixture.job().complete(COMPLETION_TIME);
-
-    // Retained success closes the task but keeps the execution readable until DELETE.
-    assertEquals(new Transition.ExecutionFinished(fixture.task(), Optional.empty()), transition);
-
-    GetResult result = fixture.job().get(COMPLETION_TIME, null);
-    assertTrue(result instanceof GetResult.Found);
-    assertEquals(
-        new SnapshotSource.Succeeded(
-            Optional.of(ID), fixture.execution(), COMPLETION_TIME - START_TIME),
-        ((GetResult.Found) result).response());
-
-    Removal removal = fixture.job().delete(COMPLETION_TIME);
-    assertTrue(removal instanceof Removal.Deleted);
-    Removal.Deleted deleted = (Removal.Deleted) removal;
-    assertEquals(Status.SUCCEEDED, deleted.responseStatus());
-    assertNull(deleted.resources().task());
-    assertSame(fixture.execution(), deleted.resources().execution());
-    assertEquals("PPL asynchronous query cancelled by user", deleted.reason());
-  }
-
-  @Test
-  public void retainedFailureDetachesExecutionAndRemainsReadable() {
-    JobFixture fixture = retainedJob();
-    Failure failure = new Failure("IllegalStateException", "query failed");
-
-    Transition transition = fixture.job().fail(failure, COMPLETION_TIME);
-
-    // Retained failure closes its execution and keeps only the failure snapshot readable.
-    assertEquals(
-        new Transition.ExecutionFinished(fixture.task(), Optional.of(fixture.execution())),
-        transition);
-
-    GetResult result = fixture.job().get(COMPLETION_TIME, null);
-    assertTrue(result instanceof GetResult.Found);
-    assertEquals(
-        new SnapshotSource.Failed(Optional.of(ID), failure, COMPLETION_TIME - START_TIME),
-        ((GetResult.Found) result).response());
-  }
-
-  @Test
-  public void getWithoutKeepAlivePreservesExpiration() {
-    JobFixture fixture = retainedJob();
-    long expirationTime = RETAINED_TIME + KEEP_ALIVE_MILLIS;
-
-    assertTrue(fixture.job().get(expirationTime - 1, null) instanceof GetResult.Found);
-    GetResult result = fixture.job().get(expirationTime, null);
-
-    // Without renewal, the lease remains valid before but not at its expiration boundary.
-    assertTrue(result instanceof GetResult.Expired);
-    Removal.Expired removal = ((GetResult.Expired) result).removal();
-    assertSame(fixture.task(), removal.resources().task());
-    assertSame(fixture.execution(), removal.resources().execution());
-    assertEquals("PPL asynchronous query expired", removal.reason());
-  }
-
-  @Test
-  public void getWithKeepAliveReplacesExpiration() {
-    JobFixture fixture = retainedJob();
-    long renewalTime = RETAINED_TIME + 500;
-    TimeValue requestedKeepAlive = TimeValue.timeValueSeconds(1);
-    long renewedExpiration = renewalTime + requestedKeepAlive.millis();
-
-    // At 2500ms, GET replaces the original lease with a one-second lease ending at 3500ms.
-    GetResult renewed = fixture.job().get(renewalTime, requestedKeepAlive);
-    // At 3499ms, the renewed lease has not expired.
-    GetResult beforeExpiration = fixture.job().get(renewedExpiration - 1, null);
-    // At 3500ms, the renewed lease expires.
-    GetResult atExpiration = fixture.job().get(renewedExpiration, null);
-
-    assertTrue(renewed instanceof GetResult.Found);
-    assertTrue(beforeExpiration instanceof GetResult.Found);
-    assertTrue(atExpiration instanceof GetResult.Expired);
-  }
-
-  @Test
-  public void expirationStartsWhenJobIsRetained() {
-    JobFixture fixture = newJob();
-    assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
-    long retainedTime = START_TIME + KEEP_ALIVE_MILLIS + 1;
-
-    // Keep-alive starts when the ID is retained, not when the unretained job is created.
-    assertNull(fixture.job().expire(retainedTime - 1));
-    fixture.job().retain(retainedTime);
-    assertNull(fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS - 1));
-    Removal removal = fixture.job().expire(retainedTime + KEEP_ALIVE_MILLIS);
-    assertTrue(removal instanceof Removal.Expired);
-    Removal.Expired expired = (Removal.Expired) removal;
-    assertSame(fixture.task(), expired.resources().task());
-    assertSame(fixture.execution(), expired.resources().execution());
-    assertEquals("PPL asynchronous query expired", expired.reason());
-  }
-
-  @Test
-  public void deleteRunningJobReturnsCancellationAndDetachesResources() {
-    JobFixture fixture = retainedJob();
-
-    Removal removal = fixture.job().delete(RETAINED_TIME + 1);
-
-    // DELETE cancels a running job, transfers its resources, and makes later DELETE return not
-    // found.
-    assertTrue(removal instanceof Removal.Deleted);
-    Removal.Deleted deleted = (Removal.Deleted) removal;
-    assertEquals(Status.CANCELLED, deleted.responseStatus());
-    assertSame(fixture.task(), deleted.resources().task());
-    assertSame(fixture.execution(), deleted.resources().execution());
-    assertEquals("PPL asynchronous query cancelled by user", deleted.reason());
-    assertTrue(deleted.resources().releasesRunningSlot());
-    assertThrows(ResourceNotFoundException.class, () -> fixture.job().delete(RETAINED_TIME + 2));
-  }
-
-  @Test
-  public void abortDetachesResourcesOnlyOnce() {
-    JobFixture fixture = newJob();
-    assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
-
-    Removal removal = fixture.job().abort();
-
-    // Abort transfers resource ownership once; repeated abort calls have no transition.
-    assertTrue(removal instanceof Removal.Discarded);
-    Removal.Discarded discarded = (Removal.Discarded) removal;
-    assertSame(fixture.task(), discarded.resources().task());
-    assertSame(fixture.execution(), discarded.resources().execution());
-    assertEquals("PPL asynchronous query startup failed", discarded.reason());
-    assertNull(fixture.job().abort());
-  }
-
-  @Test
-  public void closeDetachesResourcesOnlyOnce() {
-    JobFixture fixture = retainedJob();
-
-    Removal removal = fixture.job().close("service closing");
-
-    // Close transfers resource ownership once; repeated close calls have no transition.
-    assertTrue(removal instanceof Removal.Discarded);
-    Removal.Discarded discarded = (Removal.Discarded) removal;
-    assertSame(fixture.task(), discarded.resources().task());
-    assertSame(fixture.execution(), discarded.resources().execution());
-    assertEquals("service closing", discarded.reason());
-    assertNull(fixture.job().close("service closing"));
-  }
-
-  @Test
-  public void lateExecutionAttachmentAfterRemovalIsRejected() {
-    JobFixture fixture = newJob();
-    fixture.job().abort();
-
-    // Removal wins the race, so a late execution handle remains owned by the caller.
-    assertFalse(fixture.job().tryAttachExecution(fixture.execution()));
-  }
-
-  @Test
-  public void successfulCompletionRequiresAttachedExecution() {
-    JobFixture fixture = newJob();
-
-    // Success cannot be published until an execution handle can provide the final result.
-    assertThrows(IllegalStateException.class, () -> fixture.job().complete(COMPLETION_TIME));
-  }
-
-  private static JobFixture retainedJob() {
-    JobFixture fixture = newJob();
-    assertTrue(fixture.job().tryAttachExecution(fixture.execution()));
-    fixture.job().retain(RETAINED_TIME);
-    return fixture;
-  }
-
-  private static JobFixture newJob() {
-    JobTask task = new JobTask(mock(CancellableTask.class), () -> {});
-    AsyncQueryExecution execution = mock(AsyncQueryExecution.class);
-    return new JobFixture(
-        new PPLAsyncQueryJob(ID, PPLAsyncQueryUser.UNSECURED, START_TIME, KEEP_ALIVE_MILLIS, task),
-        task,
-        execution);
-  }
-
-  private record JobFixture(PPLAsyncQueryJob job, JobTask task, AsyncQueryExecution execution) {}
-}
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
deleted file mode 100644
index 584503bab4c..00000000000
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryServiceTest.java
+++ /dev/null
@@ -1,951 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.sql.plugin.transport.asyncquery;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.inOrder;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionStage;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Function;
-import org.junit.Test;
-import org.mockito.InOrder;
-import org.opensearch.OpenSearchSecurityException;
-import org.opensearch.OpenSearchStatusException;
-import org.opensearch.ResourceNotFoundException;
-import org.opensearch.cluster.node.DiscoveryNode;
-import org.opensearch.common.lease.Releasable;
-import org.opensearch.common.unit.TimeValue;
-import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.tasks.TaskId;
-import org.opensearch.sql.data.model.ExprValueUtils;
-import org.opensearch.sql.data.type.ExprCoreType;
-import org.opensearch.sql.executor.AsyncQueryExecution;
-import org.opensearch.sql.executor.ExecutionEngine.QueryResponse;
-import org.opensearch.sql.executor.ExecutionEngine.Schema;
-import org.opensearch.sql.executor.ExecutionEngine.Schema.Column;
-import org.opensearch.sql.legacy.metrics.BasicCounter;
-import org.opensearch.sql.legacy.metrics.MetricName;
-import org.opensearch.sql.legacy.metrics.Metrics;
-import org.opensearch.sql.legacy.metrics.NumericMetric;
-import org.opensearch.sql.plugin.transport.PPLQueryAction;
-import org.opensearch.sql.plugin.transport.PPLQueryTask;
-import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest;
-import org.opensearch.tasks.CancellableTask;
-import org.opensearch.tasks.TaskManager;
-
-public class PPLAsyncQueryServiceTest {
-  private static final PPLAsyncQueryUser OWNER = PPLAsyncQueryUser.UNSECURED;
-  private static final TimeValue KEEP_ALIVE = TimeValue.timeValueMinutes(5);
-
-  private final AtomicLong now = new AtomicLong(1_000);
-  private final AtomicReference timeoutTask = new AtomicReference<>();
-  private final AtomicBoolean timeoutCancelled = new AtomicBoolean();
-  private final PPLAsyncQueryService service = service(20, 100);
-
-  @Test
-  public void fastSuccessReturnsDirectResultWithoutRetainingJob() {
-    AsyncQueryScenario scenario =
-        scenario()
-            .withCurrentResult(response(2))
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-            .start();
-
-    scenario.completeAfterMillis(25);
-
-    scenario
-        .assertDirectSuccess(response(2), 25)
-        .assertNotRetained()
-        .assertExecutionReadOnceAndClosed()
-        .assertTimeoutCancelled();
-
-    scenario.fireTimeout().assertResponseCount(1);
-  }
-
-  @Test
-  public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() {
-    AsyncQueryScenario scenario =
-        scenario()
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-            .start()
-            .fireTimeout()
-            .assertRetainedRunning(Optional.empty())
-            .assertCapacity(1, 1);
-
-    scenario
-        .succeedAfterMillis(25, response(2))
-        .get()
-        .assertGetResponse(
-            new PPLAsyncQueryService.JobSnapshot.Succeeded(
-                Optional.of(scenario.id()), response(2), 25))
-        .assertCapacity(0, 1)
-        .assertExecutionReadCount(2)
-        .assertExecutionCloseCount(0);
-  }
-
-  @Test
-  public void retentionWaitPreventsExpiryAndLeaseStartsWhenIdIsReturned() {
-    AsyncQueryScenario scenario =
-        scenario()
-            .withKeepAlive(TimeValue.timeValueSeconds(1))
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-            .start();
-
-    scenario
-        .advanceMillis(TimeValue.timeValueSeconds(2).millis())
-        .reapExpired()
-        .assertNoInitialResponse()
-        .assertCapacity(1, 1);
-
-    scenario.fireTimeout().assertRetainedRunning(Optional.empty());
-
-    scenario
-        .advanceMillis(TimeValue.timeValueSeconds(1).millis() + 1)
-        .reapExpired()
-        .assertCapacity(0, 0);
-  }
-
-  @Test
-  public void fastFailureReturnsDirectFailureWithoutId() {
-    scenario()
-        .withCurrentResult(response(1))
-        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-        .start()
-        .fail(new IllegalStateException("boom"))
-        .assertFailure(IllegalStateException.class, "boom")
-        .assertExecutionReadCount(0)
-        .assertExecutionCloseCount(1)
-        .assertCapacity(0, 0);
-  }
-
-  @Test
-  public void expiredGetCancelsAndRemovesJob() {
-    CancellableTask task = mock(CancellableTask.class);
-    when(task.isCancelled()).thenReturn(false);
-    AsyncQueryScenario scenario = scenario().withTask(task).start();
-
-    scenario.advanceMillis(KEEP_ALIVE.millis());
-
-    ResourceNotFoundException error = assertThrows(ResourceNotFoundException.class, scenario::get);
-
-    // Expiration removes the job, so the client sees the same response as any missing job.
-    assertEquals("PPL asynchronous query not found", error.getMessage());
-    verify(task).cancel("PPL asynchronous query expired");
-    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
-  }
-
-  @Test
-  public void missingLocalJobReturnsNotFound() {
-    String remoteId = PPLAsyncQueryJobId.create("node-b").encode();
-
-    assertThrows(ResourceNotFoundException.class, () -> service.get(remoteId, OWNER, null));
-    assertThrows(ResourceNotFoundException.class, () -> service.delete(remoteId, OWNER));
-  }
-
-  @Test
-  public void deleteCancelsRunningJobAndReleasesState() {
-    CancellableTask task = mock(CancellableTask.class);
-    when(task.isCancelled()).thenReturn(false);
-    AsyncQueryScenario scenario = scenario().withTask(task).start().delete();
-
-    scenario.assertDeleteStatus(PPLAsyncQueryService.Status.CANCELLED);
-    verify(task).cancel("PPL asynchronous query cancelled by user");
-    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
-    assertThrows(ResourceNotFoundException.class, scenario::get);
-  }
-
-  @Test
-  public void deleteReturnsExistingTerminalStatus() {
-    CancellableTask task = mock(CancellableTask.class);
-    AsyncQueryScenario scenario =
-        scenario().withTask(task).withCurrentResult(response(1)).start().complete().delete();
-
-    scenario.assertDeleteStatus(PPLAsyncQueryService.Status.SUCCEEDED);
-    // DELETE removes the retained result without cancelling an already completed task.
-    verify(task, never()).cancel(org.mockito.ArgumentMatchers.anyString());
-    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
-  }
-
-  @Test
-  public void cancellationUsesTaskManagerWhenAttached() {
-    TaskManager taskManager = mock(TaskManager.class);
-    CancellableTask task = mock(CancellableTask.class);
-    when(task.isCancelled()).thenReturn(false);
-    service.attachTaskManager(taskManager);
-    scenario().withTask(task).start().delete();
-
-    verify(taskManager)
-        .cancelTaskAndDescendants(
-            org.mockito.ArgumentMatchers.eq(task),
-            org.mockito.ArgumentMatchers.eq("PPL asynchronous query cancelled by user"),
-            org.mockito.ArgumentMatchers.eq(false),
-            org.mockito.ArgumentMatchers.any());
-  }
-
-  @Test
-  public void startRegistersTaskAndCompletionReleasesIt() {
-    TaskManager taskManager = mock(TaskManager.class);
-    PPLQueryTask task = mock(PPLQueryTask.class);
-    TransportPPLQueryRequest request =
-        new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
-    RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
-    service.attachTaskManager(taskManager);
-
-    AsyncQueryScenario scenario =
-        scenario()
-            .start(request, registration.requestTask())
-            .assertRetainedRunning(Optional.empty())
-            .succeed(response(1));
-
-    InOrder registrationOrder = inOrder(taskManager);
-    registrationOrder
-        .verify(taskManager)
-        .registerChildNode(registration.requestTask().getId(), registration.localNode());
-    registrationOrder.verify(taskManager).register("transport", PPLQueryAction.NAME, request);
-    assertEquals(TaskId.EMPTY_TASK_ID, request.getParentTask());
-    verify(taskManager).unregister(task);
-    verify(registration.childNodeRegistration()).close();
-    scenario.assertCapacity(0, 1);
-  }
-
-  @Test
-  public void executionStartFailureCompletesJobAndReleasesTask() {
-    TaskManager taskManager = mock(TaskManager.class);
-    PPLQueryTask task = mock(PPLQueryTask.class);
-    TransportPPLQueryRequest request =
-        new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
-    RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
-    service.attachTaskManager(taskManager);
-
-    AsyncQueryScenario scenario =
-        scenario()
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-            .withExecutionStarter(
-                ignored -> {
-                  throw new IllegalStateException("execution did not start");
-                })
-            .start(request, registration.requestTask());
-
-    scenario
-        .assertFailure(IllegalStateException.class, "execution did not start")
-        .assertCapacity(0, 0);
-    verify(taskManager, times(1)).unregister(task);
-    verify(registration.childNodeRegistration()).close();
-  }
-
-  @Test
-  public void deleteCancelsAndReleasesServiceOwnedTask() {
-    TaskManager taskManager = mock(TaskManager.class);
-    PPLQueryTask task = mock(PPLQueryTask.class);
-    when(task.isCancelled()).thenReturn(false);
-    TransportPPLQueryRequest request =
-        new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
-    RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
-    doAnswer(
-            invocation -> {
-              ActionListener listener = invocation.getArgument(3);
-              listener.onResponse(null);
-              return null;
-            })
-        .when(taskManager)
-        .cancelTaskAndDescendants(
-            org.mockito.ArgumentMatchers.eq(task),
-            org.mockito.ArgumentMatchers.anyString(),
-            org.mockito.ArgumentMatchers.eq(false),
-            org.mockito.ArgumentMatchers.any());
-    service.attachTaskManager(taskManager);
-
-    scenario().start(request, registration.requestTask()).delete();
-
-    verify(taskManager)
-        .cancelTaskAndDescendants(
-            org.mockito.ArgumentMatchers.eq(task),
-            org.mockito.ArgumentMatchers.eq("PPL asynchronous query cancelled by user"),
-            org.mockito.ArgumentMatchers.eq(false),
-            org.mockito.ArgumentMatchers.any());
-    verify(taskManager, times(1)).unregister(task);
-    verify(registration.childNodeRegistration()).close();
-  }
-
-  @Test
-  public void startupFailureAfterJobCreationReleasesServiceOwnedTask() {
-    PPLAsyncQueryService abortingService =
-        new PPLAsyncQueryService(
-            "node-a",
-            now::get,
-            (delay, task) -> {
-              throw new IllegalStateException("scheduler unavailable");
-            },
-            () -> 20,
-            () -> 100,
-            () -> TimeValue.timeValueSeconds(60),
-            () -> TimeValue.timeValueHours(24));
-    TaskManager taskManager = mock(TaskManager.class);
-    PPLQueryTask task = mock(PPLQueryTask.class);
-    when(task.isCancelled()).thenReturn(true);
-    TransportPPLQueryRequest request =
-        new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
-    RequestTaskRegistration registration = registerRequestTask(taskManager, request, task);
-    abortingService.attachTaskManager(taskManager);
-    AsyncQueryScenario scenario =
-        scenario()
-            .withService(abortingService)
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5));
-
-    IllegalStateException failure =
-        assertThrows(
-            IllegalStateException.class, () -> scenario.start(request, registration.requestTask()));
-
-    assertEquals("scheduler unavailable", failure.getMessage());
-    verify(taskManager, times(1)).unregister(task);
-    verify(registration.childNodeRegistration()).close();
-    scenario.assertCapacity(0, 0);
-  }
-
-  @Test
-  public void childTrackingFailurePreventsRetainedTaskRegistration() {
-    TaskManager taskManager = mock(TaskManager.class);
-    PPLQueryTask requestTask = mock(PPLQueryTask.class);
-    DiscoveryNode localNode = mock(DiscoveryNode.class);
-    TransportPPLQueryRequest request =
-        new TransportPPLQueryRequest("source=t", new org.json.JSONObject(), "/_plugins/_ppl");
-    when(requestTask.getId()).thenReturn(42L);
-    when(taskManager.localNode()).thenReturn(localNode);
-    when(taskManager.registerChildNode(42L, localNode))
-        .thenThrow(new IllegalStateException("channel closed"));
-    service.attachTaskManager(taskManager);
-    AsyncQueryScenario scenario = scenario().withWaitForCompletion(TimeValue.timeValueSeconds(5));
-
-    assertThrows(IllegalStateException.class, () -> scenario.start(request, requestTask));
-
-    verify(taskManager, never()).register("transport", PPLQueryAction.NAME, request);
-    scenario.assertCapacity(0, 0);
-  }
-
-  @Test
-  public void rejectsUnauthorizedCallerWithoutRenewingOrDeleting() {
-    PPLAsyncQueryUser securedOwner = new PPLAsyncQueryUser("alice", "tenant", List.of("role-a"));
-    PPLAsyncQueryUser otherUser = new PPLAsyncQueryUser("bob", "tenant", List.of("role-a"));
-    AsyncQueryScenario scenario = scenario().withOwner(securedOwner).start();
-
-    assertThrows(
-        OpenSearchSecurityException.class, () -> service.get(scenario.id(), otherUser, null));
-    assertThrows(OpenSearchSecurityException.class, () -> service.delete(scenario.id(), otherUser));
-    assertEquals(
-        PPLAsyncQueryService.Status.RUNNING,
-        service.get(scenario.id(), securedOwner, null).status());
-  }
-
-  @Test
-  public void enforcesRunningAndRetainedCapacity() {
-    PPLAsyncQueryService limited = service(1, 1);
-    scenario().withService(limited).start();
-
-    OpenSearchStatusException exception =
-        assertThrows(
-            OpenSearchStatusException.class, () -> scenario().withService(limited).start());
-
-    assertEquals(429, exception.status().getStatus());
-  }
-
-  @Test
-  public void createFailureReleasesReservedCapacity() {
-    PPLAsyncQueryService missingOwnerNode =
-        new PPLAsyncQueryService(
-            (String) null,
-            now::get,
-            (delay, task) -> () -> {},
-            () -> 1,
-            () -> 1,
-            () -> TimeValue.timeValueSeconds(60),
-            () -> TimeValue.timeValueHours(24));
-    AsyncQueryScenario scenario = scenario().withService(missingOwnerNode);
-
-    assertThrows(NullPointerException.class, scenario::start);
-
-    scenario.assertCapacity(0, 0);
-  }
-
-  @Test
-  public void finalSnapshotDefensivelyCopiesRows() {
-    List rows = new ArrayList<>();
-    rows.add(ExprValueUtils.stringValue("first"));
-    QueryResponse response =
-        new QueryResponse(
-            new Schema(List.of(new Column("state", null, ExprCoreType.STRING))), rows, null);
-
-    AsyncQueryScenario scenario =
-        scenario()
-            .withCurrentResult(response)
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-            .start()
-            .complete();
-    rows.add(ExprValueUtils.stringValue("second"));
-
-    scenario.assertSucceededRowCount(1);
-  }
-
-  @Test
-  public void completedExecutionCanBeAttachedBeforeCompletionIsObserved() {
-    scenario()
-        .succeed(response(2))
-        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-        .start()
-        .assertDirectSuccess(response(2), 0)
-        .assertExecutionReadOnceAndClosed();
-  }
-
-  @Test
-  public void runningGetMaterializesCurrentResultOutsideJob() {
-    AsyncQueryScenario scenario = scenario().withCurrentResult(response(1)).start().get();
-
-    scenario.assertGetResponse(
-        new PPLAsyncQueryService.JobSnapshot.Running(scenario.id(), Optional.of(response(1))));
-
-    scenario
-        .withCurrentResult(response(3))
-        .get()
-        .assertGetResponse(
-            new PPLAsyncQueryService.JobSnapshot.Running(scenario.id(), Optional.of(response(3))));
-  }
-
-  @Test
-  public void failedRetainedJobReturnsNoProvisionalRowsAndClosesExecution() {
-    NumericMetric failures =
-        new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_SYS.getName(), new BasicCounter());
-    Metrics.getInstance().registerMetric(failures);
-    try {
-      AsyncQueryScenario scenario =
-          scenario()
-              .withCurrentResult(response(1))
-              .start()
-              .fail(new IllegalStateException("boom"))
-              .get();
-
-      scenario
-          .assertGetResponse(
-              new PPLAsyncQueryService.JobSnapshot.Failed(
-                  Optional.of(scenario.id()),
-                  new PPLAsyncQueryService.Failure("IllegalStateException", "boom"),
-                  0))
-          .assertExecutionReadCount(0)
-          .assertExecutionCloseCount(1);
-    } finally {
-      Metrics.getInstance().unregisterMetric(failures.getName());
-    }
-  }
-
-  @Test
-  public void failedRetainedJobRecordsFailureMetric() {
-    NumericMetric failures =
-        new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_CUS.getName(), new BasicCounter());
-    Metrics.getInstance().registerMetric(failures);
-    try {
-      scenario().start().fail(new IllegalArgumentException("invalid query"));
-
-      assertEquals(Long.valueOf(1), failures.getValue());
-    } finally {
-      Metrics.getInstance().unregisterMetric(failures.getName());
-    }
-  }
-
-  @Test
-  public void deleteBeforeExecutionAttachmentClosesLateHandle() {
-    TrackingExecution execution = new TrackingExecution(response(1));
-    AsyncQueryScenario scenario = scenario().withExecution(execution);
-    scenario
-        .withExecutionStarter(
-            ignored -> {
-              scenario.delete();
-              return execution;
-            })
-        .start();
-
-    assertEquals(1, execution.closes.get());
-    assertThrows(ResourceNotFoundException.class, scenario::get);
-  }
-
-  @Test
-  public void concurrentGetDoesNotBlockDeleteOnResultMaterialization() throws Exception {
-    BlockingExecution execution = new BlockingExecution(response(1));
-    AsyncQueryScenario scenario = scenario().withExecution(execution).start();
-
-    CompletableFuture get =
-        CompletableFuture.supplyAsync(() -> service.get(scenario.id(), OWNER, null));
-    assertTrue(execution.readStarted.await(5, TimeUnit.SECONDS));
-    CompletableFuture delete =
-        CompletableFuture.supplyAsync(() -> service.delete(scenario.id(), OWNER));
-
-    try {
-      assertEquals(PPLAsyncQueryService.Status.CANCELLED, delete.get(5, TimeUnit.SECONDS).status());
-    } finally {
-      execution.allowRead.countDown();
-    }
-
-    assertEquals(PPLAsyncQueryService.Status.RUNNING, get.get(5, TimeUnit.SECONDS).status());
-    assertEquals(1, execution.closes.get());
-    assertThrows(ResourceNotFoundException.class, scenario::get);
-  }
-
-  @Test
-  public void shutdownClosesRetainedExecutionExactlyOnce() throws Exception {
-    AsyncQueryScenario scenario = scenario().withCurrentResult(response(1)).start();
-
-    service.close();
-    service.close();
-
-    scenario.assertExecutionReadCount(0).assertExecutionCloseCount(1).assertCapacity(0, 0);
-  }
-
-  @Test
-  public void successfulCompletionRequiresFinalResultToBeVisible() {
-    scenario()
-        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-        .start()
-        .complete()
-        .assertFailure(IllegalStateException.class)
-        .assertExecutionReadOnceAndClosed()
-        .assertCapacity(0, 0);
-  }
-
-  @Test
-  public void retentionResponseMaterializationFailureAbortsUndeliverableJob() {
-    CancellableTask task = mock(CancellableTask.class);
-    when(task.isCancelled()).thenReturn(false);
-    ThrowingExecution execution = new ThrowingExecution();
-    AsyncQueryScenario scenario =
-        scenario()
-            .withTask(task)
-            .withExecution(execution)
-            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
-            .start()
-            .fireTimeout();
-
-    scenario.assertFailure(IllegalStateException.class).assertCapacity(0, 0);
-    verify(task).cancel("PPL asynchronous query startup failed");
-    assertEquals(1, execution.closes.get());
-  }
-
-  @Test
-  public void validatesDurationBounds() {
-    assertThrows(IllegalArgumentException.class, () -> service.validateKeepAlive(TimeValue.ZERO));
-    assertThrows(
-        IllegalArgumentException.class,
-        () -> service.validateKeepAlive(TimeValue.timeValueHours(25)));
-    assertThrows(
-        IllegalArgumentException.class,
-        () -> service.validateWaitForCompletion(TimeValue.timeValueSeconds(61)));
-
-    service.validateWaitForCompletion(TimeValue.ZERO);
-    service.validateWaitForCompletion(TimeValue.timeValueSeconds(60));
-    service.validateKeepAlive(TimeValue.timeValueHours(24));
-  }
-
-  private PPLAsyncQueryService service(int maxRunning, int maxRetained) {
-    return new PPLAsyncQueryService(
-        "node-a",
-        now::get,
-        (delay, task) -> {
-          timeoutTask.set(task);
-          timeoutCancelled.set(false);
-          return () -> timeoutCancelled.set(true);
-        },
-        () -> maxRunning,
-        () -> maxRetained,
-        () -> TimeValue.timeValueSeconds(60),
-        () -> TimeValue.timeValueHours(24));
-  }
-
-  private AsyncQueryScenario scenario() {
-    return new AsyncQueryScenario();
-  }
-
-  private static PPLAsyncQueryJob.JobTask jobTask(CancellableTask task) {
-    return new PPLAsyncQueryJob.JobTask(task, () -> {});
-  }
-
-  private static RequestTaskRegistration registerRequestTask(
-      TaskManager taskManager, TransportPPLQueryRequest request, PPLQueryTask retainedTask) {
-    PPLQueryTask requestTask = mock(PPLQueryTask.class);
-    DiscoveryNode localNode = mock(DiscoveryNode.class);
-    Releasable childNodeRegistration = mock(Releasable.class);
-    when(requestTask.getId()).thenReturn(42L);
-    when(localNode.getId()).thenReturn("node-a");
-    when(taskManager.localNode()).thenReturn(localNode);
-    when(taskManager.registerChildNode(42L, localNode)).thenReturn(childNodeRegistration);
-    when(taskManager.register("transport", PPLQueryAction.NAME, request))
-        .thenAnswer(
-            invocation -> {
-              assertEquals(new TaskId("node-a", 42L), request.getParentTask());
-              return retainedTask;
-            });
-    return new RequestTaskRegistration(requestTask, localNode, childNodeRegistration);
-  }
-
-  private record RequestTaskRegistration(
-      PPLQueryTask requestTask, DiscoveryNode localNode, Releasable childNodeRegistration) {}
-
-  private static QueryResponse response(int rowCount) {
-    Schema schema = new Schema(List.of(new Column("state", null, ExprCoreType.STRING)));
-    return new QueryResponse(
-        schema,
-        java.util.stream.IntStream.range(0, rowCount)
-            .mapToObj(i -> ExprValueUtils.stringValue("state-" + i))
-            .toList(),
-        null);
-  }
-
-  private final class AsyncQueryScenario {
-    private PPLAsyncQueryService targetService = service;
-    private PPLAsyncQueryUser owner = OWNER;
-    private TimeValue keepAlive = KEEP_ALIVE;
-    private TimeValue waitForCompletion = TimeValue.ZERO;
-    private CancellableTask task;
-    private TrackingExecution trackingExecution = new TrackingExecution(null);
-    private AsyncQueryExecution execution = trackingExecution;
-    private Function executionStarter = ignored -> execution;
-    private final AtomicReference initialResponse =
-        new AtomicReference<>();
-    private final AtomicReference failure = new AtomicReference<>();
-    private final AtomicInteger responseCount = new AtomicInteger();
-    private PPLAsyncQueryService.JobSnapshot getResponse;
-    private PPLAsyncQueryService.DeleteResult deleteResponse;
-
-    private AsyncQueryScenario withService(PPLAsyncQueryService service) {
-      targetService = service;
-      return this;
-    }
-
-    private AsyncQueryScenario withOwner(PPLAsyncQueryUser owner) {
-      this.owner = owner;
-      return this;
-    }
-
-    private AsyncQueryScenario withKeepAlive(TimeValue keepAlive) {
-      this.keepAlive = keepAlive;
-      return this;
-    }
-
-    private AsyncQueryScenario withWaitForCompletion(TimeValue waitForCompletion) {
-      this.waitForCompletion = waitForCompletion;
-      return this;
-    }
-
-    private AsyncQueryScenario withTask(CancellableTask task) {
-      this.task = task;
-      return this;
-    }
-
-    private AsyncQueryScenario withCurrentResult(QueryResponse response) {
-      trackingExecution.setCurrent(response);
-      return this;
-    }
-
-    private AsyncQueryScenario withExecution(AsyncQueryExecution execution) {
-      this.execution = execution;
-      trackingExecution = execution instanceof TrackingExecution tracking ? tracking : null;
-      executionStarter = ignored -> this.execution;
-      return this;
-    }
-
-    private AsyncQueryScenario withExecutionStarter(
-        Function executionStarter) {
-      this.executionStarter = executionStarter;
-      return this;
-    }
-
-    private AsyncQueryScenario start() {
-      targetService.start(
-          owner, keepAlive, waitForCompletion, jobTask(task), executionStarter, responseListener());
-      return this;
-    }
-
-    private AsyncQueryScenario start(TransportPPLQueryRequest request, PPLQueryTask requestTask) {
-      targetService.start(
-          owner,
-          keepAlive,
-          waitForCompletion,
-          request,
-          requestTask,
-          executionStarter,
-          responseListener());
-      return this;
-    }
-
-    private ActionListener responseListener() {
-      return ActionListener.wrap(
-          snapshot -> {
-            initialResponse.set(snapshot);
-            responseCount.incrementAndGet();
-          },
-          failure::set);
-    }
-
-    private AsyncQueryScenario completeAfterMillis(long elapsedMillis) {
-      now.addAndGet(elapsedMillis);
-      trackingExecution.complete();
-      return this;
-    }
-
-    private AsyncQueryScenario succeedAfterMillis(long elapsedMillis, QueryResponse response) {
-      now.addAndGet(elapsedMillis);
-      trackingExecution.succeed(response);
-      return this;
-    }
-
-    private AsyncQueryScenario complete() {
-      trackingExecution.complete();
-      return this;
-    }
-
-    private AsyncQueryScenario succeed(QueryResponse response) {
-      trackingExecution.succeed(response);
-      return this;
-    }
-
-    private AsyncQueryScenario fail(Exception failure) {
-      trackingExecution.fail(failure);
-      return this;
-    }
-
-    private AsyncQueryScenario advanceMillis(long elapsedMillis) {
-      now.addAndGet(elapsedMillis);
-      return this;
-    }
-
-    private AsyncQueryScenario reapExpired() {
-      targetService.reapExpired();
-      return this;
-    }
-
-    private AsyncQueryScenario assertDirectSuccess(
-        QueryResponse expectedResponse, long expectedTookMillis) {
-      assertEquals(
-          new PPLAsyncQueryService.JobSnapshot.Succeeded(
-              Optional.empty(), expectedResponse, expectedTookMillis),
-          initialResponse.get());
-      return this;
-    }
-
-    private AsyncQueryScenario assertSucceededRowCount(int expected) {
-      assertTrue(initialResponse.get() instanceof PPLAsyncQueryService.JobSnapshot.Succeeded);
-      PPLAsyncQueryService.JobSnapshot.Succeeded succeeded =
-          (PPLAsyncQueryService.JobSnapshot.Succeeded) initialResponse.get();
-      assertEquals(expected, succeeded.response().getResults().size());
-      return this;
-    }
-
-    private AsyncQueryScenario assertRetainedRunning(Optional expectedResponse) {
-      assertEquals(
-          new PPLAsyncQueryService.JobSnapshot.Running(id(), expectedResponse),
-          initialResponse.get());
-      return this;
-    }
-
-    private AsyncQueryScenario assertGetResponse(
-        PPLAsyncQueryService.JobSnapshot expectedResponse) {
-      assertEquals(expectedResponse, getResponse);
-      return this;
-    }
-
-    private AsyncQueryScenario assertNotRetained() {
-      return assertCapacity(0, 0);
-    }
-
-    private AsyncQueryScenario assertCapacity(int running, int retained) {
-      assertEquals(running, targetService.runningQueryCount());
-      assertEquals(retained, targetService.retainedJobCount());
-      return this;
-    }
-
-    private AsyncQueryScenario assertExecutionReadOnceAndClosed() {
-      return assertExecutionReadCount(1).assertExecutionCloseCount(1);
-    }
-
-    private AsyncQueryScenario assertExecutionReadCount(int expected) {
-      assertEquals(expected, trackingExecution.reads.get());
-      return this;
-    }
-
-    private AsyncQueryScenario assertExecutionCloseCount(int expected) {
-      assertEquals(expected, trackingExecution.closes.get());
-      return this;
-    }
-
-    private AsyncQueryScenario assertTimeoutCancelled() {
-      assertTrue(timeoutCancelled.get());
-      return this;
-    }
-
-    private AsyncQueryScenario assertNoInitialResponse() {
-      assertNull(initialResponse.get());
-      return this;
-    }
-
-    private AsyncQueryScenario assertFailure(Class type, String message) {
-      assertTrue(type.isInstance(failure.get()));
-      assertEquals(message, failure.get().getMessage());
-      return this;
-    }
-
-    private AsyncQueryScenario assertFailure(Class type) {
-      assertTrue(type.isInstance(failure.get()));
-      return this;
-    }
-
-    private AsyncQueryScenario fireTimeout() {
-      timeoutTask.get().run();
-      return this;
-    }
-
-    private AsyncQueryScenario assertResponseCount(int expected) {
-      assertEquals(expected, responseCount.get());
-      return this;
-    }
-
-    private AsyncQueryScenario get() {
-      getResponse = targetService.get(id(), owner, null);
-      return this;
-    }
-
-    private AsyncQueryScenario delete() {
-      deleteResponse = targetService.delete(id(), owner);
-      return this;
-    }
-
-    private AsyncQueryScenario assertDeleteStatus(PPLAsyncQueryService.Status status) {
-      assertEquals(new PPLAsyncQueryService.DeleteResult(id(), status), deleteResponse);
-      return this;
-    }
-
-    private String id() {
-      return initialResponse.get().id().orElseThrow();
-    }
-  }
-
-  private static final class TrackingExecution implements AsyncQueryExecution {
-    private final AtomicReference current;
-    private final CompletableFuture completion = new CompletableFuture<>();
-    private final AtomicBoolean closed = new AtomicBoolean();
-    private final AtomicInteger reads = new AtomicInteger();
-    private final AtomicInteger closes = new AtomicInteger();
-
-    private TrackingExecution(QueryResponse current) {
-      this.current = new AtomicReference<>(current);
-    }
-
-    private void setCurrent(QueryResponse response) {
-      current.set(response);
-    }
-
-    private void succeed(QueryResponse response) {
-      current.set(response);
-      completion.complete(null);
-    }
-
-    private void complete() {
-      completion.complete(null);
-    }
-
-    private void fail(Exception failure) {
-      completion.completeExceptionally(failure);
-    }
-
-    @Override
-    public Optional currentResult() {
-      reads.incrementAndGet();
-      return Optional.ofNullable(current.get());
-    }
-
-    @Override
-    public CompletionStage completion() {
-      return completion;
-    }
-
-    @Override
-    public void close() {
-      if (closed.compareAndSet(false, true)) {
-        closes.incrementAndGet();
-      }
-    }
-  }
-
-  private static final class BlockingExecution implements AsyncQueryExecution {
-    private final QueryResponse response;
-    private final CountDownLatch readStarted = new CountDownLatch(1);
-    private final CountDownLatch allowRead = new CountDownLatch(1);
-    private final AtomicBoolean closed = new AtomicBoolean();
-    private final AtomicInteger closes = new AtomicInteger();
-
-    private BlockingExecution(QueryResponse response) {
-      this.response = response;
-    }
-
-    @Override
-    public Optional currentResult() {
-      readStarted.countDown();
-      try {
-        assertTrue(allowRead.await(5, TimeUnit.SECONDS));
-      } catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
-        throw new AssertionError(e);
-      }
-      return Optional.of(response);
-    }
-
-    @Override
-    public CompletionStage completion() {
-      return new CompletableFuture<>();
-    }
-
-    @Override
-    public void close() {
-      if (closed.compareAndSet(false, true)) {
-        closes.incrementAndGet();
-      }
-    }
-  }
-
-  private static final class ThrowingExecution implements AsyncQueryExecution {
-    private final AtomicBoolean closed = new AtomicBoolean();
-    private final AtomicInteger closes = new AtomicInteger();
-
-    @Override
-    public Optional currentResult() {
-      throw new IllegalStateException("materialization failed");
-    }
-
-    @Override
-    public CompletionStage completion() {
-      return new CompletableFuture<>();
-    }
-
-    @Override
-    public void close() {
-      if (closed.compareAndSet(false, true)) {
-        closes.incrementAndGet();
-      }
-    }
-  }
-}
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobIdTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobIdTest.java
new file mode 100644
index 00000000000..d3cb6fd5e8b
--- /dev/null
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobIdTest.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.plugin.transport.asyncquery;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import org.junit.Test;
+
+public class QueryJobIdTest {
+
+  @Test
+  public void roundTripsOwnerAndContext() {
+    QueryJobId id = QueryJobId.create("node-a");
+
+    assertEquals(id, QueryJobId.parse(id.encode()));
+  }
+
+  @Test
+  public void rejectsMalformedIds() {
+    assertThrows(IllegalArgumentException.class, () -> QueryJobId.parse(""));
+    assertThrows(IllegalArgumentException.class, () -> QueryJobId.parse("not-an-id"));
+    assertThrows(IllegalArgumentException.class, () -> new QueryJobId("", "context-id"));
+  }
+}
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobOwnerTest.java
similarity index 60%
rename from plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java
rename to plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobOwnerTest.java
index fb028eea779..aa614bf7ac7 100644
--- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/PPLAsyncQueryUserTest.java
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobOwnerTest.java
@@ -16,15 +16,15 @@
 import org.opensearch.commons.ConfigConstants;
 import org.opensearch.commons.authuser.User;
 
-public class PPLAsyncQueryUserTest {
+public class QueryJobOwnerTest {
 
   @Test
   public void capturesSecurityIdentity() {
     ThreadContext context = contextWith("alice|backend-a|ppl-role|tenant-a");
 
     assertEquals(
-        new PPLAsyncQueryUser("alice", "tenant-a", List.of("backend-a")),
-        PPLAsyncQueryUser.current(context));
+        new QueryJobOwner("alice", "tenant-a", List.of("backend-a")),
+        QueryJobOwner.current(context));
   }
 
   @Test
@@ -33,51 +33,48 @@ public void acceptsUserObject() {
         contextWith(new User("alice", List.of("backend-a"), List.of("ppl-role"), null, "tenant-a"));
 
     assertEquals(
-        new PPLAsyncQueryUser("alice", "tenant-a", List.of("backend-a")),
-        PPLAsyncQueryUser.current(context));
+        new QueryJobOwner("alice", "tenant-a", List.of("backend-a")),
+        QueryJobOwner.current(context));
   }
 
   @Test
   public void rejectsUnknownSecurityContext() {
     assertThrows(
-        OpenSearchSecurityException.class,
-        () -> PPLAsyncQueryUser.current(contextWith(new Object())));
+        OpenSearchSecurityException.class, () -> QueryJobOwner.current(contextWith(new Object())));
   }
 
   @Test
   public void missingIdentityRepresentsAnUnsecuredCaller() {
     ThreadContext context = new ThreadContext(Settings.EMPTY);
 
-    assertEquals(PPLAsyncQueryUser.UNSECURED, PPLAsyncQueryUser.current(context));
+    assertEquals(QueryJobOwner.UNSECURED, QueryJobOwner.current(context));
   }
 
   @Test
   public void requiresSamePrincipalTenantAndOriginalBackendRoles() {
-    PPLAsyncQueryUser owner = new PPLAsyncQueryUser("alice", "tenant-a", List.of("role-a"));
+    QueryJobOwner owner = new QueryJobOwner("alice", "tenant-a", List.of("role-a"));
 
-    owner.authorize(
-        new PPLAsyncQueryUser("alice", "tenant-a", List.of("role-a", "newly-added-role")));
+    owner.authorize(new QueryJobOwner("alice", "tenant-a", List.of("role-a", "newly-added-role")));
 
     assertThrows(
         OpenSearchSecurityException.class,
-        () -> owner.authorize(new PPLAsyncQueryUser("bob", "tenant-a", List.of("role-a"))));
+        () -> owner.authorize(new QueryJobOwner("bob", "tenant-a", List.of("role-a"))));
     assertThrows(
         OpenSearchSecurityException.class,
-        () -> owner.authorize(new PPLAsyncQueryUser("alice", "tenant-b", List.of("role-a"))));
+        () -> owner.authorize(new QueryJobOwner("alice", "tenant-b", List.of("role-a"))));
     assertThrows(
         OpenSearchSecurityException.class,
-        () -> owner.authorize(new PPLAsyncQueryUser("alice", "tenant-a", List.of())));
+        () -> owner.authorize(new QueryJobOwner("alice", "tenant-a", List.of())));
   }
 
   @Test
   public void unsecuredModeRequiresAnUnsecuredCaller() {
-    PPLAsyncQueryUser.UNSECURED.authorize(PPLAsyncQueryUser.UNSECURED);
+    QueryJobOwner.UNSECURED.authorize(QueryJobOwner.UNSECURED);
 
     assertThrows(
         OpenSearchSecurityException.class,
         () ->
-            PPLAsyncQueryUser.UNSECURED.authorize(
-                new PPLAsyncQueryUser("alice", null, List.of("role-a"))));
+            QueryJobOwner.UNSECURED.authorize(new QueryJobOwner("alice", null, List.of("role-a"))));
   }
 
   private static ThreadContext contextWith(Object userInfo) {

From f3a125ed519344e9994fe1e0aa0710b7340eec74 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Fri, 25 Sep 2026 16:11:52 +0000
Subject: [PATCH 14/15] test(async-query): fluent-scenario tests for QueryJob
 and QueryJobRegistry

QueryJobTest drives one QueryJob through its lifecycle using a fake
clock and a fake Scheduler. Covers direct success/failure, retention
timeout, retained success/failure with metric recording, DELETE
statuses, expiry via GET boundary and via expiry timer, keep-alive
renewal, shutdown discard, authorization, and late-attach after DELETE.

QueryJobRegistryTest covers add / remove / close semantics and the
defensive jobs() snapshot.

Signed-off-by: Peng Huo 
---
 .../asyncquery/QueryJobRegistryTest.java      |  99 +++
 .../transport/asyncquery/QueryJobTest.java    | 655 ++++++++++++++++++
 2 files changed, 754 insertions(+)
 create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobRegistryTest.java
 create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobTest.java

diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobRegistryTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobRegistryTest.java
new file mode 100644
index 00000000000..e0aeec06bb4
--- /dev/null
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobRegistryTest.java
@@ -0,0 +1,99 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.plugin.transport.asyncquery;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class QueryJobRegistryTest {
+
+  @Test
+  public void addPublishesJobAndReturnsExistingOnCollision() {
+    QueryJobRegistry registry = new QueryJobRegistry();
+    QueryJobId sharedId = new QueryJobId("node-a", "context-1");
+    QueryJob first = jobWithId(sharedId);
+    QueryJob second = jobWithId(sharedId);
+
+    assertNull(registry.add(first));
+    assertSame(first, registry.add(second));
+    assertSame(first, registry.get(sharedId).orElseThrow());
+  }
+
+  @Test
+  public void removeOnlyDropsGivenJobInstance() {
+    QueryJobRegistry registry = new QueryJobRegistry();
+    QueryJobId sharedId = new QueryJobId("node-a", "context-1");
+    QueryJob first = jobWithId(sharedId);
+    QueryJob replacement = jobWithId(sharedId);
+
+    registry.add(first);
+    assertFalse(registry.remove(sharedId, replacement));
+    assertSame(first, registry.get(sharedId).orElseThrow());
+    assertTrue(registry.remove(sharedId, first));
+    assertEquals(Optional.empty(), registry.get(sharedId));
+  }
+
+  @Test
+  public void closeDiscardsEveryJobAndRejectsFurtherAdds() {
+    QueryJobRegistry registry = new QueryJobRegistry();
+    QueryJob first = jobWithId(newId());
+    QueryJob second = jobWithId(newId());
+    registry.add(first);
+    registry.add(second);
+
+    registry.close();
+
+    Mockito.verify(first).discard("PPL asynchronous query service is closing");
+    Mockito.verify(second).discard("PPL asynchronous query service is closing");
+    assertThrows(IllegalStateException.class, () -> registry.add(jobWithId(newId())));
+  }
+
+  @Test
+  public void closeIsIdempotent() {
+    QueryJobRegistry registry = new QueryJobRegistry();
+    QueryJob job = jobWithId(newId());
+    registry.add(job);
+
+    registry.close();
+    registry.close();
+
+    Mockito.verify(job, Mockito.times(1)).discard(Mockito.anyString());
+  }
+
+  @Test
+  public void jobsReturnsDefensiveSnapshot() {
+    QueryJobRegistry registry = new QueryJobRegistry();
+    QueryJob job = jobWithId(newId());
+    registry.add(job);
+
+    var snapshot = registry.jobs();
+    registry.remove(job.getJobId(), job);
+
+    assertTrue(snapshot.contains(job));
+    assertEquals(0, registry.jobs().size());
+  }
+
+  private static final AtomicInteger COUNTER = new AtomicInteger();
+
+  private static QueryJobId newId() {
+    return new QueryJobId("node-a", "context-" + COUNTER.incrementAndGet());
+  }
+
+  private static QueryJob jobWithId(QueryJobId id) {
+    QueryJob job = Mockito.mock(QueryJob.class);
+    Mockito.when(job.getJobId()).thenReturn(id);
+    return job;
+  }
+}
diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobTest.java
new file mode 100644
index 00000000000..28d317b4da8
--- /dev/null
+++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/asyncquery/QueryJobTest.java
@@ -0,0 +1,655 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.plugin.transport.asyncquery;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.function.LongSupplier;
+import org.junit.Test;
+import org.opensearch.OpenSearchSecurityException;
+import org.opensearch.ResourceNotFoundException;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.model.ExprValueUtils;
+import org.opensearch.sql.data.type.ExprCoreType;
+import org.opensearch.sql.executor.AsyncQueryExecution;
+import org.opensearch.sql.executor.ExecutionEngine.QueryResponse;
+import org.opensearch.sql.executor.ExecutionEngine.Schema;
+import org.opensearch.sql.executor.ExecutionEngine.Schema.Column;
+import org.opensearch.sql.legacy.metrics.BasicCounter;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.metrics.Metrics;
+import org.opensearch.sql.legacy.metrics.NumericMetric;
+import org.opensearch.tasks.CancellableTask;
+import org.opensearch.tasks.TaskManager;
+
+public class QueryJobTest {
+  private static final String LOCAL_NODE = "node-a";
+  private static final QueryJobOwner OWNER = QueryJobOwner.UNSECURED;
+  private static final TimeValue KEEP_ALIVE = TimeValue.timeValueMinutes(5);
+  private static final TimeValue DEFAULT_WAIT = TimeValue.timeValueSeconds(5);
+
+  @Test
+  public void fastSuccessReturnsDirectResultWithoutRetention() {
+    scenario()
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .withCurrentResult(response(2))
+        .start()
+        .completeAfterMillis(25)
+        .assertDirectSuccess(response(2), 25)
+        .assertNotRegistered()
+        .assertExecutionClosedOnce();
+  }
+
+  @Test
+  public void completedExecutionCanBeAttachedBeforeCompletionIsObserved() {
+    scenario()
+        .succeed(response(2))
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .start()
+        .assertDirectSuccess(response(2), 0)
+        .assertExecutionClosedOnce();
+  }
+
+  @Test
+  public void timeoutReturnsIdAndLaterGetReturnsCompleteResult() {
+    QueryJobScenario scenario =
+        scenario()
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start()
+            .fireRetention()
+            .assertRetainedRunning(Optional.empty());
+
+    scenario
+        .succeedAfterMillis(25, response(2))
+        .get()
+        .assertGetResponse(
+            new QueryJob.Snapshot.Succeeded(Optional.of(scenario.id()), response(2), 25));
+  }
+
+  @Test
+  public void fastFailureReturnsDirectFailureWithoutId() {
+    scenario()
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .start()
+        .fail(new IllegalStateException("boom"))
+        .assertFailure(IllegalStateException.class, "boom")
+        .assertNotRegistered()
+        .assertExecutionClosedOnce();
+  }
+
+  @Test
+  public void executionStartFailureCompletesJobAndReleasesTask() {
+    QueryJobScenario scenario =
+        scenario()
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .withExecutionStarter(
+                ignored -> {
+                  throw new IllegalStateException("execution did not start");
+                });
+    scenario.start();
+    scenario
+        .assertFailure(IllegalStateException.class, "execution did not start")
+        .assertNotRegistered()
+        .assertTaskClosed();
+  }
+
+  @Test
+  public void retainedFailureRetainsFailureSnapshotAndRecordsMetric() {
+    NumericMetric customerFailures =
+        new NumericMetric<>(MetricName.PPL_FAILED_REQ_COUNT_CUS.getName(), new BasicCounter());
+    Metrics.getInstance().registerMetric(customerFailures);
+    try {
+      QueryJobScenario scenario =
+          scenario()
+              .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+              .start()
+              .fireRetention()
+              .fail(new IllegalArgumentException("invalid query"))
+              .get();
+
+      scenario.assertGetResponse(
+          new QueryJob.Snapshot.Failed(
+              Optional.of(scenario.id()),
+              new QueryJob.Failure("IllegalArgumentException", "invalid query"),
+              0));
+      assertEquals(Long.valueOf(1), customerFailures.getValue());
+    } finally {
+      Metrics.getInstance().unregisterMetric(customerFailures.getName());
+    }
+  }
+
+  @Test
+  public void expiredGetRemovesJobAndCancelsTask() {
+    CancellableTask task = runningTask();
+    QueryJobScenario scenario =
+        scenario().withTask(task).start().fireRetention().advanceMillis(KEEP_ALIVE.millis());
+
+    ResourceNotFoundException error = assertThrows(ResourceNotFoundException.class, scenario::get);
+
+    assertEquals("PPL asynchronous query not found", error.getMessage());
+    scenario.assertNotRegistered();
+    verify(task).cancel("PPL asynchronous query expired");
+  }
+
+  @Test
+  public void expiryTimerRemovesJobAfterKeepAliveElapses() {
+    CancellableTask task = runningTask();
+    QueryJobScenario scenario =
+        scenario().withTask(task).start().fireRetention().assertRegistered();
+
+    scenario.advanceMillis(KEEP_ALIVE.millis()).fireExpiry();
+
+    scenario.assertNotRegistered();
+    verify(task).cancel("PPL asynchronous query expired");
+  }
+
+  @Test
+  public void deleteCancelsRunningJobAndReturnsCancelledStatus() {
+    CancellableTask task = runningTask();
+    QueryJobScenario scenario = scenario().withTask(task).start().fireRetention().delete();
+
+    scenario.assertDeleteStatus(QueryJob.Status.CANCELLED);
+    verify(task).cancel("PPL asynchronous query cancelled by user");
+    scenario.assertNotRegistered().assertExecutionClosedOnce();
+    assertThrows(ResourceNotFoundException.class, scenario::get);
+  }
+
+  @Test
+  public void deleteReturnsRetainedTerminalStatus() {
+    CancellableTask task = runningTask();
+    QueryJobScenario scenario =
+        scenario()
+            .withTask(task)
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .withCurrentResult(response(1))
+            .start()
+            .fireRetention()
+            .complete()
+            .delete();
+
+    scenario.assertDeleteStatus(QueryJob.Status.SUCCEEDED);
+    verify(task, never()).cancel(org.mockito.ArgumentMatchers.anyString());
+    scenario.assertNotRegistered().assertExecutionClosedOnce();
+  }
+
+  @Test
+  public void deleteReturnsNotFoundAfterExpirationBoundary() {
+    QueryJobScenario scenario = scenario().start().fireRetention();
+
+    scenario.advanceMillis(KEEP_ALIVE.millis());
+
+    assertThrows(ResourceNotFoundException.class, scenario::delete);
+    scenario.assertNotRegistered();
+  }
+
+  @Test
+  public void getRenewsLeaseAndResetsExpiryTimer() {
+    QueryJobScenario scenario =
+        scenario()
+            .withKeepAlive(TimeValue.timeValueSeconds(1))
+            .start()
+            .fireRetention()
+            .advanceMillis(500)
+            .getWithKeepAlive(TimeValue.timeValueSeconds(2))
+            .advanceMillis(1_999)
+            .assertRegistered();
+
+    // Original lease would have expired at 1000ms; renewed lease survives to 2499ms.
+    scenario.advanceMillis(1).fireExpiry().assertNotRegistered();
+  }
+
+  @Test
+  public void cancellationUsesTaskManagerWhenAttached() {
+    TaskManager taskManager = mock(TaskManager.class);
+    CancellableTask task = runningTask();
+    scenario().withTaskManager(taskManager).withTask(task).start().fireRetention().delete();
+
+    verify(taskManager)
+        .cancelTaskAndDescendants(
+            org.mockito.ArgumentMatchers.eq(task),
+            org.mockito.ArgumentMatchers.eq("PPL asynchronous query cancelled by user"),
+            org.mockito.ArgumentMatchers.eq(false),
+            org.mockito.ArgumentMatchers.any());
+  }
+
+  @Test
+  public void shutdownDiscardsRunningJobAndClosesResources() {
+    CancellableTask task = runningTask();
+    QueryJobScenario scenario =
+        scenario().withTask(task).withWaitForCompletion(TimeValue.timeValueSeconds(5)).start();
+
+    scenario.closeRegistry();
+
+    scenario.assertNotRegistered().assertExecutionClosedOnce();
+    verify(task).cancel("PPL asynchronous query service is closing");
+  }
+
+  @Test
+  public void unauthorizedCallerCannotAccessOwnedJob() {
+    QueryJobOwner alice = new QueryJobOwner("alice", "tenant-a", List.of("role-a"));
+    QueryJobOwner bob = new QueryJobOwner("bob", "tenant-a", List.of("role-a"));
+    QueryJobScenario scenario = scenario().withOwner(alice).start().fireRetention();
+
+    QueryJob job = scenario.job();
+    assertThrows(OpenSearchSecurityException.class, () -> job.getOwner().authorize(bob));
+    job.getOwner().authorize(alice);
+  }
+
+  @Test
+  public void lateExecutionAttachmentAfterDeleteIsClosed() {
+    TrackingExecution execution = new TrackingExecution(response(1));
+    QueryJobScenario scenario = scenario().withExecution(execution);
+    scenario
+        .withExecutionStarter(
+            ignored -> {
+              // Retention fires and delete removes the job before the execution attaches.
+              scenario.fireRetention();
+              QueryJobId retainedId =
+                  ((QueryJob.Snapshot.Running) scenario.initialResponse.get()).jobId();
+              scenario.registry.get(retainedId).orElseThrow().cancel("test delete");
+              return execution;
+            })
+        .start();
+
+    assertEquals(1, execution.closes.get());
+    assertTrue(scenario.registry.jobs().isEmpty());
+  }
+
+  @Test
+  public void finalSnapshotDefensivelyCopiesRows() {
+    List rows = new ArrayList<>();
+    rows.add(ExprValueUtils.stringValue("first"));
+    QueryResponse response =
+        new QueryResponse(
+            new Schema(List.of(new Column("state", null, ExprCoreType.STRING))), rows, null);
+
+    QueryJobScenario scenario =
+        scenario()
+            .withCurrentResult(response)
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start()
+            .complete();
+    rows.add(ExprValueUtils.stringValue("second"));
+
+    scenario.assertSucceededRowCount(1);
+  }
+
+  @Test
+  public void successfulCompletionRequiresFinalResultToBeVisible() {
+    scenario()
+        .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+        .start()
+        .complete()
+        .assertFailure(IllegalStateException.class)
+        .assertExecutionClosedOnce();
+  }
+
+  @Test
+  public void retentionListenerFailureAbortsUndeliverableJob() {
+    CancellableTask task = runningTask();
+    QueryJobScenario scenario =
+        scenario()
+            .withTask(task)
+            .withResponseListener(new FailingListener())
+            .withWaitForCompletion(TimeValue.timeValueSeconds(5))
+            .start();
+
+    scenario.fireRetention();
+
+    scenario.assertNotRegistered().assertExecutionClosedOnce();
+    verify(task).cancel("PPL asynchronous query startup failed");
+  }
+
+  private QueryJobScenario scenario() {
+    return new QueryJobScenario();
+  }
+
+  private static CancellableTask runningTask() {
+    CancellableTask task = mock(CancellableTask.class);
+    when(task.isCancelled()).thenReturn(false);
+    return task;
+  }
+
+  private static QueryResponse response(int rowCount) {
+    Schema schema = new Schema(List.of(new Column("state", null, ExprCoreType.STRING)));
+    return new QueryResponse(
+        schema,
+        java.util.stream.IntStream.range(0, rowCount)
+            .mapToObj(i -> ExprValueUtils.stringValue("state-" + i))
+            .toList(),
+        null);
+  }
+
+  /** Deterministic fluent scenario driving one {@link QueryJob} through its lifecycle. */
+  private final class QueryJobScenario {
+    private final AtomicLong clock = new AtomicLong(1_000);
+    private final QueryJobRegistry registry = new QueryJobRegistry();
+    private final FakeScheduler scheduler = new FakeScheduler();
+    private QueryJobOwner owner = OWNER;
+    private TimeValue keepAlive = KEEP_ALIVE;
+    private TimeValue waitForCompletion = DEFAULT_WAIT;
+    private CancellableTask task = mock(CancellableTask.class);
+    private final AtomicInteger taskReleaseCount = new AtomicInteger();
+    private TaskManager taskManager;
+    private TrackingExecution trackingExecution = new TrackingExecution(null);
+    private AsyncQueryExecution execution = trackingExecution;
+    private Function executionStarter = ignored -> execution;
+    private ActionListener responseListener;
+    private final AtomicReference initialResponse = new AtomicReference<>();
+    private final AtomicReference failure = new AtomicReference<>();
+    private final AtomicInteger responseCount = new AtomicInteger();
+    private QueryJob.Snapshot getResponse;
+    private QueryJob.Status deleteStatus;
+    private QueryJob job;
+
+    private QueryJobScenario withOwner(QueryJobOwner owner) {
+      this.owner = owner;
+      return this;
+    }
+
+    private QueryJobScenario withKeepAlive(TimeValue keepAlive) {
+      this.keepAlive = keepAlive;
+      return this;
+    }
+
+    private QueryJobScenario withWaitForCompletion(TimeValue wait) {
+      this.waitForCompletion = wait;
+      return this;
+    }
+
+    private QueryJobScenario withTask(CancellableTask task) {
+      this.task = task;
+      return this;
+    }
+
+    private QueryJobScenario withTaskManager(TaskManager taskManager) {
+      this.taskManager = taskManager;
+      return this;
+    }
+
+    private QueryJobScenario withCurrentResult(QueryResponse response) {
+      trackingExecution.setCurrent(response);
+      return this;
+    }
+
+    private QueryJobScenario withExecution(AsyncQueryExecution execution) {
+      this.execution = execution;
+      trackingExecution = execution instanceof TrackingExecution tracking ? tracking : null;
+      executionStarter = ignored -> this.execution;
+      return this;
+    }
+
+    private QueryJobScenario withExecutionStarter(
+        Function starter) {
+      this.executionStarter = starter;
+      return this;
+    }
+
+    private QueryJobScenario withResponseListener(ActionListener listener) {
+      this.responseListener = listener;
+      return this;
+    }
+
+    private QueryJobScenario start() {
+      QueryJob.JobTask jobTask = new QueryJob.JobTask(task, taskReleaseCount::incrementAndGet);
+      ActionListener listener =
+          responseListener != null ? responseListener : defaultListener();
+      job =
+          QueryJob.create(
+              LOCAL_NODE,
+              registry,
+              (LongSupplier) clock::get,
+              scheduler,
+              taskManager,
+              jobTask,
+              owner,
+              keepAlive,
+              waitForCompletion,
+              executionStarter,
+              listener);
+      return this;
+    }
+
+    private ActionListener defaultListener() {
+      return ActionListener.wrap(
+          snapshot -> {
+            initialResponse.set(snapshot);
+            responseCount.incrementAndGet();
+          },
+          failure::set);
+    }
+
+    private QueryJobScenario succeed(QueryResponse response) {
+      trackingExecution.succeed(response);
+      return this;
+    }
+
+    private QueryJobScenario complete() {
+      trackingExecution.complete();
+      return this;
+    }
+
+    private QueryJobScenario completeAfterMillis(long elapsedMillis) {
+      clock.addAndGet(elapsedMillis);
+      trackingExecution.complete();
+      return this;
+    }
+
+    private QueryJobScenario succeedAfterMillis(long elapsedMillis, QueryResponse response) {
+      clock.addAndGet(elapsedMillis);
+      trackingExecution.succeed(response);
+      return this;
+    }
+
+    private QueryJobScenario fail(Exception exception) {
+      trackingExecution.fail(exception);
+      return this;
+    }
+
+    private QueryJobScenario advanceMillis(long delta) {
+      clock.addAndGet(delta);
+      return this;
+    }
+
+    private QueryJobScenario fireRetention() {
+      scheduler.fireNext();
+      return this;
+    }
+
+    private QueryJobScenario fireExpiry() {
+      scheduler.fireNext();
+      return this;
+    }
+
+    private QueryJobScenario get() {
+      getResponse = job.get(null);
+      return this;
+    }
+
+    private QueryJobScenario getWithKeepAlive(TimeValue newKeepAlive) {
+      getResponse = job.get(newKeepAlive);
+      return this;
+    }
+
+    private QueryJobScenario delete() {
+      deleteStatus = job.cancel("PPL asynchronous query cancelled by user");
+      return this;
+    }
+
+    private QueryJobScenario closeRegistry() {
+      registry.close();
+      return this;
+    }
+
+    private QueryJobScenario assertDirectSuccess(QueryResponse expected, long tookMillis) {
+      assertEquals(
+          new QueryJob.Snapshot.Succeeded(Optional.empty(), expected, tookMillis),
+          initialResponse.get());
+      return this;
+    }
+
+    private QueryJobScenario assertRetainedRunning(Optional expected) {
+      assertEquals(new QueryJob.Snapshot.Running(id(), expected), initialResponse.get());
+      return this;
+    }
+
+    private QueryJobScenario assertGetResponse(QueryJob.Snapshot expected) {
+      assertEquals(expected, getResponse);
+      return this;
+    }
+
+    private QueryJobScenario assertDeleteStatus(QueryJob.Status expected) {
+      assertEquals(expected, deleteStatus);
+      return this;
+    }
+
+    private QueryJobScenario assertSucceededRowCount(int expected) {
+      assertTrue(initialResponse.get() instanceof QueryJob.Snapshot.Succeeded);
+      QueryJob.Snapshot.Succeeded succeeded = (QueryJob.Snapshot.Succeeded) initialResponse.get();
+      assertEquals(expected, succeeded.response().getResults().size());
+      return this;
+    }
+
+    private QueryJobScenario assertFailure(Class type) {
+      assertTrue(type.isInstance(failure.get()));
+      return this;
+    }
+
+    private QueryJobScenario assertFailure(Class type, String message) {
+      assertTrue(type.isInstance(failure.get()));
+      assertEquals(message, failure.get().getMessage());
+      return this;
+    }
+
+    private QueryJobScenario assertRegistered() {
+      assertTrue(registry.get(job.getJobId()).isPresent());
+      return this;
+    }
+
+    private QueryJobScenario assertNotRegistered() {
+      assertTrue(registry.get(job.getJobId()).isEmpty());
+      return this;
+    }
+
+    private QueryJobScenario assertExecutionClosedOnce() {
+      if (trackingExecution != null) {
+        assertEquals(1, trackingExecution.closes.get());
+      }
+      return this;
+    }
+
+    private QueryJobScenario assertTaskClosed() {
+      assertEquals(1, taskReleaseCount.get());
+      return this;
+    }
+
+    private QueryJobId id() {
+      return job.getJobId();
+    }
+
+    private QueryJob job() {
+      return job;
+    }
+  }
+
+  private static final class FakeScheduler implements QueryJob.Scheduler {
+    private final Deque pending = new ArrayDeque<>();
+
+    @Override
+    public QueryJob.Cancellable schedule(long delayMillis, Runnable task) {
+      Pending entry = new Pending(delayMillis, task);
+      pending.addLast(entry);
+      return () -> pending.remove(entry);
+    }
+
+    void fireNext() {
+      Pending entry = pending.pollFirst();
+      assertTrue("No scheduled task", entry != null);
+      entry.task.run();
+    }
+
+    private record Pending(long delayMillis, Runnable task) {}
+  }
+
+  private static final class TrackingExecution implements AsyncQueryExecution {
+    private final AtomicReference current;
+    private final CompletableFuture completion = new CompletableFuture<>();
+    private final AtomicBoolean closed = new AtomicBoolean();
+    final AtomicInteger reads = new AtomicInteger();
+    final AtomicInteger closes = new AtomicInteger();
+
+    private TrackingExecution(QueryResponse current) {
+      this.current = new AtomicReference<>(current);
+    }
+
+    private void setCurrent(QueryResponse response) {
+      current.set(response);
+    }
+
+    private void succeed(QueryResponse response) {
+      current.set(response);
+      completion.complete(null);
+    }
+
+    private void complete() {
+      completion.complete(null);
+    }
+
+    private void fail(Exception failure) {
+      completion.completeExceptionally(failure);
+    }
+
+    @Override
+    public Optional currentResult() {
+      reads.incrementAndGet();
+      return Optional.ofNullable(current.get());
+    }
+
+    @Override
+    public CompletionStage completion() {
+      return completion;
+    }
+
+    @Override
+    public void close() {
+      if (closed.compareAndSet(false, true)) {
+        closes.incrementAndGet();
+      }
+    }
+  }
+
+  private static final class FailingListener implements ActionListener {
+    @Override
+    public void onResponse(QueryJob.Snapshot snapshot) {
+      throw new IllegalStateException("listener rejected snapshot");
+    }
+
+    @Override
+    public void onFailure(Exception e) {
+      // no-op; QueryJob may deliver a subsequent failure once discard runs.
+    }
+  }
+}

From 14b9f57c62b6a3e550fee04fe986b9ad8c8bbfe6 Mon Sep 17 00:00:00 2001
From: Peng Huo 
Date: Sat, 26 Sep 2026 01:01:37 +0000
Subject: [PATCH 15/15] docs(async-query): MVP redesign proposal for
 language-neutral job lifecycle

Design-only response to review of opensearch-project/sql#5809. No code change.

- Language-neutral lifecycle layer (Livy-style control plane)
- Engine-neutral execution boundary (QueryRunner SPI for PPL/SQL/AE)
- BigQuery-style state model: retention and waiting are not state
- Split of QueryJob / QueryJobStore / QueryJobService responsibilities
- Sync-query is submit + block on completion(); async-query returns the ID

Signed-off-by: Peng Huo 
---
 docs/dev/query-job-mvp-design.md | 433 +++++++++++++++++++++++++++++++
 1 file changed, 433 insertions(+)
 create mode 100644 docs/dev/query-job-mvp-design.md

diff --git a/docs/dev/query-job-mvp-design.md b/docs/dev/query-job-mvp-design.md
new file mode 100644
index 00000000000..0a534014804
--- /dev/null
+++ b/docs/dev/query-job-mvp-design.md
@@ -0,0 +1,433 @@
+# Query Job Lifecycle — MVP Redesign
+
+> Response to review of [PR #5809](https://github.com/opensearch-project/sql/pull/5809).
+> This document is design-only; no code has been changed.
+
+## 1. What we heard from review
+
+Three concerns from @dai-chen shaped this redesign.
+
+| # | Concern | Where it landed |
+|---|---------|-----------------|
+| 1 | Make the lifecycle layer language-neutral, similar to Livy's control plane over Spark. | The `job` package no longer names PPL. It talks only to a `QueryRunner` SPI. Language modules (`ppl`, `sql`, `async-query`) implement `QueryRunnerFactory`. |
+| 2 | The execution boundary can be engine-neutral too — SQL, PPL, and AE all run behind it. | `QueryRunner` returns a neutral `QueryResult`. It knows nothing about ANTLR, `PPLService`, `SQLService`, or Spark. The same interface fits AE. |
+| 3 | Separate execution state from retention. In a BigQuery-style model the job stays `RUNNING` regardless of whether a caller is still waiting. Treat `wait_for_completion_timeout` as submit-response behavior only. | The state machine has five values (`PENDING`, `RUNNING`, `SUCCEEDED`, `FAILED`, `CANCELLED`) and no retention axis. Waiting is a *transport* concern, not a *lifecycle* concern. |
+| 3a | Extract internal responsibilities (JobStore, LeaseManager, …) from the service class. | `QueryJobService`, `QueryJobStore`, and `QueryJob` split along Single-Responsibility lines. `LeaseManager` is intentionally out of MVP; it can be added without changing the public interfaces (§7). |
+
+## 2. MVP scope
+
+**In**
+
+- One neutral job model that carries both sync and async queries.
+- Submit / get / cancel over a small, engine-agnostic control plane.
+- Owner-node routing via an opaque `QueryJobId`.
+- Caller identity captured at submit and checked at get / cancel.
+
+**Out (deferred without breaking the API)**
+
+- `wait_for_completion_timeout` — MVP treats every submission as `wait_for_completion=forever`. The transport blocks on `QueryJob#completion()` until it resolves.
+- `keep_alive` and retention leases — no expiry timer, no keep-alive renewal. Terminal jobs live only long enough for the caller to observe them.
+- Admission control (`maxRunningQueries`, `maxRetainedJobs`).
+- Persistent job store.
+- Partial-result streaming (already deferred by the current PR).
+
+Everything under **Out** is a store, service, or transport concern. None of it changes the class surface described in §4.
+
+## 3. Livy inspiration in one paragraph
+
+Livy is a REST control plane that owns *sessions* and *statements* on top of Spark. Livy has no compiler and no query engine — it accepts a language tag (`scala`, `python`, `sql`, `r`), forwards the statement to the right session, and tracks a state machine on top. The session state machine is small: waiting → running → available | error | cancelled. Language is data; the control plane is code. We adopt the same split: `QueryJob` is language-blind, and `QueryRunner` is the per-engine adapter.
+
+## 4. Class model
+
+### 4.1 Package layout
+
+```
+core/src/main/java/org/opensearch/sql/job/
+├── QueryJob.java                 // active object (state machine)
+├── QueryJobId.java               // opaque, node-routable identifier
+├── QueryJobState.java            // PENDING, RUNNING, SUCCEEDED, FAILED, CANCELLED
+├── QueryJobStatus.java           // immutable snapshot
+├── QueryJobStore.java            // registry SPI (in-memory MVP impl provided)
+├── QueryJobService.java          // control plane: submit / get / cancel
+├── QueryRunner.java              // engine SPI: run() + cancel()
+├── QueryRunnerFactory.java       // engine SPI: build a runner from a request
+├── QueryResult.java              // neutral final result
+├── QueryFailure.java             // neutral failure descriptor
+├── Principal.java                // caller identity + authorization
+└── SubmitRequest.java            // neutral submission payload
+```
+
+`core` already has no dependency on other modules, so `job` is free of language and engine imports. `ppl`, `sql`, `async-query`, and future engines depend on `job` — not the other way around.
+
+### 4.2 Class diagram
+
+```mermaid
+classDiagram
+    class QueryJobService {
+        <>
+        +submit(SubmitRequest) QueryJob
+        +get(QueryJobId, Principal) QueryJobStatus
+        +cancel(QueryJobId, Principal) QueryJobStatus
+    }
+
+    class QueryJob {
+        -QueryJobId id
+        -Principal owner
+        -QueryRunner runner
+        -QueryJobState state
+        -CompletableFuture~QueryResult~ completion
+        +id() QueryJobId
+        +owner() Principal
+        +status() QueryJobStatus
+        +completion() CompletionStage~QueryResult~
+        +cancel() void
+    }
+
+    class QueryJobStore {
+        <>
+        +register(QueryJob) QueryJob
+        +find(QueryJobId) Optional~QueryJob~
+        +remove(QueryJobId, QueryJob) boolean
+        +jobs() Collection~QueryJob~
+        +close() void
+    }
+
+    class QueryRunner {
+        <>
+        +run() CompletionStage~QueryResult~
+        +cancel() void
+    }
+
+    class QueryRunnerFactory {
+        <>
+        +create(SubmitRequest) QueryRunner
+    }
+
+    class QueryJobId {
+        <>
+        +ownerNodeId String
+        +contextId String
+        +encode() String
+        +parse(String)$ QueryJobId
+    }
+
+    class QueryJobStatus {
+        <>
+        +id QueryJobId
+        +state QueryJobState
+        +submittedAtMillis long
+        +startedAtMillis OptionalLong
+        +completedAtMillis OptionalLong
+        +failure Optional~QueryFailure~
+    }
+
+    class Principal {
+        <>
+        +name String
+        +tenant String
+        +backendRoles List~String~
+        +authorize(Principal) void
+    }
+
+    class QueryResult {
+        <>
+        +schema Schema
+        +rows List~ExprValue~
+        +tookMillis long
+    }
+
+    class QueryFailure {
+        <>
+        +type String
+        +reason String
+    }
+
+    QueryJobService --> QueryJobStore : uses
+    QueryJobService --> QueryRunnerFactory : uses
+    QueryJobService --> QueryJob : creates
+    QueryJob --> QueryRunner : delegates execution
+    QueryJob --> Principal : owner
+    QueryJob --> QueryJobId : id
+    QueryJob --> QueryJobStatus : projects
+    QueryJobStore --> QueryJob : holds
+    QueryJobStatus --> QueryJobId
+    QueryJobStatus --> QueryJobState
+    QueryJobStatus --> QueryFailure
+```
+
+### 4.3 State machine
+
+Retention is not part of state. Waiting is not part of state.
+
+```
+              submit()
+               │
+               ▼
+           ┌───────┐  runner starts   ┌────────┐
+           │PENDING│─────────────────▶│RUNNING │
+           └───────┘                   └────┬───┘
+                                            │
+                    ┌───────────────────────┼───────────────────────┐
+                    ▼                       ▼                       ▼
+               ┌─────────┐             ┌────────┐             ┌─────────┐
+               │SUCCEEDED│             │ FAILED │             │CANCELLED│
+               └─────────┘             └────────┘             └─────────┘
+
+Terminal transitions are idempotent. cancel() from any non-terminal state
+moves to CANCELLED; cancel() from a terminal state is a no-op.
+```
+
+## 5. Interface contracts
+
+Signatures below define the contract only. JavaDoc phrased for readers of Effective Java: each interface has one reason to exist (Item 20) and each record documents its invariants (Item 17).
+
+### 5.1 `QueryJobId`
+
+```java
+public record QueryJobId(String ownerNodeId, String contextId) {
+    public QueryJobId {
+        // both fields must be non-blank
+    }
+    /** Fresh ID with a random context. */
+    public static QueryJobId create(String ownerNodeId) { … }
+    /** URL-safe, versioned, opaque encoding. */
+    public String encode() { … }
+    /** Uniform error on any parse failure. */
+    public static QueryJobId parse(String encoded) { … }
+}
+```
+
+Kept intentionally identical to today's `QueryJobId`. The encoding is a public wire format; a design-level rewrite must not disturb it.
+
+### 5.2 `Principal`
+
+```java
+public record Principal(String name, String tenant, List backendRoles) {
+    public static final Principal UNSECURED = new Principal(null, null, List.of());
+    public static Principal current(ThreadContext ctx) { … }
+    public void authorize(Principal caller) { … }
+}
+```
+
+Renamed from `QueryJobOwner` for neutrality — the identity is not job-specific and will be reused by any future admin action.
+
+### 5.3 `QueryJobState`, `QueryJobStatus`, `QueryFailure`
+
+```java
+public enum QueryJobState {
+    PENDING, RUNNING, SUCCEEDED, FAILED, CANCELLED;
+    public boolean isTerminal() { … }
+}
+
+public record QueryJobStatus(
+        QueryJobId id,
+        QueryJobState state,
+        long submittedAtMillis,
+        OptionalLong startedAtMillis,
+        OptionalLong completedAtMillis,
+        Optional failure,
+        Optional result) {
+    // result present iff state == SUCCEEDED
+    // failure present iff state == FAILED
+}
+
+public record QueryFailure(String type, String reason) {
+    public static QueryFailure of(Throwable t) { … }
+}
+```
+
+`QueryJobStatus` collapses the previous `Snapshot` sealed hierarchy. The three record variants (`Running`, `Succeeded`, `Failed`) exist today only because retention returned different shapes; without retention, one record with a state discriminator is smaller and enforces the same invariants via record validation.
+
+### 5.4 `QueryResult`
+
+```java
+public record QueryResult(Schema schema, List rows, long tookMillis) {}
+```
+
+Uses the existing `core` types (`ExecutionEngine.Schema`, `ExprValue`) so all engines already implementing them can produce a `QueryResult` without introducing new dependencies.
+
+### 5.5 `SubmitRequest`
+
+```java
+public record SubmitRequest(
+        String language,             // "ppl" | "sql" | ...
+        String statement,
+        Map params,  // engine-specific, opaque to core
+        Principal submitter) {}
+```
+
+Language is data. The service uses it to pick a `QueryRunnerFactory`. Neither the service nor the job parses `statement` or reads `params`.
+
+### 5.6 `QueryRunner` and `QueryRunnerFactory`
+
+```java
+public interface QueryRunner {
+    /** Single-use. Multiple invocations must throw IllegalStateException. */
+    CompletionStage run();
+
+    /** Idempotent cooperative cancel. Safe to call before run() and after completion. */
+    void cancel();
+}
+
+public interface QueryRunnerFactory {
+    /** Language handled by this factory, matched against SubmitRequest#language. */
+    String language();
+    QueryRunner create(SubmitRequest request);
+}
+```
+
+Two-method interface (Item 21: interfaces are for use, not for reuse). PPL, SQL, and AE each provide one implementation and register it via Guice. The lifecycle package never imports any of them.
+
+### 5.7 `QueryJob`
+
+```java
+public final class QueryJob {
+
+    /** Package-private: only QueryJobService instantiates jobs. */
+    QueryJob(QueryJobId id, Principal owner, QueryRunner runner, Clock clock) { … }
+
+    public QueryJobId id() { … }
+    public Principal owner() { … }
+    public QueryJobStatus status() { … }         // snapshot, never throws
+    public CompletionStage completion() { … }
+    public void cancel() { … }                   // idempotent
+}
+```
+
+- Thread safety documented at the class level (Effective Java Item 82). All mutable fields are guarded by `this`; every side effect (runner cancel, store removal, listener notification) happens **after** the monitor is released, exactly as today's `QueryJob` does.
+- `completion()` returns a `CompletionStage` view. The internal `CompletableFuture` is never leaked, so callers cannot complete the job externally (Item 15).
+- `cancel()` is the *only* public mutator. Retention, expiry, and admission are not job concerns.
+
+### 5.8 `QueryJobStore`
+
+```java
+public interface QueryJobStore extends Closeable {
+    QueryJob register(QueryJob job);              // returns existing on duplicate ID
+    Optional find(QueryJobId id);
+    boolean remove(QueryJobId id, QueryJob job);  // conditional
+    Collection jobs();
+    @Override void close();                       // drops residual jobs
+}
+```
+
+The MVP ships one implementation: `InMemoryQueryJobStore`, a thin wrapper around `ConcurrentHashMap`. A persistent (system-index) implementation can be added later without touching `QueryJob` or `QueryJobService`.
+
+### 5.9 `QueryJobService`
+
+```java
+public interface QueryJobService {
+    QueryJob submit(SubmitRequest request);
+    QueryJobStatus get(QueryJobId id, Principal caller);
+    QueryJobStatus cancel(QueryJobId id, Principal caller);
+}
+```
+
+A single production implementation, `LocalQueryJobService`, wires:
+
+- `QueryJobStore` for the registry,
+- a `Map` keyed by language,
+- a `Clock` supplier,
+- a local-node-id supplier for job ID minting.
+
+The service is **not** an `AbstractLifecycleComponent`. Lifecycle wiring (start / stop) lives in the plugin module (`AsyncQueryLifecycle`), keeping `core` free of OpenSearch node types.
+
+## 6. Sync = async with `wait_for_timeout=forever`
+
+Only the transport layer knows whether a caller is willing to wait.
+
+```mermaid
+sequenceDiagram
+    autonumber
+    participant Client
+    participant Transport as SyncTransport
+    participant Service as QueryJobService
+    participant Job as QueryJob
+    participant Runner as QueryRunner
+
+    Client->>Transport: POST /_plugins/_ppl {query}
+    Transport->>Service: submit(request)
+    Service->>Job: new QueryJob(id, owner, runner)
+    Service->>Runner: run()
+    Service-->>Transport: QueryJob
+    Runner-->>Job: result / failure
+    Transport->>Job: completion().join()
+    Job-->>Transport: QueryResult
+    Transport-->>Client: formatted response
+```
+
+The async transport differs only after step 5:
+
+```mermaid
+sequenceDiagram
+    autonumber
+    participant Client
+    participant Transport as AsyncTransport
+    participant Service as QueryJobService
+    participant Job as QueryJob
+
+    Client->>Transport: POST /_plugins/_async_ppl {query}
+    Transport->>Service: submit(request)
+    Service-->>Transport: QueryJob
+    Transport-->>Client: 202 {jobId}
+    Note over Client,Transport: later …
+    Client->>Transport: GET /_plugins/_async_ppl/{jobId}
+    Transport->>Service: get(jobId, caller)
+    Service-->>Transport: QueryJobStatus
+    Transport-->>Client: status (+ result if SUCCEEDED)
+```
+
+Both paths use the same `QueryJob`. Nothing about the job knows which one was chosen.
+
+## 7. Extensibility without churn
+
+Every deferred feature lands behind an interface that already exists.
+
+| Feature | Change surface |
+|---------|----------------|
+| `wait_for_completion_timeout` | Transport chooses `completion().orTimeout(...)` and returns `QueryJobStatus.RUNNING` on timeout. Job unchanged. |
+| `keep_alive` | New `RetentionPolicy` composed into `QueryJobService`. Adds an `AutoCloseable` timer per terminal job; store gains `evict()`. Job unchanged. |
+| Admission control | `AdmissionController` interface, invoked before `store.register`. Job unchanged. |
+| Persistent job metadata | Alternate `QueryJobStore` implementation. Service and job unchanged. |
+| SQL and AE support | New `QueryRunnerFactory` implementations. Zero core changes. |
+| Progress / partial results | `QueryRunner` gains a second method (default returning empty). All existing engines compile without change. |
+
+Because each seam is an interface with one reason to change, the deferrals really are deferrals — not "we'll rewrite later".
+
+## 8. Effective Java / Clean Code checklist
+
+- **Item 1** — Static factory methods where they help: `QueryJobId.create`, `QueryFailure.of`, `Principal.current`. Constructors are hidden or package-private.
+- **Item 15** — `QueryJob`'s `CompletableFuture` is never exposed; only its `CompletionStage` view is.
+- **Item 17** — `QueryJobId`, `QueryJobStatus`, `QueryFailure`, `QueryResult`, `Principal`, `SubmitRequest` are `record`s with compact-constructor validation.
+- **Item 18** — `QueryJobService` composes a store, factories, and a clock. No inheritance chain.
+- **Item 20** — Every seam is an interface. Alternate engines and stores drop in.
+- **Item 22** — Each interface expresses one role; no marker or constant interfaces.
+- **Item 24** — `QueryJobStatus` is a top-level record, not a nested type of `QueryJob`.
+- **Item 55** — `Optional` on values that legitimately may be absent (`find`, `failure`, `result`, `startedAtMillis`). `OptionalLong` for the primitive fields (Item 55 forbids `Optional`).
+- **Item 82** — `QueryJob` documents thread-safety at class level and lists guarded fields.
+- **Clean Code, small classes** — the biggest class in the package is `QueryJob`; its public API is four methods.
+- **Clean Code, intent-revealing names** — `submit` / `get` / `cancel`, `register` / `find` / `remove`. No `handleXxx`, no `processYyy`.
+- **Clean Code, no duplication** — sync and async share one `QueryJob`; only the transport layer differs.
+
+## 9. What migrates from the current PR
+
+The existing branch already ships the plumbing this design keeps. The redesign is refactor-shaped, not rewrite-shaped:
+
+| Today (`feat/query-job-refactor`) | MVP redesign |
+|-----------------------------------|--------------|
+| `plugin.transport.asyncquery.QueryJob` (983 lines) | `core.job.QueryJob` (~250 lines after retention drops out). |
+| `QueryJob.Snapshot` sealed hierarchy (`Running`, `Succeeded`, `Failed`) | `QueryJobStatus` record with a `state` discriminator. |
+| `QueryJobRegistry` | `QueryJobStore` interface + `InMemoryQueryJobStore`. |
+| `QueryJobOwner` | `Principal`. |
+| `AsyncQueryExecution` (in `core.executor`) | `QueryRunner` (in `core.job`). Same shape, engine-neutral name and location. |
+| `DefaultAsyncQueryExecution` (in `ppl`) | `PPLQueryRunner` in `ppl` implementing `QueryRunner`; symmetrical `SQLQueryRunner` in `sql`. |
+| `PPLAsyncQueryService` (mentioned in dai-chen's comment #3, not yet in-tree) | `LocalQueryJobService` — no retention, no leases, no admission for MVP. |
+| `wait_for_completion_timeout` state axis | Transport-level `completion().orTimeout(...)` in a later PR. |
+| `keep_alive` expiry timer | Deferred; lands as a `RetentionPolicy` composed into the service. |
+
+## 10. Open questions
+
+1. Do we want `QueryResult` to model streaming from day one (a `Publisher` instead of `List`) so partial results are additive later? MVP says no; two engines that already produce final results would need adaptation. Flagging for review.
+2. Should `Principal.current` live on `Principal` or in a `SecurityAdapter` interface? Today it reaches into `ThreadContext`; that couples `core.job` to OpenSearch. A `SecurityAdapter` SPI is cleaner and keeps `core` engine-neutral, but adds one more file. Leaning `SecurityAdapter`.
+3. Is one job store per node acceptable in MVP, or does the first release need cluster-visible metadata (backed by a system index)?
+
+Feedback welcome on any of the above.