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

import static org.junit.Assert.assertEquals;

import alluxio.AlluxioTestDirectory;
import alluxio.ConfigurationRule;
import alluxio.ConfigurationTestUtils;
import alluxio.Process;
import alluxio.client.file.FileSystem;
import alluxio.client.file.FileSystemTestUtils;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.WritePType;
import alluxio.master.AlluxioMasterProcess;
import alluxio.master.TestUtils;
import alluxio.network.TieredIdentityFactory;
import alluxio.testutils.BaseIntegrationTest;
import alluxio.util.network.NetworkAddressUtils;
import alluxio.worker.AlluxioWorkerProcess;
import alluxio.worker.WorkerProcess;
import alluxio.worker.block.BlockWorker;

import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.powermock.reflect.Whitebox;

import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * Integration tests for functionality relating to tiered identity.
 */
public class LocalFirstPolicyIntegrationTest extends BaseIntegrationTest {
  private ExecutorService mExecutor;

  @Rule
  public ConfigurationRule mConf = new ConfigurationRule(conf(), Configuration.modifiableGlobal());

  private static Map<PropertyKey, Object> conf() {
    Map<PropertyKey, Object> map =
        ConfigurationTestUtils.testConfigurationDefaults(Configuration.global(),
        NetworkAddressUtils.getLocalHostName(
            (int) Configuration.getMs(PropertyKey.NETWORK_HOST_RESOLUTION_TIMEOUT_MS)),
        AlluxioTestDirectory.createTemporaryDirectory("tiered_identity_test").getAbsolutePath());
    map.put(PropertyKey.WORKER_RPC_PORT, 0);
    map.put(PropertyKey.WORKER_WEB_PORT, 0);

    return map;
  }

  @Before
  public void before() {
    mExecutor = Executors.newCachedThreadPool();
  }

  @After
  public void after() {
    mExecutor.shutdownNow();
  }

  @Test
  public void localNodeFirst() throws Exception {
    AlluxioMasterProcess master = AlluxioMasterProcess.Factory.create();
    WorkerProcess worker1 = AlluxioWorkerProcess.Factory
        .create(TieredIdentityFactory.fromString("node=node1,rack=rack1",
            Configuration.global()));
    WorkerProcess worker2 = AlluxioWorkerProcess.Factory
        .create(TieredIdentityFactory.fromString("node=node2,rack=rack2",
            Configuration.global()));

    runProcess(mExecutor, master);
    runProcess(mExecutor, worker1);
    runProcess(mExecutor, worker2);

    TestUtils.waitForReady(master);
    TestUtils.waitForReady(worker1);
    TestUtils.waitForReady(worker2);

    FileSystem fs = FileSystem.Factory.create();

    // Write to the worker in node1
    {
      Whitebox.setInternalState(TieredIdentityFactory.class, "sInstance",
          TieredIdentityFactory.fromString("node=node1,rack=rack1",
              Configuration.global()));
      try {
        FileSystemTestUtils.createByteFile(fs, "/file1", WritePType.MUST_CACHE, 100);
      } finally {
        Whitebox.setInternalState(TieredIdentityFactory.class, "sInstance", (Object) null);
      }
      BlockWorker blockWorker1 = worker1.getWorker(BlockWorker.class);
      BlockWorker blockWorker2 = worker2.getWorker(BlockWorker.class);
      assertEquals(100, blockWorker1.getStoreMeta().getUsedBytes());
      assertEquals(0, blockWorker2.getStoreMeta().getUsedBytes());
    }

    worker2.stop();
    worker1.stop();
    master.stop();
  }

  @Test
  public void differentNodeOfSameRack() throws Exception {
    AlluxioMasterProcess master = AlluxioMasterProcess.Factory.create();
    WorkerProcess worker1 = AlluxioWorkerProcess.Factory
        .create(TieredIdentityFactory.fromString("node=node1,rack=rack1",
            Configuration.global()));
    WorkerProcess worker2 = AlluxioWorkerProcess.Factory
        .create(TieredIdentityFactory.fromString("node=node2,rack=rack2",
            Configuration.global()));

    runProcess(mExecutor, master);
    runProcess(mExecutor, worker1);
    runProcess(mExecutor, worker2);

    TestUtils.waitForReady(master);
    TestUtils.waitForReady(worker1);
    TestUtils.waitForReady(worker2);

    FileSystem fs = FileSystem.Factory.create();
    // Write to the worker in rack2
    {
      Whitebox.setInternalState(TieredIdentityFactory.class, "sInstance",
          TieredIdentityFactory.fromString("node=node3,rack=rack2",
              Configuration.global()));
      try {
        FileSystemTestUtils.createByteFile(fs, "/file2", WritePType.MUST_CACHE, 10);
      } finally {
        Whitebox.setInternalState(TieredIdentityFactory.class, "sInstance", (Object) null);
      }
      BlockWorker blockWorker1 = worker1.getWorker(BlockWorker.class);
      BlockWorker blockWorker2 = worker2.getWorker(BlockWorker.class);
      assertEquals(0, blockWorker1.getStoreMeta().getUsedBytes());
      assertEquals(10, blockWorker2.getStoreMeta().getUsedBytes());
    }

    worker2.stop();
    worker1.stop();
    master.stop();
  }

  private void runProcess(ExecutorService e, Process p) {
    e.execute(() -> {
      try {
        p.start();
      } catch (Exception e1) {
        throw new RuntimeException(e1);
      }
    });
  }
}
