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.apache.ratis.util.LifeCycle.State.PAUSED;
  import static org.apache.ratis.util.LifeCycle.State.PAUSING;
  import static org.apache.ratis.util.LifeCycle.State.RUNNING;
  import static org.apache.ratis.util.LifeCycle.State.STARTING;
  
  import java.io.File;
  import java.io.FileNotFoundException;
  import java.io.IOException;
  import java.util.Collection;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutorService;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.base.Preconditions;
  import org.apache.ratis.io.MD5Hash;
  import org.apache.ratis.proto.RaftProtos;
  import org.apache.ratis.protocol.Message;
  import org.apache.ratis.protocol.RaftClientRequest;
  import org.apache.ratis.protocol.RaftGroupId;
  import org.apache.ratis.server.RaftServer;
  import org.apache.ratis.server.protocol.TermIndex;
  import org.apache.ratis.server.raftlog.RaftLog;
  import org.apache.ratis.server.storage.FileInfo;
  import org.apache.ratis.server.storage.RaftStorage;
  import org.apache.ratis.statemachine.SnapshotInfo;
  import org.apache.ratis.statemachine.StateMachineStorage;
  import org.apache.ratis.statemachine.TransactionContext;
  import org.apache.ratis.statemachine.impl.BaseStateMachine;
  import org.apache.ratis.statemachine.impl.SimpleStateMachineStorage;
  import org.apache.ratis.statemachine.impl.SingleFileSnapshotInfo;
  import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
  import org.apache.ratis.util.ExitUtils;
  import org.apache.ratis.util.LifeCycle;
  import org.apache.ratis.util.MD5FileUtil;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.protocol.PbMetaRequest;
  import org.apache.celeborn.common.protocol.PbMetaRequestResponse;
  import org.apache.celeborn.common.util.ThreadUtils;
  import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;
  import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.ResourceResponse;
  
  public class StateMachine extends BaseStateMachine {
    private static final Logger LOG = LoggerFactory.getLogger(StateMachine.class);
  
    private final SimpleStateMachineStorage storage =
        new SimpleStateMachineStorage() {
  
          File tmpDir = null;
  
          @Override
          public void init(RaftStorage storage) throws IOException {
            super.init(storage);
            tmpDir = storage.getStorageDir().getTmpDir();
          }
  
          @Override
          public File getTmpDir() {
            return tmpDir;
          }
        };
  
    private final HARaftServer masterRatisServer;
    private RaftGroupId raftGroupId;
    private final ExecutorService executorService;
  
    private RaftServer mServer;
  
    private final MetaHandler metaHandler;
  
    public StateMachine(HARaftServer ratisServer) {
      this.masterRatisServer = ratisServer;
      this.metaHandler = ratisServer.getMetaHandler();
  
      this.executorService =
          ThreadUtils.newDaemonSingleThreadExecutor("master-state-machine-executor");
    }
  
    /** Initializes the State Machine with the given server, group and storage. */
    @Override
    public void initialize(RaftServer server, RaftGroupId id, RaftStorage raftStorage)
        throws IOException {
      getLifeCycle()
          .startAndTransition(
              () -> {
                super.initialize(server, id, raftStorage);
                this.mServer = server;
                this.raftGroupId = id;
                storage.init(raftStorage);
              });
      loadSnapshot(storage.getLatestSnapshot());
      LOG.info("Initialized State Machine.");
    }
  
    @Override
    public void reinitialize() throws IOException {
      LOG.info("Reinitializing state machine.");
      getLifeCycle().compareAndTransition(PAUSED, STARTING);
      storage.loadLatestSnapshot();
      loadSnapshot(storage.getLatestSnapshot());
      getLifeCycle().compareAndTransition(STARTING, RUNNING);
    }
  
    @Override
    public void pause() {
      getLifeCycle().compareAndTransition(RUNNING, PAUSING);
      getLifeCycle().compareAndTransition(PAUSING, PAUSED);
    }
  
    private synchronized void loadSnapshot(SingleFileSnapshotInfo snapshot) throws IOException {
      if (snapshot == null) {
        return;
      }
      if (snapshot.getTermIndex().compareTo(getLastAppliedTermIndex()) <= 0) {
        LOG.info("obsolete snapshot provided: {}", snapshot.getTermIndex());
        return;
      }
      LOG.info("Loading Snapshot {}.", snapshot);
      final File snapshotFile = snapshot.getFile().getPath().toFile();
      if (!snapshotFile.exists()) {
        throw new FileNotFoundException(
            String.format("The snapshot file %s does not exist", snapshotFile.getPath()));
      }
      try {
        setLastAppliedTermIndex(snapshot.getTermIndex());
        install(snapshotFile);
      } catch (IOException rethrow) {
        LOG.error("Failed to load snapshot {}", snapshot);
        throw rethrow;
      }
    }
  
    private void install(File snapshotFile) throws IOException {
      try {
        metaHandler.loadSnapShot(snapshotFile);
      } catch (IOException rethrow) {
        LOG.warn("Failed to install snapshot!", rethrow);
        throw rethrow;
      }
      LOG.info("Successfully installed snapshot!");
    }
  
    @Override
    public SnapshotInfo getLatestSnapshot() {
      return storage.getLatestSnapshot();
    }
  
    @Override
    public TransactionContext startTransaction(RaftClientRequest raftClientRequest)
        throws IOException {
      Preconditions.checkArgument(raftClientRequest.getRaftGroupId().equals(raftGroupId));
      return handleStartTransactionRequests(raftClientRequest);
    }
  
    /** Apply a committed log entry to the state machine. */
    @Override
    public CompletableFuture<Message> applyTransaction(TransactionContext trx) {
      try {
        PbMetaRequest request =
            HAHelper.convertByteStringToRequest(trx.getStateMachineLogEntry().getLogData());
        long trxLogIndex = trx.getLogEntry().getIndex();
        // In the current approach we have one single global thread executor.
        // with single thread. Right now this is being done for correctness, as
        // applyTransaction will be run on multiple Master we want to execute the
        // transactions in the same order on all Master, otherwise there is a
        // chance that Master replica can be out of sync.
        // Ref: from Ozone project (OzoneManagerStateMachine)
        CompletableFuture<Message> ratisFuture = new CompletableFuture<>();
        CompletableFuture<PbMetaRequestResponse> future =
            CompletableFuture.supplyAsync(() -> runCommand(request, trxLogIndex), executorService);
        future.thenApply(
            response -> {
              if (!response.getSuccess()) {
                LOG.warn(
                    "Failed to apply log {} for this raft group {}!",
                    request.getMetaRequestType(),
                    this.raftGroupId);
              }
  
              byte[] responseBytes = response.toByteArray();
              ratisFuture.complete(Message.valueOf(ByteString.copyFrom(responseBytes)));
              return ratisFuture;
            });
        return ratisFuture;
      } catch (Exception e) {
        return completeExceptionally(e);
      }
    }
  
    /**
     * Submits write request to MetaSystem and returns the response Message.
     *
     * @param request MasterMetaRequest
     * @return response from meta system
     */
    @VisibleForTesting
    protected PbMetaRequestResponse runCommand(PbMetaRequest request, long trxLogIndex) {
      try {
        return metaHandler.handleWriteRequest(request);
      } catch (Throwable e) {
        String errorMessage = "Request " + request + "failed with exception";
        ExitUtils.terminate(1, errorMessage, e, LOG);
      }
      return null;
    }
  
    /** Query the state machine. The request must be read-only. */
    @Override
    public CompletableFuture<Message> query(Message request) {
      try {
        byte[] bytes = request.getContent().toByteArray();
        return CompletableFuture.completedFuture(
            queryCommand(ResourceProtos.ResourceRequest.parseFrom(bytes)));
      } catch (IOException e) {
        return completeExceptionally(e);
      }
    }
  
    /**
     * Submits read request to MetaSystem and returns the response Message.
     *
     * @param request MasterMetaRequest
     * @return response from meta system
     */
    private Message queryCommand(ResourceProtos.ResourceRequest request) {
      ResourceResponse response = metaHandler.handleReadRequest(request);
      return HAHelper.convertResponseToMessage(response);
    }
  
    /**
     * Store the current state as a snapshot file in the stateMachineStorage.
     *
     * @return the index of the snapshot
     */
    @Override
    public long takeSnapshot() {
      if (mServer.getLifeCycleState() != LifeCycle.State.RUNNING) {
        LOG.warn(
            "Skip taking snapshot because raft server is not in running state: "
                + "current state is {}.",
            mServer.getLifeCycleState());
        return RaftLog.INVALID_LOG_INDEX;
      }
      TermIndex lastTermIndex = getLastAppliedTermIndex();
      LOG.debug("Current Snapshot Index {}.", lastTermIndex);
      File tempFile;
      try {
        tempFile = HAHelper.createTempSnapshotFile(storage);
        metaHandler.writeToSnapShot(tempFile);
      } catch (IOException e) {
        LOG.warn("Failed to create temp snapshot file.", e);
        return RaftLog.INVALID_LOG_INDEX;
      }
      LOG.info("Taking a snapshot to file {}.", tempFile);
      final File snapshotFile =
          storage.getSnapshotFile(lastTermIndex.getTerm(), lastTermIndex.getIndex());
      try {
        final MD5Hash digest = MD5FileUtil.computeMd5ForFile(tempFile);
        LOG.info("Saving digest {} for snapshot file {}.", digest, snapshotFile);
        MD5FileUtil.saveMD5File(snapshotFile, digest);
        LOG.info("Renaming a snapshot file {} to {}.", tempFile, snapshotFile);
        if (!tempFile.renameTo(snapshotFile)) {
          tempFile.delete();
          LOG.warn("Failed to rename snapshot from {} to {}.", tempFile, snapshotFile);
          return RaftLog.INVALID_LOG_INDEX;
        }
  
        // update storage
        final FileInfo info = new FileInfo(snapshotFile.toPath(), digest);
        storage.updateLatestSnapshot(new SingleFileSnapshotInfo(info, lastTermIndex));
      } catch (Exception e) {
        tempFile.delete();
        LOG.warn("Failed to complete snapshot: {}.", snapshotFile, e);
        return RaftLog.INVALID_LOG_INDEX;
      }
      return lastTermIndex.getIndex();
    }
  
    @Override
    public void notifyLogFailed(Throwable cause, RaftProtos.LogEntryProto failedEntry) {
      LOG.warn("Log failed {}", cause.getMessage());
      if (masterRatisServer != null && masterRatisServer.isLeader()) {
        masterRatisServer.stepDown();
      }
    }
  
    /** Notifies the state machine that the raft peer is no longer leader. */
    @Override
    public void notifyNotLeader(Collection<TransactionContext> pendingEntries) throws IOException {
      masterRatisServer.updateServerRole();
    }
  
    /**
     * Handle the RaftClientRequest and return TransactionContext object.
     *
     * @param raftClientRequest
     * @return TransactionContext
     */
    private TransactionContext handleStartTransactionRequests(RaftClientRequest raftClientRequest) {
  
      return TransactionContext.newBuilder()
          .setClientRequest(raftClientRequest)
          .setStateMachine(this)
          .setServerRole(RaftProtos.RaftPeerRole.LEADER)
          .setLogData(raftClientRequest.getMessage().getContent())
          .build();
    }
  
    @VisibleForTesting
    public void setRaftGroupId(RaftGroupId raftGroupId) {
      this.raftGroupId = raftGroupId;
    }
  
    private static <T> CompletableFuture<T> completeExceptionally(Exception e) {
      final CompletableFuture<T> future = new CompletableFuture<>();
      future.completeExceptionally(e);
      return future;
    }
  
    @Override
    public StateMachineStorage getStateMachineStorage() {
      return this.storage;
    }
  }
  