private void addLeafTask(TaskAttemptToSchedulerEvent event) {
TaskAttempt taskAttempt = event.getTaskAttempt();
List<DataLocation> locations = taskAttempt.getTask().getDataLocations();
for (DataLocation location : locations) {
String host = location.getHost();
leafTaskHosts.add(host);
HostVolumeMapping hostVolumeMapping = leafTaskHostMapping.get(host);
if (hostVolumeMapping == null) {
String rack = RackResolver.resolve(host).getNetworkLocation();
hostVolumeMapping = new HostVolumeMapping(host, rack);
leafTaskHostMapping.put(host, hostVolumeMapping);
}
hostVolumeMapping.addTaskAttempt(location.getVolumeId(), taskAttempt);
if (LOG.isDebugEnabled()) {
LOG.debug("Added attempt req to host " + host);
}
HashSet<TaskAttemptId> list = leafTasksRackMapping.get(hostVolumeMapping.getRack());
if (list == null) {
list = new HashSet<>();
leafTasksRackMapping.put(hostVolumeMapping.getRack(), list);
}
list.add(taskAttempt.getId());
if (LOG.isDebugEnabled()) {
LOG.debug("Added attempt req to rack " + hostVolumeMapping.getRack());
}
}
leafTasks.add(taskAttempt.getId());
}
private synchronized TaskAttemptId getAndRemove(int volumeId) {
TaskAttemptId taskAttemptId = null;
if (!unassignedTaskForEachVolume.containsKey(volumeId)) {
if (volumeId > REMOTE) {
diskVolumeLoads.remove(volumeId);
}
return taskAttemptId;
}
LinkedHashSet<TaskAttempt> list = unassignedTaskForEachVolume.get(volumeId);
if (list != null && !list.isEmpty()) {
TaskAttempt taskAttempt;
synchronized (unassignedTaskForEachVolume) {
Iterator<TaskAttempt> iterator = list.iterator();
taskAttempt = iterator.next();
iterator.remove();
}
taskAttemptId = taskAttempt.getId();
for (DataLocation location : taskAttempt.getTask().getDataLocations()) {
HostVolumeMapping volumeMapping =
scheduledRequests.leafTaskHostMapping.get(location.getHost());
if (volumeMapping != null) {
volumeMapping.removeTaskAttempt(location.getVolumeId(), taskAttempt);
}
}
increaseConcurrency(volumeId);
}
return taskAttemptId;
}
public void releaseTaskAttempt(TaskAttempt taskAttempt) {
if (taskAttempt.isLeafTask() && taskAttempt.getWorkerConnectionInfo() != null) {
HostVolumeMapping mapping =
scheduledRequests.leafTaskHostMapping.get(
taskAttempt.getWorkerConnectionInfo().getHost());
if (mapping != null && mapping.lastAssignedVolumeId.containsKey(taskAttempt.getId())) {
mapping.decreaseConcurrency(mapping.lastAssignedVolumeId.remove(taskAttempt.getId()));
}
}
}
public void cancel(TaskAttempt taskAttempt) {
if (taskAttempt.isLeafTask()) {
releaseTaskAttempt(taskAttempt);
List<DataLocation> locations = taskAttempt.getTask().getDataLocations();
for (DataLocation location : locations) {
HostVolumeMapping volumeMapping =
scheduledRequests.leafTaskHostMapping.get(location.getHost());
volumeMapping.addTaskAttempt(location.getVolumeId(), taskAttempt);
}
scheduledRequests.leafTasks.add(taskAttempt.getId());
} else {
scheduledRequests.nonLeafTasks.add(taskAttempt.getId());
}
context
.getMasterContext()
.getEventHandler()
.handle(new TaskAttemptEvent(taskAttempt.getId(), TaskAttemptEventType.TA_ASSIGN_CANCEL));
}