alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/block/BlockSyncMasterGroup.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.worker.block;

import alluxio.ClientContext;
import alluxio.ProcessUtils;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.heartbeat.FixedIntervalSupplier;
import alluxio.heartbeat.HeartbeatContext;
import alluxio.heartbeat.HeartbeatThread;
import alluxio.master.MasterClientContext;
import alluxio.security.user.ServerUserState;
import alluxio.util.CommonUtils;
import alluxio.util.ConfigurationUtils;
import alluxio.util.WaitForOptions;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeoutException;

/**
 * An abstraction layer that manages the worker heartbeats with multiple block masters.
 * This is only active when worker.register.to.all.masters=true.
 */
public class BlockSyncMasterGroup implements Closeable {
  private static final Logger LOG = LoggerFactory.getLogger(SpecificMasterBlockSync.class);
  private volatile boolean mStarted = false;

  private final boolean mTestMode = Configuration.getBoolean(PropertyKey.TEST_MODE);

  private static BlockMasterClientFactory sBlockMasterClientFactory
      = new BlockMasterClientFactory();

  private static final long WORKER_MASTER_CONNECT_RETRY_TIMEOUT =
      Configuration.getMs(PropertyKey.WORKER_MASTER_CONNECT_RETRY_TIMEOUT);

  /**
   * Creates a block sync master group.
   * @param masterAddresses the master addresses to sync
   * @param blockWorker the block worker instance
   */
  public BlockSyncMasterGroup(
      List<InetSocketAddress> masterAddresses,
      BlockWorker blockWorker
  ) throws IOException {
    // TODO(elega): handle master membership changes
    // https://github.com/Alluxio/alluxio/issues/16898
    for (InetSocketAddress masterAddr : masterAddresses) {
      BlockMasterClient masterClient = sBlockMasterClientFactory.create(masterAddr);
      BlockHeartbeatReporter heartbeatReporter = new BlockHeartbeatReporter();

      blockWorker.getBlockStore().registerBlockStoreEventListener(heartbeatReporter);
      // Setup BlockMasterSync
      SpecificMasterBlockSync blockMasterSync = mTestMode
          ? new TestSpecificMasterBlockSync(
          blockWorker, masterClient, heartbeatReporter)
          : new SpecificMasterBlockSync(
              blockWorker, masterClient, heartbeatReporter);
      // Register each BlockMasterSync to the block events on this worker
      mMasterSyncOperators.put(masterAddr, blockMasterSync);
      LOG.info("Kick off BlockMasterSync with master {}", masterAddr);
    }
  }

  /**
   * Starts the heartbeats.
   * @param executorService the executor service to run the heartbeats
   */
  public synchronized void start(ExecutorService executorService) {
    if (!mStarted) {
      mStarted = true;
    }
    mMasterSyncOperators.values().forEach(blockMasterSync -> executorService
        .submit(new HeartbeatThread(HeartbeatContext.WORKER_BLOCK_SYNC, blockMasterSync,
            () -> new FixedIntervalSupplier(
                Configuration.getMs(PropertyKey.WORKER_BLOCK_HEARTBEAT_INTERVAL_MS)),
            Configuration.global(), ServerUserState.global())));
  }

  private final Map<InetSocketAddress, SpecificMasterBlockSync> mMasterSyncOperators =
      new HashMap<>();

  @Override
  public void close() throws IOException {
    mMasterSyncOperators.values().forEach(
        SpecificMasterBlockSync::close
    );
  }

  static void setBlockMasterClientFactory(BlockMasterClientFactory factory) {
    sBlockMasterClientFactory = factory;
  }

  /**
   * Waits until the primary master registration completes.
   * @param primaryMasterAddress the primary master address
   */
  public void waitForPrimaryMasterRegistrationComplete(InetSocketAddress primaryMasterAddress) {
    SpecificMasterBlockSync primaryMasterSync =
        mMasterSyncOperators.get(primaryMasterAddress);
    Preconditions.checkNotNull(
        primaryMasterSync, "Primary master block sync should not be null");
    try {
      CommonUtils.waitFor(this + " to start",
          primaryMasterSync::isRegistered,
          WaitForOptions.defaults().setTimeoutMs(WORKER_MASTER_CONNECT_RETRY_TIMEOUT));
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      LOG.warn("Exit the worker on interruption", e);
      throw new RuntimeException(e);
    } catch (TimeoutException e) {
      ProcessUtils.fatalError(LOG, e, "Failed to register with primary master");
    }
    LOG.info("The worker has registered with primary master, address {}", primaryMasterAddress);
  }

  /**
   * @return if the worker is registered to all masters
   */
  public boolean isRegisteredToAllMasters() {
    return mMasterSyncOperators.values().stream().allMatch(SpecificMasterBlockSync::isRegistered);
  }

  /**
   * @return the master sync operators
   */
  public Map<InetSocketAddress, SpecificMasterBlockSync> getMasterSyncOperators() {
    return mMasterSyncOperators;
  }

  /**
   * The factory class.
   */
  public static class Factory {
    /**
     * Creates a block sync master group that heartbeats to all masters.
     * @param blockWorker the block worker instance
     * @return the block sync master group instance
     */
    public static BlockSyncMasterGroup createAllMasterSync(BlockWorker blockWorker) {
      List<InetSocketAddress> masterAddresses =
          ConfigurationUtils.getMasterRpcAddresses(Configuration.global());
      try {
        return new BlockSyncMasterGroup(masterAddresses, blockWorker);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }
  }

  /**
   * A factory class for testing purpose.
   */
  @VisibleForTesting
  static class BlockMasterClientFactory {
    BlockMasterClient create(InetSocketAddress address) {
      MasterClientContext context = MasterClientContext
          .newBuilder(ClientContext.create(Configuration.global())).build();

      return new BlockMasterClient(context, address);
    }
  }
}
