[FLINK-38131][flink-runtime] Fix that JM still redistributes broadcast state when all subtask have reported their states (#26925)

Author: xiangforever2014

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

@@ -19,6 +19,7 @@
 package org.apache.flink.runtime.checkpoint;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.state.OperatorStateHandle;
 import org.apache.flink.runtime.state.OperatorStreamStateHandle;
@@ -488,13 +489,16 @@ private static final class GroupByStateNameResults {
         }
     }
 
-    private static final class StateEntry {
+    @VisibleForTesting
+    static final class StateEntry {
         final List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>> entries;
         final BitSet reportedSubtaskIndices;
+        final int parallelism;
 
         public StateEntry(int estimatedEntrySize, int parallelism) {
             this.entries = new ArrayList<>(estimatedEntrySize);
             this.reportedSubtaskIndices = new BitSet(parallelism);
+            this.parallelism = parallelism;
         }
 
         void addEntry(
@@ -506,7 +510,7 @@ void addEntry(
 
         boolean isPartiallyReported() {
             return reportedSubtaskIndices.cardinality() > 0
-                    && reportedSubtaskIndices.cardinality() < reportedSubtaskIndices.size();
+                    && reportedSubtaskIndices.cardinality() < parallelism;
         }
     }
 }

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

@@ -47,6 +47,7 @@
 import org.apache.flink.testutils.TestingUtils;
 import org.apache.flink.testutils.executor.TestExecutorExtension;
 
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -162,6 +163,23 @@ void testRepartitionUnionState() {
         verifyOneKindPartitionableStateRescale(operatorState, operatorID);
     }
 
+    @Test
+    public void testPartiallyReported() {
+        RoundRobinOperatorStateRepartitioner.StateEntry stateEntry =
+                new RoundRobinOperatorStateRepartitioner.StateEntry(0, 5);
+        stateEntry.addEntry(0, null);
+        stateEntry.addEntry(1, null);
+        stateEntry.addEntry(3, null);
+
+        // assert partially report
+        Assertions.assertTrue(stateEntry.isPartiallyReported());
+
+        // assert fully report
+        stateEntry.addEntry(2, null);
+        stateEntry.addEntry(4, null);
+        Assertions.assertFalse(stateEntry.isPartiallyReported());
+    }
+
     @Test
     void testRepartitionBroadcastState() {
         OperatorID operatorID = new OperatorID();