Fix expiration time in ES|QL async

Author: dnhatn

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/async/AsyncTaskManagementService.java

@@ -177,7 +177,6 @@ public AsyncTaskManagementService(
     public void asyncExecute(
         Request request,
         TimeValue waitForCompletionTimeout,
-        TimeValue keepAlive,
         boolean keepOnCompletion,
         ActionListener<Response> listener
     ) {
@@ -190,7 +189,7 @@ public void asyncExecute(
                 operation.execute(
                     request,
                     searchTask,
-                    wrapStoringListener(searchTask, waitForCompletionTimeout, keepAlive, keepOnCompletion, listener)
+                    wrapStoringListener(searchTask, waitForCompletionTimeout, keepOnCompletion, listener)
                 );
                 operationStarted = true;
             } finally {
@@ -205,7 +204,6 @@ public void asyncExecute(
     private ActionListener<Response> wrapStoringListener(
         T searchTask,
         TimeValue waitForCompletionTimeout,
-        TimeValue keepAlive,
         boolean keepOnCompletion,
         ActionListener<Response> listener
     ) {
@@ -227,7 +225,7 @@ private ActionListener<Response> wrapStoringListener(
                 if (keepOnCompletion) {
                     storeResults(
                         searchTask,
-                        new StoredAsyncResponse<>(response, threadPool.absoluteTimeInMillis() + keepAlive.getMillis()),
+                        new StoredAsyncResponse<>(response, searchTask.getExpirationTimeMillis()),
                         ActionListener.running(() -> acquiredListener.onResponse(response))
                     );
                 } else {
@@ -239,7 +237,7 @@ private ActionListener<Response> wrapStoringListener(
                 // We finished after timeout - saving results
                 storeResults(
                     searchTask,
-                    new StoredAsyncResponse<>(response, threadPool.absoluteTimeInMillis() + keepAlive.getMillis()),
+                    new StoredAsyncResponse<>(response, searchTask.getExpirationTimeMillis()),
                     ActionListener.running(response::decRef)
                 );
             }
@@ -251,7 +249,7 @@ private ActionListener<Response> wrapStoringListener(
                 if (keepOnCompletion) {
                     storeResults(
                         searchTask,
-                        new StoredAsyncResponse<>(e, threadPool.absoluteTimeInMillis() + keepAlive.getMillis()),
+                        new StoredAsyncResponse<>(e, searchTask.getExpirationTimeMillis()),
                         ActionListener.running(() -> acquiredListener.onFailure(e))
                     );
                 } else {
@@ -261,7 +259,7 @@ private ActionListener<Response> wrapStoringListener(
                 }
             } else {
                 // We finished after timeout - saving exception
-                storeResults(searchTask, new StoredAsyncResponse<>(e, threadPool.absoluteTimeInMillis() + keepAlive.getMillis()));
+                storeResults(searchTask, new StoredAsyncResponse<>(e, searchTask.getExpirationTimeMillis()));
             }
         });
     }

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AsyncEsqlQueryActionIT.java

@@ -312,6 +312,7 @@ public void testUpdateKeepAlive() throws Exception {
                 assertThat(resp.isRunning(), is(false));
             }
         });
+        assertThat(getExpirationFromDoc(asyncId), greaterThanOrEqualTo(nowInMillis + keepAlive.getMillis()));
         // update the keepAlive after the query has completed
         int iters = between(1, 5);
         for (int i = 0; i < iters; i++) {