@Override
public void handle(TaskSchedulerEvent event) {
if (event.getType() == EventType.T_SCHEDULE) {
if (event instanceof FragmentScheduleEvent) {
FragmentScheduleEvent castEvent = (FragmentScheduleEvent) event;
if (context.isLeafQuery()) {
TaskAttemptScheduleContext taskContext = new TaskAttemptScheduleContext();
Task task = Stage.newEmptyTask(context, taskContext, stage, nextTaskId++);
task.addFragment(castEvent.getLeftFragment(), true);
scheduledObjectNum++;
if (castEvent.hasRightFragments()) {
task.addFragments(castEvent.getRightFragments());
}
stage.getEventHandler().handle(new TaskEvent(task.getId(), TaskEventType.T_SCHEDULE));
} else {
fragmentsForNonLeafTask = new FileFragment[2];
fragmentsForNonLeafTask[0] = castEvent.getLeftFragment();
if (castEvent.hasRightFragments()) {
Collection<Fragment> var = castEvent.getRightFragments();
FileFragment[] rightFragments = var.toArray(new FileFragment[var.size()]);
fragmentsForNonLeafTask[1] = rightFragments[0];
if (rightFragments.length > 1) {
broadcastFragmentsForNonLeafTask = new FileFragment[rightFragments.length - 1];
System.arraycopy(
rightFragments,
1,
broadcastFragmentsForNonLeafTask,
0,
broadcastFragmentsForNonLeafTask.length);
} else {
broadcastFragmentsForNonLeafTask = null;
}
}
}
} else if (event instanceof FetchScheduleEvent) {
FetchScheduleEvent castEvent = (FetchScheduleEvent) event;
Map<String, List<FetchImpl>> fetches = castEvent.getFetches();
TaskAttemptScheduleContext taskScheduleContext = new TaskAttemptScheduleContext();
Task task = Stage.newEmptyTask(context, taskScheduleContext, stage, nextTaskId++);
scheduledObjectNum++;
for (Entry<String, List<FetchImpl>> eachFetch : fetches.entrySet()) {
task.addFetches(eachFetch.getKey(), eachFetch.getValue());
task.addFragment(fragmentsForNonLeafTask[0], true);
if (fragmentsForNonLeafTask[1] != null) {
task.addFragment(fragmentsForNonLeafTask[1], true);
}
}
if (broadcastFragmentsForNonLeafTask != null
&& broadcastFragmentsForNonLeafTask.length > 0) {
task.addFragments(Arrays.asList(broadcastFragmentsForNonLeafTask));
}
stage.getEventHandler().handle(new TaskEvent(task.getId(), TaskEventType.T_SCHEDULE));
} else if (event instanceof TaskAttemptToSchedulerEvent) {
TaskAttemptToSchedulerEvent castEvent = (TaskAttemptToSchedulerEvent) event;
if (context.isLeafQuery()) {
scheduledRequests.addLeafTask(castEvent);
} else {
scheduledRequests.addNonLeafTask(castEvent);
}
if (needWakeup.getAndSet(false)) {
// wake up scheduler thread after scheduled
synchronized (schedulingThread) {
schedulingThread.notifyAll();
}
}
}
} else if (event.getType() == EventType.T_SCHEDULE_CANCEL) {
// when a stage is killed, unassigned query unit attmpts are canceled from the scheduler.
// This event is triggered by TaskAttempt.
TaskAttemptToSchedulerEvent castedEvent = (TaskAttemptToSchedulerEvent) event;
scheduledRequests.leafTasks.remove(castedEvent.getTaskAttempt().getId());
LOG.info(
castedEvent.getTaskAttempt().getId()
+ " is canceled from "
+ this.getClass().getSimpleName());
((TaskAttemptToSchedulerEvent) event)
.getTaskAttempt()
.handle(
new TaskAttemptEvent(
castedEvent.getTaskAttempt().getId(), TaskAttemptEventType.TA_SCHEDULE_CANCELED));
}
}
public void schedule() throws Exception {
try {
final int incompleteTaskNum =
scheduledRequests.leafTaskNum() + scheduledRequests.nonLeafTaskNum();
if (incompleteTaskNum == 0) {
needWakeup.set(true);
// all task is done or tasks is not scheduled
synchronized (schedulingThread) {
schedulingThread.wait(1000);
}
} else {
LinkedList<TaskRequestEvent> taskRequests = createTaskRequest(incompleteTaskNum);
if (taskRequests.size() == 0) {
synchronized (schedulingThread) {
schedulingThread.wait(schedulerDelay);
}
} else {
if (LOG.isDebugEnabled()) {
LOG.debug("Get " + taskRequests.size() + " taskRequestEvents ");
}
if (isLeaf) {
scheduledRequests.assignToLeafTasks(taskRequests);
} else {
scheduledRequests.assignToNonLeafTasks(taskRequests);
}
}
}
} catch (TimeoutException e) {
LOG.error(e.getMessage());
}
}
@Override
public void stop() {
isStopped = true;
if (schedulingThread != null) {
synchronized (schedulingThread) {
schedulingThread.interrupt();
}
}
candidateWorkers.clear();
scheduledRequests.clear();
LOG.info("Task Scheduler stopped");
super.stop();
}