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

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

import alluxio.AlluxioURI;
import alluxio.client.WriteType;
import alluxio.client.file.FileInStream;
import alluxio.client.file.FileOutStream;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.runtime.UnavailableRuntimeException;
import alluxio.exception.status.UnavailableException;
import alluxio.master.MultiMasterEmbeddedJournalLocalAlluxioCluster;
import alluxio.master.block.BlockMaster;
import alluxio.master.block.DefaultBlockMaster;
import alluxio.multi.process.PortCoordination;
import alluxio.testutils.IntegrationTestUtils;
import alluxio.util.CommonUtils;
import alluxio.util.WaitForOptions;
import alluxio.wire.BlockLocationInfo;
import alluxio.worker.block.AllMasterRegistrationBlockWorker;
import alluxio.worker.block.BlockWorker;
import alluxio.worker.block.SpecificMasterBlockSync;
import alluxio.worker.block.TestSpecificMasterBlockSync;

import com.google.common.collect.Maps;
import org.apache.commons.io.IOUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;

import java.net.InetSocketAddress;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;

public class WorkerAllMasterRegistrationTest {
  private MultiMasterEmbeddedJournalLocalAlluxioCluster mCluster;

  @Rule
  public TestName mTestName = new TestName();

  private final List<DefaultBlockMaster> mBlockMasters = new ArrayList<>();
  private AllMasterRegistrationBlockWorker mWorker;
  private final int mNumMasters = 3;
  private final int mNumWorkers = 1;

  private WaitForOptions mDefaultWaitForOptions = WaitForOptions.defaults().setTimeoutMs(30000);

  @Before
  public void before() throws Exception {
    mCluster = new MultiMasterEmbeddedJournalLocalAlluxioCluster(
        mNumMasters, mNumWorkers, PortCoordination.WORKER_ALL_MASTER_REGISTRATION);
    mCluster.initConfiguration(
        IntegrationTestUtils.getTestName(getClass().getSimpleName(), mTestName.getMethodName()));
    Configuration.set(PropertyKey.MASTER_JOURNAL_LOG_SIZE_BYTES_MAX, 100);
    Configuration.set(PropertyKey.WORKER_REGISTER_TO_ALL_MASTERS, true);
    Configuration.set(PropertyKey.STANDBY_MASTER_GRPC_ENABLED, true);
    Configuration.set(PropertyKey.USER_FILE_WRITE_TYPE_DEFAULT, WriteType.MUST_CACHE);
    Configuration.set(PropertyKey.WORKER_BLOCK_HEARTBEAT_REPORT_SIZE_THRESHOLD, 5);
    Configuration.set(PropertyKey.MASTER_JOURNAL_FLUSH_TIMEOUT_MS, "30sec");
    Configuration.set(PropertyKey.MASTER_EMBEDDED_JOURNAL_WRITE_TIMEOUT, "10sec");
    Configuration.set(PropertyKey.MASTER_EMBEDDED_JOURNAL_MIN_ELECTION_TIMEOUT, "3s");
    Configuration.set(PropertyKey.MASTER_EMBEDDED_JOURNAL_MAX_ELECTION_TIMEOUT, "6s");

    mCluster.start();

    mWorker = (AllMasterRegistrationBlockWorker)
        mCluster.getWorkerProcess(0).getWorker(BlockWorker.class);
    for (int i = 0; i < mNumMasters; i++) {
      mBlockMasters.add((DefaultBlockMaster)
          mCluster.getLocalAlluxioMasterByIndex(i).getMasterProcess().getMaster(BlockMaster.class));
    }
  }

  @After
  public void after() throws Exception {
    mCluster.stop();
    mWorker.stop();
    mWorker = null;
    mBlockMasters.clear();
  }

  /**
   * Tests a happy path where added and removed blocks can all be reported to standby masters.
   */
  @Test
  public void happyPath() throws Exception {
    CommonUtils.waitFor("wait for worker registration complete", () ->
        mWorker.getBlockSyncMasterGroup().isRegisteredToAllMasters(), mDefaultWaitForOptions);

    AlluxioURI fileUri = new AlluxioURI("/foobar");
    String fileContent = "foobar";

    FileOutStream fos = mCluster.getClient().createFile(fileUri);
    fos.write(fileContent.getBytes());
    fos.close();

    FileInStream fis = mCluster.getClient().openFile(fileUri);
    assertEquals(fileContent, IOUtils.toString(fis, Charset.defaultCharset()));

    List<BlockLocationInfo> blockLocations =
        mCluster.getClient().getBlockLocations(fileUri);
    assertEquals(1, blockLocations.size());
    assertEquals(1, blockLocations.get(0).getLocations().size());
    long blockId = blockLocations.get(0).getBlockInfo().getBlockInfo().getBlockId();

    // New blocks are added by committing journals
    CommonUtils.waitFor("wait for blocks being committed to all masters", () ->
        mBlockMasters.stream().allMatch(it -> {
          try {
            return it.getBlockMetaStore().getLocations(blockId).size() == 1;
          } catch (UnavailableRuntimeException e) {
            // The RocksDB is unavailable due to events like checkpoint
            // Just retry
          }
          return false;
        }), mDefaultWaitForOptions);

    // Removed blocks are reported by worker-master heartbeats
    mWorker.removeBlock(new Random().nextLong(), blockId);
    CommonUtils.waitFor("wait for blocks being removed to all masters", () ->
        mBlockMasters.stream().allMatch(it -> {
          try {
            return it.getBlockMetaStore().getLocations(blockId).size() == 0;
          } catch (UnavailableRuntimeException e) {
            // The RocksDB is unavailable due to events like checkpoint
            // Just retry
          }
          return false;
        }), mDefaultWaitForOptions);

    assertTrue(mWorker.getBlockSyncMasterGroup().isRegisteredToAllMasters());

    fis = mCluster.getClient().openFile(fileUri);
    FileInStream finalFis = fis;

    // Make sure registration only happen once to each master
    assertTrue(getBlockSyncOperators().values().stream()
        .allMatch(it -> it.getRegistrationSuccessCount() == 1));

    assertThrows(UnavailableException.class,
        () -> IOUtils.toString(finalFis, Charset.defaultCharset()));
  }

  /**
   * Tests a scenario where the worker to master heartbeat fails.
   */
  @Test
  public void workerHeartbeatFail() throws Exception {
    CommonUtils.waitFor("wait for worker registration complete", () ->
        mWorker.getBlockSyncMasterGroup().isRegisteredToAllMasters(), mDefaultWaitForOptions);

    // Fails all heartbeats
    getBlockSyncOperators().values().forEach(TestSpecificMasterBlockSync::failHeartbeat);

    // Write a file
    AlluxioURI fileUri = new AlluxioURI("/foobar");
    String fileContent = "foobar";

    FileOutStream fos = mCluster.getClient().createFile(fileUri);
    fos.write(fileContent.getBytes());
    fos.close();

    // Committed block is on primary master even if the heartbeat is paused.
    List<BlockLocationInfo> blockLocations =
        mCluster.getClient().getBlockLocations(fileUri);
    assertEquals(1, blockLocations.size());
    assertEquals(1, blockLocations.get(0).getLocations().size());
    long blockId = blockLocations.get(0).getBlockInfo().getBlockInfo().getBlockId();

    // Added blocks are replicated to standbys by journals,
    // so even if the heartbeat fails, standby are still in sync with primary.
    CommonUtils.waitFor("wait for blocks being committed to all masters by heartbeats",
        () ->
        mBlockMasters.stream().allMatch(it -> {
          try {
            return it.getBlockMetaStore().getLocations(blockId).size() == 1;
          } catch (UnavailableRuntimeException e) {
            // The RocksDB is unavailable due to events like checkpoint
            // Just retry
          }
          return false;
        }), mDefaultWaitForOptions);

    // Remove a block
    mWorker.removeBlock(new Random().nextLong(), blockId);

    // Resume all heartbeats and the block location should be moved on standby masters,
    // by heartbeats
    getBlockSyncOperators().values().forEach(TestSpecificMasterBlockSync::restoreHeartbeat);
    CommonUtils.waitFor("wait for blocks being removed on all masters by heartbeats",
        () ->
            mBlockMasters.stream().allMatch(it -> {
              try {
                return it.getBlockMetaStore().getLocations(blockId).size() == 0;
              } catch (UnavailableRuntimeException e) {
                // The RocksDB is unavailable due to events like checkpoint
                // Just retry
              }
              return false;
            }), mDefaultWaitForOptions);

    // Make sure registration only happen once to each master
    assertTrue(getBlockSyncOperators().values().stream()
        .allMatch(it -> it.getRegistrationSuccessCount() == 1));
  }

  /**
   * Tests the master failover case and makes sure the re-registration does not happen,
   * on the new elected primary master.
   */
  @Test
  public void masterFailover() throws Exception {
    CommonUtils.waitFor("wait for worker registration complete", () ->
        mWorker.getBlockSyncMasterGroup().isRegisteredToAllMasters(), mDefaultWaitForOptions);

    AlluxioURI fileUri = new AlluxioURI("/foobar");
    String fileContent = "foobar";

    FileOutStream fos = mCluster.getClient().createFile(fileUri);
    fos.write(fileContent.getBytes());
    fos.close();

    FileInStream fis = mCluster.getClient().openFile(fileUri);
    assertEquals(fileContent, IOUtils.toString(fis, Charset.defaultCharset()));

    // Make sure registration only happen once to each master
    assertTrue(getBlockSyncOperators().values().stream()
        .allMatch(it -> it.getRegistrationSuccessCount() == 1));

    // Kill the master and let the failover happen
    int leaderId = mCluster.getLeaderIndex();
    mCluster.stopLeader();
    mCluster.waitForPrimaryMasterServing(10000);
    assertNotEquals(mCluster.getLeaderIndex(), leaderId);

    fis = mCluster.getClient().openFile(fileUri);
    FileInStream finalFis = fis;
    // The new elected primary master should be able to serve the request immediately,
    // because the added block location is replicated by journal.
    assertEquals(fileContent, IOUtils.toString(finalFis, Charset.defaultCharset()));

    // Make sure no more registration happens
    assertTrue(getBlockSyncOperators().values().stream()
        .allMatch(it -> it.getRegistrationSuccessCount() == 1));
  }

  /**
   * Tests worker being able to worker and re-register to all masters after its restart.
   */
  @Test
  public void workerRestart() throws Exception {
    CommonUtils.waitFor("wait for worker registration complete", () ->
        mWorker.getBlockSyncMasterGroup().isRegisteredToAllMasters(), mDefaultWaitForOptions);

    AlluxioURI fileUri = new AlluxioURI("/foobar");
    String fileContent = "foobar";

    FileOutStream fos = mCluster.getClient().createFile(fileUri);
    fos.write(fileContent.getBytes());
    fos.close();

    FileInStream fis = mCluster.getClient().openFile(fileUri);
    assertEquals(fileContent, IOUtils.toString(fis, Charset.defaultCharset()));

    mCluster.stopWorkers();
    mCluster.startWorkers();

    mWorker = (AllMasterRegistrationBlockWorker)
        mCluster.getWorkerProcess(0).getWorker(BlockWorker.class);

    CommonUtils.waitFor("wait for worker registration complete", () ->
        getBlockSyncOperators().values().stream().allMatch(
            SpecificMasterBlockSync::isRegistered), mDefaultWaitForOptions);

    fis = mCluster.getClient().openFile(fileUri);
    assertEquals(fileContent, IOUtils.toString(fis, Charset.defaultCharset()));
  }

  /**
   * Tests the worker can re-register with masters if its heartbeat failed too many times,
   * and the block report becomes too big and takes up too many memory.
   */
  @Test
  public void heartbeatFallsBackToRegister() throws Exception {
    CommonUtils.waitFor("wait for worker registration complete", () ->
        mWorker.getBlockSyncMasterGroup().isRegisteredToAllMasters(), mDefaultWaitForOptions);

    // Create a test file whose block will be removed later
    AlluxioURI testFileUri = new AlluxioURI("/foo");
    FileOutStream fos = mCluster.getLocalAlluxioMaster().getClient().createFile(
        new AlluxioURI("/foo"));
    fos.write("foo".getBytes());
    fos.close();
    List<BlockLocationInfo> blockLocations =
        mCluster.getClient().getBlockLocations(testFileUri);
    long testFileBlockId = blockLocations.get(0).getBlockInfo().getBlockInfo().getBlockId();

    List<Long> blockIdsToRemove = new ArrayList<>();
    int numFiles = 10;
    // Create 10 files and corresponding 10 blocks.
    // These blocks are added to standby masters by journals.
    for (int i = 0; i < numFiles; ++i) {
      AlluxioURI fileUri = new AlluxioURI("/" + i);
      fos = mCluster.getLocalAlluxioMaster().getClient().createFile(
          fileUri);
      fos.write("foo".getBytes());
      fos.close();
      blockLocations = mCluster.getClient().getBlockLocations(fileUri);
      long blockId = blockLocations.get(0).getBlockInfo().getBlockInfo().getBlockId();
      blockIdsToRemove.add(blockId);
    }

    // Make heartbeat return fail
    getBlockSyncOperators().values().forEach(TestSpecificMasterBlockSync::failHeartbeat);

    // Removing the blocks for these 10 files,
    // heartbeat report containing these blocks will be generated during the heartbeat.
    // However, the heartbeat RPC to master will not succeed because
    // we made the heartbeat fail.
    // So the heartbeat report will become larger and larger as we merge the report
    // back to the reporter on RPC failures.
    Thread fileGenerationThread = new Thread(() -> {
      for (long blockId: blockIdsToRemove) {
        try {
          mWorker.removeBlock(0, blockId);
          Thread.sleep(500);
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }
    });
    fileGenerationThread.start();

    // Registration should trigger after heartbeat throws exceptions,
    // because the heartbeat report contains too many block ids and exceeds the
    // WORKER_BLOCK_HEARTBEAT_REPORT_CAPACITY_THRESHOLD and
    // the worker will trigger re-registration.
    CommonUtils.waitFor("wait for re-registration sending the block location info",
        () -> getBlockSyncOperators().values().stream()
            .allMatch(it -> it.getRegistrationSuccessCount() >= 2),
        WaitForOptions.defaults().setTimeoutMs(60000));

    // Remove one block and resume the heartbeats,
    // and make sure all blocks are propagated to all masters after heartbeat resumes;
    fileGenerationThread.join();
    getBlockSyncOperators().values().forEach(TestSpecificMasterBlockSync::restoreHeartbeat);
    mWorker.removeBlock(1, testFileBlockId);

    // We have removed the test block id + 10 other block ids,
    // so the block meta store should not contain any block location
    blockIdsToRemove.add(testFileBlockId);
    CommonUtils.waitFor("wait for blocks propagated to masters by heartbeats",
        () -> mBlockMasters.stream().allMatch(it ->
                blockIdsToRemove.stream()
                    .allMatch(blockId -> {
                      try {
                        return it.getBlockMetaStore().getLocations(blockId).size() == 0;
                      } catch (UnavailableRuntimeException e) {
                        // The RocksDB is unavailable due to events like checkpoint
                        // Just retry
                      }
                      return false;
                    })
            ),
        mDefaultWaitForOptions);
  }

  private Map<InetSocketAddress, TestSpecificMasterBlockSync> getBlockSyncOperators() {
    return Maps.transformValues(mWorker.getBlockSyncMasterGroup().getMasterSyncOperators(),
        it -> (TestSpecificMasterBlockSync) it);
  }
}
