alluxio-master-2.x/core/server/worker/src/test/java/alluxio/worker/block/UfsIoManagerTest.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.any;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.spy;

import alluxio.AlluxioTestDirectory;
import alluxio.AlluxioURI;
import alluxio.ConfigurationRule;
import alluxio.Constants;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.UfsReadOptions;
import alluxio.underfs.UfsManager.UfsClient;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.UnderFileSystemConfiguration;
import alluxio.util.io.BufferUtils;

import com.codahale.metrics.Meter;
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.concurrent.CompletableFuture;

public final class UfsIoManagerTest {
  private static final long TEST_BLOCK_SIZE = 5 * Constants.MB;
  private static final long FIRST_BLOCK_ID = 0;
  private static final ByteBuffer TEST_BUF = ByteBuffer.allocate((int) TEST_BLOCK_SIZE);
  private UfsIOManager mUfsIOManager;

  @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());
        }
      }, Configuration.modifiableGlobal());
  private String mTestFilePath;

  @Before
  public void before() throws Exception {
    String ufsFolder = Configuration.getString(PropertyKey.MASTER_MOUNT_TABLE_ROOT_UFS);
    mTestFilePath = File.createTempFile("temp", null, new File(ufsFolder)).getAbsolutePath();
    byte[] buffer = BufferUtils.getIncreasingByteArray((int) TEST_BLOCK_SIZE * 2);
    BufferUtils.writeBufferToFile(mTestFilePath, buffer);
    UfsClient client = new UfsClient(() -> UnderFileSystem.Factory.create(mTestFilePath,
        UnderFileSystemConfiguration.defaults(Configuration.global())),
        new AlluxioURI(mTestFilePath));
    mUfsIOManager = spy(new UfsIOManager(client));
    doReturn(3d * Constants.MB).when(mUfsIOManager).getUsedThroughput(any(Meter.class));
    mUfsIOManager.start();
  }

  @Test
  public void readFullBlock() throws Exception {
    mUfsIOManager.read(TEST_BUF, 0, TEST_BLOCK_SIZE, FIRST_BLOCK_ID, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf(0, (int) TEST_BLOCK_SIZE, TEST_BUF));
    TEST_BUF.clear();
  }

  @Test
  public void readPartialBlock() throws Exception {

    mUfsIOManager.read(TEST_BUF, 0, TEST_BLOCK_SIZE - 1, FIRST_BLOCK_ID, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf(0, (int) TEST_BLOCK_SIZE - 1, TEST_BUF));
    TEST_BUF.clear();
  }

  @Test
  public void readSecondBlock() throws Exception {
    mUfsIOManager.read(TEST_BUF, TEST_BLOCK_SIZE, TEST_BLOCK_SIZE, 1, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf((int) TEST_BLOCK_SIZE, (int) TEST_BLOCK_SIZE, TEST_BUF));
    TEST_BUF.clear();
  }

  @Test
  public void offset() throws Exception {
    mUfsIOManager.read(TEST_BUF, 2, TEST_BLOCK_SIZE - 2, FIRST_BLOCK_ID, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf(2, (int) TEST_BLOCK_SIZE - 2, TEST_BUF));
    TEST_BUF.clear();
  }

  @Test
  public void readOverlap() throws Exception {
    mUfsIOManager.read(TEST_BUF, 2, TEST_BLOCK_SIZE - 2, FIRST_BLOCK_ID, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf(2, (int) TEST_BLOCK_SIZE - 2, TEST_BUF));
    TEST_BUF.clear();
    mUfsIOManager.read(TEST_BUF, TEST_BLOCK_SIZE, TEST_BLOCK_SIZE, 1, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf((int) TEST_BLOCK_SIZE, (int) TEST_BLOCK_SIZE, TEST_BUF));
    TEST_BUF.clear();
    mUfsIOManager.read(TEST_BUF, 2 + TEST_BLOCK_SIZE, TEST_BLOCK_SIZE - 2, 1, mTestFilePath,
        UfsReadOptions.getDefaultInstance()).get();
    assertTrue(checkBuf((int) TEST_BLOCK_SIZE + 2, (int) TEST_BLOCK_SIZE - 2, TEST_BUF));
    TEST_BUF.clear();
  }

  @Test
  public void readWithQuota() throws Exception {
    mUfsIOManager.setQuota("quotaTest", 2 * Constants.MB); // magic number for proper test time
    CompletableFuture<Integer> future =
        mUfsIOManager.read(TEST_BUF, 2, TEST_BLOCK_SIZE - 2, FIRST_BLOCK_ID, mTestFilePath,
            UfsReadOptions.newBuilder().setTag("quotaTest").setPositionShort(false).build());
    // sleep to make sure future is not done because of quota instead of get future too soon
    Thread.sleep(3000);
    assertFalse(future.isDone());
    mUfsIOManager.setQuota("quotaTest", 5 * Constants.MB);
    future.get();
    assertTrue(checkBuf(2, (int) TEST_BLOCK_SIZE - 2, TEST_BUF));
    TEST_BUF.clear();
  }

  private boolean checkBuf(int start, int len, ByteBuffer buf) {
    buf.rewind();
    for (int k = 0; k < len; k++) {
      if (buf.get() != (byte) (start + k)) {
        return false;
      }
    }
    return true;
  }
}
