[FLINK-27031][runtime] Assign even empty old state to the task if the…

… upstream has output states since the task should be prepared to filter these old incoming states from upstream

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java

@@ -143,7 +143,12 @@ public void assignStates() {
 
         // actually assign the state
         for (TaskStateAssignment stateAssignment : vertexAssignments.values()) {
-            if (stateAssignment.hasNonFinishedState || stateAssignment.isFullyFinished) {
+            // If upstream has output states, even the empty task state should be assigned for the
+            // current task in order to notify this task that the old states will send to it which
+            // likely should be filtered.
+            if (stateAssignment.hasNonFinishedState
+                    || stateAssignment.isFullyFinished
+                    || stateAssignment.hasUpstreamOutputStates()) {
                 assignTaskStateToExecutionJobVertices(stateAssignment);
             }
         }

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

@@ -201,6 +201,11 @@ public OperatorSubtaskState getSubtaskState(OperatorInstanceID instanceID) {
                 .build();
     }
 
+    public boolean hasUpstreamOutputStates() {
+        return Arrays.stream(getUpstreamAssignments())
+                .anyMatch(assignment -> assignment.hasOutputState);
+    }
+
     private InflightDataGateOrPartitionRescalingDescriptor log(
             InflightDataGateOrPartitionRescalingDescriptor descriptor, int subtask, int partition) {
         LOG.debug(

flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java

@@ -87,6 +87,7 @@
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.empty;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 
 /** Tests to verify state assignment operation. */
 public class StateAssignmentOperationTest extends TestLogger {
@@ -744,6 +745,46 @@ public void testChannelStateAssignmentUpscaling() throws JobException, JobExecut
                         .getInputRescalingDescriptor());
     }
 
+    @Test
+    public void testOnlyUpstreamChannelStateAssignment()
+            throws JobException, JobExecutionException {
+        // given: There is only input channel state for one subpartition.
+        List<OperatorID> operatorIds = buildOperatorIds(2);
+        Map<OperatorID, OperatorState> states = new HashMap<>();
+        Random random = new Random();
+        OperatorState upstreamState = new OperatorState(operatorIds.get(0), 2, MAX_P);
+        OperatorSubtaskState state =
+                OperatorSubtaskState.builder()
+                        .setResultSubpartitionState(
+                                new StateObjectCollection<>(
+                                        asList(
+                                                createNewResultSubpartitionStateHandle(10, random),
+                                                createNewResultSubpartitionStateHandle(
+                                                        10, random))))
+                        .build();
+        upstreamState.putState(0, state);
+
+        states.put(operatorIds.get(0), upstreamState);
+
+        Map<OperatorID, ExecutionJobVertex> vertices =
+                buildVertices(operatorIds, 3, RANGE, ROUND_ROBIN);
+
+        // when: States are assigned.
+        new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false)
+                .assignStates();
+
+        // then: All subtask have not null TaskRestore information(even if it is empty).
+        ExecutionJobVertex jobVertexWithFinishedOperator = vertices.get(operatorIds.get(0));
+        for (ExecutionVertex task : jobVertexWithFinishedOperator.getTaskVertices()) {
+            assertNotNull(task.getCurrentExecutionAttempt().getTaskRestore());
+        }
+
+        ExecutionJobVertex jobVertexWithoutFinishedOperator = vertices.get(operatorIds.get(1));
+        for (ExecutionVertex task : jobVertexWithoutFinishedOperator.getTaskVertices()) {
+            assertNotNull(task.getCurrentExecutionAttempt().getTaskRestore());
+        }
+    }
+
     @Test
     public void testStateWithFullyFinishedOperators() throws JobException, JobExecutionException {
         List<OperatorID> operatorIds = buildOperatorIds(2);