alluxio-master-2.x/job/server/src/test/java/alluxio/job/plan/replicate/SetReplicaDefinitionReplicateTest.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.job.plan.replicate;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
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.mockStatic;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import alluxio.AlluxioURI;
import alluxio.ClientContext;
import alluxio.Constants;
import alluxio.client.block.BlockStoreClient;
import alluxio.client.block.BlockWorkerInfo;
import alluxio.client.block.stream.BlockInStream;
import alluxio.client.block.stream.BlockInStream.BlockInStreamSource;
import alluxio.client.block.stream.BlockOutStream;
import alluxio.client.block.stream.TestBlockInStream;
import alluxio.client.block.stream.TestBlockOutStream;
import alluxio.client.file.FileSystem;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.URIStatus;
import alluxio.client.file.options.InStreamOptions;
import alluxio.client.file.options.OutStreamOptions;
import alluxio.collections.Pair;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.status.NotFoundException;
import alluxio.job.JobServerContext;
import alluxio.job.RunTaskContext;
import alluxio.job.SelectExecutorsContext;
import alluxio.underfs.UfsManager;
import alluxio.util.io.BufferUtils;
import alluxio.util.network.NetworkAddressUtils;
import alluxio.wire.BlockInfo;
import alluxio.wire.BlockLocation;
import alluxio.wire.FileBlockInfo;
import alluxio.wire.FileInfo;
import alluxio.wire.WorkerInfo;
import alluxio.wire.WorkerNetAddress;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.mockito.MockedStatic;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Tests replicate functionality of {@link SetReplicaDefinition}.
 */
public final class SetReplicaDefinitionReplicateTest {
  private static final long TEST_BLOCK_ID = 1L;
  private static final long TEST_BLOCK_SIZE = 512L;
  private static final int MAX_BYTES = 1000;
  private static final WorkerNetAddress ADDRESS_1 =
      new WorkerNetAddress().setHost("host1").setDataPort(10);
  private static final WorkerNetAddress ADDRESS_2 =
      new WorkerNetAddress().setHost("host2").setDataPort(10);
  private static final WorkerNetAddress ADDRESS_3 =
      new WorkerNetAddress().setHost("host3").setDataPort(10);
  private static final WorkerNetAddress LOCAL_ADDRESS =
      new WorkerNetAddress().setHost(NetworkAddressUtils
          .getLocalHostName((int) Configuration
              .getMs(PropertyKey.NETWORK_HOST_RESOLUTION_TIMEOUT_MS))).setDataPort(10);
  private static final WorkerInfo WORKER_INFO_1 = new WorkerInfo().setAddress(ADDRESS_1);
  private static final WorkerInfo WORKER_INFO_2 = new WorkerInfo().setAddress(ADDRESS_2);
  private static final WorkerInfo WORKER_INFO_3 = new WorkerInfo().setAddress(ADDRESS_3);
  private static final String TEST_PATH = "/test";

  private FileSystemContext mMockFileSystemContext;
  private BlockStoreClient mMockBlockStore;
  private FileSystem mMockFileSystem;
  private JobServerContext mMockJobServerContext;
  private UfsManager mMockUfsManager;
  private BlockInfo mTestBlockInfo;
  private URIStatus mTestStatus;
  private MockedStatic<BlockStoreClient> mMockStaticBlockStore;

  @Rule
  public final ExpectedException mThrown = ExpectedException.none();

  @Before
  public void before() throws Exception {
    mMockFileSystemContext = mock(FileSystemContext.class);
    when(mMockFileSystemContext.getClientContext())
        .thenReturn(ClientContext.create(Configuration.global()));
    when(mMockFileSystemContext.getClusterConf())
        .thenReturn(Configuration.global());
    mMockBlockStore = mock(BlockStoreClient.class);
    mMockFileSystem = mock(FileSystem.class);
    mMockUfsManager = mock(UfsManager.class);
    mMockJobServerContext =
        new JobServerContext(mMockFileSystem, mMockFileSystemContext, mMockUfsManager);
    mMockStaticBlockStore = mockStatic(BlockStoreClient.class);
    mMockStaticBlockStore.when(() ->
        BlockStoreClient.create(mMockFileSystemContext)).thenReturn(mMockBlockStore);
    mTestBlockInfo = new BlockInfo().setBlockId(TEST_BLOCK_ID).setLength(TEST_BLOCK_SIZE);
    when(mMockBlockStore.getInfo(TEST_BLOCK_ID)).thenReturn(mTestBlockInfo);
    mTestStatus = new URIStatus(
        new FileInfo().setPath(TEST_PATH).setBlockIds(Lists.newArrayList(TEST_BLOCK_ID))
            .setPersisted(true)
            .setFileBlockInfos(Lists.newArrayList(
                new FileBlockInfo().setBlockInfo(mTestBlockInfo))));
  }

  @After
  public void after() {
    mMockStaticBlockStore.close();
  }

  /**
   * Helper function to select executors.
   *
   * @param numReplicas how many replicas to replicate or evict
   * @param workerInfoList a list of current available job workers
   * @return the selection result
   */
  private Set<Pair<WorkerInfo, SetReplicaTask>> selectExecutorsTestHelper(
      int numReplicas, List<WorkerInfo> workerInfoList)
      throws Exception {
    SetReplicaConfig config = new SetReplicaConfig(TEST_PATH, TEST_BLOCK_ID, numReplicas);
    SetReplicaDefinition definition = new SetReplicaDefinition();
    return definition.selectExecutors(config, workerInfoList,
        new SelectExecutorsContext(1, mMockJobServerContext));
  }

  /**
   * Helper function to run a replicate task.
   *
   * @param blockWorkers available block workers
   * @param mockInStream mock blockInStream returned by the Block Store
   * @param mockOutStream mock blockOutStream returned by the Block Store
   */
  private void runTaskReplicateTestHelper(List<BlockWorkerInfo> blockWorkers,
      BlockInStream mockInStream, BlockOutStream mockOutStream) throws Exception {
    when(mMockFileSystem.getStatus(any(AlluxioURI.class))).thenReturn(mTestStatus);
    when(mMockFileSystemContext.getCachedWorkers()).thenReturn(blockWorkers);
    when(mMockBlockStore.getInStream(anyLong(),
            any(InStreamOptions.class))).thenReturn(mockInStream);
    when(mMockBlockStore.getInStream(any(BlockInfo.class),
        any(InStreamOptions.class), any(Map.class))).thenReturn(mockInStream);
    when(
        mMockBlockStore.getOutStream(eq(TEST_BLOCK_ID), eq(TEST_BLOCK_SIZE), eq(LOCAL_ADDRESS),
            any(OutStreamOptions.class))).thenReturn(mockOutStream);
    when(mMockBlockStore.getInfo(TEST_BLOCK_ID))
        .thenReturn(mTestBlockInfo
            .setLocations(Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1))));

    SetReplicaConfig config =
        new SetReplicaConfig(TEST_PATH, TEST_BLOCK_ID, 1 /* value not used */);
    SetReplicaDefinition definition = new SetReplicaDefinition();
    definition.runTask(config, new SetReplicaTask(Mode.REPLICATE),
        new RunTaskContext(1, 1, mMockJobServerContext));
  }

  @Test
  public void selectExecutorsOnlyOneWorkerAvailable() throws Exception {
    mTestBlockInfo.setLocations(Lists.newArrayList());
    Set<Pair<WorkerInfo, SetReplicaTask>> result =
        selectExecutorsTestHelper(2,
            Lists.newArrayList(WORKER_INFO_1));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_1, new SetReplicaTask(Mode.REPLICATE)));
    // select the only worker
    assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsOnlyOneWorkerValid() throws Exception {
    mTestBlockInfo.setLocations(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)));
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        2,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_2, new SetReplicaTask(Mode.REPLICATE)));
    // select one worker left
    assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsTwoWorkersValid() throws Exception {
    mTestBlockInfo.setLocations(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)));
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        3,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_2, new SetReplicaTask(Mode.REPLICATE)));
    expected.add(new Pair<>(WORKER_INFO_3, new SetReplicaTask(Mode.REPLICATE)));
    // select both workers left
    assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsOneOutOFTwoWorkersValid() throws Exception {
    mTestBlockInfo.setLocations(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)));
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        2,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    // select one worker out of two
    assertEquals(1, result.size());
    assertEquals(new SetReplicaTask(Mode.REPLICATE), result.iterator().next().getSecond());
  }

  @Test
  public void selectExecutorsNoWorkerValid() throws Exception {
    mTestBlockInfo.setLocations(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)));
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        2,
        Lists.newArrayList(WORKER_INFO_1));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = ImmutableSet.of();
    // select none as no choice left
    assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsInsufficientWorkerValid() throws Exception {
    mTestBlockInfo.setLocations(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)));
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        3,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_2, new SetReplicaTask(Mode.REPLICATE)));
    // select the only worker left though more copies are requested
    assertEquals(expected, result);
  }

  @Test
  public void runTaskNoBlockWorker() throws Exception {
    byte[] input = BufferUtils.getIncreasingByteArray(0, (int) TEST_BLOCK_SIZE);

    TestBlockInStream mockInStream =
        new TestBlockInStream(input, TEST_BLOCK_ID, input.length,
            BlockInStreamSource.NODE_LOCAL);
    TestBlockOutStream mockOutStream =
        new TestBlockOutStream(ByteBuffer.allocate(MAX_BYTES), TEST_BLOCK_SIZE);
    mThrown.expect(NotFoundException.class);
    mThrown.expectMessage(ExceptionMessage.NO_LOCAL_BLOCK_WORKER_LOAD_TASK
        .getMessage(TEST_BLOCK_ID));
    runTaskReplicateTestHelper(Lists.<BlockWorkerInfo>newArrayList(), mockInStream, mockOutStream);
  }

  @Test
  public void runTaskLocalBlockWorkerDifferentFileStatus() throws Exception {
    for (boolean persisted : new boolean[] {true, false}) {
      for (boolean pinned : new boolean[] {true, false}) {
        mTestStatus.getFileInfo().setPersisted(persisted)
            .setPinned(pinned)
            .setMediumTypes(pinned ? Sets.newHashSet(Constants.MEDIUM_MEM)
                : Collections.emptySet());
        byte[] input = BufferUtils.getIncreasingByteArray(0, (int) TEST_BLOCK_SIZE);
        TestBlockInStream mockInStream =
            new TestBlockInStream(input, TEST_BLOCK_ID, input.length,
                BlockInStreamSource.NODE_LOCAL);
        TestBlockOutStream mockOutStream =
            new TestBlockOutStream(ByteBuffer.allocate(MAX_BYTES), TEST_BLOCK_SIZE);
        BlockWorkerInfo localBlockWorker = new BlockWorkerInfo(LOCAL_ADDRESS, TEST_BLOCK_SIZE, 0);
        try (MockedStatic<BlockInStream> mockBlockInStream = mockStatic(BlockInStream.class)) {
          mockBlockInStream.when(() ->
              BlockInStream.create(any(FileSystemContext.class), any(BlockInfo.class),
                  any(WorkerNetAddress.class), any(BlockInStreamSource.class),
                  any(InStreamOptions.class))).thenReturn(mockInStream);
          runTaskReplicateTestHelper(Lists.newArrayList(localBlockWorker), mockInStream,
              mockOutStream);
          assertEquals(TEST_BLOCK_SIZE, mockInStream.getBytesRead());
          if (!persisted || pinned) {
            assertArrayEquals(
                String.format("input-output mismatched: pinned=%s, persisted=%s",
                    pinned, persisted), input, mockOutStream.getWrittenData());
          }
        }
      }
    }
  }

  @Test
  public void runTaskInputIOException() throws Exception {
    // file is pinned on a medium
    mTestStatus.getFileInfo().setPinned(true);
    mTestStatus.getFileInfo().setMediumTypes(Sets.newHashSet(Constants.MEDIUM_MEM));
    BlockInStream mockInStream = mock(BlockInStream.class);
    BlockOutStream mockOutStream = mock(BlockOutStream.class);

    BlockWorkerInfo localBlockWorker = new BlockWorkerInfo(LOCAL_ADDRESS, TEST_BLOCK_SIZE, 0);
    doThrow(new IOException("test")).when(mockInStream).read(any(byte[].class), anyInt(), anyInt());
    doThrow(new IOException("test")).when(mockInStream).read(any(byte[].class));
    try {
      runTaskReplicateTestHelper(Lists.newArrayList(localBlockWorker), mockInStream, mockOutStream);
      fail("Expected the task to throw and IOException");
    } catch (IOException e) {
      assertEquals("test", e.getMessage());
    }
    verify(mockOutStream).cancel();
  }
}
