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

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.spy;

import alluxio.AlluxioTestDirectory;
import alluxio.AlluxioURI;
import alluxio.ConfigurationRule;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.runtime.ResourceExhaustedRuntimeException;
import alluxio.metrics.MetricInfo;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.proto.dataserver.Protocol;
import alluxio.underfs.UfsManager;
import alluxio.underfs.UfsManager.UfsClient;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.UnderFileSystemConfiguration;
import alluxio.util.io.BufferUtils;
import alluxio.worker.block.io.BlockReader;
import alluxio.worker.block.meta.UnderFileSystemBlockMeta;

import com.codahale.metrics.Counter;
import com.codahale.metrics.Meter;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.PooledByteBufAllocator;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Optional;

public final class UnderFileSystemBlockReaderTest {
  private static final long TEST_BLOCK_SIZE = 1024;
  private static final long SESSION_ID = 1;
  private static final long BLOCK_ID = 2;

  private UnderFileSystemBlockReader mReader;
  private LocalBlockStore mAlluxioBlockStore;
  private UnderFileSystemBlockMeta mUnderFileSystemBlockMeta;
  private UfsManager.UfsClient mUfsClient;
  private UfsInputStreamCache mUfsInstreamCache;
  private Protocol.OpenUfsBlockOptions mOpenUfsBlockOptions;
  private Counter mUfsBytesRead;
  private Meter mUfsBytesReadThroughput;

  @Rule
  public TemporaryFolder mFolder = new TemporaryFolder();

  @Rule
  public ConfigurationRule mConfigurationRule =
      new ConfigurationRule(new HashMap<PropertyKey, Object>() {
        {
          put(PropertyKey.MASTER_MOUNT_TABLE_ROOT_UFS, AlluxioTestDirectory
              .createTemporaryDirectory("UnderFileSystemBlockReaderTest-RootUfs")
              .getAbsolutePath());
          // ensure tiered storage uses different tmp dir for each test case
          put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_PATH, AlluxioTestDirectory
              .createTemporaryDirectory("UnderFileSystemBlockReaderTest-WorkerDataFolder")
              .getAbsolutePath());
          put(PropertyKey.WORKER_TIERED_STORE_LEVELS, 1);
        }
      }, Configuration.modifiableGlobal());

  @Before
  public void before() throws Exception {
    String ufsFolder = Configuration.getString(PropertyKey.MASTER_MOUNT_TABLE_ROOT_UFS);
    String testFilePath = File.createTempFile("temp", null, new File(ufsFolder)).getAbsolutePath();
    byte[] buffer = BufferUtils.getIncreasingByteArray((int) TEST_BLOCK_SIZE * 2);
    BufferUtils.writeBufferToFile(testFilePath, buffer);

    mAlluxioBlockStore = new TieredBlockStore();
    mUfsInstreamCache = new UfsInputStreamCache();
    mUfsClient = new UfsClient(
        () -> UnderFileSystem.Factory.create(testFilePath,
            UnderFileSystemConfiguration.defaults(Configuration.global())),
        new AlluxioURI(testFilePath));

    mOpenUfsBlockOptions = Protocol.OpenUfsBlockOptions.newBuilder().setMaxUfsReadConcurrency(10)
        .setBlockSize(TEST_BLOCK_SIZE).setOffsetInFile(TEST_BLOCK_SIZE).setUfsPath(testFilePath)
        .build();
    mUnderFileSystemBlockMeta =
        new UnderFileSystemBlockMeta(SESSION_ID, BLOCK_ID, mOpenUfsBlockOptions);
    mUfsBytesRead = MetricsSystem.counterWithTags(
        MetricKey.WORKER_BYTES_READ_UFS.getName(),
        MetricKey.WORKER_BYTES_READ_UFS.isClusterAggregated(),
        MetricInfo.TAG_UFS, MetricsSystem.escape(mUfsClient.getUfsMountPointUri()));
    mUfsBytesReadThroughput = MetricsSystem.meterWithTags(
        MetricKey.WORKER_BYTES_READ_UFS_THROUGHPUT.getName(),
        MetricKey.WORKER_BYTES_READ_UFS_THROUGHPUT.isClusterAggregated(),
        MetricInfo.TAG_UFS,
        MetricsSystem.escape(mUfsClient.getUfsMountPointUri()));
  }

  private void checkTempBlock(long start, long length) throws Exception {
    Assert.assertTrue(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
    mAlluxioBlockStore.commitBlock(SESSION_ID, BLOCK_ID, false);
    Optional<BlockLock> lock = mAlluxioBlockStore.pinBlock(SESSION_ID, BLOCK_ID);
    BlockReader reader = mAlluxioBlockStore.createBlockReader(SESSION_ID, BLOCK_ID, 0);
    Assert.assertEquals(length, reader.getLength());
    ByteBuffer buffer = reader.read(0, length);
    assertTrue(BufferUtils.equalIncreasingByteBuffer((int) start, (int) length, buffer));
    reader.close();
    lock.get().close();
  }

  @Test
  public void readFullBlock() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(0, TEST_BLOCK_SIZE);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE, buffer));
    mReader.close();
    checkTempBlock(0, TEST_BLOCK_SIZE);
  }

  @Test
  public void readPartialBlock() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(0, TEST_BLOCK_SIZE - 1);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE - 1, buffer));
    mReader.close();
    // partial block should not be cached
    assertFalse(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
  }

  @Test
  public void offset() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(2, TEST_BLOCK_SIZE - 2);
    assertTrue(BufferUtils
        .equalIncreasingByteBuffer(2, (int) TEST_BLOCK_SIZE - 2, buffer));
    mReader.close();
    // partial block should not be cached
    assertFalse(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
  }

  @Test
  public void readOverlap() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 2, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(2, TEST_BLOCK_SIZE - 2);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(2, (int) TEST_BLOCK_SIZE - 2, buffer));
    buffer = mReader.read(0, TEST_BLOCK_SIZE - 2);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE - 2, buffer));
    buffer = mReader.read(3, TEST_BLOCK_SIZE);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(3, (int) TEST_BLOCK_SIZE - 3, buffer));
    mReader.close();
    // block should be cached as two reads covers the full block
    checkTempBlock(0, TEST_BLOCK_SIZE);
  }

  @Test
  public void readFullBlockNoCache() throws Exception {
    mUnderFileSystemBlockMeta = new UnderFileSystemBlockMeta(SESSION_ID, BLOCK_ID,
        mOpenUfsBlockOptions.toBuilder().setNoCache(true).build());
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(0, TEST_BLOCK_SIZE);
    // read should succeed even if error is thrown when caching
    assertTrue(BufferUtils.equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE, buffer));
    mReader.close();
    assertFalse(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
  }

  @Test
  public void readFullBlockRequestSpaceError() throws Exception {
    LocalBlockStore errorThrowingBlockStore = spy(mAlluxioBlockStore);
    doThrow(new ResourceExhaustedRuntimeException("Ignored", false))
        .when(errorThrowingBlockStore)
        .requestSpace(anyLong(), anyLong(), anyLong());
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        errorThrowingBlockStore, mUfsClient, mUfsInstreamCache,
        mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(0, TEST_BLOCK_SIZE);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE, buffer));
    mReader.close();
    assertFalse(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
  }

  @Test
  public void readFullBlockRequestCreateBlockError() throws Exception {
    LocalBlockStore errorThrowingBlockStore = spy(mAlluxioBlockStore);
    doThrow(new ResourceExhaustedRuntimeException("Ignored", false)).when(errorThrowingBlockStore)
        .createBlock(anyLong(), anyLong(), any(AllocateOptions.class));
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        errorThrowingBlockStore, mUfsClient, mUfsInstreamCache,
        mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuffer buffer = mReader.read(0, TEST_BLOCK_SIZE);
    assertTrue(BufferUtils.equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE, buffer));
    mReader.close();
    assertFalse(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
  }

  @Test
  public void transferFullBlock() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuf buf =
        PooledByteBufAllocator.DEFAULT.buffer((int) TEST_BLOCK_SIZE * 2, (int) TEST_BLOCK_SIZE * 2);
    try {
      while (buf.writableBytes() > 0 && mReader.transferTo(buf) != -1) {
      }
      assertTrue(BufferUtils
          .equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE, buf.nioBuffer()));
      mReader.close();
    } finally {
      buf.release();
    }
    checkTempBlock(0, TEST_BLOCK_SIZE);
  }

  @Test
  public void transferPartialBlock() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    ByteBuf buf =
        PooledByteBufAllocator.DEFAULT.buffer((int) TEST_BLOCK_SIZE / 2, (int) TEST_BLOCK_SIZE / 2);
    try {
      while (buf.writableBytes() > 0 && mReader.transferTo(buf) != -1) {
      }
      assertTrue(BufferUtils
          .equalIncreasingByteBuffer(0, (int) TEST_BLOCK_SIZE / 2, buf.nioBuffer()));
      mReader.close();
    } finally {
      buf.release();
    }
    // partial block should not be cached
    assertFalse(mAlluxioBlockStore.hasTempBlockMeta(BLOCK_ID));
  }

  @Test
  public void getLocation() throws Exception {
    mReader = UnderFileSystemBlockReader.create(mUnderFileSystemBlockMeta, 0, false,
        mAlluxioBlockStore, mUfsClient, mUfsInstreamCache, mUfsBytesRead, mUfsBytesReadThroughput);
    assertTrue(mReader.getLocation().startsWith(mOpenUfsBlockOptions.getUfsPath()));
  }
}
