content: /*
   * 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.celeborn.service.deploy.master.clustermeta.ha;
  
  import static org.junit.Assert.*;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.Mockito.when;
  
  import java.io.File;
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.util.*;
  import java.util.concurrent.atomic.AtomicLong;
  
  import scala.Tuple2;
  
  import com.google.common.collect.Lists;
  import org.junit.*;
  import org.mockito.Mockito;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.client.MasterClient;
  import org.apache.celeborn.common.client.MasterNotLeaderException;
  import org.apache.celeborn.common.exception.CelebornRuntimeException;
  import org.apache.celeborn.common.identity.UserIdentifier;
  import org.apache.celeborn.common.meta.DiskInfo;
  import org.apache.celeborn.common.meta.WorkerInfo;
  import org.apache.celeborn.common.meta.WorkerStatus;
  import org.apache.celeborn.common.network.sasl.SaslCredentials;
  import org.apache.celeborn.common.network.sasl.SaslTestBase;
  import org.apache.celeborn.common.network.sasl.SecretRegistry;
  import org.apache.celeborn.common.network.sasl.registration.RegistrationClientBootstrap;
  import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo;
  import org.apache.celeborn.common.network.sasl.registration.RegistrationServerBootstrap;
  import org.apache.celeborn.common.network.util.TransportConf;
  import org.apache.celeborn.common.protocol.WorkerEventType;
  import org.apache.celeborn.common.quota.ResourceConsumption;
  import org.apache.celeborn.common.rpc.RpcAddress;
  import org.apache.celeborn.common.rpc.RpcEndpointAddress;
  import org.apache.celeborn.common.rpc.RpcEndpointRef;
  import org.apache.celeborn.common.rpc.RpcEnv;
  import org.apache.celeborn.common.rpc.netty.NettyRpcEndpointRef;
  import org.apache.celeborn.common.util.Utils;
  import org.apache.celeborn.common.util.Utils$;
  import org.apache.celeborn.service.deploy.master.MasterSecretRegistryImpl;
  import org.apache.celeborn.service.deploy.master.clustermeta.AbstractMetaManager;
  import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;
  
  public class RatisMasterStatusSystemSuiteJ {
    protected static HARaftServer RATISSERVER1 = null;
    protected static HARaftServer RATISSERVER2 = null;
    protected static HARaftServer RATISSERVER3 = null;
    protected static HAMasterMetaManager STATUSSYSTEM1 = null;
    protected static HAMasterMetaManager STATUSSYSTEM2 = null;
    protected static HAMasterMetaManager STATUSSYSTEM3 = null;
  
    private static final RpcEndpointRef dummyRef =
        new NettyRpcEndpointRef(
            new CelebornConf(), RpcEndpointAddress.apply("localhost", 111, "dummy"), null);
  
    protected static RpcEnv mockRpcEnv = Mockito.mock(RpcEnv.class);
    protected static RpcEndpointRef mockRpcEndpoint = Mockito.mock(RpcEndpointRef.class);
  
    @BeforeClass
    public static void init() throws Exception {
      resetRaftServer(
          configureServerConf(new CelebornConf(), 1),
          configureServerConf(new CelebornConf(), 2),
          configureServerConf(new CelebornConf(), 3),
          false);
    }
  
    private static void stopAllRaftServers() {
      if (RATISSERVER1 != null) {
        RATISSERVER1.stop();
      }
      if (RATISSERVER2 != null) {
        RATISSERVER2.stop();
      }
      if (RATISSERVER3 != null) {
        RATISSERVER3.stop();
      }
    }
  
    static CelebornConf configureServerConf(CelebornConf conf, int id) throws IOException {
      File tmpDir = File.createTempFile("celeborn-ratis" + id, "for-test-only");
      tmpDir.delete();
      tmpDir.mkdirs();
      conf.set(CelebornConf.HA_MASTER_RATIS_STORAGE_DIR().key(), tmpDir.getAbsolutePath());
      return conf;
    }
  
    public static void resetRaftServer(
        CelebornConf conf1, CelebornConf conf2, CelebornConf conf3, boolean sslEnabled)
        throws IOException, InterruptedException {
      Mockito.when(mockRpcEnv.setupEndpointRef(Mockito.any(), Mockito.any()))
          .thenReturn(mockRpcEndpoint);
      when(mockRpcEnv.setupEndpointRef(any(), any())).thenReturn(dummyRef);
  
      stopAllRaftServers();
  
      int retryCount = 0;
      boolean serversStarted = false;
  
      while (!serversStarted) {
        try {
          STATUSSYSTEM1 = new HAMasterMetaManager(mockRpcEnv, conf1);
          STATUSSYSTEM2 = new HAMasterMetaManager(mockRpcEnv, conf2);
          STATUSSYSTEM3 = new HAMasterMetaManager(mockRpcEnv, conf3);
  
          MetaHandler handler1 = new MetaHandler(STATUSSYSTEM1);
          MetaHandler handler2 = new MetaHandler(STATUSSYSTEM2);
          MetaHandler handler3 = new MetaHandler(STATUSSYSTEM3);
  
          String id1 = UUID.randomUUID().toString();
          String id2 = UUID.randomUUID().toString();
          String id3 = UUID.randomUUID().toString();
  
          int ratisPort1 = Utils$.MODULE$.selectRandomInt(1024, 65535);
          int ratisPort2 = ratisPort1 + 1;
          int ratisPort3 = ratisPort2 + 1;
  
          MasterNode masterNode1 =
              new MasterNode.Builder()
                  .setHost(Utils.localHostName(conf1))
                  .setRatisPort(ratisPort1)
                  .setRpcPort(ratisPort1)
                  .setInternalRpcPort(ratisPort1)
                  .setSslEnabled(sslEnabled)
                  .setNodeId(id1)
                  .build();
          MasterNode masterNode2 =
              new MasterNode.Builder()
                  .setHost(Utils.localHostName(conf2))
                  .setRatisPort(ratisPort2)
                  .setRpcPort(ratisPort2)
                  .setInternalRpcPort(ratisPort2)
                  .setSslEnabled(sslEnabled)
                  .setNodeId(id2)
                  .build();
          MasterNode masterNode3 =
              new MasterNode.Builder()
                  .setHost(Utils.localHostName(conf3))
                  .setRatisPort(ratisPort3)
                  .setRpcPort(ratisPort3)
                  .setInternalRpcPort(ratisPort3)
                  .setSslEnabled(sslEnabled)
                  .setNodeId(id3)
                  .build();
  
          List<MasterNode> peersForNode1 = Arrays.asList(masterNode2, masterNode3);
          List<MasterNode> peersForNode2 = Arrays.asList(masterNode1, masterNode3);
          List<MasterNode> peersForNode3 = Arrays.asList(masterNode1, masterNode2);
  
          RATISSERVER1 =
              HARaftServer.newMasterRatisServer(handler1, conf1, masterNode1, peersForNode1);
          RATISSERVER2 =
              HARaftServer.newMasterRatisServer(handler2, conf2, masterNode2, peersForNode2);
          RATISSERVER3 =
              HARaftServer.newMasterRatisServer(handler3, conf3, masterNode3, peersForNode3);
  
          STATUSSYSTEM1.setRatisServer(RATISSERVER1);
          STATUSSYSTEM2.setRatisServer(RATISSERVER2);
          STATUSSYSTEM3.setRatisServer(RATISSERVER3);
  
          RATISSERVER1.start();
          RATISSERVER2.start();
          RATISSERVER3.start();
          Thread.sleep(15 * 1000);
          serversStarted = true;
        } catch (Exception e) {
          stopAllRaftServers();
          retryCount += 1;
          if (retryCount == 3) {
            throw e;
          }
        }
      }
    }
  
    @Test
    public void testLeaderAvailable() {
      boolean hasLeader =
          RATISSERVER1.isLeader() || RATISSERVER2.isLeader() || RATISSERVER3.isLeader();
      Assert.assertTrue(hasLeader);
  
      // Check if the rpc endpoint and internal rpc endpoint of the leader is as expected.
  
      HARaftServer leader =
          RATISSERVER1.isLeader()
              ? RATISSERVER1
              : (RATISSERVER2.isLeader() ? RATISSERVER2 : RATISSERVER3);
      // one of them must be the follower given the three servers we have
      HARaftServer follower = RATISSERVER1.isLeader() ? RATISSERVER2 : RATISSERVER1;
  
      // This is expected to be false, but as a side effect, updates getCachedLeaderPeerRpcEndpoint
      boolean isFollowerCurrentLeader = follower.isLeader();
      Assert.assertFalse(isFollowerCurrentLeader);
  
      Optional<HARaftServer.LeaderPeerEndpoints> cachedLeaderPeerRpcEndpoint =
          follower.getCachedLeaderPeerRpcEndpoint();
  
      Assert.assertTrue(cachedLeaderPeerRpcEndpoint.isPresent());
  
      Tuple2<String, String> rpcEndpointsPair = cachedLeaderPeerRpcEndpoint.get().rpcEndpoints;
      Tuple2<String, String> rpcInternalEndpointsPair =
          cachedLeaderPeerRpcEndpoint.get().rpcInternalEndpoints;
  
      // rpc endpoint may use custom host name then this ut need check ever ip/host
      Assert.assertTrue(
          leader.getRpcEndpoint().equals(rpcEndpointsPair._1)
              || leader.getRpcEndpoint().equals(rpcEndpointsPair._2));
      Assert.assertTrue(
          leader.getInternalRpcEndpoint().equals(rpcInternalEndpointsPair._1)
              || leader.getRpcEndpoint().equals(rpcInternalEndpointsPair._2));
    }
  
    private static final String HOSTNAME1 = "host1";
    private static final int RPCPORT1 = 1111;
    private static final int PUSHPORT1 = 1112;
    private static final int FETCHPORT1 = 1113;
    private static final int REPLICATEPORT1 = 1114;
    private static final int INTERNALPORT1 = 1115;
    private static final String NETWORK_LOCATION1 = "networkLocation1";
    private static final Map<String, DiskInfo> disks1 = new HashMap<>();
    private static final Map<UserIdentifier, ResourceConsumption> userResourceConsumption1 =
        new HashMap<>();
  
    private static final String HOSTNAME2 = "host2";
    private static final int RPCPORT2 = 2111;
    private static final int PUSHPORT2 = 2112;
    private static final int FETCHPORT2 = 2113;
    private static final int REPLICATEPORT2 = 2114;
    private static final int INTERNALPORT2 = 2115;
    private static final String NETWORK_LOCATION2 = "networkLocation2";
    private static final Map<String, DiskInfo> disks2 = new HashMap<>();
    private static final Map<UserIdentifier, ResourceConsumption> userResourceConsumption2 =
        new HashMap<>();
  
    private static final String HOSTNAME3 = "host3";
    private static final int RPCPORT3 = 3111;
    private static final int PUSHPORT3 = 3112;
    private static final int FETCHPORT3 = 3113;
    private static final int REPLICATEPORT3 = 3114;
    private static final int INTERNALPORT3 = 3115;
    private static final String NETWORK_LOCATION3 = "networkLocation3";
    private static final Map<String, DiskInfo> disks3 = new HashMap<>();
    private static final Map<UserIdentifier, ResourceConsumption> userResourceConsumption3 =
        new HashMap<>();
  
    private final AtomicLong callerId = new AtomicLong();
    private static final String APPID1 = "appId1";
    private static final int SHUFFLEID1 = 1;
    private static final String SHUFFLEKEY1 = APPID1 + "-" + SHUFFLEID1;
    private static final String POLICY1 = "ShufflePartitionsFallbackPolicy";
    private static final String POLICY2 = "QuotaFallbackPolicy";
    private static final WorkerStatus workerStatus = WorkerStatus.normalWorkerStatus();
  
    private String getNewReqeustId() {
      return MasterClient.encodeRequestId(UUID.randomUUID().toString(), callerId.incrementAndGet());
    }
  
    public HAMasterMetaManager pickLeaderStatusSystem() {
      if (RATISSERVER1.isLeader()) {
        return STATUSSYSTEM1;
      }
      if (RATISSERVER2.isLeader()) {
        return STATUSSYSTEM2;
      }
      if (RATISSERVER3.isLeader()) {
        return STATUSSYSTEM3;
      }
      return null;
    }
  
    private void stopNoneLeaderRaftServer(HARaftServer... raftServers) {
      for (HARaftServer raftServer : raftServers) {
        if (!raftServer.isLeader()) {
          raftServer.stop();
        }
      }
    }
  
    @Test
    public void testRaftSystemException() throws Exception {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
      try {
        stopNoneLeaderRaftServer(RATISSERVER1, RATISSERVER2, RATISSERVER3);
        statusSystem.handleRegisterWorker(
            HOSTNAME1,
            RPCPORT1,
            PUSHPORT1,
            FETCHPORT1,
            REPLICATEPORT1,
            INTERNALPORT1,
            NETWORK_LOCATION1,
            disks1,
            userResourceConsumption1,
            getNewReqeustId());
        fail();
      } catch (CelebornRuntimeException e) {
        Assert.assertTrue(true);
      } finally {
        init();
      }
    }
  
    @Test
    public void testHandleRegisterWorker() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(3, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM3.workersMap.size());
  
      Assert.assertEquals(3, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(3, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(3, STATUSSYSTEM3.availableWorkers.size());
  
      assertWorkers(STATUSSYSTEM1.workersMap.values());
      assertWorkers(STATUSSYSTEM2.workersMap.values());
      assertWorkers(STATUSSYSTEM3.workersMap.values());
    }
  
    private void assertWorkers(Collection<WorkerInfo> workerInfos) {
      for (WorkerInfo workerInfo : workerInfos) {
        assertWorker(workerInfo);
      }
    }
  
    private void assertWorker(WorkerInfo workerInfo) {
      Map<String, DiskInfo> diskInfos = workerInfo.diskInfos();
      Assert.assertEquals(96 * 1024 * 1024 * 1024L, diskInfos.get("disk1").totalSpace());
      Assert.assertEquals(96 * 1024 * 1024 * 1024L, diskInfos.get("disk2").totalSpace());
      Assert.assertEquals(96 * 1024 * 1024 * 1024L, diskInfos.get("disk3").totalSpace());
      Assert.assertEquals(96 * 1024 * 1024 * 1024L, diskInfos.get("disk4").totalSpace());
    }
  
    @Test
    public void testHandleWorkerExclude() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      WorkerInfo workerInfo1 =
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1);
      WorkerInfo workerInfo2 =
          new WorkerInfo(
              HOSTNAME2,
              RPCPORT2,
              PUSHPORT2,
              FETCHPORT2,
              REPLICATEPORT2,
              INTERNALPORT2,
              disks2,
              userResourceConsumption2);
  
      statusSystem.handleRegisterWorker(
          workerInfo1.host(),
          workerInfo1.rpcPort(),
          workerInfo1.pushPort(),
          workerInfo1.fetchPort(),
          workerInfo1.replicatePort(),
          workerInfo1.internalPort(),
          workerInfo1.networkLocation(),
          workerInfo1.diskInfos(),
          workerInfo1.userResourceConsumption(),
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          workerInfo2.host(),
          workerInfo2.rpcPort(),
          workerInfo2.pushPort(),
          workerInfo2.fetchPort(),
          workerInfo2.replicatePort(),
          workerInfo2.internalPort(),
          workerInfo2.networkLocation(),
          workerInfo2.diskInfos(),
          workerInfo2.userResourceConsumption(),
          getNewReqeustId());
  
      statusSystem.handleWorkerExclude(
          Arrays.asList(workerInfo1, workerInfo2), Collections.emptyList(), getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(2, STATUSSYSTEM1.manuallyExcludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.manuallyExcludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.manuallyExcludedWorkers.size());
  
      Assert.assertEquals(0, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM3.availableWorkers.size());
  
      statusSystem.handleWorkerExclude(
          Collections.emptyList(), Collections.singletonList(workerInfo1), getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(1, STATUSSYSTEM1.manuallyExcludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.manuallyExcludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.manuallyExcludedWorkers.size());
  
      Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testHandleWorkerLost() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      statusSystem.handleWorkerLost(
          HOSTNAME1, RPCPORT1, PUSHPORT1, FETCHPORT1, REPLICATEPORT1, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(2, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(2, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(2, STATUSSYSTEM3.workersMap.size());
  
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testHandleRequestSlots() {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      WorkerInfo workerInfo1 =
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1);
      WorkerInfo workerInfo2 =
          new WorkerInfo(
              HOSTNAME2,
              RPCPORT2,
              PUSHPORT2,
              FETCHPORT2,
              REPLICATEPORT2,
              INTERNALPORT2,
              disks2,
              userResourceConsumption2);
      WorkerInfo workerInfo3 =
          new WorkerInfo(
              HOSTNAME3,
              RPCPORT3,
              PUSHPORT3,
              FETCHPORT3,
              REPLICATEPORT3,
              INTERNALPORT3,
              disks3,
              userResourceConsumption3);
  
      Map<String, Map<String, Integer>> workersToAllocate = new HashMap<>();
      Map<String, Integer> allocation1 = new HashMap<>();
      allocation1.put("disk1", 15);
      Map<String, Integer> allocation2 = new HashMap<>();
      allocation2.put("disk2", 25);
      Map<String, Integer> allocation3 = new HashMap<>();
      allocation3.put("disk3", 35);
      workersToAllocate.put(workerInfo1.toUniqueId(), allocation1);
      workersToAllocate.put(workerInfo2.toUniqueId(), allocation2);
      workersToAllocate.put(workerInfo3.toUniqueId(), allocation3);
  
      statusSystem.handleRequestSlots(SHUFFLEKEY1, HOSTNAME1, workersToAllocate, getNewReqeustId());
  
      // Do not update diskinfo's activeslots
  
      Assert.assertEquals(
          0,
          statusSystem.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME1))
              .findFirst()
              .get()
              .usedSlots());
      Assert.assertEquals(
          0,
          statusSystem.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME2))
              .findFirst()
              .get()
              .usedSlots());
      Assert.assertEquals(
          0,
          statusSystem.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME3))
              .findFirst()
              .get()
              .usedSlots());
    }
  
    @Test
    public void testHandleReleaseSlots() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION3,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(3, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM3.workersMap.size());
  
      Map<String, Map<String, Integer>> workersToAllocate = new HashMap<>();
      Map<String, Integer> allocations = new HashMap<>();
      allocations.put("disk1", 5);
      workersToAllocate.put(
          statusSystem.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME1))
              .findFirst()
              .get()
              .toUniqueId(),
          allocations);
      workersToAllocate.put(
          statusSystem.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME2))
              .findFirst()
              .get()
              .toUniqueId(),
          allocations);
  
      statusSystem.handleRequestSlots(SHUFFLEKEY1, HOSTNAME1, workersToAllocate, getNewReqeustId());
      Thread.sleep(3000L);
  
      // Do not update diskinfo's activeslots
  
      Assert.assertEquals(
          0,
          STATUSSYSTEM1.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME1))
              .findFirst()
              .get()
              .usedSlots());
      Assert.assertEquals(
          0,
          STATUSSYSTEM2.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME1))
              .findFirst()
              .get()
              .usedSlots());
      Assert.assertEquals(
          0,
          STATUSSYSTEM3.workersMap.values().stream()
              .filter(w -> w.host().equals(HOSTNAME1))
              .findFirst()
              .get()
              .usedSlots());
    }
  
    @Test
    public void testHandleAppLost() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      Thread.sleep(3000L);
      WorkerInfo workerInfo1 =
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1);
      WorkerInfo workerInfo2 =
          new WorkerInfo(
              HOSTNAME2,
              RPCPORT2,
              PUSHPORT2,
              FETCHPORT2,
              REPLICATEPORT2,
              INTERNALPORT2,
              disks2,
              userResourceConsumption2);
      Map<String, Map<String, Integer>> workersToAllocate = new HashMap<>();
      Map<String, Integer> allocations = new HashMap<>();
      allocations.put("disk1", 5);
      workersToAllocate.put(workerInfo1.toUniqueId(), allocations);
      workersToAllocate.put(workerInfo2.toUniqueId(), allocations);
  
      statusSystem.handleRequestSlots(SHUFFLEKEY1, HOSTNAME1, workersToAllocate, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(1, STATUSSYSTEM1.registeredAppAndShuffles.size());
      Assert.assertEquals(1, STATUSSYSTEM2.registeredAppAndShuffles.size());
      Assert.assertEquals(1, STATUSSYSTEM3.registeredAppAndShuffles.size());
  
      statusSystem.handleAppLost(APPID1, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertTrue(STATUSSYSTEM1.registeredAppAndShuffles.isEmpty());
      Assert.assertTrue(STATUSSYSTEM2.registeredAppAndShuffles.isEmpty());
      Assert.assertTrue(STATUSSYSTEM3.registeredAppAndShuffles.isEmpty());
    }
  
    @Test
    public void testHandleUnRegisterShuffle() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      WorkerInfo workerInfo1 =
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1);
      WorkerInfo workerInfo2 =
          new WorkerInfo(
              HOSTNAME2,
              RPCPORT2,
              PUSHPORT2,
              FETCHPORT2,
              REPLICATEPORT2,
              INTERNALPORT2,
              disks2,
              userResourceConsumption2);
  
      Map<String, Map<String, Integer>> workersToAllocate = new HashMap<>();
      Map<String, Integer> allocations = new HashMap<>();
      allocations.put("disk1", 5);
      workersToAllocate.put(workerInfo1.toUniqueId(), allocations);
      workersToAllocate.put(workerInfo2.toUniqueId(), allocations);
  
      statusSystem.handleRequestSlots(SHUFFLEKEY1, HOSTNAME1, workersToAllocate, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(1, STATUSSYSTEM1.registeredAppAndShuffles.size());
      Assert.assertEquals(1, STATUSSYSTEM2.registeredAppAndShuffles.size());
      Assert.assertEquals(1, STATUSSYSTEM3.registeredAppAndShuffles.size());
  
      statusSystem.handleUnRegisterShuffle(SHUFFLEKEY1, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertTrue(STATUSSYSTEM1.registeredAppAndShuffles.isEmpty());
      Assert.assertTrue(STATUSSYSTEM2.registeredAppAndShuffles.isEmpty());
      Assert.assertTrue(STATUSSYSTEM3.registeredAppAndShuffles.isEmpty());
    }
  
    @Test
    public void testHandleBatchUnRegisterShuffle() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      WorkerInfo workerInfo1 =
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1);
      WorkerInfo workerInfo2 =
          new WorkerInfo(
              HOSTNAME2,
              RPCPORT2,
              PUSHPORT2,
              FETCHPORT2,
              REPLICATEPORT2,
              INTERNALPORT2,
              disks2,
              userResourceConsumption2);
  
      Map<String, Map<String, Integer>> workersToAllocate = new HashMap<>();
      Map<String, Integer> allocations = new HashMap<>();
      allocations.put("disk1", 5);
      workersToAllocate.put(workerInfo1.toUniqueId(), allocations);
      workersToAllocate.put(workerInfo2.toUniqueId(), allocations);
  
      List<String> shuffleKeysAll = new ArrayList<>();
      for (int i = 1; i <= 4; i++) {
        String shuffleKey = APPID1 + "-" + i;
        shuffleKeysAll.add(shuffleKey);
        statusSystem.handleRequestSlots(shuffleKey, HOSTNAME1, workersToAllocate, getNewReqeustId());
      }
  
      Thread.sleep(3000L);
  
      Assert.assertEquals(4, STATUSSYSTEM1.registeredShuffleCount());
      Assert.assertEquals(4, STATUSSYSTEM2.registeredShuffleCount());
      Assert.assertEquals(4, STATUSSYSTEM3.registeredShuffleCount());
  
      List<String> shuffleKeys1 = new ArrayList<>();
      shuffleKeys1.add(shuffleKeysAll.get(0));
  
      statusSystem.handleBatchUnRegisterShuffles(shuffleKeys1, getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(3, STATUSSYSTEM1.registeredShuffleCount());
      Assert.assertEquals(3, STATUSSYSTEM2.registeredShuffleCount());
      Assert.assertEquals(3, STATUSSYSTEM3.registeredShuffleCount());
  
      statusSystem.handleBatchUnRegisterShuffles(shuffleKeysAll, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertTrue(STATUSSYSTEM1.registeredShuffleCount() == 0);
      Assert.assertTrue(STATUSSYSTEM2.registeredShuffleCount() == 0);
      Assert.assertTrue(STATUSSYSTEM3.registeredShuffleCount() == 0);
    }
  
    @Test
    public void testHandleAppHeartbeat() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      long dummy = 1235L;
      statusSystem.handleAppHeartbeat(
          APPID1, 1, 1, 1, 1, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(Long.valueOf(dummy), STATUSSYSTEM1.appHeartbeatTime.get(APPID1));
      Assert.assertEquals(Long.valueOf(dummy), STATUSSYSTEM2.appHeartbeatTime.get(APPID1));
      Assert.assertEquals(Long.valueOf(dummy), STATUSSYSTEM3.appHeartbeatTime.get(APPID1));
  
      String appId2 = "app02";
      statusSystem.handleAppHeartbeat(
          appId2, 1, 1, 2, 2, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(Long.valueOf(dummy), STATUSSYSTEM1.appHeartbeatTime.get(appId2));
      Assert.assertEquals(Long.valueOf(dummy), STATUSSYSTEM2.appHeartbeatTime.get(appId2));
      Assert.assertEquals(Long.valueOf(dummy), STATUSSYSTEM3.appHeartbeatTime.get(appId2));
  
      Assert.assertEquals(2, STATUSSYSTEM1.appHeartbeatTime.size());
      Assert.assertEquals(2, STATUSSYSTEM2.appHeartbeatTime.size());
      Assert.assertEquals(2, STATUSSYSTEM3.appHeartbeatTime.size());
    }
  
    @Test
    public void testHandleWorkerHeartbeat() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      statusSystem.handleWorkerHeartbeat(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          new HashMap<>(),
          userResourceConsumption1,
          1,
          false,
          workerStatus,
          getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
  
      Assert.assertEquals(3, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM3.workersMap.size());
  
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
      statusSystem.handleWorkerHeartbeat(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          new HashMap<>(),
          userResourceConsumption2,
          1,
          false,
          workerStatus,
          getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(2, statusSystem.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM1.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.excludedWorkers.size());
  
      Assert.assertEquals(3, statusSystem.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM3.workersMap.size());
  
      Assert.assertEquals(1, statusSystem.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
  
      statusSystem.handleWorkerHeartbeat(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          disks1,
          userResourceConsumption1,
          1,
          false,
          workerStatus,
          getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(1, statusSystem.excludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
  
      Assert.assertEquals(3, statusSystem.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM3.workersMap.size());
  
      Assert.assertEquals(2, statusSystem.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
      statusSystem.handleWorkerHeartbeat(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          disks1,
          userResourceConsumption1,
          1,
          true,
          workerStatus,
          getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(2, statusSystem.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM1.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.excludedWorkers.size());
  
      Assert.assertEquals(3, statusSystem.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM1.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM2.workersMap.size());
      Assert.assertEquals(3, STATUSSYSTEM3.workersMap.size());
  
      Assert.assertEquals(1, statusSystem.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testAutoReleaseHighWorkLoadWorkers() throws InterruptedException, IOException {
      CelebornConf conf1 = new CelebornConf();
      conf1.set(CelebornConf.MASTER_AUTO_RELEASE_HIGH_WORKLOAD_WORKER_ENABLED(), true);
      conf1.set(CelebornConf.MASTER_AUTO_RELEASE_HIGH_WORKLOAD_WORKER_RATIO_THRESHOLD(), 0.8);
      CelebornConf conf2 = new CelebornConf();
      conf2.set(CelebornConf.MASTER_AUTO_RELEASE_HIGH_WORKLOAD_WORKER_ENABLED(), true);
      conf2.set(CelebornConf.MASTER_AUTO_RELEASE_HIGH_WORKLOAD_WORKER_RATIO_THRESHOLD(), 0.8);
      CelebornConf conf3 = new CelebornConf();
      conf3.set(CelebornConf.MASTER_AUTO_RELEASE_HIGH_WORKLOAD_WORKER_ENABLED(), true);
      conf3.set(CelebornConf.MASTER_AUTO_RELEASE_HIGH_WORKLOAD_WORKER_RATIO_THRESHOLD(), 0.8);
  
      try {
        resetRaftServer(
            configureServerConf(conf1, 1),
            configureServerConf(conf2, 2),
            configureServerConf(conf3, 3),
            false);
  
        AbstractMetaManager statusSystem = pickLeaderStatusSystem();
        Assert.assertNotNull(statusSystem);
  
        statusSystem.handleRegisterWorker(
            HOSTNAME1,
            RPCPORT1,
            PUSHPORT1,
            FETCHPORT1,
            REPLICATEPORT1,
            INTERNALPORT1,
            NETWORK_LOCATION1,
            disks1,
            userResourceConsumption1,
            getNewReqeustId());
        statusSystem.handleRegisterWorker(
            HOSTNAME2,
            RPCPORT2,
            PUSHPORT2,
            FETCHPORT2,
            REPLICATEPORT2,
            INTERNALPORT2,
            NETWORK_LOCATION2,
            disks2,
            userResourceConsumption2,
            getNewReqeustId());
        statusSystem.handleRegisterWorker(
            HOSTNAME3,
            RPCPORT3,
            PUSHPORT3,
            FETCHPORT3,
            REPLICATEPORT3,
            INTERNALPORT3,
            NETWORK_LOCATION3,
            disks3,
            userResourceConsumption3,
            getNewReqeustId());
  
        Thread.sleep(3000L);
  
        // worker2 and work3 are unhealthy
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME2,
            RPCPORT2,
            PUSHPORT2,
            FETCHPORT2,
            REPLICATEPORT2,
            new HashMap<>(),
            userResourceConsumption2,
            1,
            false,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(1, statusSystem.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(2, statusSystem.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME3,
            RPCPORT3,
            PUSHPORT3,
            FETCHPORT3,
            REPLICATEPORT3,
            new HashMap<>(),
            userResourceConsumption3,
            1,
            false,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(2, statusSystem.excludedWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(1, statusSystem.availableWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
  
        // worker2 and work3 have high workload
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME2,
            RPCPORT2,
            PUSHPORT2,
            FETCHPORT2,
            REPLICATEPORT2,
            disks2,
            userResourceConsumption2,
            1,
            false,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(1, statusSystem.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(2, statusSystem.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME3,
            RPCPORT3,
            PUSHPORT3,
            FETCHPORT3,
            REPLICATEPORT3,
            disks3,
            userResourceConsumption3,
            1,
            false,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(0, statusSystem.excludedWorkers.size());
        Assert.assertEquals(0, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(0, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(0, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(3, statusSystem.availableWorkers.size());
        Assert.assertEquals(3, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(3, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(3, STATUSSYSTEM3.availableWorkers.size());
  
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME2,
            RPCPORT2,
            PUSHPORT2,
            FETCHPORT2,
            REPLICATEPORT2,
            disks2,
            userResourceConsumption2,
            1,
            true,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(1, statusSystem.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(2, statusSystem.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME3,
            RPCPORT3,
            PUSHPORT3,
            FETCHPORT3,
            REPLICATEPORT3,
            disks3,
            userResourceConsumption3,
            1,
            true,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        // release 2 workers with high workload
        Assert.assertEquals(0, statusSystem.excludedWorkers.size());
        Assert.assertEquals(0, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(0, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(0, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(3, statusSystem.availableWorkers.size());
        Assert.assertEquals(3, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(3, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(3, STATUSSYSTEM3.availableWorkers.size());
  
        // work2 has high workload and work3 is unhealthy
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME2,
            RPCPORT2,
            PUSHPORT2,
            FETCHPORT2,
            REPLICATEPORT2,
            disks2,
            userResourceConsumption2,
            1,
            true,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(1, statusSystem.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(2, statusSystem.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME3,
            RPCPORT3,
            PUSHPORT3,
            FETCHPORT3,
            REPLICATEPORT3,
            new HashMap<>(),
            userResourceConsumption3,
            1,
            false,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        // release worker2
        Assert.assertEquals(1, statusSystem.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(2, statusSystem.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
        statusSystem.handleWorkerHeartbeat(
            HOSTNAME2,
            RPCPORT2,
            PUSHPORT2,
            FETCHPORT2,
            REPLICATEPORT2,
            new HashMap<>(),
            userResourceConsumption2,
            1,
            true,
            workerStatus,
            getNewReqeustId());
        Thread.sleep(3000L);
        Assert.assertEquals(2, statusSystem.excludedWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM1.excludedWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM2.excludedWorkers.size());
        Assert.assertEquals(2, STATUSSYSTEM3.excludedWorkers.size());
        Assert.assertEquals(1, statusSystem.availableWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
        Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
      } catch (Exception e) {
        fail(e.getMessage());
      } finally {
        resetRaftServer(
            configureServerConf(new CelebornConf(), 1),
            configureServerConf(new CelebornConf(), 2),
            configureServerConf(new CelebornConf(), 3),
            false);
      }
    }
  
    @Before
    public void resetStatus() {
      STATUSSYSTEM1.registeredAppAndShuffles.clear();
      STATUSSYSTEM1.hostnameSet.clear();
      STATUSSYSTEM1.workersMap.clear();
      STATUSSYSTEM1.availableWorkers.clear();
      STATUSSYSTEM1.appHeartbeatTime.clear();
      STATUSSYSTEM1.excludedWorkers.clear();
      STATUSSYSTEM1.workerLostEvents.clear();
  
      STATUSSYSTEM2.registeredAppAndShuffles.clear();
      STATUSSYSTEM2.hostnameSet.clear();
      STATUSSYSTEM2.workersMap.clear();
      STATUSSYSTEM2.availableWorkers.clear();
      STATUSSYSTEM2.appHeartbeatTime.clear();
      STATUSSYSTEM2.excludedWorkers.clear();
      STATUSSYSTEM2.workerLostEvents.clear();
  
      STATUSSYSTEM3.registeredAppAndShuffles.clear();
      STATUSSYSTEM3.hostnameSet.clear();
      STATUSSYSTEM3.workersMap.clear();
      STATUSSYSTEM3.availableWorkers.clear();
      STATUSSYSTEM3.appHeartbeatTime.clear();
      STATUSSYSTEM3.excludedWorkers.clear();
      STATUSSYSTEM3.workerLostEvents.clear();
  
      disks1.clear();
      disks1.put(
          "disk1",
          new DiskInfo("disk1", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks1.put(
          "disk2",
          new DiskInfo("disk2", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks1.put(
          "disk3",
          new DiskInfo("disk3", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks1.put(
          "disk4",
          new DiskInfo("disk4", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
  
      disks2.clear();
      disks2.put(
          "disk1",
          new DiskInfo("disk1", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks2.put(
          "disk2",
          new DiskInfo("disk2", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks2.put(
          "disk3",
          new DiskInfo("disk3", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks2.put(
          "disk4",
          new DiskInfo("disk4", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
  
      disks3.clear();
      disks3.put(
          "disk1",
          new DiskInfo("disk1", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks3.put(
          "disk2",
          new DiskInfo("disk2", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks3.put(
          "disk3",
          new DiskInfo("disk3", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
      disks3.put(
          "disk4",
          new DiskInfo("disk4", 64 * 1024 * 1024 * 1024L, 100, 100, 0)
              .setTotalSpace(96 * 1024 * 1024 * 1024L));
    }
  
    @Test
    public void testHandleReportWorkerFailure() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      List<WorkerInfo> failedWorkers = new ArrayList<>();
      failedWorkers.add(
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1));
  
      statusSystem.handleReportWorkerUnavailable(failedWorkers, getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(1, STATUSSYSTEM1.shutdownWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.shutdownWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.shutdownWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM1.excludedWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM2.excludedWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM3.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testHandleRemoveWorkersUnavailableInfo() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      WorkerInfo workerInfo1 =
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1);
  
      List<WorkerInfo> unavailableWorkers = new ArrayList<>();
      unavailableWorkers.add(workerInfo1);
  
      statusSystem.handleWorkerLost(
          HOSTNAME1, RPCPORT1, PUSHPORT1, FETCHPORT1, REPLICATEPORT1, getNewReqeustId());
      statusSystem.handleReportWorkerUnavailable(unavailableWorkers, getNewReqeustId());
  
      Thread.sleep(3000L);
      Assert.assertEquals(2, STATUSSYSTEM1.workersMap.size());
  
      Assert.assertEquals(1, STATUSSYSTEM1.shutdownWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.shutdownWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.shutdownWorkers.size());
  
      Assert.assertEquals(1, STATUSSYSTEM1.lostWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.lostWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.lostWorkers.size());
  
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
  
      statusSystem.handleRemoveWorkersUnavailableInfo(unavailableWorkers, getNewReqeustId());
      Thread.sleep(3000L);
  
      Assert.assertEquals(0, STATUSSYSTEM1.shutdownWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM2.shutdownWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM3.shutdownWorkers.size());
  
      Assert.assertEquals(0, STATUSSYSTEM1.lostWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM2.lostWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM3.lostWorkers.size());
  
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testHandleUpdatePartitionSize() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      CelebornConf conf = new CelebornConf();
      statusSystem.partitionTotalWritten.reset();
      statusSystem.partitionTotalFileCount.reset();
  
      statusSystem.handleUpdatePartitionSize();
      Assert.assertEquals(statusSystem.estimatedPartitionSize, conf.initialEstimatedPartitionSize());
  
      Long dummy = 1235L;
      statusSystem.handleAppHeartbeat(
          APPID1, 10000000000l, 1, 1, 1, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
      String appId2 = "app02";
      statusSystem.handleAppHeartbeat(
          appId2, 1, 1, 2, 2, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
  
      // Max size
      statusSystem.handleUpdatePartitionSize();
      Assert.assertEquals(statusSystem.estimatedPartitionSize, conf.maxPartitionSizeToEstimate());
  
      statusSystem.handleAppHeartbeat(
          APPID1, 1000000000l, 1, 1, 1, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
      statusSystem.handleAppHeartbeat(
          appId2, 1, 1, 2, 2, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
  
      // Size between minEstimateSize -> maxEstimateSize
      statusSystem.handleUpdatePartitionSize();
      Assert.assertEquals(statusSystem.estimatedPartitionSize, 500000000);
  
      statusSystem.handleAppHeartbeat(
          APPID1, 1000l, 1, 1, 1, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
      statusSystem.handleAppHeartbeat(
          appId2, 1000l, 1, 2, 2, new HashMap<>(), new HashMap<>(), dummy, getNewReqeustId());
  
      // Min size
      statusSystem.handleUpdatePartitionSize();
      Assert.assertEquals(statusSystem.estimatedPartitionSize, conf.minPartitionSizeToEstimate());
    }
  
    @Test
    public void testHandleWorkerEvent() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      WorkerInfo workerInfo1 =
          WorkerInfo.fromUniqueId(
              HOSTNAME1 + ":" + RPCPORT1 + ":" + PUSHPORT1 + ":" + FETCHPORT1 + ":" + REPLICATEPORT1);
      WorkerInfo workerInfo2 =
          WorkerInfo.fromUniqueId(
              HOSTNAME2 + ":" + RPCPORT2 + ":" + PUSHPORT2 + ":" + FETCHPORT2 + ":" + REPLICATEPORT2);
      statusSystem.handleWorkerEvent(
          ResourceProtos.WorkerEventType.Decommission_VALUE,
          Lists.newArrayList(workerInfo1, workerInfo2),
          getNewReqeustId());
  
      Thread.sleep(3000L);
      Assert.assertEquals(2, STATUSSYSTEM1.workerEventInfos.size());
      Assert.assertEquals(2, STATUSSYSTEM2.workerEventInfos.size());
      Assert.assertEquals(2, STATUSSYSTEM3.workerEventInfos.size());
  
      Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
  
      Assert.assertTrue(STATUSSYSTEM1.workerEventInfos.containsKey(workerInfo1));
      Assert.assertTrue(STATUSSYSTEM1.workerEventInfos.containsKey(workerInfo2));
  
      Assert.assertEquals(
          WorkerEventType.Decommission,
          STATUSSYSTEM1.workerEventInfos.get(workerInfo1).getEventType());
  
      statusSystem.handleWorkerEvent(
          ResourceProtos.WorkerEventType.None_VALUE,
          Lists.newArrayList(workerInfo1),
          getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(1, STATUSSYSTEM1.workerEventInfos.size());
      Assert.assertEquals(1, STATUSSYSTEM2.workerEventInfos.size());
      Assert.assertEquals(1, STATUSSYSTEM3.workerEventInfos.size());
      Assert.assertTrue(STATUSSYSTEM1.workerEventInfos.containsKey(workerInfo2));
      Assert.assertEquals(
          WorkerEventType.Decommission,
          STATUSSYSTEM1.workerEventInfos.get(workerInfo2).getEventType());
  
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testHandleWorkerEventByTransportMessage() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      WorkerInfo workerInfo1 =
          WorkerInfo.fromUniqueId(
              HOSTNAME1 + ":" + RPCPORT1 + ":" + PUSHPORT1 + ":" + FETCHPORT1 + ":" + REPLICATEPORT1);
      WorkerInfo workerInfo2 =
          WorkerInfo.fromUniqueId(
              HOSTNAME2 + ":" + RPCPORT2 + ":" + PUSHPORT2 + ":" + FETCHPORT2 + ":" + REPLICATEPORT2);
      statusSystem.handleWorkerEvent(
          WorkerEventType.Decommission_VALUE,
          Lists.newArrayList(workerInfo1, workerInfo2),
          getNewReqeustId());
  
      Thread.sleep(3000L);
      Assert.assertEquals(2, STATUSSYSTEM1.workerEventInfos.size());
      Assert.assertEquals(2, STATUSSYSTEM2.workerEventInfos.size());
      Assert.assertEquals(2, STATUSSYSTEM3.workerEventInfos.size());
  
      Assert.assertEquals(1, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.availableWorkers.size());
  
      Assert.assertTrue(STATUSSYSTEM1.workerEventInfos.containsKey(workerInfo1));
      Assert.assertTrue(STATUSSYSTEM1.workerEventInfos.containsKey(workerInfo2));
  
      Assert.assertEquals(
          WorkerEventType.Decommission,
          STATUSSYSTEM1.workerEventInfos.get(workerInfo1).getEventType());
  
      statusSystem.handleWorkerEvent(
          WorkerEventType.None_VALUE, Lists.newArrayList(workerInfo1), getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(1, STATUSSYSTEM1.workerEventInfos.size());
      Assert.assertEquals(1, STATUSSYSTEM2.workerEventInfos.size());
      Assert.assertEquals(1, STATUSSYSTEM3.workerEventInfos.size());
      Assert.assertTrue(STATUSSYSTEM1.workerEventInfos.containsKey(workerInfo2));
      Assert.assertEquals(
          WorkerEventType.Decommission,
          STATUSSYSTEM1.workerEventInfos.get(workerInfo2).getEventType());
  
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
    }
  
    @Test
    public void testReviseShuffles() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleReviseLostShuffles(
          "app-1", Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8), getNewReqeustId());
      Thread.sleep(1000l);
      Assert.assertEquals(STATUSSYSTEM1.registeredShuffleCount(), 8);
      Assert.assertEquals(STATUSSYSTEM2.registeredShuffleCount(), 8);
      Assert.assertEquals(STATUSSYSTEM3.registeredShuffleCount(), 8);
    }
  
    @AfterClass
    public static void testNotifyLogFailed() {
      List<HARaftServer> list = Arrays.asList(RATISSERVER1, RATISSERVER2, RATISSERVER3);
      for (HARaftServer haRaftServer : list) {
        if (haRaftServer.isLeader()) {
          haRaftServer
              .getMasterStateMachine()
              .notifyLogFailed(new Exception("test leader step down"), null);
          Assert.assertFalse(haRaftServer.isLeader());
          break;
        }
      }
    }
  
    @Test
    public void testHandleReportWorkerDecommission() throws InterruptedException {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
  
      statusSystem.handleRegisterWorker(
          HOSTNAME1,
          RPCPORT1,
          PUSHPORT1,
          FETCHPORT1,
          REPLICATEPORT1,
          INTERNALPORT1,
          NETWORK_LOCATION1,
          disks1,
          userResourceConsumption1,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME2,
          RPCPORT2,
          PUSHPORT2,
          FETCHPORT2,
          REPLICATEPORT2,
          INTERNALPORT2,
          NETWORK_LOCATION2,
          disks2,
          userResourceConsumption2,
          getNewReqeustId());
      statusSystem.handleRegisterWorker(
          HOSTNAME3,
          RPCPORT3,
          PUSHPORT3,
          FETCHPORT3,
          REPLICATEPORT3,
          INTERNALPORT3,
          NETWORK_LOCATION3,
          disks3,
          userResourceConsumption3,
          getNewReqeustId());
  
      List<WorkerInfo> workers = new ArrayList<>();
      workers.add(
          new WorkerInfo(
              HOSTNAME1,
              RPCPORT1,
              PUSHPORT1,
              FETCHPORT1,
              REPLICATEPORT1,
              INTERNALPORT1,
              disks1,
              userResourceConsumption1));
  
      statusSystem.handleReportWorkerDecommission(workers, getNewReqeustId());
      Thread.sleep(3000L);
      Assert.assertEquals(1, STATUSSYSTEM1.decommissionWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM2.decommissionWorkers.size());
      Assert.assertEquals(1, STATUSSYSTEM3.decommissionWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM1.excludedWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM2.excludedWorkers.size());
      Assert.assertEquals(0, STATUSSYSTEM3.excludedWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM1.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM2.availableWorkers.size());
      Assert.assertEquals(2, STATUSSYSTEM3.availableWorkers.size());
    }
  
    @Test
    public void testRedirectToLeaderForRegistration() throws Throwable {
      CelebornConf celebornConf = new CelebornConf();
      celebornConf.set("celeborn.shuffle.io.maxRetries", "1");
      celebornConf.set("celeborn.network.bind.preferIpAddress", "false");
      TransportConf conf = new TransportConf("shuffle", celebornConf);
  
      // if server1 is the leader, pick server2, else server1 - so that it is always a follower
      AbstractMetaManager followerStatusSystem =
          RATISSERVER1.isLeader() ? STATUSSYSTEM2 : STATUSSYSTEM1;
  
      SecretRegistry registry =
          new SecretRegistry() {
            // Only canRegisterElseReturnFailure is relevant for this test, rest dont matter
            public String getSecretKey(String appId) {
              return null;
            }
  
            public boolean isRegistered(String appId) {
              return false;
            }
  
            public void register(String appId, String secret) {}
  
            public void unregister(String appId) {}
  
            public boolean registrationEnabled() {
              // what MasterSecretRegistryImpl does - except we purposefully pick a follower
              return MasterSecretRegistryImpl.registrationEnabledImpl(followerStatusSystem, false);
            }
          };
  
      RegistrationServerBootstrap serverBootstrap = new RegistrationServerBootstrap(conf, registry);
      RegistrationClientBootstrap clientBootstrap =
          new RegistrationClientBootstrap(
              conf, "app", new SaslCredentials("user", "password"), new RegistrationInfo());
      try {
        SaslTestBase.authHelper(conf, serverBootstrap, clientBootstrap);
        fail("Should have redirected to leader");
      } catch (Exception e) {
        assertTrue(e.getCause() instanceof MasterNotLeaderException);
        MasterNotLeaderException notLeaderException = (MasterNotLeaderException) e.getCause();
  
        HARaftServer leader =
            RATISSERVER1.isLeader()
                ? RATISSERVER1
                : RATISSERVER2.isLeader() ? RATISSERVER2 : RATISSERVER3;
  
        // By default, prefer ip is true - which is used while the config above returns hostnames.
        // Resolve the host in order to match
        RpcAddress leaderEndpointAddress = RpcAddress.fromHostAndPort(leader.getRpcEndpoint());
        RpcAddress suggestedEndpointAddress =
            RpcAddress.fromHostAndPort(notLeaderException.getSuggestedLeaderAddress());
  
        assertEquals(
            new InetSocketAddress(leaderEndpointAddress.host(), leaderEndpointAddress.port()),
            new InetSocketAddress(suggestedEndpointAddress.host(), suggestedEndpointAddress.port()));
      }
    }
  
    @Test
    public void testShuffleAndApplicationCountWithFallback() {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
      statusSystem.shuffleTotalCount.reset();
      statusSystem.applicationTotalCount.reset();
      statusSystem.shuffleFallbackCounts.clear();
      statusSystem.applicationFallbackCounts.clear();
  
      Long dummy = 1235L;
      Map<String, Long> shuffleFallbackCounts = new HashMap<>();
      Map<String, Long> applicationFallbackCounts = new HashMap<>();
      shuffleFallbackCounts.put(POLICY1, 2L);
      applicationFallbackCounts.put(POLICY1, 1L);
      statusSystem.handleAppHeartbeat(
          APPID1,
          1000l,
          1,
          2,
          1,
          shuffleFallbackCounts,
          applicationFallbackCounts,
          dummy,
          getNewReqeustId());
      shuffleFallbackCounts.put(POLICY1, 1L);
      shuffleFallbackCounts.put(POLICY2, 2L);
      applicationFallbackCounts.put(POLICY1, 1L);
      applicationFallbackCounts.put(POLICY2, 1L);
      statusSystem.handleAppHeartbeat(
          APPID1,
          1000l,
          1,
          3,
          2,
          shuffleFallbackCounts,
          applicationFallbackCounts,
          dummy,
          getNewReqeustId());
  
      assertEquals(statusSystem.shuffleTotalCount.longValue(), 5);
      assertEquals(statusSystem.applicationTotalCount.longValue(), 3);
      assertEquals(statusSystem.shuffleFallbackCounts.get(POLICY1).longValue(), 3);
      assertEquals(statusSystem.shuffleFallbackCounts.get(POLICY2).longValue(), 2);
      assertEquals(statusSystem.applicationFallbackCounts.get(POLICY1).longValue(), 2);
      assertEquals(statusSystem.applicationFallbackCounts.get(POLICY2).longValue(), 1);
    }
  
    @Test
    public void testRegisterApplicationInfo() {
      AbstractMetaManager statusSystem = pickLeaderStatusSystem();
      Assert.assertNotNull(statusSystem);
      statusSystem.applicationInfos.clear();
      UserIdentifier userIdentifier = new UserIdentifier("tenant", "celeborn");
  
      String appId = "app1";
      Map<String, String> extraInfo = Collections.singletonMap("k1", "v1");
      statusSystem.handleRegisterApplicationInfo(appId, userIdentifier, extraInfo, getNewReqeustId());
  
      assertEquals(statusSystem.applicationInfos.get(appId).userIdentifier(), userIdentifier);
      assertEquals(statusSystem.applicationInfos.get(appId).extraInfo(), extraInfo);
    }
  }
  