add unit test & fix code style

Author: huwh

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

@@ -64,9 +64,7 @@ public class TaskExecutorRegistration implements Serializable {
      */
     private final String nodeId;
 
-    /**
-     * Number of slots in static slot allocation.
-     */
+    /** Number of slots in static slot allocation. */
     private final int numberSlots;
 
     public TaskExecutorRegistration(

flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorRecoveryTest.java

@@ -26,12 +26,15 @@
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.clusterframework.types.SlotID;
+import org.apache.flink.runtime.entrypoint.ClusterInformation;
 import org.apache.flink.runtime.entrypoint.WorkingDirectory;
 import org.apache.flink.runtime.highavailability.TestingHighAvailabilityServices;
+import org.apache.flink.runtime.instance.InstanceID;
 import org.apache.flink.runtime.jobmaster.utils.TestingJobMasterGateway;
 import org.apache.flink.runtime.jobmaster.utils.TestingJobMasterGatewayBuilder;
 import org.apache.flink.runtime.leaderretrieval.SettableLeaderRetrievalService;
 import org.apache.flink.runtime.messages.Acknowledge;
+import org.apache.flink.runtime.resourcemanager.TaskExecutorRegistration;
 import org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway;
 import org.apache.flink.runtime.rpc.TestingRpcService;
 import org.apache.flink.runtime.rpc.TestingRpcServiceExtension;
@@ -119,7 +122,8 @@ void testRecoveredTaskExecutorWillRestoreAllocationState(@TempDir File tempDir)
         assertThat(slotReport.getNumSlotStatus(), is(2));
 
         final SlotStatus slotStatus = slotReport.iterator().next();
-        final SlotID allocatedSlotID = slotStatus.getSlotID();
+        final SlotID allocatedSlotID =
+                SlotID.getDynamicSlotID(slotStatus.getSlotID().getResourceID());
 
         final AllocationID allocationId = new AllocationID();
         taskExecutorGateway
@@ -178,6 +182,145 @@ void testRecoveredTaskExecutorWillRestoreAllocationState(@TempDir File tempDir)
         assertThat(offeredSlot.getAllocationId(), is(allocationId));
     }
 
+    @Test
+    void testRecoveredTaskExecutorWillRestoreAllocationStateWithDynamicSlotRequest(
+            @TempDir File tempDir) throws Exception {
+        final ResourceID resourceId = ResourceID.generate();
+
+        final Configuration configuration = new Configuration();
+        configuration.set(TaskManagerOptions.NUM_TASK_SLOTS, 2);
+        configuration.set(StateRecoveryOptions.LOCAL_RECOVERY, true);
+
+        final TestingResourceManagerGateway testingResourceManagerGateway =
+                new TestingResourceManagerGateway();
+        final ArrayBlockingQueue<TaskExecutorSlotReport> queue = new ArrayBlockingQueue<>(2);
+        testingResourceManagerGateway.setSendSlotReportFunction(
+                slotReportInformation -> {
+                    queue.offer(
+                            TaskExecutorSlotReport.create(
+                                    slotReportInformation.f0, slotReportInformation.f2));
+                    return CompletableFuture.completedFuture(Acknowledge.get());
+                });
+
+        final ArrayBlockingQueue<TaskExecutorRegistration> taskExecutorRegistrations =
+                new ArrayBlockingQueue<>(2);
+
+        testingResourceManagerGateway.setRegisterTaskExecutorFunction(
+                taskExecutorRegistration -> {
+                    taskExecutorRegistrations.offer(taskExecutorRegistration);
+                    return CompletableFuture.completedFuture(
+                            new TaskExecutorRegistrationSuccess(
+                                    new InstanceID(),
+                                    taskExecutorRegistration.getResourceId(),
+                                    new ClusterInformation("localhost", 1234),
+                                    null));
+                });
+
+        final TestingRpcService rpcService = rpcServiceExtension.getTestingRpcService();
+        rpcService.registerGateway(
+                testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
+
+        final JobID jobId = new JobID();
+
+        final TestingHighAvailabilityServices highAvailabilityServices =
+                new TestingHighAvailabilityServices();
+
+        highAvailabilityServices.setResourceManagerLeaderRetriever(
+                new SettableLeaderRetrievalService(
+                        testingResourceManagerGateway.getAddress(),
+                        testingResourceManagerGateway.getFencingToken().toUUID()));
+        final SettableLeaderRetrievalService jobMasterLeaderRetriever =
+                new SettableLeaderRetrievalService();
+        highAvailabilityServices.setJobMasterLeaderRetriever(jobId, jobMasterLeaderRetriever);
+
+        final WorkingDirectory workingDirectory = WorkingDirectory.create(tempDir);
+        final TaskExecutor taskExecutor =
+                TaskExecutorBuilder.newBuilder(
+                                rpcService, highAvailabilityServices, workingDirectory)
+                        .setConfiguration(configuration)
+                        .setResourceId(resourceId)
+                        .build();
+
+        taskExecutor.start();
+
+        final TaskExecutorGateway taskExecutorGateway =
+                taskExecutor.getSelfGateway(TaskExecutorGateway.class);
+
+        final TaskExecutorSlotReport taskExecutorSlotReport = queue.take();
+
+        final SlotReport slotReport = taskExecutorSlotReport.getSlotReport();
+
+        assertThat(slotReport.getNumSlotStatus(), is(2));
+
+        final TaskExecutorRegistration taskExecutorRegistration = taskExecutorRegistrations.take();
+        assertThat(taskExecutorRegistration.getNumberSlots(), is(2));
+
+        final SlotStatus slotStatus = slotReport.iterator().next();
+        final SlotID allocatedSlotID =
+                SlotID.getDynamicSlotID(slotStatus.getSlotID().getResourceID());
+
+        final AllocationID allocationId = new AllocationID();
+        taskExecutorGateway
+                .requestSlot(
+                        allocatedSlotID,
+                        jobId,
+                        allocationId,
+                        slotStatus.getResourceProfile(),
+                        "localhost",
+                        testingResourceManagerGateway.getFencingToken(),
+                        Duration.ofSeconds(10L))
+                .join();
+
+        taskExecutor.close();
+
+        final BlockingQueue<Collection<SlotOffer>> offeredSlots = new ArrayBlockingQueue<>(1);
+
+        final TestingJobMasterGateway jobMasterGateway =
+                new TestingJobMasterGatewayBuilder()
+                        .setOfferSlotsFunction(
+                                (resourceID, slotOffers) -> {
+                                    offeredSlots.offer(new HashSet<>(slotOffers));
+                                    return CompletableFuture.completedFuture(slotOffers);
+                                })
+                        .build();
+        rpcService.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);
+        jobMasterLeaderRetriever.notifyListener(
+                jobMasterGateway.getAddress(), jobMasterGateway.getFencingToken().toUUID());
+
+        // recover the TaskExecutor
+        final TaskExecutor recoveredTaskExecutor =
+                TaskExecutorBuilder.newBuilder(
+                                rpcService, highAvailabilityServices, workingDirectory)
+                        .setConfiguration(configuration)
+                        .setResourceId(resourceId)
+                        .build();
+
+        recoveredTaskExecutor.start();
+
+        final TaskExecutorSlotReport recoveredSlotReport = queue.take();
+        assertThat(recoveredSlotReport.getSlotReport().getNumSlotStatus(), is(3));
+        for (SlotStatus status : recoveredSlotReport.getSlotReport()) {
+            // dynamic slot request
+            if (status.getSlotID().getSlotNumber() == 2) {
+                assertThat(status.getJobID(), is(jobId));
+                assertThat(status.getAllocationID(), is(allocationId));
+            } else {
+                assertThat(status.getJobID(), is(nullValue()));
+            }
+        }
+
+        final TaskExecutorRegistration recoveredTaskExecutorRegistration =
+                taskExecutorRegistrations.take();
+        assertThat(recoveredTaskExecutorRegistration.getNumberSlots(), is(2));
+
+        final Collection<SlotOffer> take = offeredSlots.take();
+
+        assertThat(take, hasSize(1));
+        final SlotOffer offeredSlot = take.iterator().next();
+
+        assertThat(offeredSlot.getAllocationId(), is(allocationId));
+    }
+
     private static final class TaskExecutorSlotReport {
         private final ResourceID taskExecutorResourceId;
         private final SlotReport slotReport;