alluxio-master-2.x/core/server/worker/src/test/java/alluxio/worker/grpc/BlockReadHandlerTest.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.worker.grpc;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.verify;

import alluxio.Constants;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.ReadRequest;
import alluxio.grpc.ReadResponse;
import alluxio.util.CommonUtils;
import alluxio.util.WaitForOptions;
import alluxio.util.io.BufferUtils;
import alluxio.worker.block.DefaultBlockWorker;
import alluxio.worker.block.io.BlockReader;
import alluxio.worker.block.io.LocalFileBlockReader;

import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.StatusException;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import io.netty.util.ResourceLeakDetector;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.mockito.stubbing.Answer;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.channels.FileChannel;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.TimeoutException;

/**
 * Unit tests for {@link BlockReadHandler}.
 */
public class BlockReadHandlerTest {
  private static final long CHUNK_SIZE =
      Configuration.getBytes(PropertyKey.USER_STREAMING_READER_CHUNK_SIZE_BYTES);
  private final Random mRandom = new Random();

  private BlockReadHandler mReadHandler;
  private ServerCallStreamObserver<ReadResponse> mResponseObserver;
  private final List<ReadResponse> mResponses = new ArrayList<>();
  private boolean mResponseCompleted;
  private Throwable mError;
  private BlockReader mBlockReader;
  private File mFile;

  @Rule
  public TemporaryFolder mTestFolder = new TemporaryFolder();

  @Before
  public void before() throws Exception {
    ResourceLeakDetector.setLevel(ResourceLeakDetector.Level.ADVANCED);

    mFile = mTestFolder.newFile();
    mBlockReader = new LocalFileBlockReader(mFile.getPath());
    mResponseObserver = Mockito.mock(ServerCallStreamObserver.class);
    Mockito.when(mResponseObserver.isReady()).thenReturn(true);
    doAnswer(args -> {
      mResponseCompleted = true;
      return null;
    }).when(mResponseObserver).onCompleted();
    doAnswer(args -> {
      mResponseCompleted = true;
      mError = args.getArgument(0, Throwable.class);
      return null;
    }).when(mResponseObserver).onError(any(Throwable.class));
    doAnswer((args) -> {
      // make a copy of response data before it is released
      mResponses.add(ReadResponse.parseFrom(
          args.getArgument(0, ReadResponse.class).toByteString()));
      return null;
    }).when(mResponseObserver).onNext(any(ReadResponse.class));
    DefaultBlockWorker blockWorker = Mockito.mock(DefaultBlockWorker.class);
    Mockito.when(blockWorker
            .createBlockReader(anyLong(), anyLong(), anyLong(), anyBoolean(), any()))
        .thenAnswer((Answer<BlockReader>) invocationOnMock -> {
          long offset = invocationOnMock.getArgument(2);
          ((FileChannel) mBlockReader.getChannel()).position(offset);
          return mBlockReader;
        });
    mReadHandler = new BlockReadHandler(GrpcExecutors.BLOCK_READER_EXECUTOR, blockWorker,
        mResponseObserver, false);
  }

  /**
   * Reads all bytes of a file.
   */
  @Test
  public void readFullFile() throws Exception {
    long checksumExpected = populateInputFile(CHUNK_SIZE * 10, 0, CHUNK_SIZE * 10 - 1);
    mReadHandler.onNext(buildReadRequest(0, CHUNK_SIZE * 10));
    checkAllReadResponses(mResponses, checksumExpected);
  }

  /**
   * Reads a sub-region of a file.
   */
  @Test
  public void readPartialFile() throws Exception {
    long start = 3;
    long end = CHUNK_SIZE * 10 - 99;
    long checksumExpected = populateInputFile(CHUNK_SIZE * 10, start, end);
    mReadHandler.onNext(buildReadRequest(start, end + 1 - start));
    checkAllReadResponses(mResponses, checksumExpected);
  }

  /**
   * Fails if the read request tries to read an empty file.
   */
  @Test
  public void readEmptyFile() throws Exception {
    populateInputFile(0, 0, 0);
    mReadHandler.onNext(buildReadRequest(0, 0));
    checkErrorCode(mResponseObserver, Status.Code.INVALID_ARGUMENT);
  }

  /**
   * Cancels the read request immediately after the read request is sent.
   */
  @Test
  public void cancelRequest() throws Exception {
    long fileSize = CHUNK_SIZE * 100 + 1;
    populateInputFile(fileSize, 0, fileSize - 1);
    mReadHandler.onNext(buildReadRequest(0, fileSize));
    mReadHandler.onCompleted();

    CommonUtils.waitFor("response", () -> mResponseCompleted || mError != null,
        WaitForOptions.defaults().setTimeoutMs(Constants.MINUTE_MS));
    verify((StreamObserver<ReadResponse>) mResponseObserver).onCompleted();
  }

  @Test
  public void ErrorReceived() {
    mReadHandler.onError(new IOException("test error"));
  }

  @Test
  public void ErrorReceivedAfterRequest() throws Exception {
    populateInputFile(CHUNK_SIZE * 10, 0, CHUNK_SIZE * 10 - 1);
    mReadHandler.onNext(buildReadRequest(0, CHUNK_SIZE * 10));
    mReadHandler.onError(new IOException("test error"));
  }

  @Test
  public void readFailure() throws Exception {
    long fileSize = CHUNK_SIZE * 10 + 1;
    populateInputFile(0, 0, fileSize - 1);
    mBlockReader.close();
    mReadHandler.onNext(buildReadRequest(0, fileSize));
    checkErrorCode(mResponseObserver, Status.Code.FAILED_PRECONDITION);
  }

  /**
   * Populates the input file, also computes the checksum for part of the file.
   *
   * @param length the length of the file
   * @param start the start position to compute the checksum
   * @param end the last position to compute the checksum
   * @return the checksum
   */
  private long populateInputFile(long length, long start, long end) throws Exception {
    long checksum = 0;
    long pos = 0;
    if (length > 0) {
      try (FileOutputStream fileOutputStream = new FileOutputStream(mFile)) {
        while (length > 0) {
          byte[] buffer = new byte[(int) Math.min(length, Constants.MB)];
          mRandom.nextBytes(buffer);
          for (byte b : buffer) {
            if (pos >= start && pos <= end) {
              checksum += BufferUtils.byteToInt(b);
            }
            pos++;
          }
          fileOutputStream.write(buffer);
          length -= buffer.length;
        }
      }
    }
    return checksum;
  }

  /**
   * Checks all the read responses.
   */
  private void checkAllReadResponses(List<ReadResponse> responses,
      long checksumExpected) throws Exception {
    long checksumActual = 0;
    CommonUtils.waitFor("response", () -> mResponseCompleted || mError != null,
        WaitForOptions.defaults().setTimeoutMs(Constants.MINUTE_MS));
    for (ReadResponse readResponse : responses) {
      if (readResponse == null) {
        Assert.fail();
        break;
      }
      assertTrue(readResponse.hasChunk());
      assertTrue(readResponse.getChunk().hasData());
      ByteString buffer = readResponse.getChunk().getData();
      if (buffer != null) {
        for (byte b : buffer) {
          checksumActual += BufferUtils.byteToInt(b);
        }
      }
    }
    assertEquals(checksumExpected, checksumActual);
  }

  /**
   * Checks an error is returned with given code.
   *
   * @param responseObserver the response stream observer
   * @param code the expected error code
   */
  private void checkErrorCode(final StreamObserver<ReadResponse> responseObserver,
      Status.Code code) throws TimeoutException, InterruptedException {
    CommonUtils.waitFor("response", () -> mResponseCompleted || mError != null,
        WaitForOptions.defaults().setTimeoutMs(Constants.MINUTE_MS));
    ArgumentCaptor<Throwable> captor = ArgumentCaptor.forClass(Throwable.class);
    verify(responseObserver).onError(captor.capture());
    Throwable t = captor.getValue();
    assertTrue(t instanceof StatusException);
    assertEquals(code, ((StatusException) t).getStatus().getCode());
  }

  private ReadRequest buildReadRequest(long offset, long len) {
    return ReadRequest.newBuilder().setBlockId(1L).setOffset(offset).setLength(len)
        .setChunkSize(CHUNK_SIZE).build();
  }
}
