alluxio-master-2.x/core/server/common/src/main/java/alluxio/master/journal/raft/RaftJournalServiceHandler.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.master.journal.raft;

import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.LatestSnapshotInfoPRequest;
import alluxio.grpc.RaftJournalServiceGrpc;
import alluxio.grpc.SnapshotData;
import alluxio.grpc.SnapshotMetadata;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.util.compression.DirectoryMarshaller;

import com.google.protobuf.ByteString;
import com.google.protobuf.UnsafeByteOperations;
import io.grpc.Context;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import org.apache.ratis.server.protocol.TermIndex;
import org.apache.ratis.statemachine.SnapshotInfo;
import org.apache.ratis.statemachine.StateMachineStorage;
import org.apache.ratis.statemachine.impl.SimpleStateMachineStorage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.OutputStream;
import java.nio.file.Path;
import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.TimeUnit;

/**
 * RPC handler for raft journal service.
 */
public class RaftJournalServiceHandler extends RaftJournalServiceGrpc.RaftJournalServiceImplBase {
  private static final Logger LOG = LoggerFactory.getLogger(RaftJournalServiceHandler.class);

  private final StateMachineStorage mStateMachineStorage;
  private volatile long mLastSnapshotUploadDurationMs = -1;
  private volatile long mLastSnapshotUploadSize = -1;
  private volatile long mLastSnapshotUploadDiskSize = -1;

  /**
   * @param storage the storage that the state machine uses for its snapshots
   */
  public RaftJournalServiceHandler(StateMachineStorage storage) {
    mStateMachineStorage = storage;

    MetricsSystem.registerGaugeIfAbsent(
        MetricKey.MASTER_EMBEDDED_JOURNAL_LAST_SNAPSHOT_UPLOAD_DURATION_MS.getName(),
        () -> mLastSnapshotUploadDurationMs);
    MetricsSystem.registerGaugeIfAbsent(
        MetricKey.MASTER_EMBEDDED_JOURNAL_LAST_SNAPSHOT_UPLOAD_SIZE.getName(),
        () -> mLastSnapshotUploadSize);
    MetricsSystem.registerGaugeIfAbsent(
        MetricKey.MASTER_EMBEDDED_JOURNAL_LAST_SNAPSHOT_UPLOAD_DISK_SIZE.getName(),
        () -> mLastSnapshotUploadDiskSize);
  }

  @Override
  public void requestLatestSnapshotInfo(LatestSnapshotInfoPRequest request,
                                        StreamObserver<SnapshotMetadata> responseObserver) {
    LOG.info("Received request for latest snapshot info");
    if (Context.current().isCancelled()) {
      responseObserver.onError(
          Status.CANCELLED.withDescription("Cancelled by client").asRuntimeException());
      return;
    }
    SnapshotInfo snapshot = mStateMachineStorage.getLatestSnapshot();
    SnapshotMetadata.Builder metadata = SnapshotMetadata.newBuilder();
    if (snapshot == null) {
      LOG.info("No snapshot to send");
      metadata.setExists(false);
    } else {
      LOG.info("Found snapshot {}", snapshot.getTermIndex());
      metadata.setExists(true)
          .setSnapshotTerm(snapshot.getTerm())
          .setSnapshotIndex(snapshot.getIndex());
    }
    responseObserver.onNext(metadata.build());
    responseObserver.onCompleted();
  }

  @Override
  public void requestLatestSnapshotData(SnapshotMetadata request,
                                        StreamObserver<SnapshotData> responseObserver) {
    TermIndex index = TermIndex.valueOf(request.getSnapshotTerm(), request.getSnapshotIndex());
    LOG.info("Received request for snapshot data {}", index);
    if (Context.current().isCancelled()) {
      responseObserver.onError(
          Status.CANCELLED.withDescription("Cancelled by client").asRuntimeException());
      return;
    }

    String snapshotDirName = SimpleStateMachineStorage
        .getSnapshotFileName(request.getSnapshotTerm(), request.getSnapshotIndex());
    Path snapshotPath = new File(mStateMachineStorage.getSnapshotDir(), snapshotDirName).toPath();
    long totalBytesSent;
    long diskSize;
    LOG.info("Begin snapshot upload of {}", index);
    Instant start = Instant.now();
    try (SnapshotGrpcOutputStream stream = new SnapshotGrpcOutputStream(responseObserver)) {
      DirectoryMarshaller marshaller = DirectoryMarshaller.Factory.create();
      diskSize = marshaller.write(snapshotPath, stream);
      totalBytesSent = stream.totalBytes();
    } catch (Exception e) {
      LOG.warn("Failed to upload snapshot {}", index, e);
      responseObserver.onError(Status.INTERNAL.withCause(e).asRuntimeException());
      return;
    }
    responseObserver.onCompleted();
    // update last duration and duration timer metrics
    mLastSnapshotUploadDurationMs = Duration.between(start, Instant.now()).toMillis();
    MetricsSystem.timer(MetricKey.MASTER_EMBEDDED_JOURNAL_SNAPSHOT_UPLOAD_TIMER.getName())
        .update(mLastSnapshotUploadDurationMs, TimeUnit.MILLISECONDS);
    LOG.info("Total milliseconds to upload {}: {}", index, mLastSnapshotUploadDurationMs);
    // update uncompressed snapshot size metric
    mLastSnapshotUploadDiskSize = diskSize;
    MetricsSystem.histogram(
            MetricKey.MASTER_EMBEDDED_JOURNAL_SNAPSHOT_UPLOAD_DISK_HISTOGRAM.getName())
        .update(mLastSnapshotUploadDiskSize);
    LOG.info("Total snapshot uncompressed bytes for {}: {}", index, mLastSnapshotUploadDiskSize);
    // update compressed snapshot size (aka size sent over the network)
    mLastSnapshotUploadSize = totalBytesSent;
    MetricsSystem.histogram(MetricKey.MASTER_EMBEDDED_JOURNAL_SNAPSHOT_UPLOAD_HISTOGRAM.getName())
        .update(mLastSnapshotUploadSize);
    LOG.info("Total bytes sent for {}: {}", index, mLastSnapshotUploadSize);
    LOG.info("Uploaded snapshot {}", index);
  }

  static class SnapshotGrpcOutputStream extends OutputStream {
    private final int mSnapshotReplicationChunkSize = (int) Configuration.getBytes(
        PropertyKey.MASTER_EMBEDDED_JOURNAL_SNAPSHOT_REPLICATION_CHUNK_SIZE);
    private final StreamObserver<SnapshotData> mObserver;
    private long mTotalBytesSent = 0;
    private byte[] mBuffer = new byte[mSnapshotReplicationChunkSize];
    private int mBufferPosition = 0;

    public SnapshotGrpcOutputStream(StreamObserver<SnapshotData> responseObserver) {
      mObserver = responseObserver;
    }

    @Override
    public void write(int b) {
      mBuffer[mBufferPosition++] = (byte) b;
      if (mBufferPosition == mBuffer.length) {
        flushBuffer();
      }
    }

    @Override
    public void close() {
      if (mBufferPosition > 0) {
        flushBuffer();
      }
    }

    private void flushBuffer() {
      // avoids copy
      ByteString bytes = UnsafeByteOperations.unsafeWrap(mBuffer, 0, mBufferPosition);
      mBuffer = new byte[mSnapshotReplicationChunkSize];
      LOG.debug("Sending chunk of size {}: {}", mBufferPosition, bytes);
      mObserver.onNext(SnapshotData.newBuilder().setChunk(bytes).build());
      mTotalBytesSent += mBufferPosition;
      mBufferPosition = 0;
    }

    public long totalBytes() {
      return mTotalBytesSent + mBufferPosition;
    }
  }
}
