com.tencent.angel.master.task.AMTaskManager Maven / Gradle / Ivy

Go to download
/*
 * Tencent is pleased to support the open source community by making Angel available.
 *
 * Copyright (C) 2017-2018 THL A29 Limited, a Tencent company. All rights reserved.
 *
 * Licensed 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
 *
 * https://opensource.org/licenses/Apache-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 com.tencent.angel.master.task;

import com.tencent.angel.worker.task.TaskId;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Task manager in master, it manages all tasks {@link com.tencent.angel.master.task.AMTask} in the application.
 */
public class AMTaskManager {
  private static final Log LOG = LogFactory.getLog(AMTaskManager.class);
  /**
   * task id to task map
   */
  private final HashMap idToTaskMap;
  private final Lock readLock;
  private final Lock writeLock;

  public AMTaskManager() {
    idToTaskMap = new HashMap();
    ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
    readLock = readWriteLock.readLock();
    writeLock = readWriteLock.writeLock();
  }

  /**
   * add the task to task map
   *
   * @param taskId the task id
   * @param task   the task
   */
  public void putTask(TaskId taskId, AMTask task) {
    try {
      writeLock.lock();
      idToTaskMap.put(taskId, task);
    } finally {
      writeLock.unlock();
    }
  }

  /**
   * get the task corresponding to specified task id
   *
   * @param taskId task id
   * @return AMTask the task corresponding to specified task id
   */
  public AMTask getTask(TaskId taskId) {
    try {
      readLock.lock();
      return idToTaskMap.get(taskId);
    } finally {
      readLock.unlock();
    }
  }

  /**
   * get all tasks
   *
   * @return Collection all tasks
   */
  public Collection getTasks() {
    try {
      readLock.lock();
      Set cloneTasks = new HashSet();
      cloneTasks.addAll(idToTaskMap.values());
      return cloneTasks;
    } finally {
      readLock.unlock();
    }
  }

  /**
   * write all tasks state to a output stream
   *
   * @param output the output stream
   * @throws IOException
   */
  public void serialize(DataOutputStream output) throws IOException {
    try {
      readLock.lock();
      output.writeInt(idToTaskMap.size());
      for (Entry entry : idToTaskMap.entrySet()) {
        output.writeInt(entry.getKey().getIndex());
        entry.getValue().serialize(output);
      }
    } finally {
      readLock.unlock();
    }
  }

  /**
   * read all tasks state from a input stream
   *
   * @param input the input stream
   * @throws IOException
   */
  public void deserialize(DataInputStream input) throws IOException {
    try {
      writeLock.lock();
      int size = input.readInt();
      for (int i = 0; i < size; i++) {
        TaskId taskId = new TaskId(input.readInt());
        AMTask task = new AMTask(taskId, null);
        task.deserialize(input);
        LOG.info("task deserialize=" + task);
        idToTaskMap.put(taskId, task);
      }
    } finally {
      writeLock.unlock();
    }
  }
}