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

import static org.mockito.AdditionalMatchers.or;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.startsWith;

import alluxio.ClientContext;
import alluxio.Constants;
import alluxio.client.block.BlockMasterClient;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.runtime.UnknownRuntimeException;
import alluxio.grpc.StorageList;
import alluxio.grpc.WorkerLostStorageInfo;
import alluxio.master.LocalAlluxioCluster;
import alluxio.master.MasterClientContext;
import alluxio.testutils.BaseIntegrationTest;
import alluxio.testutils.LocalAlluxioClusterResource;
import alluxio.worker.block.meta.DefaultStorageDir;
import alluxio.worker.block.meta.DefaultStorageTier;

import com.google.common.io.Files;
import org.apache.commons.io.FileUtils;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentMatchers;
import org.mockito.Mockito;
import org.powermock.api.mockito.PowerMockito;
import org.powermock.core.classloader.annotations.PowerMockIgnore;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;

/**
 * Tests for getting worker lost storage information.
 */
@RunWith(PowerMockRunner.class)
@PrepareForTest({DefaultStorageDir.class})
@PowerMockIgnore({"javax.*.*", "com.sun.*", "org.xml.*"})
public class LostStorageIntegrationTest extends BaseIntegrationTest {
  private static final int CAPACITY_BYTES = Constants.KB;
  private static final String SSD_TIER = Constants.MEDIUM_SSD;
  private static final String HDD_TIER = Constants.MEDIUM_HDD;
  private static final String WORKER_STORAGE_SUFFIX = "/alluxioworker";

  @Rule
  public LocalAlluxioClusterResource mLocalAlluxioClusterResource =
      new LocalAlluxioClusterResource.Builder()
          .setStartCluster(false).build();

  private LocalAlluxioCluster mLocalAlluxioCluster = null;

  private BlockMasterClient mBlockMasterClient = null;

  @Test
  public void reportLostStorageInWorkerRegister() throws Exception {
    File ssdDir = Files.createTempDir();
    String ssdPath = ssdDir.getAbsolutePath();
    File hddDir = Files.createTempDir();
    String hddPath = hddDir.getAbsolutePath();

    // Mock no write permission so worker storage paths cannot be initialize
    PowerMockito.mockStatic(DefaultStorageDir.class);
    Mockito.when(DefaultStorageDir.newStorageDir(any(DefaultStorageTier.class),
        anyInt(),
        anyLong(),
        anyLong(),
        anyString(),
        anyString())).thenCallRealMethod();
    Mockito.when(DefaultStorageDir.newStorageDir(any(DefaultStorageTier.class),
        anyInt(),
        anyLong(),
        anyLong(),
        or(startsWith(ssdPath), startsWith(hddPath)),
        anyString())).thenThrow(
            new UnknownRuntimeException(new IOException("mock no write permission exception")));

    startClusterWithWorkerStorage(ssdPath, hddPath);
    checkLostStorageResults(ssdPath, hddPath);
  }

  @Test
  public void reportLostStorageInHeartbeat() throws Exception {
    File ssdDir = Files.createTempDir();
    String ssdPath = ssdDir.getAbsolutePath();
    File hddDir = Files.createTempDir();
    String hddPath = hddDir.getAbsolutePath();

    startClusterWithWorkerStorage(ssdPath, hddPath);

    FileUtils.deleteDirectory(ssdDir);
    FileUtils.deleteDirectory(hddDir);

    // Make sure worker lost storage is detected and heartbeat with the master
    Thread.sleep(10 * Configuration.getMs(PropertyKey.WORKER_BLOCK_HEARTBEAT_INTERVAL_MS));
    checkLostStorageResults(ssdPath, hddPath);
  }

  @Test
  public void lostStorageWhenRestart() throws Exception {
    File ssdDir = Files.createTempDir();
    String ssdPath = ssdDir.getAbsolutePath();
    File hddDir = Files.createTempDir();
    String hddPath = hddDir.getAbsolutePath();

    // Mock no write permission so worker storage paths cannot be initialize
    PowerMockito.mockStatic(DefaultStorageDir.class);
    Mockito.when(DefaultStorageDir.newStorageDir(ArgumentMatchers.any(DefaultStorageTier.class),
        ArgumentMatchers.anyInt(),
        ArgumentMatchers.anyLong(),
        ArgumentMatchers.anyLong(),
        ArgumentMatchers.anyString(),
        ArgumentMatchers.anyString())).thenCallRealMethod();
    Mockito.when(DefaultStorageDir.newStorageDir(ArgumentMatchers.any(DefaultStorageTier.class),
        ArgumentMatchers.anyInt(),
        ArgumentMatchers.anyLong(),
        ArgumentMatchers.anyLong(),
        ArgumentMatchers.startsWith(ssdPath),
        ArgumentMatchers.anyString())).thenThrow(
            new UnknownRuntimeException(new IOException("mock no write permission exception")));

    startClusterWithWorkerStorage(ssdPath, hddPath);

    FileUtils.deleteDirectory(hddDir);
    // Make sure lost storage is detected and reported to master
    Thread.sleep(10 * Configuration.getMs(PropertyKey.WORKER_BLOCK_HEARTBEAT_INTERVAL_MS));
    checkLostStorageResults(ssdPath, hddPath);

    mLocalAlluxioCluster.restartMasters();
    mLocalAlluxioCluster.waitForWorkersRegistered(6 * Constants.SECOND_MS);
    checkLostStorageResults(ssdPath, hddPath);
  }

  /**
   * Starts the {@link LocalAlluxioCluster}.
   *
   * @param ssdPath the local path representing a worker storage in ssd tier
   * @param hddPath the local path representing a worker storage in hdd tier
   */
  private void startClusterWithWorkerStorage(String ssdPath, String hddPath) throws Exception {
    mLocalAlluxioClusterResource
        .setProperty(PropertyKey.USER_BLOCK_SIZE_BYTES_DEFAULT, "1KB")
        .setProperty(PropertyKey.WORKER_TIERED_STORE_LEVELS, 3)
        .setProperty(PropertyKey.WORKER_RAMDISK_SIZE, CAPACITY_BYTES)
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_ALIAS.format(1), SSD_TIER)
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_ALIAS.format(2), HDD_TIER)
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_PATH.format(0),
            Files.createTempDir().getAbsolutePath())
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_PATH.format(1),
            ssdPath)
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_QUOTA.format(1),
            String.valueOf(CAPACITY_BYTES))
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_PATH.format(2),
            hddPath)
        .setProperty(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_QUOTA.format(2),
            String.valueOf(CAPACITY_BYTES));
    mLocalAlluxioClusterResource.start();
    mLocalAlluxioCluster = mLocalAlluxioClusterResource.get();
    mBlockMasterClient =
        BlockMasterClient.Factory.create(MasterClientContext
            .newBuilder(ClientContext.create(Configuration.global())).build());
    mBlockMasterClient.connect();
  }

  /**
   * Checks if the worker lost storage results are as expected.
   *
   * @param ssdPath the local path representing a lost worker storage in ssd tier
   * @param hddPath the local path representing a lost worker storage in hdd tier
   */
  private void checkLostStorageResults(String ssdPath, String hddPath) throws IOException {
    List<WorkerLostStorageInfo>  infoList = mBlockMasterClient.getWorkerLostStorage();
    Assert.assertEquals(1, infoList.size());
    Map<String, StorageList> lostStorageMapTwo = infoList.get(0).getLostStorageMap();
    Assert.assertEquals(ssdPath + WORKER_STORAGE_SUFFIX,
        lostStorageMapTwo.get(SSD_TIER).getStorage(0));
    Assert.assertEquals(hddPath + WORKER_STORAGE_SUFFIX,
        lostStorageMapTwo.get(HDD_TIER).getStorage(0));
  }
}
