[FLINK-28145][runtime] Let ResourceManager support blocklist

This closes #20218.

Author: wanglijie95

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManager.java

@@ -24,6 +24,9 @@
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.blob.TransientBlobKey;
+import org.apache.flink.runtime.blocklist.BlockedNode;
+import org.apache.flink.runtime.blocklist.BlocklistContext;
+import org.apache.flink.runtime.blocklist.BlocklistHandler;
 import org.apache.flink.runtime.clusterframework.ApplicationStatus;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
@@ -96,6 +99,7 @@
 import java.util.stream.Collectors;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
 
 /**
  * ResourceManager implementation. The resource manager is responsible for resource de-/allocation
@@ -157,6 +161,8 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 
     private final DelegationTokenManager delegationTokenManager;
 
+    protected final BlocklistHandler blocklistHandler;
+
     public ResourceManager(
             RpcService rpcService,
             UUID leaderSessionId,
@@ -165,6 +171,7 @@ public ResourceManager(
             DelegationTokenManager delegationTokenManager,
             SlotManager slotManager,
             ResourceManagerPartitionTrackerFactory clusterPartitionTrackerFactory,
+            BlocklistHandler.Factory blocklistHandlerFactory,
             JobLeaderIdService jobLeaderIdService,
             ClusterInformation clusterInformation,
             FatalErrorHandler fatalErrorHandler,
@@ -189,6 +196,12 @@ public ResourceManager(
         this.jmResourceIdRegistrations = new HashMap<>(4);
         this.taskExecutors = new HashMap<>(8);
         this.taskExecutorGatewayFutures = new HashMap<>(8);
+        this.blocklistHandler =
+                blocklistHandlerFactory.create(
+                        new ResourceManagerBlocklistContext(),
+                        this::getNodeIdOfTaskManager,
+                        getMainThreadExecutor(),
+                        log);
 
         this.jobManagerHeartbeatManager = NoOpHeartbeatManager.getInstance();
         this.taskManagerHeartbeatManager = NoOpHeartbeatManager.getInstance();
@@ -246,7 +259,10 @@ private void startResourceManagerServices() throws Exception {
             startHeartbeatServices();
 
             slotManager.start(
-                    getFencingToken(), getMainThreadExecutor(), new ResourceActionsImpl());
+                    getFencingToken(),
+                    getMainThreadExecutor(),
+                    new ResourceActionsImpl(),
+                    blocklistHandler::isBlockedTaskManager);
 
             delegationTokenManager.start();
 
@@ -833,10 +849,21 @@ public CompletableFuture<TaskExecutorThreadInfoGateway> requestTaskExecutorThrea
         }
     }
 
+    @Override
+    public CompletableFuture<Acknowledge> notifyNewBlockedNodes(Collection<BlockedNode> newNodes) {
+        blocklistHandler.addNewBlockedNodes(newNodes);
+        return CompletableFuture.completedFuture(Acknowledge.get());
+    }
+
     // ------------------------------------------------------------------------
     //  Internal methods
     // ------------------------------------------------------------------------
 
+    private String getNodeIdOfTaskManager(ResourceID taskManagerId) {
+        checkState(taskExecutors.containsKey(taskManagerId));
+        return taskExecutors.get(taskManagerId).getNodeId();
+    }
+
     /**
      * Registers a new JobMaster.
      *
@@ -1444,6 +1471,18 @@ public Void retrievePayload(ResourceID resourceID) {
         }
     }
 
+    private class ResourceManagerBlocklistContext implements BlocklistContext {
+        @Override
+        public void blockResources(Collection<BlockedNode> blockedNodes) {}
+
+        @Override
+        public void unblockResources(Collection<BlockedNode> unBlockedNodes) {
+            // when a node is unblocked, we should trigger the resource requirements because the
+            // slots on this node become available again.
+            slotManager.triggerResourceRequirementsCheck();
+        }
+    }
+
     // ------------------------------------------------------------------------
     //  Resource Management
     // ------------------------------------------------------------------------

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManagerGateway.java

@@ -24,6 +24,7 @@
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.blob.BlobServer;
 import org.apache.flink.runtime.blob.TransientBlobKey;
+import org.apache.flink.runtime.blocklist.BlocklistListener;
 import org.apache.flink.runtime.clusterframework.ApplicationStatus;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
@@ -55,7 +56,7 @@
 
 /** The {@link ResourceManager}'s RPC gateway interface. */
 public interface ResourceManagerGateway
-        extends FencedRpcGateway<ResourceManagerId>, ClusterPartitionManager {
+        extends FencedRpcGateway<ResourceManagerId>, ClusterPartitionManager, BlocklistListener {
 
     /**
      * Register a {@link JobMaster} at the resource manager.

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/StandaloneResourceManager.java

@@ -19,6 +19,7 @@
 package org.apache.flink.runtime.resourcemanager;
 
 import org.apache.flink.api.common.time.Time;
+import org.apache.flink.runtime.blocklist.BlocklistHandler;
 import org.apache.flink.runtime.clusterframework.ApplicationStatus;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.entrypoint.ClusterInformation;
@@ -57,6 +58,7 @@ public StandaloneResourceManager(
             DelegationTokenManager delegationTokenManager,
             SlotManager slotManager,
             ResourceManagerPartitionTrackerFactory clusterPartitionTrackerFactory,
+            BlocklistHandler.Factory blocklistHandlerFactory,
             JobLeaderIdService jobLeaderIdService,
             ClusterInformation clusterInformation,
             FatalErrorHandler fatalErrorHandler,
@@ -72,6 +74,7 @@ public StandaloneResourceManager(
                 delegationTokenManager,
                 slotManager,
                 clusterPartitionTrackerFactory,
+                blocklistHandlerFactory,
                 jobLeaderIdService,
                 clusterInformation,
                 fatalErrorHandler,

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/StandaloneResourceManagerFactory.java

@@ -24,6 +24,7 @@
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.ConfigurationUtils;
 import org.apache.flink.configuration.ResourceManagerOptions;
+import org.apache.flink.runtime.blocklist.BlocklistUtils;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.entrypoint.ClusterInformation;
 import org.apache.flink.runtime.heartbeat.HeartbeatServices;
@@ -82,6 +83,7 @@ protected ResourceManager<ResourceID> createResourceManager(
                 delegationTokenManager,
                 resourceManagerRuntimeServices.getSlotManager(),
                 ResourceManagerPartitionTrackerImpl::new,
+                BlocklistUtils.loadBlocklistHandlerFactory(configuration),
                 resourceManagerRuntimeServices.getJobLeaderIdService(),
                 clusterInformation,
                 fatalErrorHandler,

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/active/ActiveResourceManager.java

@@ -22,6 +22,7 @@
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.configuration.AkkaOptions;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.blocklist.BlocklistHandler;
 import org.apache.flink.runtime.clusterframework.ApplicationStatus;
 import org.apache.flink.runtime.clusterframework.TaskExecutorProcessSpec;
 import org.apache.flink.runtime.clusterframework.TaskExecutorProcessUtils;
@@ -50,7 +51,6 @@
 
 import java.time.Duration;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -113,6 +113,7 @@ public ActiveResourceManager(
             DelegationTokenManager delegationTokenManager,
             SlotManager slotManager,
             ResourceManagerPartitionTrackerFactory clusterPartitionTrackerFactory,
+            BlocklistHandler.Factory blocklistHandlerFactory,
             JobLeaderIdService jobLeaderIdService,
             ClusterInformation clusterInformation,
             FatalErrorHandler fatalErrorHandler,
@@ -129,6 +130,7 @@ public ActiveResourceManager(
                 delegationTokenManager,
                 slotManager,
                 clusterPartitionTrackerFactory,
+                blocklistHandlerFactory,
                 jobLeaderIdService,
                 clusterInformation,
                 fatalErrorHandler,
@@ -157,12 +159,11 @@ public ActiveResourceManager(
     @Override
     protected void initialize() throws ResourceManagerException {
         try {
-            // TODO: pass the blocked node retriever
             resourceManagerDriver.initialize(
                     this,
                     new GatewayMainThreadExecutor(),
                     ioExecutor,
-                    () -> Collections.emptySet());
+                    blocklistHandler::getAllBlockedNodeIds);
         } catch (Exception e) {
             throw new ResourceManagerException("Cannot initialize resource provider.", e);
         }

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/active/ActiveResourceManagerFactory.java

@@ -25,6 +25,7 @@
 import org.apache.flink.configuration.MemorySize;
 import org.apache.flink.configuration.ResourceManagerOptions;
 import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.runtime.blocklist.BlocklistUtils;
 import org.apache.flink.runtime.clusterframework.TaskExecutorProcessUtils;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.clusterframework.types.ResourceIDRetrievable;
@@ -119,6 +120,7 @@ public ResourceManager<WorkerType> createResourceManager(
                 delegationTokenManager,
                 resourceManagerRuntimeServices.getSlotManager(),
                 ResourceManagerPartitionTrackerImpl::new,
+                BlocklistUtils.loadBlocklistHandlerFactory(configuration),
                 resourceManagerRuntimeServices.getJobLeaderIdService(),
                 clusterInformation,
                 fatalErrorHandler,

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/slotmanager/DeclarativeSlotManager.java

@@ -185,6 +185,11 @@ public void setFailUnfulfillableRequest(boolean failUnfulfillableRequest) {
         }
     }
 
+    @Override
+    public void triggerResourceRequirementsCheck() {
+        checkResourceRequirementsWithDelay();
+    }
+
     // ---------------------------------------------------------------------------------------------
     // Component lifecycle methods
     // ---------------------------------------------------------------------------------------------

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/slotmanager/FineGrainedSlotManager.java

@@ -172,6 +172,11 @@ public void setFailUnfulfillableRequest(boolean failUnfulfillableRequest) {
         }
     }
 
+    @Override
+    public void triggerResourceRequirementsCheck() {
+        checkResourceRequirementsWithDelay();
+    }
+
     // ---------------------------------------------------------------------------------------------
     // Component lifecycle methods
     // ---------------------------------------------------------------------------------------------

flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/slotmanager/SlotManager.java

@@ -152,4 +152,10 @@ boolean registerTaskManager(
     void freeSlot(SlotID slotId, AllocationID allocationId);
 
     void setFailUnfulfillableRequest(boolean failUnfulfillableRequest);
+
+    /**
+     * Trigger the resource requirement check. This method will be called when some slot statuses
+     * changed.
+     */
+    void triggerResourceRequirementsCheck();
 }

flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTest.java

@@ -22,6 +22,10 @@
 import org.apache.flink.api.common.JobStatus;
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.core.testutils.OneShotLatch;
+import org.apache.flink.runtime.blocklist.BlockedNode;
+import org.apache.flink.runtime.blocklist.BlocklistHandler;
+import org.apache.flink.runtime.blocklist.DefaultBlocklistHandler;
+import org.apache.flink.runtime.blocklist.NoOpBlocklistHandler;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
 import org.apache.flink.runtime.heartbeat.HeartbeatServices;
@@ -62,6 +66,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -511,6 +516,33 @@ void testDisconnectTaskManager() throws Exception {
         assertThat(stopWorkerFuture.get()).isEqualTo(taskExecutorId);
     }
 
+    @Test
+    void testUnblockResourcesWillTriggerResourceRequirementsCheck() throws Exception {
+
+        final CompletableFuture<Void> triggerRequirementsCheckFuture = new CompletableFuture<>();
+
+        final SlotManager slotManager =
+                new TestingSlotManagerBuilder()
+                        .setTriggerRequirementsCheckConsumer(
+                                triggerRequirementsCheckFuture::complete)
+                        .createSlotManager();
+        resourceManager =
+                new ResourceManagerBuilder()
+                        .withSlotManager(slotManager)
+                        .withBlocklistHandlerFactory(
+                                new DefaultBlocklistHandler.Factory(Duration.ofMillis(100L)))
+                        .buildAndStart();
+
+        final ResourceManagerGateway resourceManagerGateway =
+                resourceManager.getSelfGateway(ResourceManagerGateway.class);
+
+        resourceManagerGateway.notifyNewBlockedNodes(
+                Collections.singleton(
+                        new BlockedNode("node", "Test cause", System.currentTimeMillis())));
+
+        triggerRequirementsCheckFuture.get();
+    }
+
     private void testDisconnectJobManager(JobStatus jobStatus) throws Exception {
         final TestingJobMasterGateway jobMasterGateway =
                 new TestingJobMasterGatewayBuilder()
@@ -598,6 +630,8 @@ private class ResourceManagerBuilder {
         private HeartbeatServices heartbeatServices = null;
         private JobLeaderIdService jobLeaderIdService = null;
         private SlotManager slotManager = null;
+        private BlocklistHandler.Factory blocklistHandlerFactory =
+                new NoOpBlocklistHandler.Factory();
         private Function<ResourceID, Boolean> stopWorkerFunction = null;
 
         private ResourceManagerBuilder withHeartbeatServices(HeartbeatServices heartbeatServices) {
@@ -616,6 +650,12 @@ private ResourceManagerBuilder withSlotManager(SlotManager slotManager) {
             return this;
         }
 
+        private ResourceManagerBuilder withBlocklistHandlerFactory(
+                BlocklistHandler.Factory blocklistHandlerFactory) {
+            this.blocklistHandlerFactory = blocklistHandlerFactory;
+            return this;
+        }
+
         private ResourceManagerBuilder withStopWorkerFunction(
                 Function<ResourceID, Boolean> stopWorkerFunction) {
             this.stopWorkerFunction = stopWorkerFunction;
@@ -655,6 +695,7 @@ private TestingResourceManager buildAndStart() throws Exception {
                             new NoOpDelegationTokenManager(),
                             slotManager,
                             NoOpResourceManagerPartitionTracker::get,
+                            blocklistHandlerFactory,
                             jobLeaderIdService,
                             testingFatalErrorHandler,
                             UnregisteredMetricGroups.createUnregisteredResourceManagerMetricGroup(),