Align transitionToFailed and transitionToCanceled in QueryStateMachine

Before commit 8a89c29, the only
difference between `transitionToCanceled` and `transitionToFailed`
seemed to be logging the failure thru QUERY_STATE_LOG.

After that commit, one of those methods sets `finalQueryInfo` to ensure
it's set. Unify the handling to ensure `finalQueryInfo` is set for
`transitionToCanceled` as well.

core/trino-main/src/main/java/io/trino/execution/QueryStateMachine.java

@@ -1109,7 +1109,17 @@ private void transitionToFinishedIfReady()
         queryState.setIf(FINISHED, currentState -> !currentState.isDone());
     }
 
+    public boolean transitionToCanceled()
+    {
+        return transitionToFailed(new TrinoException(USER_CANCELED, "Query was canceled"), false);
+    }
+
     public boolean transitionToFailed(Throwable throwable)
+    {
+        return transitionToFailed(throwable, true);
+    }
+
+    private boolean transitionToFailed(Throwable throwable, boolean log)
     {
         queryStateTimer.endQuery();
 
@@ -1123,12 +1133,16 @@ public boolean transitionToFailed(Throwable throwable)
 
         QueryState oldState = queryState.trySet(FAILED);
         if (oldState.isDone()) {
-            QUERY_STATE_LOG.debug(throwable, "Failure after query %s finished", queryId);
+            if (log) {
+                QUERY_STATE_LOG.debug(throwable, "Failure after query %s finished", queryId);
+            }
             return false;
         }
 
         try {
-            QUERY_STATE_LOG.debug(throwable, "Query %s failed", queryId);
+            if (log) {
+                QUERY_STATE_LOG.debug(throwable, "Query %s failed", queryId);
+            }
             session.getTransactionId().flatMap(transactionManager::getTransactionInfoIfExist).ifPresent(transaction -> {
                 try {
                     if (transaction.isAutoCommitContext()) {
@@ -1140,7 +1154,9 @@ public boolean transitionToFailed(Throwable throwable)
                 }
                 catch (RuntimeException e) {
                     // This shouldn't happen but be safe and just fail the transaction directly
-                    QUERY_STATE_LOG.error(e, "Error aborting transaction for failed query. Transaction will be failed directly");
+                    if (log) {
+                        QUERY_STATE_LOG.error(e, "Error aborting transaction for failed query. Transaction will be failed directly");
+                    }
                 }
             });
         }
@@ -1154,32 +1170,6 @@ public boolean transitionToFailed(Throwable throwable)
         return true;
     }
 
-    public boolean transitionToCanceled()
-    {
-        queryStateTimer.endQuery();
-
-        // NOTE: The failure cause must be set before triggering the state change, so
-        // listeners can observe the exception. This is safe because the failure cause
-        // can only be observed if the transition to FAILED is successful.
-        failureCause.compareAndSet(null, toFailure(new TrinoException(USER_CANCELED, "Query was canceled")));
-
-        cleanupQueryQuietly();
-
-        boolean canceled = queryState.setIf(FAILED, currentState -> !currentState.isDone());
-        if (canceled) {
-            session.getTransactionId().flatMap(transactionManager::getTransactionInfoIfExist).ifPresent(transaction -> {
-                if (transaction.isAutoCommitContext()) {
-                    transactionManager.asyncAbort(transaction.getTransactionId());
-                }
-                else {
-                    transactionManager.fail(transaction.getTransactionId());
-                }
-            });
-        }
-
-        return canceled;
-    }
-
     private void cleanupQuery()
     {
         // only execute cleanup once

core/trino-main/src/test/java/io/trino/execution/TestQueryStateMachine.java

@@ -482,7 +482,11 @@ public void cleanupQuery(Session session)
         cancellingThread.get(10, SECONDS);
         anotherThread.get(10, SECONDS);
 
-        // TODO queryStateMachine.getFinalQueryInfo() does not exist for cancelled queries, but may be created by anotherThread due to a race
+        ExecutionFailureInfo failureInfo = queryStateMachine.getFinalQueryInfo().orElseThrow().getFailureInfo();
+        assertThat(failureInfo).isNotNull();
+        assertThat(failureInfo.getErrorCode()).isEqualTo(USER_CANCELED.toErrorCode());
+        assertThat(failureInfo.getMessage()).isEqualTo("Query was canceled");
+
         BasicQueryInfo basicQueryInfo = queryStateMachine.getBasicQueryInfo(Optional.empty());
         assertThat(basicQueryInfo.getErrorCode()).isEqualTo(USER_CANCELED.toErrorCode());
     }