alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/metrics/MetricsHeartbeatContext.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.metrics;

import alluxio.ClientContext;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.master.MasterInquireClient;
import alluxio.util.IdUtils;
import alluxio.util.ThreadFactoryUtils;

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

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import javax.annotation.concurrent.ThreadSafe;

/**
 * A class used to track metrics heartbeats to a master.
 *
 * The class should be instantiated when a new FileSystemContext is created with a
 * configuration that points to a given master. As new FileSystemContexts are created, if they
 * utilize the same connection details, then they can simply be added to this context so
 * that their information is included in the metrics heartbeat. To add them, one should simply
 * call {@link #addHeartbeat(ClientContext, MasterInquireClient)} with the necessary arguments.
 *
 * For each separate set of connection details, a new instance of this class is created. As
 * FileSystemContexts are closed, they remove themselves from the internal metrics heartbeat.
 * When a context reaches 0 tracked contexts it will automatically close and remove itself from
 * the internal MASTER_METRICS_HEARTBEAT map.
 *
 * When the final FileSystemContext closes and removes its heartbeat from metrics it will also
 * shutdown and close the executor service until a new {@link alluxio.client.file.FileSystemContext}
 * is created.
 */
@ThreadSafe
public class MetricsHeartbeatContext {
  private static final Logger LOG = LoggerFactory.getLogger(MetricsHeartbeatContext.class);

  /** A map from master connection details to heartbeat context instances. */
  private static final Map<MasterInquireClient.ConnectDetails, MetricsHeartbeatContext>
      MASTER_METRICS_HEARTBEAT = new ConcurrentHashMap<>(2);

  /** A value that tracks whether or not we've registered the shutdown hook for the JVM. */
  private static boolean sAddedShudownHook = false;

  /** Application ID for the JVM. Initialized lazily once the first heartbeat is added. */
  private static String sAppId = null;

  /** The service which executes metrics heartbeat RPCs. */
  private static ScheduledExecutorService sExecutorService;

  private final MasterInquireClient.ConnectDetails mConnectDetails;
  private final ClientMasterSync mClientMasterSync;
  private final AlluxioConfiguration mConf;

  // This can only be a primitive if all accesses are synchronized
  private int mCtxCount;
  private ScheduledFuture<?> mMetricsMasterHeartbeatTask;

  private MetricsHeartbeatContext(ClientContext ctx, MasterInquireClient inquireClient) {
    mCtxCount = 0;
    mConnectDetails = inquireClient.getConnectDetails();
    mConf = ctx.getClusterConf();
    mClientMasterSync = new ClientMasterSync(sAppId, ctx, inquireClient);
  }

  private synchronized void addContext() {
    // increment and lazily schedule the new heartbeat task if it is the first one
    if (mCtxCount++ == 0) {
      mMetricsMasterHeartbeatTask =
          sExecutorService.scheduleWithFixedDelay(mClientMasterSync::heartbeat,
              mConf.getMs(PropertyKey.USER_METRICS_HEARTBEAT_INTERVAL_MS),
              mConf.getMs(PropertyKey.USER_METRICS_HEARTBEAT_INTERVAL_MS),
              TimeUnit.MILLISECONDS);
    }
  }

  private synchronized void heartbeat() {
    mClientMasterSync.heartbeat();
  }

  /**
   * Remove an application from this metrics heartbeat.
   *
   * A user who calls this method should assume the reference to this context is invalid
   * afterwards. It will automatically close and remove itself from all tracking if the number
   * of open contexts for this heartbeat reaches 0. Never attempt to add another context with
   * the same reference after removing.
   */
  private synchronized void removeContext() {
    if (--mCtxCount <= 0) {
      close();
    }
  }

  /**
   * When closed, this method will remove its task from the scheduled executor.
   *
   * It will also remove itself from being tracked in the MASTER_METRICS_HEARTBEAT. It should
   * only ever be called in {@link #removeContext()} when the context count reaches 0. Afterwards,
   * this reference should be discarded.
   */
  private synchronized void close() {
    if (mMetricsMasterHeartbeatTask != null) {
      mMetricsMasterHeartbeatTask.cancel(false);
    }
    MASTER_METRICS_HEARTBEAT.remove(mConnectDetails);
    // Trigger the last heartbeat to preserve the client side metrics changes
    heartbeat();
    mClientMasterSync.close();
  }

  /**
   * Sets up a new metrics heartbeat with the given client information.
   *
   * This will instantiate a new executor service if it is the first heartbeat to be added,
   * otherwise the application Id is simply included in an already existing heartbeat. This helps
   * to consolidate RPCs and utilize less resources on the client.
   *
   * @param ctx The application's client context
   * @param inquireClient the master inquire client used to connect to the master
   */
  public static synchronized void addHeartbeat(ClientContext ctx,
      MasterInquireClient inquireClient) {
    Preconditions.checkNotNull(ctx);
    Preconditions.checkNotNull(inquireClient);

    // Lazily initializing the executor service for first heartbeat
    // Relies on the method being synchronized
    if (sExecutorService == null) {
      sExecutorService = Executors.newSingleThreadScheduledExecutor(
          ThreadFactoryUtils.build("metrics-master-heartbeat-%d", true));
    }

    // register the shutdown hook if it hasn't been set up already
    if (!sAddedShudownHook) {
      try {
        Runtime.getRuntime().addShutdownHook(new MetricsMasterSyncShutDownHook());
        sAddedShudownHook = true;
      } catch (IllegalStateException e) {
        // this exception is thrown when the system is already in the process of shutting down. In
        // such a situation, we are about to shutdown anyway and there is no need to register this
        // shutdown hook
      } catch (SecurityException e) {
        LOG.info("Not registering metrics shutdown hook due to security exception. Regular "
            + "heartbeats will still be performed to collect metrics data, but no final "
            + "heartbeat will be performed on JVM exit. Security exception: {}", e.toString());
      }
    }

    if (sAppId == null) {
      sAppId = IdUtils.createOrGetAppIdFromConfig(ctx.getClusterConf());
      LOG.info("Created metrics heartbeat with ID {}. This ID will be used for identifying info "
              + "from the client. It can be set manually through the {} property",
          sAppId, PropertyKey.Name.USER_APP_ID);
    }

    MetricsHeartbeatContext heartbeatCtx = MASTER_METRICS_HEARTBEAT.computeIfAbsent(
        inquireClient.getConnectDetails(),
        (addr) -> new MetricsHeartbeatContext(ctx, inquireClient));
    heartbeatCtx.addContext();
    LOG.debug("Registered metrics heartbeat with appId: {}", sAppId);
  }

  /**
   * Removes an application from the metrics heartbeat.
   *
   * If this is the last application to be removed for a given master then it will cancel the
   * execution of the metrics RPC for that master.
   *
   * If this is the last application to be removed for the whole JVM, it will shutdown the
   * executor threadpool until another application needs to heartbeat. If the appId which is to
   * be removed isn't found, the application will silently continue.
   *
   * @param ctx The client context used to register the heartbeat
   */
  public static synchronized void removeHeartbeat(ClientContext ctx) {
    MasterInquireClient.ConnectDetails connectDetails =
        MasterInquireClient.Factory.getConnectDetails(ctx.getClusterConf());
    MetricsHeartbeatContext heartbeatCtx = MASTER_METRICS_HEARTBEAT.get(connectDetails);
    if (heartbeatCtx != null) {
      heartbeatCtx.removeContext();
      LOG.debug("De-registered metrics heartbeat with appId: {}", sAppId);
    }

    if (MASTER_METRICS_HEARTBEAT.isEmpty()) {
      sExecutorService.shutdown();
      try {
        sExecutorService.awaitTermination(5000, TimeUnit.MILLISECONDS);
      } catch (InterruptedException e) {
        LOG.warn("Metrics heartbeat executor did not shut down in a timely manner: {}",
            e.toString());
      }
      sExecutorService = null;
      sAppId = null;
    }
  }

  /**
   * Class that heartbeats to the metrics master before exit. The heartbeat is performed in a second
   * thread so that we can exit early if the heartbeat is taking too long.
   */
  private static final class MetricsMasterSyncShutDownHook extends Thread {
    private final Thread mLastHeartbeatThread;

    /**
     * Creates a new metrics master shutdown hook.
     */
    public MetricsMasterSyncShutDownHook() {
      mLastHeartbeatThread = new Thread(() -> {
        if (sExecutorService != null) {
          MASTER_METRICS_HEARTBEAT.forEach((key, value) -> value.heartbeat());
        }
      });
      mLastHeartbeatThread.setDaemon(true);
    }

    @Override
    public void run() {
      mLastHeartbeatThread.start();
      try {
        // Shutdown hooks should run quickly, so we limit the wait time to 500ms. It isn't the end
        // of the world if the final heartbeat fails.
        mLastHeartbeatThread.join(500);
      } catch (InterruptedException e) {
        return;
      } finally {
        if (mLastHeartbeatThread.isAlive()) {
          LOG.warn("Failed to heartbeat to the metrics master before exit");
        }
      }
    }
  }
}
