Fix graceful shutdown

Graceful shutdown logic compared previous nodestate instead of new SHUTTING_DOWN
nodestate and in effect it never finished the waitActiveTasksToFinish.
This behavior has been recreated and is covered by the new test.

Node State is now correctly assigned before starting the shutdown sequence.

Author: brybacki

core/trino-main/src/main/java/io/trino/server/NodeStateManager.java

@@ -216,15 +216,22 @@ public synchronized void transitionState(NodeState state)
                 }
             }
             case SHUTTING_DOWN -> {
-                if (currState.state() == DRAINED && nodeState.compareAndSet(currState, currState.toShuttingDown())) {
+                if (isCoordinator) {
+                    throw new UnsupportedOperationException("Cannot shutdown coordinator");
+                }
+                VersionedState shuttingDown = currState.toShuttingDown();
+                if (currState.state() == DRAINED && nodeState.compareAndSet(currState, shuttingDown)) {
                     requestTerminate();
                     return;
                 }
+                nodeState.set(shuttingDown);
                 requestGracefulShutdown();
-                nodeState.set(currState.toShuttingDown());
                 return;
             }
             case DRAINING -> {
+                if (isCoordinator) {
+                    throw new UnsupportedOperationException("Cannot drain coordinator");
+                }
                 if (currState.state() == ACTIVE && nodeState.compareAndSet(currState, currState.toDraining())) {
                     requestDrain();
                     return;
@@ -246,9 +253,6 @@ private long nextStateVersion()
     private synchronized void requestDrain()
     {
         log.debug("Drain requested, NodeState: %s", getServerState());
-        if (isCoordinator) {
-            throw new UnsupportedOperationException("Cannot drain coordinator");
-        }
 
         // wait for a grace period (so that draining state is observed by the coordinator) before starting draining
         // when coordinator observes draining no new tasks are assigned to this worker
@@ -259,22 +263,16 @@ private synchronized void requestDrain()
     private void requestTerminate()
     {
         log.info("Immediate Shutdown requested");
-        if (isCoordinator) {
-            throw new UnsupportedOperationException("Cannot shutdown coordinator");
-        }
 
         shutdownHandler.schedule(this::terminate, 0, MILLISECONDS);
     }
 
     private void requestGracefulShutdown()
     {
         log.info("Shutdown requested");
-        if (isCoordinator) {
-            throw new UnsupportedOperationException("Cannot shutdown coordinator");
-        }
 
-        // wait for a grace period (so that shutting down state is observed by the coordinator) to start the shutdown sequence
         VersionedState expectedState = nodeState.get();
+        // wait for a grace period (so that shutting down state is observed by the coordinator) to start the shutdown sequence
         shutdownHandler.schedule(() -> shutdown(expectedState), gracePeriod.toMillis(), MILLISECONDS);
     }
 

core/trino-main/src/test/java/io/trino/server/TestNodeStateManager.java

@@ -50,6 +50,7 @@
 import static io.trino.metadata.NodeState.ACTIVE;
 import static io.trino.metadata.NodeState.DRAINED;
 import static io.trino.metadata.NodeState.DRAINING;
+import static io.trino.metadata.NodeState.SHUTTING_DOWN;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -148,7 +149,41 @@ void testImmediateTransitionToShuttingDownWhenDrained()
     }
 
     @Test
-    void testWaitActiveTasksToFinish()
+    void testWaitActiveTasksToFinishDuringShutdown()
+            throws URISyntaxException
+    {
+        List<TaskInfo> taskInfos = new ArrayList<>();
+        TaskInfo task = TaskInfo.createInitialTask(
+                new TaskId(new StageId("query1", 1), 1, 1),
+                new URI(""),
+                "1",
+                false,
+                Optional.empty(),
+                new TaskStats(DateTime.now(), null));
+        taskInfos.add(task);
+        tasks.set(taskInfos);
+
+        // Draining - will wait for tasks to finish
+        nodeStateManager.transitionState(SHUTTING_DOWN);
+        assertThat(nodeStateManager.getServerState()).isEqualTo(SHUTTING_DOWN);
+
+        // make sure that nodeStateManager registered a listener for tasks to finish
+        ticker.increment(1, SECONDS);
+        executor.run();
+        await().atMost(1, SECONDS).until(() -> sqlTasksObservable.getTasks().size() == 1);
+
+        // simulate task completion after some time
+        tasks.set(Collections.emptyList());
+        sqlTasksObservable.getTasks().get(task.taskStatus().getTaskId())
+                .stateChanged(TaskState.FINISHED);
+
+        // when NodeStateManager sees task finished - it will drain after another drain period
+        await().atMost(1, SECONDS)
+                .untilAsserted(() -> assertThat(nodeStateManager.getServerState()).isEqualTo(SHUTTING_DOWN));
+    }
+
+    @Test
+    void testWaitActiveTasksToFinishDuringDraining()
             throws URISyntaxException
     {
         List<TaskInfo> taskInfos = new ArrayList<>();