content: /*
   * 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.master.block;
  
  import static org.hamcrest.CoreMatchers.containsString;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThrows;
  import static org.junit.Assert.assertTrue;
  
  import alluxio.clock.ManualClock;
  import alluxio.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.grpc.BlockHeartbeatPRequest;
  import alluxio.grpc.BlockHeartbeatPResponse;
  import alluxio.grpc.BlockIdList;
  import alluxio.grpc.BlockStoreLocationProto;
  import alluxio.grpc.GetRegisterLeasePRequest;
  import alluxio.grpc.LocationBlockIdListEntry;
  import alluxio.grpc.RegisterWorkerPOptions;
  import alluxio.grpc.RegisterWorkerPRequest;
  import alluxio.grpc.RegisterWorkerPResponse;
  import alluxio.master.AlwaysPrimaryPrimarySelector;
  import alluxio.master.CoreMasterContext;
  import alluxio.master.MasterRegistry;
  import alluxio.master.MasterTestUtils;
  import alluxio.master.journal.noop.NoopJournalSystem;
  import alluxio.master.metrics.MetricsMaster;
  import alluxio.master.metrics.MetricsMasterFactory;
  import alluxio.util.SleepUtils;
  import alluxio.util.ThreadFactoryUtils;
  import alluxio.util.executor.ExecutorServiceFactories;
  import alluxio.wire.RegisterLease;
  import alluxio.wire.WorkerNetAddress;
  import alluxio.worker.block.BlockStoreLocation;
  
  import com.google.common.collect.ImmutableList;
  import io.grpc.stub.StreamObserver;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  
  import java.util.Optional;
  import java.util.Queue;
  import java.util.concurrent.ConcurrentLinkedDeque;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  
  public class BlockMasterWorkerServiceHandlerTest {
    private static final WorkerNetAddress NET_ADDRESS_1 = new WorkerNetAddress().setHost("localhost")
        .setRpcPort(80).setDataPort(81).setWebPort(82);
  
    private BlockMaster mBlockMaster;
    private MasterRegistry mRegistry;
    private ManualClock mClock;
    private ExecutorService mExecutorService;
    private MetricsMaster mMetricsMaster;
    private BlockMasterWorkerServiceHandler mHandler;
  
    /**
     * Sets up the dependencies before a test runs.
     */
    @Before
    public void before() throws Exception {
      initServiceHandler(true);
    }
  
    public void initServiceHandler(boolean leaseEnabled) throws Exception {
      if (leaseEnabled) {
        Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_ENABLED, true);
        Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_TTL, "3s");
        // Tests on the JVM check logic will be done separately
        Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_RESPECT_JVM_SPACE, false);
        Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_COUNT, 1);
      } else {
        Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_ENABLED, false);
      }
  
      mRegistry = new MasterRegistry();
      CoreMasterContext masterContext = MasterTestUtils.testMasterContext(new NoopJournalSystem(),
          null, new AlwaysPrimaryPrimarySelector());
      mMetricsMaster = new MetricsMasterFactory().create(mRegistry, masterContext);
      mClock = new ManualClock();
      mExecutorService =
          Executors.newFixedThreadPool(2, ThreadFactoryUtils.build("TestBlockMaster-%d", true));
      mBlockMaster = new DefaultBlockMaster(mMetricsMaster, masterContext, mClock,
          ExecutorServiceFactories.constantExecutorServiceFactory(mExecutorService));
      mRegistry.add(BlockMaster.class, mBlockMaster);
      mRegistry.start(true);
      mHandler = new BlockMasterWorkerServiceHandler(mBlockMaster);
    }
  
    @After
    public void after() throws Exception {
      mRegistry.stop();
    }
  
    @Test
    public void registerWithNoLeaseIsRejected() {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
  
      // Prepare LocationBlockIdListEntry objects
      BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
      BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder()
          .setTierAlias(loc.tierAlias())
          .setMediumType(loc.mediumType())
          .build();
  
      BlockIdList blockIdList1 = BlockIdList.newBuilder()
          .addAllBlockId(ImmutableList.of(1L, 2L)).build();
      LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder()
          .setKey(locationProto).setValue(blockIdList1).build();
  
      RegisterWorkerPRequest request = RegisterWorkerPRequest.newBuilder()
          .setWorkerId(workerId)
          .addStorageTiers("MEM")
          .putTotalBytesOnTiers("MEM", 1000L).putUsedBytesOnTiers("MEM", 0L)
          .setOptions(RegisterWorkerPOptions.getDefaultInstance())
          .addCurrentBlocks(listEntry1)
          .build();
  
      Queue<Throwable> errors = new ConcurrentLinkedDeque<>();
      StreamObserver<RegisterWorkerPResponse> noopResponseObserver =
          new StreamObserver<RegisterWorkerPResponse>() {
            @Override
            public void onNext(RegisterWorkerPResponse response) {}
  
            @Override
            public void onError(Throwable t) {
              errors.offer(t);
            }
  
            @Override
            public void onCompleted() {}
          };
  
      // The responseObserver should see an error
      mHandler.registerWorker(request, noopResponseObserver);
      assertEquals(1, errors.size());
      Throwable t = errors.poll();
      Assert.assertThat(t.getMessage(),
          containsString("does not have a lease or the lease has expired."));
    }
  
    @Test
    public void registerWorkerFailsOnDuplicateBlockLocation() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
  
      // Prepare LocationBlockIdListEntry objects
      BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
      BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder()
          .setTierAlias(loc.tierAlias())
          .setMediumType(loc.mediumType())
          .build();
  
      BlockIdList blockIdList1 = BlockIdList.newBuilder()
          .addAllBlockId(ImmutableList.of(1L, 2L)).build();
      LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder()
          .setKey(locationProto).setValue(blockIdList1).build();
  
      BlockIdList blockIdList2 = BlockIdList.newBuilder()
          .addAllBlockId(ImmutableList.of(3L, 4L)).build();
      LocationBlockIdListEntry listEntry2 = LocationBlockIdListEntry.newBuilder()
          .setKey(locationProto).setValue(blockIdList2).build();
  
      // Prepare a lease
      GetRegisterLeasePRequest leaseRequest = GetRegisterLeasePRequest.newBuilder()
          .setWorkerId(workerId)
          .setBlockCount(blockIdList1.getBlockIdCount() + blockIdList2.getBlockIdCount())
          .build();
      Optional<RegisterLease> lease = mBlockMaster.tryAcquireRegisterLease(leaseRequest);
      assertTrue(lease.isPresent());
  
      // The request is not deduplicated
      RegisterWorkerPRequest request = RegisterWorkerPRequest.newBuilder()
          .setWorkerId(workerId)
          .addStorageTiers("MEM")
          .putTotalBytesOnTiers("MEM", 1000L).putUsedBytesOnTiers("MEM", 0L)
          .setOptions(RegisterWorkerPOptions.getDefaultInstance())
          .addCurrentBlocks(listEntry1)
          .addCurrentBlocks(listEntry2)
          .build();
  
      // Noop response observer
      StreamObserver<RegisterWorkerPResponse> noopResponseObserver =
          new StreamObserver<RegisterWorkerPResponse>() {
            @Override
            public void onNext(RegisterWorkerPResponse response) {}
  
            @Override
            public void onError(Throwable t) {}
  
            @Override
            public void onCompleted() {}
          };
  
      assertThrows(AssertionError.class, () -> {
        mHandler.registerWorker(request, noopResponseObserver);
      });
  
      mBlockMaster.releaseRegisterLease(workerId);
    }
  
    @Test
    public void registerLeaseExpired() {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
  
      // Prepare LocationBlockIdListEntry objects
      BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
      BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder()
          .setTierAlias(loc.tierAlias())
          .setMediumType(loc.mediumType())
          .build();
  
      BlockIdList blockIdList1 = BlockIdList.newBuilder()
          .addAllBlockId(ImmutableList.of(1L, 2L)).build();
      LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder()
          .setKey(locationProto).setValue(blockIdList1).build();
  
      // Prepare a lease
      GetRegisterLeasePRequest leaseRequest = GetRegisterLeasePRequest.newBuilder()
          .setWorkerId(workerId)
          .setBlockCount(blockIdList1.getBlockIdCount())
          .build();
      Optional<RegisterLease> lease = mBlockMaster.tryAcquireRegisterLease(leaseRequest);
      assertTrue(lease.isPresent());
  
      // Sleep for a while so that the lease expires
      SleepUtils.sleepMs(5000);
  
      // The lease is recycled and taken away
      GetRegisterLeasePRequest newLeaseRequest = GetRegisterLeasePRequest.newBuilder()
          .setWorkerId(workerId + 1)
          .setBlockCount(blockIdList1.getBlockIdCount())
          .build();
      Optional<RegisterLease> newLease = mBlockMaster.tryAcquireRegisterLease(newLeaseRequest);
      assertTrue(newLease.isPresent());
  
      RegisterWorkerPRequest request = RegisterWorkerPRequest.newBuilder()
          .setWorkerId(workerId)
          .addStorageTiers("MEM")
          .putTotalBytesOnTiers("MEM", 1000L).putUsedBytesOnTiers("MEM", 0L)
          .setOptions(RegisterWorkerPOptions.getDefaultInstance())
          .addCurrentBlocks(listEntry1)
          .build();
  
      // Noop response observer
      Queue<Throwable> errors = new ConcurrentLinkedDeque<>();
      StreamObserver<RegisterWorkerPResponse> noopResponseObserver =
          new StreamObserver<RegisterWorkerPResponse>() {
            @Override
            public void onNext(RegisterWorkerPResponse response) {}
  
            @Override
            public void onError(Throwable t) {
              errors.offer(t);
            }
  
            @Override
            public void onCompleted() {}
          };
  
      mHandler.registerWorker(request, noopResponseObserver);
      assertEquals(1, errors.size());
      Throwable t = errors.poll();
      Assert.assertThat(t.getMessage(),
          containsString("does not have a lease or the lease has expired."));
  
      mBlockMaster.releaseRegisterLease(workerId + 1);
    }
  
    @Test
    public void registerLeaseTurnedOff() throws Exception {
      initServiceHandler(false);
  
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
  
      // Prepare LocationBlockIdListEntry objects
      BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
      BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder()
          .setTierAlias(loc.tierAlias())
          .setMediumType(loc.mediumType())
          .build();
  
      BlockIdList blockIdList1 = BlockIdList.newBuilder()
          .addAllBlockId(ImmutableList.of(1L, 2L)).build();
      LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder()
          .setKey(locationProto).setValue(blockIdList1).build();
  
      // No lease is acquired
      RegisterWorkerPRequest request = RegisterWorkerPRequest.newBuilder()
          .setWorkerId(workerId)
          .addStorageTiers("MEM")
          .putTotalBytesOnTiers("MEM", 1000L).putUsedBytesOnTiers("MEM", 0L)
          .setOptions(RegisterWorkerPOptions.getDefaultInstance())
          .addCurrentBlocks(listEntry1)
          .build();
  
      Queue<Throwable> errors = new ConcurrentLinkedDeque<>();
      StreamObserver<RegisterWorkerPResponse> noopResponseObserver =
          new StreamObserver<RegisterWorkerPResponse>() {
            @Override
            public void onNext(RegisterWorkerPResponse response) {}
  
            @Override
            public void onError(Throwable t) {
              errors.offer(t);
            }
  
            @Override
            public void onCompleted() {}
          };
  
      mHandler.registerWorker(request, noopResponseObserver);
      assertEquals(0, errors.size());
    }
  
    @Test
    public void workerHeartbeatFailsOnDuplicateBlockLocation() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
  
      // Prepare LocationBlockIdListEntry objects
      BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
      BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder()
              .setTierAlias(loc.tierAlias())
              .setMediumType(loc.mediumType())
              .build();
  
      BlockIdList blockIdList1 = BlockIdList.newBuilder()
              .addAllBlockId(ImmutableList.of(1L, 2L)).build();
      LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder()
              .setKey(locationProto).setValue(blockIdList1).build();
  
      BlockIdList blockIdList2 = BlockIdList.newBuilder()
              .addAllBlockId(ImmutableList.of(3L, 4L)).build();
      LocationBlockIdListEntry listEntry2 = LocationBlockIdListEntry.newBuilder()
              .setKey(locationProto).setValue(blockIdList2).build();
  
      BlockHeartbeatPRequest request = BlockHeartbeatPRequest.newBuilder()
              .setWorkerId(workerId)
              .addAddedBlocks(listEntry1)
              .addAddedBlocks(listEntry2)
              .build();
  
      // Noop response observer
      StreamObserver<BlockHeartbeatPResponse> noopResponseObserver =
          new StreamObserver<BlockHeartbeatPResponse>() {
            @Override
            public void onNext(BlockHeartbeatPResponse response) {}
  
            @Override
            public void onError(Throwable t) {}
  
            @Override
            public void onCompleted() {}
          };
  
      assertThrows(AssertionError.class, () -> {
        mHandler.blockHeartbeat(request, noopResponseObserver);
      });
    }
  }
  