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

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import alluxio.AlluxioURI;
import alluxio.ClientContext;
import alluxio.Constants;
import alluxio.client.UnderStorageType;
import alluxio.client.WriteType;
import alluxio.client.block.BlockStoreClient;
import alluxio.client.block.BlockWorkerInfo;
import alluxio.client.block.policy.BlockLocationPolicy;
import alluxio.client.block.stream.BlockOutStream;
import alluxio.client.block.stream.TestBlockOutStream;
import alluxio.client.block.stream.TestUnderFileSystemFileOutStream;
import alluxio.client.block.stream.UnderFileSystemFileOutStream;
import alluxio.client.file.options.OutStreamOptions;
import alluxio.client.util.ClientTestUtils;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.Configuration;
import alluxio.conf.InstancedConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.PreconditionMessage;
import alluxio.exception.status.UnavailableException;
import alluxio.grpc.CompleteFilePOptions;
import alluxio.grpc.CreateFilePOptions;
import alluxio.grpc.FileSystemMasterCommonPOptions;
import alluxio.grpc.GetStatusPOptions;
import alluxio.grpc.TtlAction;
import alluxio.grpc.WritePType;
import alluxio.network.TieredIdentityFactory;
import alluxio.resource.DummyCloseableResource;
import alluxio.security.GroupMappingServiceTestUtils;
import alluxio.util.io.BufferUtils;
import alluxio.wire.FileInfo;
import alluxio.wire.WorkerNetAddress;

import com.google.common.collect.Lists;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.powermock.api.mockito.PowerMockito;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Tests for the {@link FileOutStream} class.
 */
@RunWith(PowerMockRunner.class)
@PrepareForTest({FileSystemContext.class, FileSystemMasterClient.class, BlockStoreClient.class,
    UnderFileSystemFileOutStream.class})
public class FileOutStreamTest {

  private static InstancedConfiguration sConf = Configuration.copyGlobal();

  private static final long BLOCK_LENGTH = 100L;
  private static final AlluxioURI FILE_NAME = new AlluxioURI("/file");

  private FileSystemContext mFileSystemContext;
  private BlockStoreClient mBlockStore;
  private FileSystemMasterClient mFileSystemMasterClient;

  private Map<Long, TestBlockOutStream> mAlluxioOutStreamMap;
  private TestUnderFileSystemFileOutStream mUnderStorageOutputStream;
  private AtomicBoolean mUnderStorageFlushed;

  private AlluxioFileOutStream mTestStream;
  private ClientContext mClientContext;

  /**
   * Sets up the different contexts and clients before a test runs.
   */
  @Before
  public void before() throws Exception {
    GroupMappingServiceTestUtils.resetCache();
    ClientTestUtils.setSmallBufferSizes(sConf);

    mClientContext = ClientContext.create(sConf);

    // PowerMock enums and final classes
    mFileSystemContext = PowerMockito.mock(FileSystemContext.class);
    when(mFileSystemContext.getClientContext()).thenReturn(mClientContext);
    when(mFileSystemContext.getClusterConf()).thenReturn(sConf);
    when(mFileSystemContext.getPathConf(any(AlluxioURI.class))).thenReturn(sConf);
    when(mFileSystemContext.getWriteBlockLocationPolicy(any(AlluxioConfiguration.class)))
        .thenAnswer((Answer) invocation -> {
          AlluxioConfiguration conf =
              invocation.getArgument(0, AlluxioConfiguration.class);
          return BlockLocationPolicy.Factory.create(
              conf.getClass(PropertyKey.USER_BLOCK_WRITE_LOCATION_POLICY), conf);
        });
    mBlockStore = PowerMockito.mock(BlockStoreClient.class);
    mFileSystemMasterClient = PowerMockito.mock(FileSystemMasterClient.class);

    PowerMockito.mockStatic(BlockStoreClient.class);
    PowerMockito.when(BlockStoreClient.create(mFileSystemContext)).thenReturn(mBlockStore);

    when(mFileSystemContext.acquireMasterClientResource())
        .thenReturn(new DummyCloseableResource<>(mFileSystemMasterClient));
    when(mFileSystemMasterClient.getStatus(any(AlluxioURI.class), any(GetStatusPOptions.class)))
        .thenReturn(new URIStatus(new FileInfo()));

    // Return sequentially increasing numbers for new block ids
    when(mFileSystemMasterClient.getNewBlockIdForFile(FILE_NAME))
        .thenAnswer(new Answer<Long>() {
          private long mCount = 0;

          @Override
          public Long answer(InvocationOnMock invocation) throws Throwable {
            return mCount++;
          }
        });

    // Set up out streams. When they are created, add them to outStreamMap
    final Map<Long, TestBlockOutStream> outStreamMap = new HashMap<>();
    when(mBlockStore.getOutStream(anyLong(), eq(BLOCK_LENGTH),
        any(OutStreamOptions.class))).thenAnswer(new Answer<TestBlockOutStream>() {
          @Override
          public TestBlockOutStream answer(InvocationOnMock invocation) throws Throwable {
            Long blockId = invocation.getArgument(0, Long.class);
            if (!outStreamMap.containsKey(blockId)) {
              TestBlockOutStream newStream =
                  new TestBlockOutStream(ByteBuffer.allocate(1000), BLOCK_LENGTH);
              outStreamMap.put(blockId, newStream);
            }
            return outStreamMap.get(blockId);
          }
        });
    BlockWorkerInfo workerInfo =
        new BlockWorkerInfo(new WorkerNetAddress().setHost("localhost")
            .setTieredIdentity(TieredIdentityFactory.fromString("node=localhost", sConf))
            .setRpcPort(1).setDataPort(2).setWebPort(3), Constants.GB, 0);
    when(mFileSystemContext.getCachedWorkers()).thenReturn(Lists.newArrayList(workerInfo));
    mAlluxioOutStreamMap = outStreamMap;

    // Create an under storage stream so that we can check whether it has been flushed
    final AtomicBoolean underStorageFlushed = new AtomicBoolean(false);
    mUnderStorageOutputStream =
        new TestUnderFileSystemFileOutStream(ByteBuffer.allocate(5000)) {
          @Override
          public void flush() throws IOException {
            super.flush();
            underStorageFlushed.set(true);
          }
        };
    mUnderStorageFlushed = underStorageFlushed;

    PowerMockito.mockStatic(UnderFileSystemFileOutStream.class);
    PowerMockito.when(
        UnderFileSystemFileOutStream.create(any(FileSystemContext.class),
            any(WorkerNetAddress.class), any(OutStreamOptions.class))).thenReturn(
        mUnderStorageOutputStream);

    OutStreamOptions options =
        OutStreamOptions.defaults(mFileSystemContext).setBlockSizeBytes(BLOCK_LENGTH)
            .setWriteType(WriteType.CACHE_THROUGH).setUfsPath(FILE_NAME.getPath());
    mTestStream = createTestStream(FILE_NAME, options);
  }

  @After
  public void after() {
    ClientTestUtils.resetClient(sConf);
  }

  /**
   * Tests that a single byte is written to the out stream correctly.
   */
  @Test
  public void singleByteWrite() throws Exception {
    mTestStream.write(5);
    mTestStream.close();
    assertArrayEquals(new byte[] {5}, mAlluxioOutStreamMap.get(0L).getWrittenData());
  }

  /**
   * Tests that many bytes, written one at a time, are written to the out streams correctly.
   */
  @Test
  public void manyBytesWrite() throws IOException {
    int bytesToWrite = (int) ((BLOCK_LENGTH * 5) + (BLOCK_LENGTH / 2));
    for (int i = 0; i < bytesToWrite; i++) {
      mTestStream.write(i);
    }
    mTestStream.close();
    verifyIncreasingBytesWritten(bytesToWrite);
  }

  /**
   * Tests that writing a buffer all at once will write bytes to the out streams correctly.
   */
  @Test
  public void writeBuffer() throws IOException {
    int bytesToWrite = (int) ((BLOCK_LENGTH * 5) + (BLOCK_LENGTH / 2));
    mTestStream.write(BufferUtils.getIncreasingByteArray(bytesToWrite));
    mTestStream.close();
    verifyIncreasingBytesWritten(bytesToWrite);
  }

  /**
   * Tests writing a buffer at an offset.
   */
  @Test
  public void writeOffset() throws IOException {
    int bytesToWrite = (int) ((BLOCK_LENGTH * 5) + (BLOCK_LENGTH / 2));
    int offset = (int) (BLOCK_LENGTH / 3);
    mTestStream.write(BufferUtils.getIncreasingByteArray(bytesToWrite + offset), offset,
        bytesToWrite);
    mTestStream.close();
    verifyIncreasingBytesWritten(offset, bytesToWrite);
  }

  /**
   * Tests that {@link FileOutStream#close()} will close but not cancel the underlying out streams.
   * Also checks that {@link FileOutStream#close()} persists and completes the file.
   */
  @Test
  public void close() throws Exception {
    mTestStream.write(BufferUtils.getIncreasingByteArray((int) (BLOCK_LENGTH * 1.5)));
    mTestStream.close();
    for (long streamIndex = 0; streamIndex < 2; streamIndex++) {
      assertFalse(mAlluxioOutStreamMap.get(streamIndex).isCanceled());
      assertTrue(mAlluxioOutStreamMap.get(streamIndex).isClosed());
    }
    verify(mFileSystemMasterClient).completeFile(eq(FILE_NAME), any(CompleteFilePOptions.class));
  }

  /**
   * Tests that {@link FileOutStream#cancel()} will cancel and close the underlying out streams, and
   * delete from the under file system when the delegation flag is set. Also makes sure that
   * cancel() doesn't persist or complete the file.
   */
  @Test
  public void cancelWithDelegation() throws Exception {
    mTestStream.write(BufferUtils.getIncreasingByteArray((int) (BLOCK_LENGTH * 1.5)));
    mTestStream.cancel();
    for (long streamIndex = 0; streamIndex < 2; streamIndex++) {
      assertTrue(mAlluxioOutStreamMap.get(streamIndex).isClosed());
      assertTrue(mAlluxioOutStreamMap.get(streamIndex).isCanceled());
    }
    // Don't complete the file if the stream was canceled
    verify(mFileSystemMasterClient, times(0)).completeFile(any(AlluxioURI.class),
        any(CompleteFilePOptions.class));
  }

  /**
   * Tests that {@link FileOutStream#flush()} will flush the under store stream.
   */
  @Test
  public void flush() throws IOException {
    assertFalse(mUnderStorageFlushed.get());
    mTestStream.flush();
    assertTrue(mUnderStorageFlushed.get());
  }

  /**
   * Tests that if an exception is thrown by the underlying out stream, and the user is using
   * {@link UnderStorageType#NO_PERSIST} for their under storage type, the correct exception
   * message will be thrown.
   */
  @Test
  public void cacheWriteExceptionNonSyncPersist() throws IOException {
    OutStreamOptions options =
        OutStreamOptions.defaults(mFileSystemContext).setBlockSizeBytes(BLOCK_LENGTH)
            .setWriteType(WriteType.MUST_CACHE);
    BlockOutStream stream = mock(BlockOutStream.class);
    when(mBlockStore.getOutStream(anyLong(), anyLong(), any(OutStreamOptions.class)))
        .thenReturn(stream);
    mTestStream = createTestStream(FILE_NAME, options);

    when(stream.remaining()).thenReturn(BLOCK_LENGTH);
    doThrow(new IOException("test error")).when(stream).write(7);

    try {
      mTestStream.write(7);
      fail("the test should fail");
    } catch (IOException e) {
      assertEquals(ExceptionMessage.FAILED_CACHE.getMessage("test error"), e.getMessage());
    }
  }

  /**
   * Tests that if an exception is thrown by the underlying out stream, and the user is using
   * {@link UnderStorageType#SYNC_PERSIST} for their under storage type, the error is recovered
   * from by writing the data to the under storage out stream.
   */
  @Test
  public void cacheWriteExceptionSyncPersist() throws IOException {
    BlockOutStream stream = mock(BlockOutStream.class);
    when(mBlockStore.getOutStream(anyLong(), anyLong(), any(OutStreamOptions.class)))
        .thenReturn(stream);

    when(stream.remaining()).thenReturn(BLOCK_LENGTH);
    doThrow(new IOException("test error")).when(stream).write((byte) 7);
    mTestStream.write(7);
    mTestStream.write(8);
    assertArrayEquals(new byte[] {7, 8}, mUnderStorageOutputStream.getWrittenData());
    // The cache stream is written to only once - the FileInStream gives up on it after it throws
    // the first exception.
    verify(stream, times(1)).write(anyInt());
  }

  /**
   * Tests that write only writes a byte.
   */
  @Test
  public void truncateWrite() throws IOException {
    // Only writes the lowest byte
    mTestStream.write(0x1fffff00);
    mTestStream.write(0x1fffff01);
    mTestStream.close();
    verifyIncreasingBytesWritten(2);
  }

  /**
   * Tests that the correct exception is thrown when a buffer is written with invalid offset/length.
   */
  @Test
  public void writeBadBufferOffset() throws IOException {
    try {
      mTestStream.write(new byte[10], 5, 6);
      fail("buffer write with invalid offset/length should fail");
    } catch (IllegalArgumentException e) {
      assertEquals(String.format(PreconditionMessage.ERR_BUFFER_STATE.toString(), 10, 5, 6),
          e.getMessage());
    }
  }

  /**
   * Tests that writing a null buffer throws the correct exception.
   */
  @Test
  public void writeNullBuffer() throws IOException {
    try {
      mTestStream.write(null);
      fail("writing null should fail");
    } catch (IllegalArgumentException e) {
      assertEquals(PreconditionMessage.ERR_WRITE_BUFFER_NULL.toString(), e.getMessage());
    }
  }

  /**
   * Tests that writing a null buffer with offset/length information throws the correct exception.
   */
  @Test
  public void writeNullBufferOffset() throws IOException {
    try {
      mTestStream.write(null, 0, 0);
      fail("writing null should fail");
    } catch (IllegalArgumentException e) {
      assertEquals(PreconditionMessage.ERR_WRITE_BUFFER_NULL.toString(), e.getMessage());
    }
  }

  /**
   * Tests that the async write invokes the expected client APIs.
   */
  @Test
  public void asyncWrite() throws Exception {
    OutStreamOptions options =
        OutStreamOptions.defaults(mFileSystemContext).setBlockSizeBytes(BLOCK_LENGTH)
            .setWriteType(WriteType.ASYNC_THROUGH);
    mTestStream = createTestStream(FILE_NAME, options);

    mTestStream.write(BufferUtils.getIncreasingByteArray((int) (BLOCK_LENGTH * 1.5)));
    mTestStream.close();

    // Verify that async persist request is sent with complete file request.
    ArgumentCaptor<CompleteFilePOptions> parameterCaptor =
            ArgumentCaptor.forClass(CompleteFilePOptions.class);
    verify(mFileSystemMasterClient).completeFile(eq(FILE_NAME), parameterCaptor.capture());
    Assert.assertTrue(parameterCaptor.getValue().hasAsyncPersistOptions());
  }

  /**
   * Tests that common options are propagated to async write request.
   */
  @Test
  public void asyncWriteOptionPropagation() throws Exception {
    Random rand = new Random();
    FileSystemMasterCommonPOptions commonOptions =
        FileSystemMasterCommonPOptions.newBuilder().setTtl(rand.nextLong())
            .setTtlAction(TtlAction.values()[rand.nextInt(TtlAction.values().length)])
            .setSyncIntervalMs(rand.nextLong()).build();

    OutStreamOptions options =
        new OutStreamOptions(CreateFilePOptions.newBuilder().setWriteType(WritePType.ASYNC_THROUGH)
            .setBlockSizeBytes(BLOCK_LENGTH).setCommonOptions(commonOptions).build(),
            mFileSystemContext, sConf);

    // Verify that OutStreamOptions have captured the common options properly.
    assertEquals(options.getCommonOptions(), commonOptions);

    mTestStream = createTestStream(FILE_NAME, options);
    mTestStream.write(BufferUtils.getIncreasingByteArray((int) (BLOCK_LENGTH * 1.5)));
    mTestStream.close();
    verify(mFileSystemMasterClient).completeFile(eq(FILE_NAME), any(CompleteFilePOptions.class));

    // Verify that common options for OutStreamOptions are propagated to ScheduleAsyncPersistence.
    ArgumentCaptor<CompleteFilePOptions> parameterCaptor =
            ArgumentCaptor.forClass(CompleteFilePOptions.class);

    verify(mFileSystemMasterClient).completeFile(eq(FILE_NAME), parameterCaptor.capture());
    assertEquals(parameterCaptor.getValue().getAsyncPersistOptions().getCommonOptions(),
        options.getCommonOptions());
  }

  /**
   * Tests that the number of bytes written is correct when the stream is created with different
   * under storage types.
   */
  @Test
  public void getBytesWrittenWithDifferentUnderStorageType() throws IOException {
    for (WriteType type : WriteType.values()) {
      OutStreamOptions options =
          OutStreamOptions.defaults(mFileSystemContext).setBlockSizeBytes(BLOCK_LENGTH)
              .setWriteType(type).setUfsPath(FILE_NAME.getPath());
      mTestStream = createTestStream(FILE_NAME, options);
      mTestStream.write(BufferUtils.getIncreasingByteArray((int) BLOCK_LENGTH));
      mTestStream.flush();
      assertEquals(BLOCK_LENGTH, mTestStream.getBytesWritten());
    }
  }

  @Test
  public void createWithNoWorker()  {
    // The default 2 minutes is too long.
    sConf.set(PropertyKey.USER_FILE_WRITE_INIT_MAX_DURATION, "10sec");
    mClientContext = ClientContext.create(sConf);
    OutStreamOptions options = OutStreamOptions.defaults(mFileSystemContext)
        .setLocationPolicy((getWorkerOptions) -> Optional.empty())
        .setWriteType(WriteType.CACHE_THROUGH);
    Exception e = assertThrows(UnavailableException.class,
        () -> mTestStream = createTestStream(FILE_NAME, options));
    assertTrue(e.getMessage().contains(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage()));
  }

  private void verifyIncreasingBytesWritten(int len) {
    verifyIncreasingBytesWritten(0, len);
  }

  /**
   * Verifies that the out streams have had exactly `len` increasing bytes written to them, with the
   * first byte starting at `start`. Also verifies that the same bytes have been written to the
   * under storage file stream.
   */
  private void verifyIncreasingBytesWritten(int start, int len) {
    long filledStreams = len / BLOCK_LENGTH;
    for (long streamIndex = 0; streamIndex < filledStreams; streamIndex++) {
      assertTrue("stream " + streamIndex + " was never written",
          mAlluxioOutStreamMap.containsKey(streamIndex));
      assertArrayEquals(BufferUtils
          .getIncreasingByteArray((int) (streamIndex * BLOCK_LENGTH + start), (int) BLOCK_LENGTH),
          mAlluxioOutStreamMap.get(streamIndex).getWrittenData());
    }
    long lastStreamBytes = len - filledStreams * BLOCK_LENGTH;
    assertArrayEquals(
        BufferUtils.getIncreasingByteArray((int) (filledStreams * BLOCK_LENGTH + start),
            (int) lastStreamBytes),
        mAlluxioOutStreamMap.get(filledStreams).getWrittenData());

    assertArrayEquals(BufferUtils.getIncreasingByteArray(start, len),
        mUnderStorageOutputStream.getWrittenData());
  }

  /**
   * Creates a {@link FileOutStream} for test.
   *
   * @param path the file path
   * @param options the set of options specific to this operation
   * @return a {@link FileOutStream}
   */
  private AlluxioFileOutStream createTestStream(AlluxioURI path, OutStreamOptions options)
      throws IOException {
    return new AlluxioFileOutStream(path, options, mFileSystemContext);
  }
}
