From af73e70e7fb232974f4447e2187aa32469337080 Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Thu, 24 Sep 2026 19:09:44 +0000 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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);
   }