alluxio-master-2.x/core/client/fs/src/test/java/alluxio/client/block/stream/SharedGrpcDataReaderTest.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.client.block.stream;

import alluxio.Constants;
import alluxio.collections.ConcurrentHashSet;
import alluxio.grpc.ReadRequest;
import alluxio.grpc.ReadResponse;
import alluxio.wire.WorkerNetAddress;

import com.google.common.base.Throwables;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;

/**
 * Tests {@link SharedGrpcDataReaderTest}.
 */
public final class SharedGrpcDataReaderTest {
  private static final long BLOCK_ID = 1L;
  private static final int CHUNK_SIZE = 1024;
  private static final long TIMEOUT = 10 * Constants.SECOND_MS;

  private int mBlockSize = CHUNK_SIZE * 7 + CHUNK_SIZE / 3;
  private TestBufferCachingGrpcDataReader mBufferCachingDataReader;
  private ReadRequest mReadRequest;

  @Before
  public void before() {
    WorkerNetAddress address = new WorkerNetAddress();
    BlockWorkerClient client = Mockito.mock(BlockWorkerClient.class);
    GrpcBlockingStream<ReadRequest, ReadResponse> unusedStream
        = new GrpcBlockingStream<>(client::readBlock, 5, "test message");
    mReadRequest = ReadRequest.newBuilder().setOffset(0).setLength(mBlockSize)
        .setChunkSize(CHUNK_SIZE).setBlockId(BLOCK_ID).build();
    mBufferCachingDataReader = new TestBufferCachingGrpcDataReader(address,
        new NoopClosableResource<>(client), TIMEOUT,
        mReadRequest, unusedStream, CHUNK_SIZE, mBlockSize);
  }

  @Test
  public void singleThreadReadFullFile() throws Exception {
    SharedGrpcDataReader sharedReader
        = new SharedGrpcDataReader(mReadRequest, mBufferCachingDataReader);
    for (int i = 0; i < getChunkNum(mBlockSize); i++) {
      Assert.assertTrue(mBufferCachingDataReader.validateBuffer(i, sharedReader.readChunk()));
    }
    Assert.assertNull(sharedReader.readChunk());
  }

  @Test(timeout = 1000 * 60)
  public void singleThreadReadPartialFile() throws Exception {
    int[][] moves = {{10, 2438}, {1024, 2096}, {2042, 1222}, {3057, 2}};

    for (int[] offsetAndLen : moves) {
      int offset = offsetAndLen[0];
      int len = offsetAndLen[1];
      ReadRequest partialReadRequest = ReadRequest.newBuilder().setOffset(offset).setLength(len)
          .setBlockId(mReadRequest.getBlockId()).setChunkSize(mReadRequest.getChunkSize()).build();
      SharedGrpcDataReader sharedReader
          = new SharedGrpcDataReader(partialReadRequest, mBufferCachingDataReader);
      while (offset != -1) {
        offset = validateRead(sharedReader, offset, getChunkNum(len));
      }
    }
  }

  @Test(timeout = 1000 * 60)
  public void twoThreadSequentialRead() throws Exception {
    int readerOneOffset = 10;
    int readerOneLen = CHUNK_SIZE * 6 + CHUNK_SIZE / 3;
    int readerTwoOffset = 80;
    int readerTwoLen = CHUNK_SIZE * 5 + CHUNK_SIZE / 2;
    ReadRequest readRequestOne = ReadRequest.newBuilder()
        .setBlockId(BLOCK_ID).setChunkSize(CHUNK_SIZE)
        .setOffset(readerOneOffset).setLength(readerOneLen).build();
    ReadRequest readRequestTwo = ReadRequest.newBuilder()
        .setBlockId(BLOCK_ID).setChunkSize(CHUNK_SIZE)
        .setOffset(readerTwoOffset).setLength(readerTwoLen).build();
    SharedGrpcDataReader sharedReaderOne
        = new SharedGrpcDataReader(readRequestOne, mBufferCachingDataReader);
    SharedGrpcDataReader sharedReaderTwo
        = new SharedGrpcDataReader(readRequestTwo, mBufferCachingDataReader);
    boolean readerOne = true;
    while (readerOneOffset != -1 && readerTwoOffset != -1) {
      if (readerOne) {
        readerOneOffset = validateRead(sharedReaderOne, readerOneOffset, getChunkNum(readerOneLen));
      } else {
        readerTwoOffset = validateRead(sharedReaderTwo, readerTwoOffset, getChunkNum(readerTwoLen));
      }
      readerOne = !readerOne;
    }
  }

  /**
   * 10 threads read from the same file with the shared cache concurrently.
   */
  @Test(timeout = 1000 * 60)
  public void MultiThreadConcurrentRead() throws Exception {
    int concurrency = 10;
    List<Thread> threads = new ArrayList<>(concurrency);
    // If there are exceptions, we will store them here
    final ConcurrentHashSet<Throwable> errors = new ConcurrentHashSet<>();
    Thread.UncaughtExceptionHandler exceptionHandler = (th, ex) -> errors.add(ex);
    for (int i = 0; i < concurrency; i++) {
      Thread t = new Thread(() -> {
        Random random = new Random();
        int offset = random.nextInt(mBlockSize);
        int len = random.nextInt(mBlockSize - offset);
        ReadRequest partialReadRequest = ReadRequest.newBuilder()
            .setOffset(offset).setLength(len).setBlockId(mReadRequest.getBlockId())
            .setChunkSize(mReadRequest.getChunkSize()).build();
        try (SharedGrpcDataReader reader
                 = new SharedGrpcDataReader(partialReadRequest, mBufferCachingDataReader)) {
          while (offset != -1) {
            offset = validateRead(reader, offset, getChunkNum(len));
          }
        } catch (Exception e) {
          Throwables.throwIfUnchecked(e);
        }
      });
      t.setUncaughtExceptionHandler(exceptionHandler);
      threads.add(t);
    }
    Collections.shuffle(threads);
    for (Thread t : threads) {
      t.start();
    }
    for (Thread t : threads) {
      t.join();
    }
  }

  private int validateRead(SharedGrpcDataReader reader, int offset, int chunkNum)
      throws Exception {
    int index = offset / CHUNK_SIZE;
    Assert.assertTrue(mBufferCachingDataReader.validateBuffer(offset,
        (index + 1) * CHUNK_SIZE - offset, reader.readChunk()));
    return index >= chunkNum ? -1 : (index + 1) * CHUNK_SIZE;
  }

  private int getChunkNum(int len) {
    if (len % CHUNK_SIZE == 0) {
      return len / CHUNK_SIZE;
    }
    return len / CHUNK_SIZE + 1;
  }
}
