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.heartbeat;
  
  import alluxio.resource.LockResource;
  
  import com.google.common.base.Preconditions;
  
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.locks.Condition;
  import java.util.concurrent.locks.Lock;
  import java.util.concurrent.locks.ReentrantLock;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * This class can be used for controlling heartbeat execution of a thread.
   *
   * In particular, the {@link #await(String)} method can be used to wait for a thread to start
   * waiting to be scheduled and the {@link #schedule(String)} method can be used to schedule a thread
   * that is waiting to be scheduled.
   *
   * The contract of this class is that the {@link #schedule(String)} method should only be called for
   * threads that are in fact waiting to be scheduled. If this contract is not met, concurrent
   * execution of a thread and the heartbeat scheduler logic to schedule the thread can lead to a
   * deadlock.
   *
   * For an example of how to use the {@link HeartbeatScheduler}, see unit test of
   * {@link HeartbeatThread}.
   */
  @ThreadSafe
  public final class HeartbeatScheduler {
    /**
     * A map from thread name to active timer for that thread. A timer is active when its thread is
     * waiting to be scheduled.
     */
    private static Map<String, ScheduledTimer> sTimers = new HashMap<>();
    private static Lock sLock = new ReentrantLock();
    private static Condition sCondition = sLock.newCondition();
  
    private HeartbeatScheduler() {} // to prevent initialization
  
    /**
     * @param timer a timer to add to the scheduler
     */
    public static void addTimer(ScheduledTimer timer) {
      Preconditions.checkNotNull(timer, "timer");
      try (LockResource r = new LockResource(sLock)) {
        Preconditions.checkState(!sTimers.containsKey(timer.getThreadName()),
            "The timer for thread %s is already waiting to be scheduled", timer.getThreadName());
        sTimers.put(timer.getThreadName(), timer);
        sCondition.signalAll();
      }
    }
  
    /**
     * Removes a timer name from the scheduler if it exists.
     *
     * @param name the name to clear
     */
    public static void clearTimer(String name) {
      try (LockResource r = new LockResource(sLock)) {
        sTimers.remove(name);
      }
    }
  
    /**
     * Removes a timer from the scheduler.
     *
     * This method will fail if the timer is not in the scheduler.
     *
     * @param timer the timer to remove
     */
    public static void removeTimer(ScheduledTimer timer) {
      Preconditions.checkNotNull(timer, "timer");
      try (LockResource r = new LockResource(sLock)) {
        ScheduledTimer removedTimer = sTimers.remove(timer.getThreadName());
        Preconditions.checkNotNull(removedTimer, "sTimers should contain %s", timer.getThreadName());
        Preconditions.checkState(removedTimer == timer,
            "sTimers should contain the timer being removed");
      }
    }
  
    /**
     * @return the set of threads present in the scheduler
     */
    public static Set<String> getThreadNames() {
      try (LockResource r = new LockResource(sLock)) {
        return sTimers.keySet();
      }
    }
  
    /**
     * Schedules execution of a heartbeat for the given thread.
     *
     * @param threadName a name of the thread for which heartbeat is to be executed
     */
    public static void schedule(String threadName) {
      try (LockResource r = new LockResource(sLock)) {
        ScheduledTimer timer = sTimers.get(threadName);
        if (timer == null) {
          throw new RuntimeException("Timer for thread " + threadName + " not found.");
        }
        timer.schedule();
      }
    }
  
    /**
     * Waits for the given thread to be ready to be scheduled.
     *
     * @param name a name of the thread to wait for
     * @throws InterruptedException if the waiting thread is interrupted
     */
    public static void await(String name) throws InterruptedException {
      try (LockResource r = new LockResource(sLock)) {
        while (!sTimers.containsKey(name)) {
          sCondition.await();
        }
      }
    }
  
    /**
     * Waits until the given thread can be executed, throwing an unchecked exception of the given
     * timeout expires.
     *
     * @param name a name of the thread to wait for
     * @param time the maximum time to wait
     * @param unit the time unit of the {@code time} argument
     * @throws InterruptedException if the waiting thread is interrupted
     */
    public static void await(String name, long time, TimeUnit unit) throws InterruptedException {
      try (LockResource r = new LockResource(sLock)) {
        while (!sTimers.containsKey(name)) {
          if (!sCondition.await(time, unit)) {
            throw new RuntimeException(
                "Timed out waiting for thread " + name + " to be ready for scheduling");
          }
        }
      }
    }
  
    /**
     * Convenience method for executing a heartbeat and waiting for it to complete.
     *
     * @param name the name of the heartbeat to execute
     * @throws InterruptedException if the waiting thread is interrupted
     */
    public static void execute(String name) throws InterruptedException {
      await(name);
      schedule(name);
      await(name);
    }
  }
  