alluxio.heartbeat.HeartbeatScheduler.java Source code

Introduction

Here is the source code for alluxio.heartbeat.HeartbeatScheduler.java

Source

/*
 * 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);
    }
}