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

import static junit.framework.TestCase.assertTrue;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

import alluxio.AlluxioTestDirectory;
import alluxio.AlluxioURI;
import alluxio.ConfigurationRule;
import alluxio.Constants;
import alluxio.Sessions;
import alluxio.client.block.stream.BlockWorkerDataReader;
import alluxio.client.block.stream.DataReader;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.URIStatus;
import alluxio.client.file.options.InStreamOptions;
import alluxio.conf.Configuration;
import alluxio.conf.InstancedConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.OpenFilePOptions;
import alluxio.grpc.ReadPType;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.underfs.UfsManager;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.UnderFileSystemConfiguration;
import alluxio.util.FileSystemOptionsUtils;
import alluxio.util.io.BufferUtils;
import alluxio.wire.BlockInfo;
import alluxio.wire.FileBlockInfo;
import alluxio.wire.FileInfo;
import alluxio.worker.block.BlockMasterClient;
import alluxio.worker.block.BlockMasterClientPool;
import alluxio.worker.block.BlockWorker;
import alluxio.worker.block.CreateBlockOptions;
import alluxio.worker.block.DefaultBlockWorker;
import alluxio.worker.block.MonoBlockStore;
import alluxio.worker.block.TieredBlockStore;
import alluxio.worker.block.io.BlockWriter;
import alluxio.worker.file.FileSystemMasterClient;

import com.google.common.collect.ImmutableMap;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

import java.io.File;
import java.io.IOException;
import java.util.Collections;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Unit tests for {@link BlockWorkerDataReader}.
 */
public class BlockWorkerDataReaderTest {
  // Use 2L block id here since block id includes the special sequence number info
  private static final long BLOCK_ID = 2L;
  private static final int CHUNK_SIZE = 128;
  private static final long SESSION_ID = 10L;
  private static final int LOCK_NUM = 5;

  private final InstancedConfiguration mConf = Configuration.modifiableGlobal();
  private final String mMemDir =
      AlluxioTestDirectory.createTemporaryDirectory(Constants.MEDIUM_MEM).getAbsolutePath();

  private BlockWorker mBlockWorker;
  private BlockWorkerDataReader.Factory mDataReaderFactory;
  private String mRootUfs;

  @Rule
  public ConfigurationRule mConfigurationRule =
      new ConfigurationRule(new ImmutableMap.Builder<PropertyKey, Object>()
          .put(PropertyKey.WORKER_TIERED_STORE_LEVELS, 1)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_ALIAS, Constants.MEDIUM_MEM)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_MEDIUMTYPE, Constants.MEDIUM_MEM)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_QUOTA, "1GB")
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_PATH, mMemDir)
          .put(PropertyKey.WORKER_TIERED_STORE_BLOCK_LOCKS, LOCK_NUM)
          .put(PropertyKey.WORKER_RPC_PORT, 0)
          .put(PropertyKey.MASTER_MOUNT_TABLE_ROOT_UFS,
              AlluxioTestDirectory.createTemporaryDirectory("BlockWorkerDataReaderTest")
                  .getAbsolutePath()).build(), Configuration.modifiableGlobal());

  @Before
  public void before() throws Exception {
    BlockMasterClient blockMasterClient = mock(BlockMasterClient.class);
    BlockMasterClientPool blockMasterClientPool = spy(new BlockMasterClientPool());
    when(blockMasterClientPool.createNewResource()).thenReturn(blockMasterClient);

    FileSystemMasterClient fileSystemMasterClient = mock(FileSystemMasterClient.class);
    Sessions sessions = mock(Sessions.class);

    // Connect to the real UFS for UFS read testing
    UfsManager ufsManager = mock(UfsManager.class);
    mRootUfs = Configuration.getString(PropertyKey.MASTER_MOUNT_TABLE_ROOT_UFS);
    UfsManager.UfsClient ufsClient = new UfsManager.UfsClient(
        () -> UnderFileSystem.Factory.create(mRootUfs,
            UnderFileSystemConfiguration.defaults(Configuration.global())),
        new AlluxioURI(mRootUfs));
    when(ufsManager.get(anyLong())).thenReturn(ufsClient);

    TieredBlockStore tieredBlockStore = new TieredBlockStore();
    AtomicReference<Long> workerId = new AtomicReference<>(-1L);
    MonoBlockStore blockStore =
        new MonoBlockStore(tieredBlockStore, blockMasterClientPool, ufsManager, workerId);
    mBlockWorker = new DefaultBlockWorker(blockMasterClientPool, fileSystemMasterClient,
        sessions, blockStore, workerId);

    URIStatus dummyStatus =
        new URIStatus(new FileInfo().setBlockIds(Collections.singletonList(BLOCK_ID)));
    InStreamOptions options =
        new InStreamOptions(dummyStatus, FileSystemOptionsUtils.openFileDefaults(mConf), mConf,
            FileSystemContext.create(mConf));
    mDataReaderFactory =
        new BlockWorkerDataReader.Factory(mBlockWorker, BLOCK_ID, CHUNK_SIZE, options);
  }

  @Test
  public void createWithBlockNotExists() {
    assertThrows(IOException.class, () -> mDataReaderFactory.create(BLOCK_ID, 100));
  }

  @Test
  public void create() throws Exception {
    mBlockWorker.createBlock(SESSION_ID, BLOCK_ID, 0,
        new CreateBlockOptions(null, Constants.MEDIUM_MEM, 1));
    mBlockWorker.commitBlock(SESSION_ID, BLOCK_ID, true);
    DataReader dataReader = mDataReaderFactory.create(100, 200);
    assertEquals(100, dataReader.pos());
  }

  // See https://github.com/Alluxio/alluxio/issues/13255
  @Test
  public void createAndCloseManyReader() throws Exception {
    for (int i = 0; i < LOCK_NUM * 10; i++) {
      long blockId = i;
      mBlockWorker.createBlock(SESSION_ID, blockId, 0,
          new CreateBlockOptions(null, Constants.MEDIUM_MEM, 1));
      mBlockWorker.commitBlock(SESSION_ID, blockId, true);
      InStreamOptions inStreamOptions = new InStreamOptions(
          new URIStatus(new FileInfo().setBlockIds(Collections.singletonList(blockId))),
          FileSystemOptionsUtils.openFileDefaults(mConf), mConf, FileSystemContext.create(mConf));
      mDataReaderFactory =
          new BlockWorkerDataReader.Factory(mBlockWorker, blockId, CHUNK_SIZE, inStreamOptions);
      DataReader dataReader = mDataReaderFactory.create(0, 100);
      dataReader.close();
    }
  }

  @Test
  public void readChunkUfs() throws Exception {
    // Write an actual file to UFS
    String testFilePath = File.createTempFile("temp", null, new File(mRootUfs)).getAbsolutePath();
    byte[] buffer = BufferUtils.getIncreasingByteArray(CHUNK_SIZE * 5);
    BufferUtils.writeBufferToFile(testFilePath, buffer);

    BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID).setLength(CHUNK_SIZE  * 5);
    URIStatus dummyStatus = new URIStatus(new FileInfo().setPersisted(true)
        .setUfsPath(testFilePath)
        .setBlockIds(Collections.singletonList(BLOCK_ID))
        .setLength(CHUNK_SIZE * 5)
        .setBlockSizeBytes(CHUNK_SIZE)
        .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info))));
    OpenFilePOptions readOptions = OpenFilePOptions.newBuilder()
        .setReadType(ReadPType.NO_CACHE).build();
    InStreamOptions options = new InStreamOptions(dummyStatus, readOptions, mConf,
        FileSystemContext.create(mConf));

    BlockWorkerDataReader.Factory factory = new BlockWorkerDataReader
        .Factory(mBlockWorker, BLOCK_ID, CHUNK_SIZE, options);
    int len = CHUNK_SIZE * 3 / 2;
    try (DataReader dataReader = factory.create(0, len)) {
      validateBuffer(dataReader.readChunk(), 0, CHUNK_SIZE);
      assertEquals(CHUNK_SIZE, dataReader.pos());
      validateBuffer(dataReader.readChunk(), CHUNK_SIZE, len - CHUNK_SIZE);
      assertEquals(len, dataReader.pos());
    }
  }

  @Test
  public void readChunkFullFile() throws Exception {
    int len = CHUNK_SIZE * 2;
    mBlockWorker.createBlock(SESSION_ID, BLOCK_ID, 0,
        new CreateBlockOptions(null, Constants.MEDIUM_MEM, 1));
    try (BlockWriter writer = mBlockWorker.createBlockWriter(SESSION_ID, BLOCK_ID)) {
      writer.append(BufferUtils.getIncreasingByteBuffer(len));
    }
    mBlockWorker.commitBlock(SESSION_ID, BLOCK_ID, true);
    DataReader dataReader = mDataReaderFactory.create(0, len);
    validateBuffer(dataReader.readChunk(), 0, CHUNK_SIZE);
    assertEquals(CHUNK_SIZE, dataReader.pos());
    validateBuffer(dataReader.readChunk(), CHUNK_SIZE, CHUNK_SIZE);
    assertEquals(len, dataReader.pos());
    dataReader.close();
  }

  @Test
  public void readChunkPartial() throws Exception {
    int len = CHUNK_SIZE * 5;
    mBlockWorker.createBlock(SESSION_ID, BLOCK_ID, 0,
        new CreateBlockOptions(null, Constants.MEDIUM_MEM, 1));
    try (BlockWriter writer = mBlockWorker.createBlockWriter(SESSION_ID, BLOCK_ID)) {
      writer.append(BufferUtils.getIncreasingByteBuffer(len));
    }
    mBlockWorker.commitBlock(SESSION_ID, BLOCK_ID, true);
    int start = len / 5 * 2;
    int end = len / 5 * 4;
    DataReader dataReader = mDataReaderFactory.create(start, end);
    for (int s = start; s < end; s += CHUNK_SIZE) {
      int currentLen = Math.min(CHUNK_SIZE, end - s);
      validateBuffer(dataReader.readChunk(), s, currentLen);
    }
  }

  private void validateBuffer(DataBuffer buffer, int start, int len) {
    byte[] bytes = new byte[buffer.readableBytes()];
    buffer.readBytes(bytes, 0, buffer.readableBytes());
    assertTrue(BufferUtils.equalIncreasingByteArray(start, len, bytes));
  }
}
