server/src/test/java/org/apache/uniffle/server/TopNShuffleDataSizeOfAppCalcTaskTest.java - incubator-uniffle - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *    http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 package org.apache.uniffle.server;

 import java.io.File;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;

 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.io.TempDir;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcClient;
 import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcNettyClient;
 import org.apache.uniffle.client.request.RssRegisterShuffleRequest;
 import org.apache.uniffle.client.request.RssSendShuffleDataRequest;
 import org.apache.uniffle.client.response.RssRegisterShuffleResponse;
 import org.apache.uniffle.client.response.RssSendShuffleDataResponse;
 import org.apache.uniffle.common.PartitionRange;
 import org.apache.uniffle.common.ShuffleBlockInfo;
 import org.apache.uniffle.common.metrics.TestUtils;
 import org.apache.uniffle.common.rpc.ServerType;
 import org.apache.uniffle.common.rpc.StatusCode;
 import org.apache.uniffle.common.util.RssUtils;
 import org.apache.uniffle.storage.util.StorageType;

 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;

 public class TopNShuffleDataSizeOfAppCalcTaskTest {
   private static final Logger LOG =
       LoggerFactory.getLogger(TopNShuffleDataSizeOfAppCalcTaskTest.class);

   protected static List<ShuffleServer> shuffleServers = Lists.newArrayList();
   private static ShuffleServerConf grpcShuffleServerConfig;
   private static ShuffleServerConf nettyShuffleServerConfig;
   private static final Long EVENT_THRESHOLD_SIZE = 2048L;
   protected static final int SHUFFLE_SERVER_PORT = 20001;

   private static AtomicInteger serverRpcPortCounter = new AtomicInteger();
   private static AtomicInteger nettyPortCounter = new AtomicInteger();
   private static AtomicInteger jettyPortCounter = new AtomicInteger();

   static @TempDir File tempDir;

   protected static final String LOCALHOST;

   static {
     try {
       LOCALHOST = RssUtils.getHostIp();
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
   }

   protected static final int COORDINATOR_PORT_1 = 19999;
   protected static final int NETTY_PORT = 21000;
   protected static final String COORDINATOR_QUORUM = LOCALHOST + ":" + COORDINATOR_PORT_1;

   @BeforeAll
   public static void setupServers(@TempDir File tmpDir) throws Exception {
     File dataDir1 = new File(tmpDir, "data1");
     String grpcBasePath = dataDir1.getAbsolutePath();
     ShuffleServerConf grpcShuffleServerConf = buildShuffleServerConf(ServerType.GRPC, grpcBasePath);
     createShuffleServer(grpcShuffleServerConf);

     File dataDir2 = new File(tmpDir, "data2");
     String nettyBasePath = dataDir2.getAbsolutePath();
     ShuffleServerConf nettyShuffleServerConf =
         buildShuffleServerConf(ServerType.GRPC_NETTY, nettyBasePath);
     createShuffleServer(nettyShuffleServerConf);

     startServers();
     grpcShuffleServerConfig = grpcShuffleServerConf;
     nettyShuffleServerConfig = nettyShuffleServerConf;
   }

   private static ShuffleServerConf buildShuffleServerConf(ServerType serverType, String basePath)
       throws Exception {
     ShuffleServerConf shuffleServerConf = buildShuffleServerConf(serverType);
     shuffleServerConf.setString(
         ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE_HDFS.name());
     shuffleServerConf.set(
         ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, EVENT_THRESHOLD_SIZE);
     shuffleServerConf.set(ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(basePath));
     shuffleServerConf.set(ShuffleServerConf.RPC_METRICS_ENABLED, true);
     shuffleServerConf.set(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 2000L);
     shuffleServerConf.set(ShuffleServerConf.SERVER_PRE_ALLOCATION_EXPIRED, 5000L);
     shuffleServerConf.set(ShuffleServerConf.TOP_N_APP_SHUFFLE_DATA_REFRESH_INTERVAL, 700);
     shuffleServerConf.set(ShuffleServerConf.TOP_N_APP_SHUFFLE_DATA_SIZE_NUMBER, 5);
     return shuffleServerConf;
   }

   protected static ShuffleServerConf buildShuffleServerConf(ServerType serverType)
       throws Exception {
     ShuffleServerConf serverConf = new ShuffleServerConf();
     serverConf.setInteger(
         "rss.rpc.server.port", SHUFFLE_SERVER_PORT + serverRpcPortCounter.getAndIncrement());
     serverConf.setString("rss.storage.type", StorageType.MEMORY_LOCALFILE_HDFS.name());
     serverConf.setString("rss.storage.basePath", tempDir.getAbsolutePath());
     serverConf.setString("rss.server.buffer.capacity", "671088640");
     serverConf.setString("rss.server.memory.shuffle.highWaterMark", "50.0");
     serverConf.setString("rss.server.memory.shuffle.lowWaterMark", "0.0");
     serverConf.setString("rss.server.read.buffer.capacity", "335544320");
     serverConf.setString("rss.coordinator.quorum", COORDINATOR_QUORUM);
     serverConf.setString("rss.server.heartbeat.delay", "1000");
     serverConf.setString("rss.server.heartbeat.interval", "1000");
     serverConf.setInteger("rss.jetty.http.port", 18080 + jettyPortCounter.getAndIncrement());
     serverConf.setInteger("rss.jetty.corePool.size", 64);
     serverConf.setInteger("rss.rpc.executor.size", 10);
     serverConf.setString("rss.server.hadoop.dfs.replication", "2");
     serverConf.setLong("rss.server.disk.capacity", 10L * 1024L * 1024L * 1024L);
     serverConf.setBoolean("rss.server.health.check.enable", false);
     serverConf.setBoolean(ShuffleServerConf.RSS_TEST_MODE_ENABLE, true);
     serverConf.set(ShuffleServerConf.SERVER_TRIGGER_FLUSH_CHECK_INTERVAL, 500L);
     serverConf.set(ShuffleServerConf.RPC_SERVER_TYPE, serverType);
     if (serverType == ServerType.GRPC_NETTY) {
       serverConf.setInteger(
           ShuffleServerConf.NETTY_SERVER_PORT, NETTY_PORT + nettyPortCounter.getAndIncrement());
     }
     return serverConf;
   }

   protected static void createShuffleServer(ShuffleServerConf serverConf) throws Exception {
     shuffleServers.add(new ShuffleServer(serverConf));
   }

   public static void startServers() throws Exception {
     for (ShuffleServer shuffleServer : shuffleServers) {
       shuffleServer.start();
     }
   }

   private void registerAndRequireBuffer(String appId, int length, boolean isNettyMode)
       throws Exception {
     ShuffleServerGrpcClient shuffleServerClient =
         isNettyMode
             ? new ShuffleServerGrpcNettyClient(
                 LOCALHOST,
                 nettyShuffleServerConfig.getInteger(ShuffleServerConf.RPC_SERVER_PORT),
                 nettyShuffleServerConfig.getInteger(ShuffleServerConf.NETTY_SERVER_PORT))
             : new ShuffleServerGrpcClient(
                 LOCALHOST, grpcShuffleServerConfig.getInteger(ShuffleServerConf.RPC_SERVER_PORT));
     int shuffleId = 0;
     int partitionId = 0;
     List<PartitionRange> partitionIds = Lists.newArrayList(new PartitionRange(0, 3));

     RssRegisterShuffleRequest registerShuffleRequest =
         new RssRegisterShuffleRequest(appId, shuffleId, partitionIds, "");
     RssRegisterShuffleResponse registerResponse =
         shuffleServerClient.registerShuffle(registerShuffleRequest);
     assertSame(StatusCode.SUCCESS, registerResponse.getStatusCode());

     List<ShuffleBlockInfo> blockInfos =
         Lists.newArrayList(
             new ShuffleBlockInfo(
                 shuffleId,
                 partitionId,
                 0,
                 length,
                 0,
                 new byte[length],
                 Lists.newArrayList(),
                 0,
                 100,
                 0));

     Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
     partitionToBlocks.put(partitionId, blockInfos);
     Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
     shuffleToBlocks.put(shuffleId, partitionToBlocks);

     RssSendShuffleDataRequest sendShuffleDataRequest =
         new RssSendShuffleDataRequest(appId, 1, 1000, shuffleToBlocks);
     RssSendShuffleDataResponse response =
         shuffleServerClient.sendShuffleData(sendShuffleDataRequest);
     assertSame(StatusCode.SUCCESS, response.getStatusCode());
     shuffleServerClient.close();
   }

   private static Stream<Arguments> testTopNShuffleDataSizeOfAppCalcTaskProvider() {
     return Stream.of(Arguments.of(true), Arguments.of(false));
   }

   @ParameterizedTest
   @MethodSource("testTopNShuffleDataSizeOfAppCalcTaskProvider")
   private void testTopNShuffleDataSizeOfAppCalcTask(boolean isNettyMode) throws Exception {
     // Here is 6 app, but config max top n number is 5
     registerAndRequireBuffer("application_id_1", 1000, isNettyMode);
     registerAndRequireBuffer("application_id_2", 2000, isNettyMode);
     registerAndRequireBuffer("application_id_3", 3000, isNettyMode);
     registerAndRequireBuffer("application_id_4", 4000, isNettyMode);
     registerAndRequireBuffer("application_id_5", 5000, isNettyMode);
     registerAndRequireBuffer("application_id_6", 6000, isNettyMode);

     Thread.sleep(500);
     int jettyPort =
         isNettyMode
             ? nettyShuffleServerConfig.getInteger(ShuffleServerConf.JETTY_HTTP_PORT)
             : grpcShuffleServerConfig.getInteger(ShuffleServerConf.JETTY_HTTP_PORT);
     String content =
         TestUtils.httpGet(String.format("http://127.0.0.1:%s/metrics/server", jettyPort));
     LOG.info(content);
     ObjectMapper mapper = new ObjectMapper();
     JsonNode actualObj = mapper.readTree(content);
     JsonNode metricsNode = actualObj.get("metrics");
     Set<String> topNTotalDataSizeApps = new HashSet<>();
     Set<String> topNInMemoryDataSizeApps = new HashSet<>();
     for (int i = 0; i < metricsNode.size(); i++) {
       JsonNode metricsName = metricsNode.get(i).get("name");
       if (ShuffleServerMetrics.TOPN_OF_TOTAL_DATA_SIZE_FOR_APP.equals(metricsName.textValue())) {
         Iterator<Map.Entry<String, JsonNode>> it = metricsNode.get(i).fields();
         while (it.hasNext()) {
           Map.Entry<String, JsonNode> entry = it.next();
           if ("labelValues".equalsIgnoreCase(entry.getKey())) {
             topNTotalDataSizeApps.add(entry.getValue().toString());
           }
         }
       }
       if (ShuffleServerMetrics.TOPN_OF_IN_MEMORY_DATA_SIZE_FOR_APP.equals(
           metricsName.textValue())) {
         Iterator<Map.Entry<String, JsonNode>> it = metricsNode.get(i).fields();
         while (it.hasNext()) {
           Map.Entry<String, JsonNode> entry = it.next();
           if ("labelValues".equalsIgnoreCase(entry.getKey())) {
             topNInMemoryDataSizeApps.add(entry.getValue().toString());
           }
         }
       }
     }

     Set<String> expectedTopNApps =
         Sets.newHashSet(
             "[\"application_id_6\"]",
             "[\"application_id_5\"]",
             "[\"application_id_4\"]",
             "[\"application_id_3\"]",
             "[\"application_id_2\"]");
     assertTrue(
         expectedTopNApps.containsAll(topNTotalDataSizeApps)
             && expectedTopNApps.size() == topNTotalDataSizeApps.size());
     assertTrue(
         expectedTopNApps.containsAll(topNInMemoryDataSizeApps)
             && expectedTopNApps.size() == topNInMemoryDataSizeApps.size());
   }

   @AfterEach
   public void cleanMetrics() throws Exception {
     ShuffleServerMetrics.clear();
   }

   @AfterAll
   public static void shutdownServers() throws Exception {
     for (ShuffleServer shuffleServer : shuffleServers) {
       shuffleServer.stopServer();
     }
     shuffleServers = Lists.newArrayList();
     ShuffleServerMetrics.clear();
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	package org.apache.uniffle.server;

	import java.io.File;
	import java.util.Arrays;
	import java.util.HashSet;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.stream.Stream;

	import com.fasterxml.jackson.databind.JsonNode;
	import com.fasterxml.jackson.databind.ObjectMapper;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Maps;
	import com.google.common.collect.Sets;
	import org.junit.jupiter.api.AfterAll;
	import org.junit.jupiter.api.AfterEach;
	import org.junit.jupiter.api.BeforeAll;
	import org.junit.jupiter.api.io.TempDir;
	import org.junit.jupiter.params.ParameterizedTest;
	import org.junit.jupiter.params.provider.Arguments;
	import org.junit.jupiter.params.provider.MethodSource;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcClient;
	import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcNettyClient;
	import org.apache.uniffle.client.request.RssRegisterShuffleRequest;
	import org.apache.uniffle.client.request.RssSendShuffleDataRequest;
	import org.apache.uniffle.client.response.RssRegisterShuffleResponse;
	import org.apache.uniffle.client.response.RssSendShuffleDataResponse;
	import org.apache.uniffle.common.PartitionRange;
	import org.apache.uniffle.common.ShuffleBlockInfo;
	import org.apache.uniffle.common.metrics.TestUtils;
	import org.apache.uniffle.common.rpc.ServerType;
	import org.apache.uniffle.common.rpc.StatusCode;
	import org.apache.uniffle.common.util.RssUtils;
	import org.apache.uniffle.storage.util.StorageType;

	import static org.junit.jupiter.api.Assertions.assertSame;
	import static org.junit.jupiter.api.Assertions.assertTrue;

	public class TopNShuffleDataSizeOfAppCalcTaskTest {
	private static final Logger LOG =
	LoggerFactory.getLogger(TopNShuffleDataSizeOfAppCalcTaskTest.class);

	protected static List<ShuffleServer> shuffleServers = Lists.newArrayList();
	private static ShuffleServerConf grpcShuffleServerConfig;
	private static ShuffleServerConf nettyShuffleServerConfig;
	private static final Long EVENT_THRESHOLD_SIZE = 2048L;
	protected static final int SHUFFLE_SERVER_PORT = 20001;

	private static AtomicInteger serverRpcPortCounter = new AtomicInteger();
	private static AtomicInteger nettyPortCounter = new AtomicInteger();
	private static AtomicInteger jettyPortCounter = new AtomicInteger();

	static @TempDir File tempDir;

	protected static final String LOCALHOST;

	static {
	try {
	LOCALHOST = RssUtils.getHostIp();
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	}

	protected static final int COORDINATOR_PORT_1 = 19999;
	protected static final int NETTY_PORT = 21000;
	protected static final String COORDINATOR_QUORUM = LOCALHOST + ":" + COORDINATOR_PORT_1;

	@BeforeAll
	public static void setupServers(@TempDir File tmpDir) throws Exception {
	File dataDir1 = new File(tmpDir, "data1");
	String grpcBasePath = dataDir1.getAbsolutePath();
	ShuffleServerConf grpcShuffleServerConf = buildShuffleServerConf(ServerType.GRPC, grpcBasePath);
	createShuffleServer(grpcShuffleServerConf);

	File dataDir2 = new File(tmpDir, "data2");
	String nettyBasePath = dataDir2.getAbsolutePath();
	ShuffleServerConf nettyShuffleServerConf =
	buildShuffleServerConf(ServerType.GRPC_NETTY, nettyBasePath);
	createShuffleServer(nettyShuffleServerConf);

	startServers();
	grpcShuffleServerConfig = grpcShuffleServerConf;
	nettyShuffleServerConfig = nettyShuffleServerConf;
	}

	private static ShuffleServerConf buildShuffleServerConf(ServerType serverType, String basePath)
	throws Exception {
	ShuffleServerConf shuffleServerConf = buildShuffleServerConf(serverType);
	shuffleServerConf.setString(
	ShuffleServerConf.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE_HDFS.name());
	shuffleServerConf.set(
	ShuffleServerConf.FLUSH_COLD_STORAGE_THRESHOLD_SIZE, EVENT_THRESHOLD_SIZE);
	shuffleServerConf.set(ShuffleServerConf.RSS_STORAGE_BASE_PATH, Arrays.asList(basePath));
	shuffleServerConf.set(ShuffleServerConf.RPC_METRICS_ENABLED, true);
	shuffleServerConf.set(ShuffleServerConf.SERVER_APP_EXPIRED_WITHOUT_HEARTBEAT, 2000L);
	shuffleServerConf.set(ShuffleServerConf.SERVER_PRE_ALLOCATION_EXPIRED, 5000L);
	shuffleServerConf.set(ShuffleServerConf.TOP_N_APP_SHUFFLE_DATA_REFRESH_INTERVAL, 700);
	shuffleServerConf.set(ShuffleServerConf.TOP_N_APP_SHUFFLE_DATA_SIZE_NUMBER, 5);
	return shuffleServerConf;
	}

	protected static ShuffleServerConf buildShuffleServerConf(ServerType serverType)
	throws Exception {
	ShuffleServerConf serverConf = new ShuffleServerConf();
	serverConf.setInteger(
	"rss.rpc.server.port", SHUFFLE_SERVER_PORT + serverRpcPortCounter.getAndIncrement());
	serverConf.setString("rss.storage.type", StorageType.MEMORY_LOCALFILE_HDFS.name());
	serverConf.setString("rss.storage.basePath", tempDir.getAbsolutePath());
	serverConf.setString("rss.server.buffer.capacity", "671088640");
	serverConf.setString("rss.server.memory.shuffle.highWaterMark", "50.0");
	serverConf.setString("rss.server.memory.shuffle.lowWaterMark", "0.0");
	serverConf.setString("rss.server.read.buffer.capacity", "335544320");
	serverConf.setString("rss.coordinator.quorum", COORDINATOR_QUORUM);
	serverConf.setString("rss.server.heartbeat.delay", "1000");
	serverConf.setString("rss.server.heartbeat.interval", "1000");
	serverConf.setInteger("rss.jetty.http.port", 18080 + jettyPortCounter.getAndIncrement());
	serverConf.setInteger("rss.jetty.corePool.size", 64);
	serverConf.setInteger("rss.rpc.executor.size", 10);
	serverConf.setString("rss.server.hadoop.dfs.replication", "2");
	serverConf.setLong("rss.server.disk.capacity", 10L * 1024L * 1024L * 1024L);
	serverConf.setBoolean("rss.server.health.check.enable", false);
	serverConf.setBoolean(ShuffleServerConf.RSS_TEST_MODE_ENABLE, true);
	serverConf.set(ShuffleServerConf.SERVER_TRIGGER_FLUSH_CHECK_INTERVAL, 500L);
	serverConf.set(ShuffleServerConf.RPC_SERVER_TYPE, serverType);
	if (serverType == ServerType.GRPC_NETTY) {
	serverConf.setInteger(
	ShuffleServerConf.NETTY_SERVER_PORT, NETTY_PORT + nettyPortCounter.getAndIncrement());
	}
	return serverConf;
	}

	protected static void createShuffleServer(ShuffleServerConf serverConf) throws Exception {
	shuffleServers.add(new ShuffleServer(serverConf));
	}

	public static void startServers() throws Exception {
	for (ShuffleServer shuffleServer : shuffleServers) {
	shuffleServer.start();
	}
	}

	private void registerAndRequireBuffer(String appId, int length, boolean isNettyMode)
	throws Exception {
	ShuffleServerGrpcClient shuffleServerClient =
	isNettyMode
	? new ShuffleServerGrpcNettyClient(
	LOCALHOST,
	nettyShuffleServerConfig.getInteger(ShuffleServerConf.RPC_SERVER_PORT),
	nettyShuffleServerConfig.getInteger(ShuffleServerConf.NETTY_SERVER_PORT))
	: new ShuffleServerGrpcClient(
	LOCALHOST, grpcShuffleServerConfig.getInteger(ShuffleServerConf.RPC_SERVER_PORT));
	int shuffleId = 0;
	int partitionId = 0;
	List<PartitionRange> partitionIds = Lists.newArrayList(new PartitionRange(0, 3));

	RssRegisterShuffleRequest registerShuffleRequest =
	new RssRegisterShuffleRequest(appId, shuffleId, partitionIds, "");
	RssRegisterShuffleResponse registerResponse =
	shuffleServerClient.registerShuffle(registerShuffleRequest);
	assertSame(StatusCode.SUCCESS, registerResponse.getStatusCode());

	List<ShuffleBlockInfo> blockInfos =
	Lists.newArrayList(
	new ShuffleBlockInfo(
	shuffleId,
	partitionId,
	0,
	length,
	0,
	new byte[length],
	Lists.newArrayList(),
	0,
	100,
	0));

	Map<Integer, List<ShuffleBlockInfo>> partitionToBlocks = Maps.newHashMap();
	partitionToBlocks.put(partitionId, blockInfos);
	Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleToBlocks = Maps.newHashMap();
	shuffleToBlocks.put(shuffleId, partitionToBlocks);

	RssSendShuffleDataRequest sendShuffleDataRequest =
	new RssSendShuffleDataRequest(appId, 1, 1000, shuffleToBlocks);
	RssSendShuffleDataResponse response =
	shuffleServerClient.sendShuffleData(sendShuffleDataRequest);
	assertSame(StatusCode.SUCCESS, response.getStatusCode());
	shuffleServerClient.close();
	}

	private static Stream<Arguments> testTopNShuffleDataSizeOfAppCalcTaskProvider() {
	return Stream.of(Arguments.of(true), Arguments.of(false));
	}

	@ParameterizedTest
	@MethodSource("testTopNShuffleDataSizeOfAppCalcTaskProvider")
	private void testTopNShuffleDataSizeOfAppCalcTask(boolean isNettyMode) throws Exception {
	// Here is 6 app, but config max top n number is 5
	registerAndRequireBuffer("application_id_1", 1000, isNettyMode);
	registerAndRequireBuffer("application_id_2", 2000, isNettyMode);
	registerAndRequireBuffer("application_id_3", 3000, isNettyMode);
	registerAndRequireBuffer("application_id_4", 4000, isNettyMode);
	registerAndRequireBuffer("application_id_5", 5000, isNettyMode);
	registerAndRequireBuffer("application_id_6", 6000, isNettyMode);

	Thread.sleep(500);
	int jettyPort =
	isNettyMode
	? nettyShuffleServerConfig.getInteger(ShuffleServerConf.JETTY_HTTP_PORT)
	: grpcShuffleServerConfig.getInteger(ShuffleServerConf.JETTY_HTTP_PORT);
	String content =
	TestUtils.httpGet(String.format("http://127.0.0.1:%s/metrics/server", jettyPort));
	LOG.info(content);
	ObjectMapper mapper = new ObjectMapper();
	JsonNode actualObj = mapper.readTree(content);
	JsonNode metricsNode = actualObj.get("metrics");
	Set<String> topNTotalDataSizeApps = new HashSet<>();
	Set<String> topNInMemoryDataSizeApps = new HashSet<>();
	for (int i = 0; i < metricsNode.size(); i++) {
	JsonNode metricsName = metricsNode.get(i).get("name");
	if (ShuffleServerMetrics.TOPN_OF_TOTAL_DATA_SIZE_FOR_APP.equals(metricsName.textValue())) {
	Iterator<Map.Entry<String, JsonNode>> it = metricsNode.get(i).fields();
	while (it.hasNext()) {
	Map.Entry<String, JsonNode> entry = it.next();
	if ("labelValues".equalsIgnoreCase(entry.getKey())) {
	topNTotalDataSizeApps.add(entry.getValue().toString());
	}
	}
	}
	if (ShuffleServerMetrics.TOPN_OF_IN_MEMORY_DATA_SIZE_FOR_APP.equals(
	metricsName.textValue())) {
	Iterator<Map.Entry<String, JsonNode>> it = metricsNode.get(i).fields();
	while (it.hasNext()) {
	Map.Entry<String, JsonNode> entry = it.next();
	if ("labelValues".equalsIgnoreCase(entry.getKey())) {
	topNInMemoryDataSizeApps.add(entry.getValue().toString());
	}
	}
	}
	}

	Set<String> expectedTopNApps =
	Sets.newHashSet(
	"[\"application_id_6\"]",
	"[\"application_id_5\"]",
	"[\"application_id_4\"]",
	"[\"application_id_3\"]",
	"[\"application_id_2\"]");
	assertTrue(
	expectedTopNApps.containsAll(topNTotalDataSizeApps)
	&& expectedTopNApps.size() == topNTotalDataSizeApps.size());
	assertTrue(
	expectedTopNApps.containsAll(topNInMemoryDataSizeApps)
	&& expectedTopNApps.size() == topNInMemoryDataSizeApps.size());
	}

	@AfterEach
	public void cleanMetrics() throws Exception {
	ShuffleServerMetrics.clear();
	}

	@AfterAll
	public static void shutdownServers() throws Exception {
	for (ShuffleServer shuffleServer : shuffleServers) {
	shuffleServer.stopServer();
	}
	shuffleServers = Lists.newArrayList();
	ShuffleServerMetrics.clear();
	}
	}