@Override
public void killFragment(String fragmentId) {
synchronized (lock) {
TaskWrapper taskWrapper = knownTasks.remove(fragmentId);
// Can be null since the task may have completed meanwhile.
if (taskWrapper != null) {
if (taskWrapper.inWaitQueue) {
if (isDebugEnabled) {
LOG.debug("Removing {} from waitQueue", fragmentId);
}
taskWrapper.setIsInWaitQueue(false);
waitQueue.remove(taskWrapper);
}
if (taskWrapper.inPreemptionQueue) {
if (isDebugEnabled) {
LOG.debug("Removing {} from preemptionQueue", fragmentId);
}
taskWrapper.setIsInPreemptableQueue(false);
preemptionQueue.remove(taskWrapper);
}
taskWrapper.getTaskRunnerCallable().killTask();
} else {
LOG.info("Ignoring killFragment request for {} since it isn't known", fragmentId);
}
lock.notify();
}
}
private void finishableStateUpdated(TaskWrapper taskWrapper, boolean newFinishableState) {
synchronized (lock) {
if (taskWrapper.isInWaitQueue()) {
// Re-order the wait queue
LOG.debug(
"Re-ordering the wait queue since {} finishable state moved to {}",
taskWrapper.getRequestId(),
newFinishableState);
if (waitQueue.remove(taskWrapper)) {
// Put element back only if it existed.
waitQueue.offer(taskWrapper);
} else {
LOG.warn(
"Failed to remove {} from waitQueue",
taskWrapper.getTaskRunnerCallable().getRequestId());
}
}
if (newFinishableState == true && taskWrapper.isInPreemptionQueue()) {
LOG.debug(
"Removing {} from preemption queue because it's state changed to {}",
taskWrapper.getRequestId(),
newFinishableState);
preemptionQueue.remove(taskWrapper.getTaskRunnerCallable());
} else if (newFinishableState == false
&& !taskWrapper.isInPreemptionQueue()
&& !taskWrapper.isInWaitQueue()) {
LOG.debug(
"Adding {} to preemption queue since finishable state changed to {}",
taskWrapper.getRequestId(),
newFinishableState);
preemptionQueue.offer(taskWrapper);
}
lock.notify();
}
}
@Override
public void schedule(TaskRunnerCallable task) throws RejectedExecutionException {
TaskWrapper taskWrapper = new TaskWrapper(task, this);
TaskWrapper evictedTask;
synchronized (lock) {
// If the queue does not have capacity, it does not throw a Rejection. Instead it will
// return the task with the lowest priority, which could be the task which is currently being
// processed.
// TODO HIVE-11687 It's possible for a bunch of tasks to come in around the same time, without
// the
// actual executor threads picking up any work. This will lead to unnecessary rejection of
// tasks.
// The wait queue should be able to fit at least (waitQueue + currentFreeExecutor slots)
evictedTask = waitQueue.offer(taskWrapper);
if (evictedTask != taskWrapper) {
knownTasks.put(taskWrapper.getRequestId(), taskWrapper);
taskWrapper.setIsInWaitQueue(true);
if (isDebugEnabled) {
LOG.debug(
"{} added to wait queue. Current wait queue size={}",
task.getRequestId(),
waitQueue.size());
}
} else {
if (isInfoEnabled) {
LOG.info("wait queue full, size={}. {} not added", waitQueue.size(), task.getRequestId());
}
evictedTask.getTaskRunnerCallable().killTask();
throw new RejectedExecutionException("Wait queue full");
}
}
// At this point, the task has been added into the queue. It may have caused an eviction for
// some other task.
// This registration has to be done after knownTasks has been populated.
// Register for state change notifications so that the waitQueue can be re-ordered correctly
// if the fragment moves in or out of the finishable state.
boolean canFinish = taskWrapper.getTaskRunnerCallable().canFinish();
// It's safe to register outside of the lock since the stateChangeTracker ensures that updates
// and registrations are mutually exclusive.
taskWrapper.maybeRegisterForFinishedStateNotifications(canFinish);
if (isDebugEnabled) {
LOG.debug("Wait Queue: {}", waitQueue);
}
if (evictedTask != null) {
knownTasks.remove(evictedTask.getRequestId());
evictedTask.maybeUnregisterForFinishedStateNotifications();
evictedTask.setIsInWaitQueue(false);
evictedTask.getTaskRunnerCallable().killTask();
if (isInfoEnabled) {
LOG.info(
"{} evicted from wait queue in favor of {} because of lower priority",
evictedTask.getRequestId(),
task.getRequestId());
}
}
synchronized (lock) {
lock.notify();
}
}
