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 java.io.File;
  import java.io.IOException;
  import java.net.InetAddress;
  import java.net.InetSocketAddress;
  import java.nio.charset.StandardCharsets;
  import java.util.*;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.locks.ReentrantReadWriteLock;
  
  import javax.net.ssl.KeyManager;
  import javax.net.ssl.TrustManager;
  
  import scala.Tuple2;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.protobuf.InvalidProtocolBufferException;
  import org.apache.ratis.RaftConfigKeys;
  import org.apache.ratis.client.RaftClientConfigKeys;
  import org.apache.ratis.conf.Parameters;
  import org.apache.ratis.conf.RaftProperties;
  import org.apache.ratis.grpc.GrpcConfigKeys;
  import org.apache.ratis.grpc.GrpcTlsConfig;
  import org.apache.ratis.netty.NettyConfigKeys;
  import org.apache.ratis.proto.RaftProtos;
  import org.apache.ratis.protocol.*;
  import org.apache.ratis.protocol.exceptions.RaftException;
  import org.apache.ratis.rpc.CallId;
  import org.apache.ratis.rpc.RpcType;
  import org.apache.ratis.rpc.SupportedRpcType;
  import org.apache.ratis.server.RaftServer;
  import org.apache.ratis.server.RaftServerConfigKeys;
  import org.apache.ratis.server.storage.RaftStorage.StartupOption;
  import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
  import org.apache.ratis.util.LifeCycle;
  import org.apache.ratis.util.SizeInBytes;
  import org.apache.ratis.util.TimeDuration;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.client.MasterClient;
  import org.apache.celeborn.common.exception.CelebornRuntimeException;
  import org.apache.celeborn.common.network.ssl.SSLFactory;
  import org.apache.celeborn.common.protocol.TransportModuleConstants;
  import org.apache.celeborn.common.util.ThreadUtils;
  import org.apache.celeborn.common.util.Utils;
  import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;
  import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.ResourceResponse;
  
  public class HARaftServer {
    public static final Logger LOG = LoggerFactory.getLogger(HARaftServer.class);
  
    public static final UUID CELEBORN_UUID =
        UUID.nameUUIDFromBytes("CELEBORN".getBytes(StandardCharsets.UTF_8));
  
    public static final RaftGroupId RAFT_GROUP_ID = RaftGroupId.valueOf(CELEBORN_UUID);
  
    public static final long REQUEST_TIMEOUT_MS = 60_000;
  
    private final MasterNode localNode;
    private final InetSocketAddress ratisAddr;
    private final String rpcEndpoint;
    private final String internalRpcEndpoint;
    private final RaftServer server;
    private final RaftGroup raftGroup;
    private final RaftPeerId raftPeerId;
  
    private final ClientId clientId = ClientId.randomId();
  
    private final MetaHandler metaHandler;
  
    private final StateMachine masterStateMachine;
  
    private final ScheduledExecutorService scheduledRoleChecker =
        ThreadUtils.newDaemonSingleThreadScheduledExecutor("master-ratis-role-checker");
    private long roleCheckIntervalMs;
    private final ReentrantReadWriteLock roleCheckLock = new ReentrantReadWriteLock();
    private Optional<RaftProtos.RaftPeerRole> cachedPeerRole = Optional.empty();
    private Optional<LeaderPeerEndpoints> cachedLeaderPeerRpcEndpoints = Optional.empty();
  
    private final CelebornConf conf;
    private long workerTimeoutDeadline;
    private long appTimeoutDeadline;
  
    /**
     * Returns a Master Ratis server.
     *
     * @param conf configuration
     * @param localRaftPeerId raft peer id of this Ratis server
     * @param localNode local node of this Ratis server
     * @param raftPeers peer nodes in the raft ring
     * @throws IOException
     */
    private HARaftServer(
        MetaHandler metaHandler,
        CelebornConf conf,
        RaftPeerId localRaftPeerId,
        MasterNode localNode,
        List<RaftPeer> raftPeers)
        throws IOException {
      this.metaHandler = metaHandler;
      this.localNode = localNode;
      this.ratisAddr = localNode.ratisAddr();
      this.rpcEndpoint = localNode.rpcEndpoint();
      this.internalRpcEndpoint = localNode.internalRpcEndpoint();
      this.raftPeerId = localRaftPeerId;
      this.raftGroup = RaftGroup.valueOf(RAFT_GROUP_ID, raftPeers);
      this.masterStateMachine = getStateMachine();
      this.conf = conf;
  
      final RpcType rpc = SupportedRpcType.valueOfIgnoreCase(conf.haMasterRatisRpcType());
      RaftProperties serverProperties = newRaftProperties(conf, rpc);
      Parameters sslParameters =
          localNode.sslEnabled() ? configureSsl(conf, serverProperties, rpc) : null;
      setDeadlineTime(Integer.MAX_VALUE, Integer.MAX_VALUE); // for default
      this.server =
          RaftServer.newBuilder()
              .setServerId(this.raftPeerId)
              .setGroup(this.raftGroup)
              .setProperties(serverProperties)
              .setParameters(sslParameters)
              .setStateMachine(masterStateMachine)
              // RATIS-1677. Do not auto format RaftStorage in RECOVER.
              .setOption(StartupOption.valueOf(conf.haMasterRatisStorageStartupOption()))
              .build();
  
      StringBuilder raftPeersStr = new StringBuilder();
      for (RaftPeer peer : raftPeers) {
        raftPeersStr.append(", ").append(peer.getAddress());
      }
      LOG.info(
          "Ratis server started with GroupID: {} and Raft Peers: {}.",
          RAFT_GROUP_ID,
          raftPeersStr.substring(2));
  
      // Run a scheduler to check and update the server role on the leader periodically
      this.scheduledRoleChecker.scheduleWithFixedDelay(
          () -> {
            // Run this check only on the leader OM
            if (cachedPeerRole.isPresent()
                && cachedPeerRole.get() == RaftProtos.RaftPeerRole.LEADER) {
              updateServerRole();
            }
          },
          roleCheckIntervalMs,
          roleCheckIntervalMs,
          TimeUnit.MILLISECONDS);
    }
  
    public static HARaftServer newMasterRatisServer(
        MetaHandler metaHandler, CelebornConf conf, MasterNode localNode, List<MasterNode> peerNodes)
        throws IOException {
      String nodeId = localNode.nodeId();
      RaftPeerId localRaftPeerId = RaftPeerId.getRaftPeerId(nodeId);
  
      InetSocketAddress ratisAddr = localNode.ratisAddr();
      RaftPeer localRaftPeer =
          RaftPeer.newBuilder()
              .setId(localRaftPeerId)
              .setAddress(ratisAddr)
              .setClientAddress(localNode.rpcEndpoint())
              // We use admin address to host the internal rpc address
              .setAdminAddress(localNode.internalRpcEndpoint())
              .build();
      List<RaftPeer> raftPeers = new ArrayList<>();
      // Add this Ratis server to the Ratis ring
      raftPeers.add(localRaftPeer);
      peerNodes.forEach(
          peer -> {
            String peerNodeId = peer.nodeId();
            RaftPeerId raftPeerId = RaftPeerId.valueOf(peerNodeId);
            RaftPeer raftPeer;
            if (peer.isRatisHostUnresolved()) {
              raftPeer =
                  RaftPeer.newBuilder()
                      .setId(raftPeerId)
                      .setAddress(peer.ratisEndpoint())
                      .setClientAddress(peer.rpcEndpoint())
                      // We use admin address to host the internal rpc address
                      .setAdminAddress(peer.internalRpcEndpoint())
                      .build();
            } else {
              InetSocketAddress peerRatisAddr = peer.ratisAddr();
              raftPeer =
                  RaftPeer.newBuilder()
                      .setId(raftPeerId)
                      .setAddress(peerRatisAddr)
                      .setClientAddress(peer.rpcEndpoint())
                      // We use admin address to host the internal rpc address
                      .setAdminAddress(peer.internalRpcEndpoint())
                      .build();
            }
  
            // Add other nodes belonging to the same service to the Ratis ring
            raftPeers.add(raftPeer);
          });
  
      return new HARaftServer(metaHandler, conf, localRaftPeerId, localNode, raftPeers);
    }
  
    public ResourceResponse submitRequest(ResourceProtos.ResourceRequest request)
        throws CelebornRuntimeException {
      String requestId = request.getRequestId();
      Tuple2<String, Long> decoded = MasterClient.decodeRequestId(requestId);
      if (decoded == null) {
        throw new CelebornRuntimeException(
            "RequestId:" + requestId + " invalid, should be: uuid#callId.");
      }
      ClientId clientId = ClientId.valueOf(UUID.fromString(decoded._1));
      long callId = decoded._2;
      RaftClientRequest raftClientRequest =
          new RaftClientRequest.Builder()
              .setClientId(clientId)
              .setServerId(server.getId())
              .setGroupId(RAFT_GROUP_ID)
              .setCallId(callId)
              .setType(RaftClientRequest.writeRequestType())
              .setMessage(Message.valueOf(HAHelper.convertRequestToByteString(request)))
              .build();
  
      RaftClientReply raftClientReply;
      try {
        raftClientReply = server.submitClientRequestAsync(raftClientRequest).get();
      } catch (Exception ex) {
        throw new CelebornRuntimeException(ex.getMessage(), ex);
      }
  
      if (!raftClientReply.isSuccess()) {
        RaftException exception = raftClientReply.getException();
        throw new CelebornRuntimeException(
            exception == null ? "Encounter raft exception." : exception.getMessage());
      }
  
      try {
        byte[] bytes = raftClientReply.getMessage().getContent().toByteArray();
        return ResourceResponse.newBuilder(ResourceResponse.parseFrom(bytes)).build();
      } catch (InvalidProtocolBufferException ex) {
        throw new CelebornRuntimeException(ex.getMessage(), ex);
      }
    }
  
    /**
     * Start the Ratis server.
     *
     * @throws IOException
     */
    public void start() throws IOException {
      LOG.info(
          "Starting {} {} at port {}",
          getClass().getSimpleName(),
          server.getId(),
          ratisAddr.getPort());
      server.start();
    }
  
    public void stop() {
      try {
        server.close();
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }
  
    private RaftProperties newRaftProperties(CelebornConf conf, RpcType rpc) {
      final RaftProperties properties = new RaftProperties();
      // Set RPC type
      RaftConfigKeys.Rpc.setType(properties, rpc);
  
      // Set the ratis port number
      if (rpc == SupportedRpcType.GRPC) {
        GrpcConfigKeys.Server.setPort(properties, ratisAddr.getPort());
      } else if (rpc == SupportedRpcType.NETTY) {
        NettyConfigKeys.Server.setPort(properties, ratisAddr.getPort());
      }
  
      // Set Ratis storage directory
      String storageDir = conf.haMasterRatisStorageDir();
      RaftServerConfigKeys.setStorageDir(properties, Collections.singletonList(new File(storageDir)));
  
      // Set RAFT segment size
      long raftSegmentSize = conf.haMasterRatisLogSegmentSizeMax();
      RaftServerConfigKeys.Log.setSegmentSizeMax(properties, SizeInBytes.valueOf(raftSegmentSize));
      RaftServerConfigKeys.Log.setPurgeUptoSnapshotIndex(properties, true);
  
      // Set RAFT segment pre-allocated size
      long raftSegmentPreallocatedSize = conf.haMasterRatisLogPreallocatedSize();
      long raftSegmentWriteBufferSize = conf.haMasterRatisLogWriteBufferSize();
      int logAppenderQueueNumElements = conf.haMasterRatisLogAppenderQueueNumElements();
      long logAppenderQueueByteLimit = conf.haMasterRatisLogAppenderQueueBytesLimit();
      // RATIS-589. Eliminate buffer copying in SegmentedRaftLogOutputStream.
      // 4 bytes (serialized size) + logAppenderQueueByteLimit + 4 bytes (checksum)
      if (raftSegmentWriteBufferSize < logAppenderQueueByteLimit + 8) {
        throw new IllegalArgumentException(
            CelebornConf.HA_MASTER_RATIS_LOG_WRITE_BUFFER_SIZE().key()
                + " (= "
                + raftSegmentWriteBufferSize
                + ") is less than "
                + CelebornConf.HA_MASTER_RATIS_LOG_APPENDER_QUEUE_BYTE_LIMIT().key()
                + " + 8 (= "
                + (logAppenderQueueByteLimit + 8)
                + ")");
      }
      boolean shouldInstallSnapshot = conf.haMasterRatisLogInstallSnapshotEnabled();
      RaftServerConfigKeys.Log.Appender.setBufferElementLimit(
          properties, logAppenderQueueNumElements);
      RaftServerConfigKeys.Log.Appender.setBufferByteLimit(
          properties, SizeInBytes.valueOf(logAppenderQueueByteLimit));
      RaftServerConfigKeys.Log.setPreallocatedSize(
          properties, SizeInBytes.valueOf(raftSegmentPreallocatedSize));
      RaftServerConfigKeys.Log.setWriteBufferSize(
          properties, SizeInBytes.valueOf(raftSegmentWriteBufferSize));
      RaftServerConfigKeys.Log.Appender.setInstallSnapshotEnabled(properties, shouldInstallSnapshot);
      int logPurgeGap = conf.haMasterRatisLogPurgeGap();
      RaftServerConfigKeys.Log.setPurgeGap(properties, logPurgeGap);
  
      // For grpc set the maximum message size
      // RATIS-2135. The leader keeps sending inconsistent entries repeatedly to followers.
      // raft.grpc.message.size.max must be 1m larger than raft.server.log.appender.buffer.byte-limit
      GrpcConfigKeys.setMessageSizeMax(
          properties, SizeInBytes.valueOf(logAppenderQueueByteLimit + SizeInBytes.ONE_MB.getSize()));
  
      // Set the server request timeout
      TimeDuration serverRequestTimeout =
          TimeDuration.valueOf(conf.haMasterRatisRpcRequestTimeout(), TimeUnit.SECONDS);
      RaftServerConfigKeys.Rpc.setRequestTimeout(properties, serverRequestTimeout);
  
      // Set timeout for server retry cache entry
      TimeDuration retryCacheExpiryTime =
          TimeDuration.valueOf(conf.haMasterRatisRetryCacheExpiryTime(), TimeUnit.SECONDS);
      RaftServerConfigKeys.RetryCache.setExpiryTime(properties, retryCacheExpiryTime);
  
      // Set the server min and max timeout
      TimeDuration rpcTimeoutMin =
          TimeDuration.valueOf(conf.haMasterRatisRpcTimeoutMin(), TimeUnit.SECONDS);
      TimeDuration rpcTimeoutMax =
          TimeDuration.valueOf(conf.haMasterRatisRpcTimeoutMax(), TimeUnit.SECONDS);
      RaftServerConfigKeys.Rpc.setTimeoutMin(properties, rpcTimeoutMin);
      RaftServerConfigKeys.Rpc.setTimeoutMax(properties, rpcTimeoutMax);
  
      TimeDuration firstElectionTimeoutMin =
          TimeDuration.valueOf(conf.haMasterRatisFirstElectionTimeoutMin(), TimeUnit.SECONDS);
      TimeDuration firstElectionTimeoutMax =
          TimeDuration.valueOf(conf.haMasterRatisFirstElectionTimeoutMax(), TimeUnit.SECONDS);
      RaftServerConfigKeys.Rpc.setFirstElectionTimeoutMin(properties, firstElectionTimeoutMin);
      RaftServerConfigKeys.Rpc.setFirstElectionTimeoutMax(properties, firstElectionTimeoutMax);
  
      boolean leaderElectionMemberMajorityAdd = conf.hasMasterRatisLeaderElectionMemberMajorityAdd();
      RaftServerConfigKeys.LeaderElection.setMemberMajorityAdd(
          properties, leaderElectionMemberMajorityAdd);
  
      // Set the rpc client timeout
      TimeDuration clientRpcTimeout =
          TimeDuration.valueOf(conf.haMasterRatisClientRpcTimeout(), TimeUnit.SECONDS);
      TimeDuration clientRpcWatchTimeout =
          TimeDuration.valueOf(conf.haMasterRatisClientRpcWatchTimeout(), TimeUnit.SECONDS);
      RaftClientConfigKeys.Rpc.setRequestTimeout(properties, clientRpcTimeout);
      RaftClientConfigKeys.Rpc.setWatchRequestTimeout(properties, clientRpcWatchTimeout);
  
      // Set the number of maximum cached segments
      RaftServerConfigKeys.Log.setSegmentCacheNumMax(properties, 2);
  
      TimeDuration noLeaderTimeout =
          TimeDuration.valueOf(conf.haMasterRatisNotificationNoLeaderTimeout(), TimeUnit.SECONDS);
      RaftServerConfigKeys.Notification.setNoLeaderTimeout(properties, noLeaderTimeout);
      TimeDuration slownessTimeout =
          TimeDuration.valueOf(conf.haMasterRatisRpcSlownessTimeout(), TimeUnit.SECONDS);
      RaftServerConfigKeys.Rpc.setSlownessTimeout(properties, slownessTimeout);
  
      // Set role checker time
      this.roleCheckIntervalMs = conf.haMasterRatisRoleCheckInterval();
  
      // snapshot retention
      int numSnapshotRetentionFileNum = conf.haMasterRatisSnapshotRetentionFileNum();
      RaftServerConfigKeys.Snapshot.setRetentionFileNum(properties, numSnapshotRetentionFileNum);
  
      // snapshot interval
      RaftServerConfigKeys.Snapshot.setAutoTriggerEnabled(
          properties, conf.haMasterRatisSnapshotAutoTriggerEnabled());
  
      long snapshotAutoTriggerThreshold = conf.haMasterRatisSnapshotAutoTriggerThreshold();
      RaftServerConfigKeys.Snapshot.setAutoTriggerThreshold(properties, snapshotAutoTriggerThreshold);
  
      for (Map.Entry<String, String> ratisEntry : conf.haRatisCustomConfigs().entrySet()) {
        properties.set(ratisEntry.getKey().replace("celeborn.ratis.", ""), ratisEntry.getValue());
      }
  
      return properties;
    }
  
    private Parameters configureSsl(CelebornConf conf, RaftProperties properties, RpcType rpc) {
  
      if (rpc != SupportedRpcType.GRPC) {
        LOG.error(
            "SSL has been disabled for Raft communication between masters. "
                + "This is only supported when ratis is configured with GRPC");
        return null;
      }
  
      // This is used only for querying state after initialization - not actual SSL
      // also why nThreads does not matter
      SSLFactory factory =
          SSLFactory.createSslFactory(
              Utils.fromCelebornConf(conf, TransportModuleConstants.RPC_SERVICE_MODULE, 1));
  
      assert (null != factory);
      assert (factory.hasKeyManagers());
      assert (!factory.getTrustManagers().isEmpty());
  
      TrustManager trustManager = factory.getTrustManagers().get(0);
      KeyManager keyManager = factory.getKeyManagers().get(0);
  
      Parameters params = new Parameters();
      GrpcConfigKeys.TLS.setEnabled(properties, true);
      GrpcConfigKeys.TLS.setConf(params, new GrpcTlsConfig(keyManager, trustManager, true));
  
      LOG.info("SSL enabled for ratis communication between masters");
  
      return params;
    }
  
    private StateMachine getStateMachine() {
      StateMachine stateMachine = new StateMachine(this);
      stateMachine.setRaftGroupId(RAFT_GROUP_ID);
      return stateMachine;
    }
  
    public MetaHandler getMetaHandler() {
      return metaHandler;
    }
  
    @VisibleForTesting
    public LifeCycle.State getServerState() {
      return server.getLifeCycleState();
    }
  
    public RaftGroup getRaftGroup() {
      return this.raftGroup;
    }
  
    public StateMachine getMasterStateMachine() {
      return this.masterStateMachine;
    }
  
    /**
     * Check the cached leader status.
     *
     * @return true if cached role is Leader, false otherwise.
     */
    private boolean checkCachedPeerRoleIsLeader() {
      this.roleCheckLock.readLock().lock();
      try {
        return cachedPeerRole.isPresent() && cachedPeerRole.get() == RaftProtos.RaftPeerRole.LEADER;
      } finally {
        this.roleCheckLock.readLock().unlock();
      }
    }
  
    /**
     * Check if the current node is the leader node.
     *
     * @return true if Leader, false otherwise.
     */
    public boolean isLeader() {
      if (checkCachedPeerRoleIsLeader()) {
        return true;
      }
  
      // Get the server role from ratis server and update the cached values.
      updateServerRole();
  
      // After updating the server role, check and return if leader or not.
      return checkCachedPeerRoleIsLeader();
    }
  
    /**
     * Get the suggested leader peer id.
     *
     * @return RaftPeerId of the suggested leader node - {@code Optional<LeaderPeerEndpoints>}
     */
    public Optional<LeaderPeerEndpoints> getCachedLeaderPeerRpcEndpoint() {
      this.roleCheckLock.readLock().lock();
      try {
        return cachedLeaderPeerRpcEndpoints;
      } finally {
        this.roleCheckLock.readLock().unlock();
      }
    }
  
    /**
     * Get the group info (peer role and leader peer id) from Ratis server and update the server role.
     */
    public void updateServerRole() {
      try {
        GroupInfoReply groupInfo = getGroupInfo();
        RaftProtos.RoleInfoProto roleInfoProto = groupInfo.getRoleInfoProto();
        RaftProtos.RaftPeerRole thisNodeRole = roleInfoProto.getRole();
        Tuple2<String, String> leaderPeerRpcEndpoint = null;
        Tuple2<String, String> leaderPeerInternalRpcEndpoint = null;
        if (thisNodeRole.equals(RaftProtos.RaftPeerRole.LEADER)) {
          // Current Node always uses original rpcEndpoint/internalRpcEndpoint, as if something wrong
          // they would never return to client.
          setServerRole(
              thisNodeRole,
              Tuple2.apply(this.rpcEndpoint, this.rpcEndpoint),
              Tuple2.apply(this.internalRpcEndpoint, this.internalRpcEndpoint));
        } else if (thisNodeRole.equals(RaftProtos.RaftPeerRole.FOLLOWER)) {
          ByteString leaderNodeId = roleInfoProto.getFollowerInfo().getLeaderInfo().getId().getId();
          // There may be a chance, here we get leaderNodeId as null. For
          // example, in 3 node Ratis, if 2 nodes are down, there will
          // be no leader.
          if (leaderNodeId != null && !leaderNodeId.isEmpty()) {
            String clientAddress =
                roleInfoProto.getFollowerInfo().getLeaderInfo().getId().getClientAddress();
            leaderPeerRpcEndpoint = Utils.addressToIpHostAddressPair(clientAddress);
            // We use admin address to host the internal rpc address
            if (conf.internalPortEnabled()) {
              String adminAddress =
                  roleInfoProto.getFollowerInfo().getLeaderInfo().getId().getAdminAddress();
              leaderPeerInternalRpcEndpoint = Utils.addressToIpHostAddressPair(adminAddress);
            } else {
              leaderPeerInternalRpcEndpoint = leaderPeerRpcEndpoint;
            }
          }
  
          setServerRole(thisNodeRole, leaderPeerRpcEndpoint, leaderPeerInternalRpcEndpoint);
  
        } else {
          setServerRole(thisNodeRole, null, null);
        }
      } catch (IOException e) {
        LOG.error(
            "Failed to retrieve RaftPeerRole. Setting cached role to "
                + "{} and resetting leader info.",
            RaftProtos.RaftPeerRole.UNRECOGNIZED,
            e);
        setServerRole(null, null, null);
      }
    }
  
    /** Set the current server role and the leader peer rpc endpoint. */
    private void setServerRole(
        RaftProtos.RaftPeerRole currentRole,
        Tuple2<String, String> leaderPeerRpcEndpoint,
        Tuple2<String, String> leaderPeerInternalRpcEndpoint) {
      this.roleCheckLock.writeLock().lock();
      try {
        boolean leaderChanged = false;
        if (RaftProtos.RaftPeerRole.LEADER == currentRole && !checkCachedPeerRoleIsLeader()) {
          leaderChanged = true;
          setDeadlineTime(conf.workerHeartbeatTimeout(), conf.appHeartbeatTimeoutMs());
        } else if (RaftProtos.RaftPeerRole.LEADER != currentRole && checkCachedPeerRoleIsLeader()) {
          leaderChanged = true;
          setDeadlineTime(Integer.MAX_VALUE, Integer.MAX_VALUE); // for revoke
        }
  
        if (leaderChanged) {
          LOG.warn(
              "Raft Role changed, CurrentNode Role: {}, Leader: {}",
              currentRole,
              leaderPeerRpcEndpoint);
        }
  
        this.cachedPeerRole = Optional.ofNullable(currentRole);
        if (null != leaderPeerRpcEndpoint) {
          this.cachedLeaderPeerRpcEndpoints =
              Optional.of(
                  new LeaderPeerEndpoints(leaderPeerRpcEndpoint, leaderPeerInternalRpcEndpoint));
        } else {
          this.cachedLeaderPeerRpcEndpoints = Optional.empty();
        }
      } finally {
        this.roleCheckLock.writeLock().unlock();
      }
    }
  
    public GroupInfoReply getGroupInfo() throws IOException {
      GroupInfoRequest groupInfoRequest =
          new GroupInfoRequest(clientId, raftPeerId, RAFT_GROUP_ID, CallId.getAndIncrement());
      return server.getGroupInfo(groupInfoRequest);
    }
  
    // Exposed for testing
    public InetAddress getRaftAddress() {
      return this.ratisAddr.getAddress();
    }
  
    public int getRaftPort() {
      return this.ratisAddr.getPort();
    }
  
    public String getRpcEndpoint() {
      return this.rpcEndpoint;
    }
  
    public String getInternalRpcEndpoint() {
      return this.internalRpcEndpoint;
    }
  
    void stepDown() {
      try {
        TransferLeadershipRequest request =
            new TransferLeadershipRequest(
                clientId,
                server.getId(),
                raftGroup.getGroupId(),
                CallId.getAndIncrement(),
                null,
                REQUEST_TIMEOUT_MS);
        RaftClientReply reply = server.transferLeadership(request);
        if (reply.isSuccess()) {
          LOG.info("Successfully step down leader {}.", server.getId());
        } else {
          LOG.warn("Step down leader failed!");
        }
      } catch (Exception e) {
        LOG.warn("Step down leader failed!", e);
      }
    }
  
    public void setDeadlineTime(long increaseWorkerTime, long increaseAppTime) {
      this.workerTimeoutDeadline = System.currentTimeMillis() + increaseWorkerTime;
      this.appTimeoutDeadline = System.currentTimeMillis() + increaseAppTime;
    }
  
    public long getWorkerTimeoutDeadline() {
      return workerTimeoutDeadline;
    }
  
    public long getAppTimeoutDeadline() {
      return appTimeoutDeadline;
    }
  
    public ClientId getClientId() {
      return clientId;
    }
  
    public RaftServer getServer() {
      return server;
    }
  
    public RaftGroupId getGroupId() {
      return raftGroup.getGroupId();
    }
  
    public String getLocalAddress() {
      return localNode.ratisEndpoint();
    }
  
    public static class LeaderPeerEndpoints {
      // the rpcEndpoints Tuple2 (ip:port, host:port)
      public final Tuple2<String, String> rpcEndpoints;
  
      // the rpcInternalEndpoints Tuple2 (ip:port, host:port)
      public final Tuple2<String, String> rpcInternalEndpoints;
  
      public LeaderPeerEndpoints(
          Tuple2<String, String> rpcEndpoints, Tuple2<String, String> rpcInternalEndpoints) {
        this.rpcEndpoints = rpcEndpoints;
        this.rpcInternalEndpoints = rpcInternalEndpoints;
      }
    }
  }
  