Coordinator-driven graceful decommission/recommission of worker nodes

core/trino-main/src/main/java/io/trino/metadata/AllNodes.java

@@ -23,23 +23,49 @@
 public class AllNodes
 {
     private final Set<InternalNode> activeNodes;
+    private final Set<InternalNode> decommissionedNodes;
+    private final Set<InternalNode> decommissioningNodes;
     private final Set<InternalNode> inactiveNodes;
     private final Set<InternalNode> shuttingDownNodes;
     private final Set<InternalNode> activeCoordinators;
+    private final Set<InternalNode> aliveNodes;
 
-    public AllNodes(Set<InternalNode> activeNodes, Set<InternalNode> inactiveNodes, Set<InternalNode> shuttingDownNodes, Set<InternalNode> activeCoordinators)
+    public AllNodes(Set<InternalNode> activeNodes,
+            Set<InternalNode> decommissionedNodes,
+            Set<InternalNode> decommissioningNodes,
+            Set<InternalNode> inactiveNodes,
+            Set<InternalNode> shuttingDownNodes,
+            Set<InternalNode> activeCoordinators)
     {
         this.activeNodes = ImmutableSet.copyOf(requireNonNull(activeNodes, "activeNodes is null"));
+        this.decommissionedNodes = ImmutableSet.copyOf(requireNonNull(decommissionedNodes, "decommissionedNodes is null"));
+        this.decommissioningNodes = ImmutableSet.copyOf(requireNonNull(decommissioningNodes, "decommissioningNodes is null"));
         this.inactiveNodes = ImmutableSet.copyOf(requireNonNull(inactiveNodes, "inactiveNodes is null"));
         this.shuttingDownNodes = ImmutableSet.copyOf(requireNonNull(shuttingDownNodes, "shuttingDownNodes is null"));
         this.activeCoordinators = ImmutableSet.copyOf(requireNonNull(activeCoordinators, "activeCoordinators is null"));
+        this.aliveNodes = ImmutableSet.<InternalNode>builder()
+                .addAll(activeNodes)
+                .addAll(decommissionedNodes)
+                .addAll(decommissioningNodes)
+                .addAll(shuttingDownNodes)
+                .build();
     }
 
     public Set<InternalNode> getActiveNodes()
     {
         return activeNodes;
     }
 
+    public Set<InternalNode> getDecommissionedNodes()
+    {
+        return decommissionedNodes;
+    }
+
+    public Set<InternalNode> getDecommissioningNodes()
+    {
+        return decommissioningNodes;
+    }
+
     public Set<InternalNode> getInactiveNodes()
     {
         return inactiveNodes;
@@ -50,6 +76,11 @@ public Set<InternalNode> getShuttingDownNodes()
         return shuttingDownNodes;
     }
 
+    public Set<InternalNode> getAliveNodes()
+    {
+        return aliveNodes;
+    }
+
     public Set<InternalNode> getActiveCoordinators()
     {
         return activeCoordinators;
@@ -66,6 +97,8 @@ public boolean equals(Object o)
         }
         AllNodes allNodes = (AllNodes) o;
         return Objects.equals(activeNodes, allNodes.activeNodes) &&
+                Objects.equals(decommissionedNodes, allNodes.decommissionedNodes) &&
+                Objects.equals(decommissioningNodes, allNodes.decommissioningNodes) &&
                 Objects.equals(inactiveNodes, allNodes.inactiveNodes) &&
                 Objects.equals(shuttingDownNodes, allNodes.shuttingDownNodes) &&
                 Objects.equals(activeCoordinators, allNodes.activeCoordinators);
@@ -74,6 +107,12 @@ public boolean equals(Object o)
     @Override
     public int hashCode()
     {
-        return Objects.hash(activeNodes, inactiveNodes, shuttingDownNodes, activeCoordinators);
+        return Objects.hash(
+                activeNodes,
+                inactiveNodes,
+                shuttingDownNodes,
+                decommissioningNodes,
+                decommissionedNodes,
+                activeCoordinators);
     }
 }

core/trino-main/src/main/java/io/trino/metadata/DiscoveryNodeManager.java

@@ -18,7 +18,6 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.SetMultimap;
-import com.google.common.collect.Sets;
 import com.google.common.collect.Sets.SetView;
 import io.airlift.discovery.client.ServiceDescriptor;
 import io.airlift.discovery.client.ServiceSelector;
@@ -43,6 +42,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -59,6 +59,8 @@
 import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
 import static io.trino.connector.system.GlobalSystemConnector.CATALOG_HANDLE;
 import static io.trino.metadata.NodeState.ACTIVE;
+import static io.trino.metadata.NodeState.DECOMMISSIONED;
+import static io.trino.metadata.NodeState.DECOMMISSIONING;
 import static io.trino.metadata.NodeState.INACTIVE;
 import static io.trino.metadata.NodeState.SHUTTING_DOWN;
 import static java.util.Locale.ENGLISH;
@@ -93,6 +95,9 @@ public final class DiscoveryNodeManager
     @GuardedBy("this")
     private Set<InternalNode> coordinators;
 
+    @GuardedBy("this")
+    private Set<String> nodesToBeDecommissioned = new HashSet<>();
+
     @GuardedBy("this")
     private final List<Consumer<AllNodes>> listeners = new ArrayList<>();
 
@@ -169,10 +174,7 @@ public void destroy()
     private void pollWorkers()
     {
         AllNodes allNodes = getAllNodes();
-        Set<InternalNode> aliveNodes = ImmutableSet.<InternalNode>builder()
-                .addAll(allNodes.getActiveNodes())
-                .addAll(allNodes.getShuttingDownNodes())
-                .build();
+        Set<InternalNode> aliveNodes = allNodes.getAliveNodes();
 
         ImmutableSet<String> aliveNodeIds = aliveNodes.stream()
                 .map(InternalNode::getNodeIdentifier)
@@ -215,9 +217,7 @@ private synchronized void refreshNodesInternal()
                 .filter(service -> !failureDetector.getFailed().contains(service))
                 .collect(toImmutableSet());
 
-        ImmutableSet.Builder<InternalNode> activeNodesBuilder = ImmutableSet.builder();
-        ImmutableSet.Builder<InternalNode> inactiveNodesBuilder = ImmutableSet.builder();
-        ImmutableSet.Builder<InternalNode> shuttingDownNodesBuilder = ImmutableSet.builder();
+        ImmutableSetMultimap.Builder<NodeState, InternalNode> nodeStateMapBuilder = ImmutableSetMultimap.builder();
         ImmutableSet.Builder<InternalNode> coordinatorsBuilder = ImmutableSet.builder();
         ImmutableSetMultimap.Builder<CatalogHandle, InternalNode> byCatalogHandleBuilder = ImmutableSetMultimap.builder();
 
@@ -229,55 +229,65 @@ private synchronized void refreshNodesInternal()
                 InternalNode node = new InternalNode(service.getNodeId(), uri, nodeVersion, coordinator);
                 NodeState nodeState = getNodeState(node);
 
-                switch (nodeState) {
-                    case ACTIVE:
-                        activeNodesBuilder.add(node);
-                        if (coordinator) {
-                            coordinatorsBuilder.add(node);
-                        }
+                // nodesToBeDecommissioned is the authoritative list of node to be decommissioned
+                // from coordinator perspective. Once a worker appears in the list,
+                // its state become DECOMMISSIONING even if worker has yet confirmed such
+                // so that no new tasks will be scheduled on it.
+                if (!coordinator && nodesToBeDecommissioned.contains(node.getNodeIdentifier())
+                        && nodeState == ACTIVE) {
+                    log.debug("Treat " + node.getNodeIdentifier() + " as DECOMMISSIONING");
+                    nodeState = DECOMMISSIONING;
+                }
 
-                        // record available active nodes organized by catalog handle
-                        String catalogHandleIds = service.getProperties().get("catalogHandleIds");
-                        if (catalogHandleIds != null) {
-                            catalogHandleIds = catalogHandleIds.toLowerCase(ENGLISH);
-                            for (String catalogHandleId : CATALOG_HANDLE_ID_SPLITTER.split(catalogHandleIds)) {
-                                byCatalogHandleBuilder.put(CatalogHandle.fromId(catalogHandleId), node);
-                            }
+                // Add node to node state map.
+                nodeStateMapBuilder.put(nodeState, node);
+
+                if (nodeState == ACTIVE) {
+                    if (coordinator) {
+                        coordinatorsBuilder.add(node);
+                    }
+
+                    // record available active nodes organized by catalog handle
+                    String catalogHandleIds = service.getProperties().get("catalogHandleIds");
+                    if (catalogHandleIds != null) {
+                        catalogHandleIds = catalogHandleIds.toLowerCase(ENGLISH);
+                        for (String catalogHandleId : CATALOG_HANDLE_ID_SPLITTER.split(catalogHandleIds)) {
+                            byCatalogHandleBuilder.put(CatalogHandle.fromId(catalogHandleId), node);
                         }
+                    }
 
-                        // always add system connector
-                        byCatalogHandleBuilder.put(CATALOG_HANDLE, node);
-                        break;
-                    case INACTIVE:
-                        inactiveNodesBuilder.add(node);
-                        break;
-                    case SHUTTING_DOWN:
-                        shuttingDownNodesBuilder.add(node);
-                        break;
-                    default:
-                        log.error("Unknown state %s for node %s", nodeState, node);
+                    // always add system connector
+                    byCatalogHandleBuilder.put(CATALOG_HANDLE, node);
                 }
             }
         }
 
-        if (allNodes != null) {
+        // nodes by catalog handle changes anytime a node adds or removes a catalog (note: this is not part of the listener system)
+        if (!allCatalogsOnAllNodes) {
+            activeNodesByCatalogHandle = Optional.of(byCatalogHandleBuilder.build());
+        }
+
+        SetMultimap<NodeState, InternalNode> nodeStateMap = nodeStateMapBuilder.build();
+        AllNodes currAllNodes = new AllNodes(
+                nodeStateMap.get(ACTIVE),
+                nodeStateMap.get(DECOMMISSIONED),
+                nodeStateMap.get(DECOMMISSIONING),
+                nodeStateMap.get(INACTIVE),
+                nodeStateMap.get(SHUTTING_DOWN),
+                coordinatorsBuilder.build());
+
+        if (this.allNodes != null) {
             // log node that are no longer active (but not shutting down)
-            SetView<InternalNode> missingNodes = difference(allNodes.getActiveNodes(), Sets.union(activeNodesBuilder.build(), shuttingDownNodesBuilder.build()));
+            SetView<InternalNode> missingNodes = difference(allNodes.getActiveNodes(), currAllNodes.getAliveNodes());
             for (InternalNode missingNode : missingNodes) {
                 log.info("Previously active node is missing: %s (last seen at %s)", missingNode.getNodeIdentifier(), missingNode.getHost());
             }
         }
 
-        // nodes by catalog handle changes anytime a node adds or removes a catalog (note: this is not part of the listener system)
-        if (!allCatalogsOnAllNodes) {
-            activeNodesByCatalogHandle = Optional.of(byCatalogHandleBuilder.build());
-        }
-
-        AllNodes allNodes = new AllNodes(activeNodesBuilder.build(), inactiveNodesBuilder.build(), shuttingDownNodesBuilder.build(), coordinatorsBuilder.build());
         // only update if all nodes actually changed (note: this does not include the connectors registered with the nodes)
-        if (!allNodes.equals(this.allNodes)) {
+        if (!currAllNodes.equals(this.allNodes)) {
             // assign allNodes to a local variable for use in the callback below
-            this.allNodes = allNodes;
+            this.allNodes = currAllNodes;
             coordinators = coordinatorsBuilder.build();
 
             // notify listeners
@@ -289,22 +299,19 @@ private synchronized void refreshNodesInternal()
     private NodeState getNodeState(InternalNode node)
     {
         if (expectedNodeVersion.equals(node.getNodeVersion())) {
-            if (isNodeShuttingDown(node.getNodeIdentifier())) {
-                return SHUTTING_DOWN;
+            String nodeId = node.getNodeIdentifier();
+            Optional<NodeState> remoteNodeState = nodeStates.containsKey(nodeId)
+                    ? nodeStates.get(nodeId).getNodeState()
+                    : Optional.empty();
+            if (remoteNodeState.isPresent()) {
+                return remoteNodeState.get();
             }
+            // no remote node state
             return ACTIVE;
         }
         return INACTIVE;
     }
 
-    private boolean isNodeShuttingDown(String nodeId)
-    {
-        Optional<NodeState> remoteNodeState = nodeStates.containsKey(nodeId)
-                ? nodeStates.get(nodeId).getNodeState()
-                : Optional.empty();
-        return remoteNodeState.isPresent() && remoteNodeState.get() == SHUTTING_DOWN;
-    }
-
     @Override
     public synchronized AllNodes getAllNodes()
     {
@@ -317,6 +324,18 @@ public int getActiveNodeCount()
         return getAllNodes().getActiveNodes().size();
     }
 
+    @Managed
+    public int getDecommissionedNodeCount()
+    {
+        return getAllNodes().getDecommissionedNodes().size();
+    }
+
+    @Managed
+    public int getDecommissioningNodeCount()
+    {
+        return getAllNodes().getDecommissioningNodes().size();
+    }
+
     @Managed
     public int getInactiveNodeCount()
     {
@@ -335,6 +354,10 @@ public Set<InternalNode> getNodes(NodeState state)
         switch (state) {
             case ACTIVE:
                 return getAllNodes().getActiveNodes();
+            case DECOMMISSIONED:
+                return getAllNodes().getDecommissionedNodes();
+            case DECOMMISSIONING:
+                return getAllNodes().getDecommissioningNodes();
             case INACTIVE:
                 return getAllNodes().getInactiveNodes();
             case SHUTTING_DOWN:
@@ -407,4 +430,9 @@ private static boolean isCoordinator(ServiceDescriptor service)
     {
         return Boolean.parseBoolean(service.getProperties().get("coordinator"));
     }
+
+    public synchronized void setNodesToExclude(Set<String> nodesToExclude)
+    {
+        this.nodesToBeDecommissioned = nodesToExclude;
+    }
 }

core/trino-main/src/main/java/io/trino/metadata/InMemoryNodeManager.java

@@ -59,7 +59,11 @@ public Set<InternalNode> getNodes(NodeState state)
     {
         switch (state) {
             case ACTIVE:
-                return allNodes;
+                return getAllNodes().getActiveNodes();
+            case DECOMMISSIONED:
+                return getAllNodes().getDecommissionedNodes();
+            case DECOMMISSIONING:
+                return getAllNodes().getDecommissioningNodes();
             case INACTIVE:
             case SHUTTING_DOWN:
                 return ImmutableSet.of();
@@ -86,6 +90,8 @@ public AllNodes getAllNodes()
                 allNodes,
                 ImmutableSet.of(),
                 ImmutableSet.of(),
+                ImmutableSet.of(),
+                ImmutableSet.of(),
                 ImmutableSet.of(CURRENT_NODE));
     }
 

core/trino-main/src/main/java/io/trino/metadata/NodeState.java

@@ -16,6 +16,8 @@
 public enum NodeState
 {
     ACTIVE,
+    DECOMMISSIONED,
+    DECOMMISSIONING,
     INACTIVE,
     SHUTTING_DOWN
 }

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

@@ -341,6 +341,15 @@ protected void setup(Binder binder)
 
         install(new QueryExecutionFactoryModule());
 
+        // nodes and queries for monitoring and auto-scaling
+        jaxrsBinder(binder).bind(NodesResource.class);
+        httpClientBinder(binder).bindHttpClient("nodes", ForNodes.class)
+                .withTracing()
+                .withConfigDefaults(config -> {
+                    config.setIdleTimeout(new Duration(30, SECONDS));
+                    config.setRequestTimeout(new Duration(10, SECONDS));
+                });
+
         // cleanup
         binder.bind(ExecutorCleanup.class).asEagerSingleton();
     }

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

@@ -0,0 +1,31 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.server;
+
+import javax.inject.Qualifier;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@Qualifier
+public @interface ForNodes
+{
+}