alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/block/SpecificMasterBlockSync.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.ProcessUtils;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.ConnectionFailedException;
import alluxio.exception.FailedToAcquireRegisterLeaseException;
import alluxio.grpc.Command;
import alluxio.heartbeat.HeartbeatExecutor;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.retry.ExponentialBackoffRetry;
import alluxio.retry.RetryPolicy;
import alluxio.util.CommonUtils;
import alluxio.wire.WorkerNetAddress;

import com.codahale.metrics.Counter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.net.SocketAddress;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * The block master sync thread when workers are registered to all masters.
 * With respect to behaviors, this implementation differs from {@link BlockMasterSync} in:
 * 1. The registration takes place asynchronously and the caller can poll the registration state.
 *  We need to make the process async because when standby master read is enabled, workers have to
 *  register to all masters and these registrations can happen concurrently to speed up the process.
 * 2. A registration failure doesn't throw a fatal exception. Instead, it retries endlessly.
 *  This is because a standby master registration failure
 *  should be a soft failure and can be retried later.
 */
@NotThreadSafe
public class SpecificMasterBlockSync implements HeartbeatExecutor, Closeable {
  private static final Logger LOG = LoggerFactory.getLogger(SpecificMasterBlockSync.class);
  private static final long ACQUIRE_LEASE_WAIT_MAX_DURATION =
      Configuration.getMs(PropertyKey.WORKER_REGISTER_LEASE_RETRY_MAX_DURATION);

  private final long mWorkerBlockHeartbeatReportSizeThreshold =
      Configuration.getInt(PropertyKey.WORKER_BLOCK_HEARTBEAT_REPORT_SIZE_THRESHOLD);

  private final SocketAddress mMasterAddress;

  /**
   * The worker registration state.
   * If the state is NOT_REGISTERED, heartbeat will trigger a registration.
   * During the registration process, the state will be set to REGISTERING.
   * When the registration is done, the state will be set to REGISTERED.
   * When the sync receives a registration command from the master during the heartbeat,
   * the state will be reset to NOT_REGISTERED and the sync will attempt to register it again
   * in the next heartbeat.
   */
  private volatile WorkerMasterRegistrationState mWorkerState =
      WorkerMasterRegistrationState.NOT_REGISTERED;

  /**
   * An async service to remove block.
   */
  private final AsyncBlockRemover mAsyncBlockRemover;

  /**
   * The worker ID for the worker. This may change if the master asks the worker to re-register.
   */
  private final AtomicReference<Long> mWorkerId;

  /**
   * Client for all master communication.
   */
  private final BlockMasterClient mMasterClient;

  /**
   * The net address of the worker.
   */
  private final WorkerNetAddress mWorkerAddress;

  /**
   * The helper instance for sync related methods.
   */
  private final BlockMasterSyncHelper mBlockMasterSyncHelper;

  /**
   * The block worker responsible for interacting with Alluxio and UFS storage.
   */
  private final BlockWorker mBlockWorker;
  /**
   * Last System.currentTimeMillis() timestamp when a heartbeat successfully completed.
   */
  private long mLastSuccessfulHeartbeatMs = 0;

  private final BlockHeartbeatReporter mBlockHeartbeatReporter;

  /**
   * Creates a new instance of {@link SpecificMasterBlockSync}.
   *
   * @param blockWorker the {@link BlockWorker} this syncer is updating to
   * @param masterClient the block master client
   * @param heartbeatReporter the heartbeat reporter
   */
  public SpecificMasterBlockSync(
      BlockWorker blockWorker,
      BlockMasterClient masterClient, BlockHeartbeatReporter heartbeatReporter)
      throws IOException {
    mBlockWorker = blockWorker;
    mWorkerId = blockWorker.getWorkerId();
    mWorkerAddress = blockWorker.getWorkerAddress();
    mMasterClient = masterClient;
    mAsyncBlockRemover = new AsyncBlockRemover(mBlockWorker);
    mBlockMasterSyncHelper = new BlockMasterSyncHelper(mMasterClient);
    mMasterAddress = masterClient.getRemoteSockAddress();
    mBlockHeartbeatReporter = heartbeatReporter;
  }

  private void registerWithMaster() {
    RetryPolicy retry = createEndlessRetry();
    while (retry.attempt()) {
      try {
        LOG.info("Registering with master {}", mMasterAddress);
        // The content in the report can be cleared because registration will
        // report these block information anyways.
        mBlockHeartbeatReporter.clear();
        registerWithMasterInternal();
        LOG.info("Finished registration with {}", mMasterAddress);
        return;
      } catch (Exception e) {
        LOG.error("Failed to register with master {}, error {}, retry count {} Will retry...",
            mMasterAddress, e, retry.getAttemptCount());
        mWorkerState = WorkerMasterRegistrationState.NOT_REGISTERED;
      }
    }
    // Should not reach here because the retry is indefinite
    ProcessUtils.fatalError(LOG, new RuntimeException(),
        "Failed to register with master %s", mMasterAddress);
  }

  protected void registerWithMasterInternal()
      throws IOException, FailedToAcquireRegisterLeaseException {
    // The target master is not necessarily the one that allocated the workerID
    LOG.info("Notify the master {} about the workerID {}", mMasterAddress, mWorkerId);
    mMasterClient.notifyWorkerId(mWorkerId.get(), mWorkerAddress);
    // TODO(elega) If worker registration to all masters happens at the same time,
    // this might cause worker OOM issues because each block sync thread will hold a BlockStoreMeta
    // instance during the registration.
    // If this happens, consider limiting the worker registration concurrency,
    // e.g. register the worker to masters one by one.
    BlockStoreMeta storeMeta = mBlockWorker.getStoreMetaFull();

    try {
      mBlockMasterSyncHelper.tryAcquireLease(mWorkerId.get(), storeMeta);
    } catch (FailedToAcquireRegisterLeaseException e) {
      if (Configuration.getBoolean(PropertyKey.TEST_MODE)) {
        throw new RuntimeException(String.format("Master register lease timeout exceeded: %dms",
            ACQUIRE_LEASE_WAIT_MAX_DURATION));
      }
      throw e;
    }
    mWorkerState = WorkerMasterRegistrationState.REGISTERING;
    mBlockMasterSyncHelper.registerToMaster(mWorkerId.get(), storeMeta);

    mWorkerState = WorkerMasterRegistrationState.REGISTERED;
    Metrics.WORKER_MASTER_REGISTRATION_SUCCESS_COUNT.inc();
    mLastSuccessfulHeartbeatMs = CommonUtils.getCurrentMs();
  }

  private RetryPolicy createEndlessRetry() {
    return new ExponentialBackoffRetry(
        1000, 60 * 1000, Integer.MAX_VALUE);
  }

  @Override
  public synchronized void heartbeat(long runLimit) throws InterruptedException {
    if (mWorkerState == WorkerMasterRegistrationState.NOT_REGISTERED) {
      // Not registered because:
      // 1. The worker just started, we kick off the 1st registration here.
      // 2. Master sends a registration command during
      // the heartbeat and resets the registration state. (e.g. master restarted)
      // 3. The heartbeat message becomes too big that we decide to fall back to a full re-register
      LOG.info("The worker needs to register with master {}", mMasterAddress);
      // This will retry indefinitely and essentially block here if the master is not ready
      registerWithMaster();
      LOG.info("BlockMasterSync to master {} has started", mMasterAddress);
    }
    if (mWorkerState == WorkerMasterRegistrationState.REGISTERING) {
      return;
    }

    RetryPolicy endlessRetry = createEndlessRetry();
    while (endlessRetry.attempt()) {
      BlockHeartbeatReport report = mBlockHeartbeatReporter.generateReportAndClear();
      boolean success = false;
      try {
        beforeHeartbeat();
        success = mBlockMasterSyncHelper.heartbeat(
            mWorkerId.get(), report,
            mBlockWorker.getStoreMeta(), this::handleMasterCommand);
      } catch (Exception e) {
        LOG.error("Failed to receive master heartbeat command. worker id {}", mWorkerId, e);
      }
      if (success) {
        mLastSuccessfulHeartbeatMs = CommonUtils.getCurrentMs();
        break;
      } else {
        LOG.warn(
            "Heartbeat failed, worker id {}, worker host {} # of attempts {}, last success ts {}",
            mWorkerId.get(), mWorkerAddress.getHost(), endlessRetry.getAttemptCount(),
            mLastSuccessfulHeartbeatMs);
        if (report.getBlockChangeCount() >= mWorkerBlockHeartbeatReportSizeThreshold) {
          // If the report becomes too big, merging it back to the reporter might cause OOM issue.
          // We throw away the result and let the worker re-register with the master.
          mWorkerState = WorkerMasterRegistrationState.NOT_REGISTERED;
          return;
        } else {
          mBlockHeartbeatReporter.mergeBack(report);
        }
      }
    }
  }

  protected void beforeHeartbeat() {
  }

  @Override
  public void close() {
    mAsyncBlockRemover.shutDown();
    mMasterClient.close();
  }

  /**
   * @return if the worker has registered with the master successfully
   */
  public boolean isRegistered() {
    return mWorkerState == WorkerMasterRegistrationState.REGISTERED;
  }

  /**
   * Handles a master command. The command is one of Unknown, Nothing, Register, Free, or Delete.
   * This call will block until the command is complete.
   *
   * @param cmd the command to execute
   * @throws IOException               if I/O errors occur
   * @throws ConnectionFailedException if connection fails
   */
  private void handleMasterCommand(Command cmd) throws IOException, ConnectionFailedException {
    if (cmd == null) {
      return;
    }
    switch (cmd.getCommandType()) {
      // Currently unused
      case Delete:
        break;
      // Master requests blocks to be removed from Alluxio managed space.
      case Free:
        mAsyncBlockRemover.addBlocksToDelete(cmd.getDataList());
        break;
      // No action required
      case Nothing:
        break;
      // Master requests re-registration
      case Register:
        mWorkerState = WorkerMasterRegistrationState.NOT_REGISTERED;
        break;
      // Unknown request
      case Unknown:
        LOG.error("Master heartbeat sends unknown command {}", cmd);
        break;
      default:
        throw new RuntimeException("Un-recognized command from master " + cmd);
    }
  }

  /**
   * Metrics.
   */
  public static final class Metrics {
    private static final Counter WORKER_MASTER_REGISTRATION_SUCCESS_COUNT
        = MetricsSystem.counter(MetricKey.WORKER_MASTER_REGISTRATION_SUCCESS_COUNT.getName());
  }
}
