@Override
public void onSuccess(TaskRunner2Result result) {
knownTasks.remove(taskWrapper.getRequestId());
taskWrapper.setIsInPreemptableQueue(false);
taskWrapper.maybeUnregisterForFinishedStateNotifications();
taskWrapper.getTaskRunnerCallable().getCallback().onSuccess(result);
updatePreemptionListAndNotify(result.getEndReason());
}
@Override
public void onFailure(Throwable t) {
knownTasks.remove(taskWrapper.getRequestId());
taskWrapper.setIsInPreemptableQueue(false);
taskWrapper.maybeUnregisterForFinishedStateNotifications();
taskWrapper.getTaskRunnerCallable().getCallback().onFailure(t);
updatePreemptionListAndNotify(null);
LOG.error("Failed notification received: Stacktrace: " + ExceptionUtils.getStackTrace(t));
}
@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();
}
}