mesos/src/main/java/org/apache/hama/bsp/TaskDelegator.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.hama.bsp;

import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.commons.collections.map.MultiValueMap;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.commons.math3.util.Pair;
import org.apache.hama.ipc.GroomProtocol;
import org.apache.mesos.Protos;
import org.apache.mesos.SchedulerDriver;

public class TaskDelegator implements GroomStatusListener {
  public static final Log LOG = LogFactory.getLog(MesosScheduler.class);

  private Set<TaskInProgress> recoveryTasks;

  /**
   * Map to hold assignments from groomServerNames to TasksInProgress
   */
  private MultiValueMap assignments = new MultiValueMap();

  private MultiValueMap jobAssignments = new MultiValueMap();

  private AtomicReference<GroomServerManager> groomServerManager;

  /**
   * Map from Pair of groomServerName and port number to the latest
   * GroomServerStatus
   */
  private Map<Pair<String, Integer>, GroomServerStatus> groomServers = new HashMap<Pair<String, Integer>, GroomServerStatus>();

  private Map<Pair<String, Integer>, Protos.TaskID> groomTaskIDs = new HashMap<Pair<String, Integer>, Protos.TaskID>();

  private SchedulerDriver driver;

  public TaskDelegator(AtomicReference<GroomServerManager> groomServerManager,
      SchedulerDriver driver, Set<TaskInProgress> recoveryTasks) {
    this.groomServerManager = groomServerManager;
    groomServerManager.get().addJobInProgressListener(
        new TaskDelegatorJobListener());
    this.driver = driver;
    this.recoveryTasks = recoveryTasks;
  }

  @Override
  public void groomServerRegistered(GroomServerStatus status) {
    Pair<String, Integer> key = new Pair<String, Integer>(
        status.getGroomHostName(), BSPMaster.resolveWorkerAddress(
            status.rpcServer).getPort());
    LOG.debug("Received Groom From: " + key.getKey() + ":" + key.getValue());

    if (assignments.containsKey(key)) {
      for (Object tip : assignments.getCollection(key)) {
        execute((TaskInProgress) tip, status);
      }
    } else {
      LOG.error("Unexpected host found: " + key.getKey() + ":" + key.getValue());
    }
    groomServers.put(key, status);
  }

  /**
   * Add a task for execution when the groom server becomes available
   * 
   * @param tip The TaskInProgress to execute
   * @param hostName The hostname where the resource reservation was made
   */
  public void addTask(TaskInProgress tip, Protos.TaskID taskId,
      String hostName, Integer port) {
    LOG.trace("Adding Host: " + hostName + ":" + port + " for Task:"
        + tip.getTaskId());
    Pair<String, Integer> key = new Pair<String, Integer>(hostName, port);

    groomTaskIDs.put(key, taskId);

    if (groomServers.containsKey(key)) {
      execute(tip, groomServers.get(key));
    } else {
      assignments.put(key, tip);
      jobAssignments.put(tip.getJob(), new Pair<Object, Object>(key, tip));
    }
  }

  private void execute(TaskInProgress tip, GroomServerStatus status) {
    Task task = tip.constructTask(status);

    GroomServerAction[] actions;
    GroomProtocol worker = groomServerManager.get().findGroomServer(status);

    if (!recoveryTasks.contains(tip)) {
      actions = new GroomServerAction[1];
      actions[0] = new LaunchTaskAction(task);
    } else {
      LOG.trace("Executing a recovery task");
      recoveryTasks.remove(tip);
      HashMap<String, GroomServerStatus> groomStatuses = new HashMap<String, GroomServerStatus>(
          1);
      groomStatuses.put(status.hostName, status);
      Map<GroomServerStatus, List<GroomServerAction>> actionMap = new HashMap<GroomServerStatus, List<GroomServerAction>>(
          2 * groomStatuses.size());
      try {
        tip.getJob().recoverTasks(groomStatuses, actionMap);
      } catch (IOException e) {
        LOG.warn("Task recovery failed", e);
      }

      List<GroomServerAction> actionList = actionMap.get(status);
      actions = new GroomServerAction[actionList.size()];
      actionList.toArray(actions);
    }
    Directive d1 = new DispatchTasksDirective(actions);
    try {
      worker.dispatch(d1);
    } catch (IOException ioe) {
      LOG.error(
          "Fail to dispatch tasks to GroomServer " + status.getGroomName(), ioe);
    }
  }

  @Override
  public void taskComplete(GroomServerStatus status, TaskInProgress task) {
    Pair<String, Integer> key = new Pair<String, Integer>(
        status.getGroomHostName(), BSPMaster.resolveWorkerAddress(
            status.rpcServer).getPort());
    groomServers.put(key, status);
    assignments.remove(key, task);
    jobAssignments.remove(task.getJob(), new Pair<Object, Object>(key, task));

    if (assignments.getCollection(key) == null) {
      groomServers.remove(key);
      driver.killTask(groomTaskIDs.get(key));
    }
  }

  private class TaskDelegatorJobListener extends JobInProgressListener {

    @Override
    public void jobAdded(JobInProgress job) throws IOException {

    }

    @Override
    public void jobRemoved(JobInProgress job) throws IOException {
      @SuppressWarnings("unchecked")
      Collection<Pair<Object, Object>> remainingTasks = jobAssignments
          .getCollection(job);
      if (remainingTasks != null) {
        for (Pair<Object, Object> taskToRemove : remainingTasks) {
          assignments.remove(taskToRemove.getKey(), taskToRemove.getValue());
          if (assignments.getCollection(taskToRemove.getKey()) == null) {
            groomServers.remove(taskToRemove.getKey());
            driver.killTask(groomTaskIDs.get(taskToRemove.getKey()));
          }
        }
        jobAssignments.remove(job);
      }
    }

    @Override
    public void recoverTaskInJob(JobInProgress job) throws IOException {

    }
  }
}