Fix missing QueryCompletedEvent race condition

trinodb · May 10, 2024 · f891fdd · f891fdd
1 parent 4fd1b0c
commit f891fdd
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/dispatcher/LocalDispatchQuery.java b/core/trino-main/src/main/java/io/trino/dispatcher/LocalDispatchQuery.java
@@ -93,8 +93,13 @@ public LocalDispatchQuery(
 
         stateMachine.addStateChangeListener(state -> {
             if (state == QueryState.FAILED) {
+                // notificationSentOrGuaranteed.compareAndSet(false, true) ensures the queryCompletedEvent is only fired once.
+                // Either via an immediateFailureEvent or a finalQueryInfoListener.
+                // In cases when finalQueryInfoListener wins the race AND finalQueryInfo is not set the listener isn't triggered.
+                // getFullQueryInfo() force sets finalQueryInfo so the finalQueryInfoListener condition is met.
+                ExecutionFailureInfo failureInfo = getFullQueryInfo().getFailureInfo();
                 if (notificationSentOrGuaranteed.compareAndSet(false, true)) {
-                    queryMonitor.queryImmediateFailureEvent(getBasicQueryInfo(), getFullQueryInfo().getFailureInfo());
+                    queryMonitor.queryImmediateFailureEvent(getBasicQueryInfo(), failureInfo);
                 }
             }
             // any PLANNING or later state means the query has been submitted for execution

diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java b/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java
@@ -1501,6 +1501,19 @@ public void testAnonymizedJsonPlan()
                 .isEqualTo(Optional.of(ANONYMIZED_PLAN_JSON_CODEC.toJson(anonymizedPlan)));
     }
 
+    @Test
+    public void testAllImmediateFailureEventsPresent()
+            throws Exception
+    {
+        String immediatelyFailingQuery = "grant select on fake_catalog_%s.fake_schema.fake_table to fake_role";
+        String expectedFailure = "line 1:1: Table 'fake_catalog_%s.fake_schema.fake_table' does not exist";
+        int queryCount = 100;
+
+        for (int i = 0; i < queryCount; i++) {
+            assertFailedQuery(immediatelyFailingQuery.formatted(i), expectedFailure.formatted(i));
+        }
+    }
+
     private void assertLineage(String baseQuery, Set<String> inputTables, OutputColumnMetadata... outputColumnMetadata)
             throws Exception
     {