PoolManager.java

package org.apache.hadoop.corona;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * Manages pools for a given type. Needs to be thread safe because addSession()
 * and getSortedPools() can be called from different threads.
 */
public class PoolManager {

  public static final Log LOG = LogFactory.getLog(PoolManager.class);
  public final static String DEFAULT_POOL_NAME = "default";
  private final String type;
  private final ConcurrentHashMap<String, PoolSchedulable> nameToPool;
  private final ConfigManager configManager;
  private CoronaConf conf;
  private Collection<PoolSchedulable> snapshotPools;
  private Queue<PoolSchedulable> scheduleQueue;
  private Queue<PoolSchedulable> preemptQueue;

  public PoolManager(String type, ConfigManager configManager) {
    this.type = type;
    this.configManager = configManager;

    // This needs to be thread safe because addSession() and getSortedPools()
    // are called from different threads
    this.nameToPool = new ConcurrentHashMap<String, PoolSchedulable>();
  }

  /**
   * Take snapshots for all pools and sessions
   */
  public void snapshot() {
    snapshotPools = new ArrayList<PoolSchedulable>(nameToPool.values());
    for (PoolSchedulable pool : snapshotPools) {
      pool.snapshot();
    }
    scheduleQueue = null;
    preemptQueue = null;
  }

  public Queue<PoolSchedulable> getScheduleQueue() {
    if (scheduleQueue == null) {
      scheduleQueue = createPoolQueue(ScheduleComparator.FAIR);
    }
    return scheduleQueue;
  }

  public Queue<PoolSchedulable> getPreemptQueue() {
    if (preemptQueue == null) {
      preemptQueue = createPoolQueue(ScheduleComparator.FAIR_PREEMPT);
    }
    return preemptQueue;
  }

  public Queue<PoolSchedulable> createPoolQueue(ScheduleComparator comparator) {
    int initCapacity = snapshotPools.size() == 0 ? 1 : snapshotPools.size();
    Queue<PoolSchedulable> poolQueue = new PriorityQueue<PoolSchedulable>(
        initCapacity, comparator);
    poolQueue.addAll(snapshotPools);
    return poolQueue;
  }

  public void addSession(String id, Session session) {
    String poolName = getPoolName(session);
    getPoolSchedulable(poolName).addSession(id, session);
  }

  public static String getPoolName(Session session) {
    String poolName = session.getPoolId();
    // If there is no explicit pool name, take user name.
    if (poolName == null || poolName.equals("")) {
      poolName = session.getUserId();
    }
    // If pool name still is bad, use default.
    if (poolName == null || poolName.equals("")) {
      return DEFAULT_POOL_NAME;
    }
    if (!isLegalPoolName(poolName)) {
      LOG.warn("Bad poolName:" + poolName +
          " from session:" + session.sessionId);
      return DEFAULT_POOL_NAME;
    }
    return poolName;
  }

  /**
   * Returns whether or not the given pool name is legal.
   * 
   * Legal pool names are of nonzero length and are formed only of alphanumeric 
   * characters, underscores (_), and hyphens (-).
   */
  private static boolean isLegalPoolName(String poolName) {
    return !poolName.matches(".*[^0-9a-zA-Z\\-\\_].*") 
            && (poolName.length() > 0);

  }

  private PoolSchedulable getPoolSchedulable(String name) {
    PoolSchedulable poolSchedulable = nameToPool.get(name);
    if (poolSchedulable == null) {
      poolSchedulable = new PoolSchedulable(name, type, configManager);
      // Note that we never remove pools from this map
      nameToPool.put(name, poolSchedulable);
    }
    return poolSchedulable;
  }

  public void distributeShare(int total) {
    Schedulable.distributeShare(
        total, snapshotPools, ScheduleComparator.FAIR);
  }

  public Collection<PoolSchedulable> getPools() {
    return Collections.unmodifiableCollection(snapshotPools);
  }

  public void setConf(CoronaConf conf) {
    this.conf = conf;
  }

  public CoronaConf getConf() {
    return conf;
  }
}