HDFS-17558. RBF: Make maxSizeOfFederatedStateToPropagate work on setResponseHeaderState. (#6902)

Co-authored-by: fuchaohong <fuchaohong@chinatelecom.cn>

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/router/RouterStateIdContext.java

@@ -104,8 +104,10 @@ public void setResponseHeaderState(RpcResponseHeaderProto.Builder headerBuilder)
         builder.putNamespaceStateIds(k, v.get());
       }
     });
+    if (builder.getNamespaceStateIdsCount() <= maxSizeOfFederatedStateToPropagate) {
       headerBuilder.setRouterFederatedState(builder.build().toByteString());
     }
+  }
 
   public LongAccumulator getNamespaceStateId(String nsId) {
     return namespaceIdMap.computeIfAbsent(nsId, key -> new LongAccumulator(Math::max, Long.MIN_VALUE));
@@ -155,10 +157,8 @@ public static long getClientStateIdFromCurrentCall(String nsId) {
 
   @Override
   public void updateResponseState(RpcResponseHeaderProto.Builder header) {
-    if (namespaceIdMap.size() <= maxSizeOfFederatedStateToPropagate) {
     setResponseHeaderState(header);
   }
-  }
 
   @Override
   public void receiveResponseState(RpcResponseHeaderProto header) {

hadoop-hdfs-project/hadoop-hdfs-rbf/src/test/java/org/apache/hadoop/hdfs/server/federation/router/TestObserverWithRouter.java

@@ -641,6 +641,46 @@ public void testRouterResponseHeaderState() {
     Assertions.assertEquals(10L, latestFederateState.get("ns0"));
   }
 
+  @Test
+  @Tag(SKIP_BEFORE_EACH_CLUSTER_STARTUP)
+  public void testRouterResponseHeaderStateMaxSizeLimit() {
+    Configuration conf = new Configuration();
+    conf.setBoolean(RBFConfigKeys.DFS_ROUTER_OBSERVER_READ_DEFAULT_KEY, true);
+    conf.setInt(RBFConfigKeys.DFS_ROUTER_OBSERVER_FEDERATED_STATE_PROPAGATION_MAXSIZE, 1);
+
+    RouterStateIdContext routerStateIdContext = new RouterStateIdContext(conf);
+
+    ConcurrentHashMap<String, LongAccumulator> namespaceIdMap =
+        routerStateIdContext.getNamespaceIdMap();
+    namespaceIdMap.put("ns0", new LongAccumulator(Math::max, 10));
+    namespaceIdMap.put("ns1", new LongAccumulator(Math::max, Long.MIN_VALUE));
+
+    RpcHeaderProtos.RpcResponseHeaderProto.Builder responseHeaderBuilder =
+        RpcHeaderProtos.RpcResponseHeaderProto
+            .newBuilder()
+            .setCallId(1)
+            .setStatus(RpcHeaderProtos.RpcResponseHeaderProto.RpcStatusProto.SUCCESS);
+    routerStateIdContext.updateResponseState(responseHeaderBuilder);
+
+    Map<String, Long> latestFederateState = RouterStateIdContext.getRouterFederatedStateMap(
+        responseHeaderBuilder.build().getRouterFederatedState());
+    // Validate that ns0 is still part of the header
+    Assertions.assertEquals(1, latestFederateState.size());
+
+    namespaceIdMap.put("ns2", new LongAccumulator(Math::max, 20));
+    // Rebuild header
+    responseHeaderBuilder =
+        RpcHeaderProtos.RpcResponseHeaderProto
+            .newBuilder()
+            .setCallId(1)
+            .setStatus(RpcHeaderProtos.RpcResponseHeaderProto.RpcStatusProto.SUCCESS);
+    routerStateIdContext.updateResponseState(responseHeaderBuilder);
+    latestFederateState = RouterStateIdContext.getRouterFederatedStateMap(
+        responseHeaderBuilder.build().getRouterFederatedState());
+    // Validate that ns0 is still part of the header
+    Assertions.assertEquals(0, latestFederateState.size());
+  }
+
   @EnumSource(ConfigSetting.class)
   @ParameterizedTest
   public void testStateIdProgressionInRouter(ConfigSetting configSetting) throws Exception {