elasticsearch中的优先级线程池

es中通过一个优先级的线程池PrioritizedEsThreadPoolExecutor来根据线程的优先级来控制优先顺序。

类的继承关系非常的清楚，继承自EsThreadPoolExecutor。

PrioritizedEsThreadPoolExecutor中的线程任务优先级先后关系通过一个队列来实现，当两个线程的优先级一样的时候根据FIFO的原则来确定线程任务的前后顺序。

 

当任务通过execute()方法来执行的时候，普通的线程任务会被包装成TieBreakingPrioritizedRunnable，其继承自PrioritizedRunnable，实现了Runnable接口的同时也实现了Comparable接口来方便线程任务先后顺序的确定。

public void execute(Runnable command, final TimeValue timeout, final Runnable timeoutCallback) {
    command = wrapRunnable(command);
    doExecute(command);
    if (timeout.nanos() >= 0) {
        if (command instanceof TieBreakingPrioritizedRunnable) {
            ((TieBreakingPrioritizedRunnable) command).scheduleTimeout(timer, timeoutCallback, timeout);
        } else {
            // We really shouldn't be here. The only way we can get here if somebody created PrioritizedFutureTask
            // and passed it to execute, which doesn't make much sense
            throw new UnsupportedOperationException("Execute with timeout is not supported for future tasks");
        }
    }
}
@Override
protected Runnable wrapRunnable(Runnable command) {
    if (command instanceof PrioritizedRunnable) {
        if ((command instanceof TieBreakingPrioritizedRunnable)) {
            return command;
        }
        Priority priority = ((PrioritizedRunnable) command).priority();
        return new TieBreakingPrioritizedRunnable(super.wrapRunnable(command), priority, insertionOrder.incrementAndGet());
    } else if (command instanceof PrioritizedFutureTask) {
        return command;
    } else { // it might be a callable wrapper...
        if (command instanceof TieBreakingPrioritizedRunnable) {
            return command;
        }
        return new TieBreakingPrioritizedRunnable(super.wrapRunnable(command), Priority.NORMAL, insertionOrder.incrementAndGet());
    }
}

TieBreakingPrioritizedRunnable相比PrioritizedRunnable，主要是多了一个insertionOrder，insertionOrder是有一个AtomicLong，每次一个新的线程任务要被执行的时候都会加一，并将其作为任务进入线程池的依据，以便作为判断先后。

@Override
public int compareTo(PrioritizedRunnable pr) {
    int res = super.compareTo(pr);
    if (res != 0 || !(pr instanceof TieBreakingPrioritizedRunnable)) {
        return res;
    }
    return insertionOrder < ((TieBreakingPrioritizedRunnable) pr).insertionOrder ? -1 : 1;
}

TieBreakingPrioritizedRunnable的compareTo()方法先调用父类的compareTo()方法来比较优先级，如果相等，则通过比较insertionOrder的大小来确定先后。

 

实现优先级的核心在于线程池的等待队列，也就是PriorityBlockingQueue。

PrioritizedEsThreadPoolExecutor(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                ThreadFactory threadFactory, ThreadContext contextHolder, ScheduledExecutorService timer) {
    super(name, corePoolSize, maximumPoolSize, keepAliveTime, unit, new PriorityBlockingQueue<>(), threadFactory, contextHolder);
    this.timer = timer;
}

PriorityBlockingQueue的核心在于siftUpComparable()方法，这个方法将在add()将任务加入队列的时候 被调用，用来确认任务在队列当中的任务位置。

private static <T> void siftUpComparable(int k, T x, Object[] array) {
    Comparable<? super T> key = (Comparable<? super T>) x;
    while (k > 0) {
        int parent = (k - 1) >>> 1;
        Object e = array[parent];
        if (key.compareTo((T) e) >= 0)
            break;
        array[k] = e;
        k = parent;
    }
    array[k] = key;
}

可以看到，在这个队列的存储更类似树，第一个元素存储在下标0的位置，1,2位置则与0进行比较，3,4位置则与1进行比较以此类推，更类似一个大根堆。

再取出来的时候则是通过dequeue()方法，虽然只是简单的取数组下标0位置的，但是也要相应的通过siftDownComparable()方法重新使得当前队列最先被调用的任务在下标0的位置，同时确保别的任务下标往前移。

private static <T> void siftDownComparable(int k, T x, Object[] array,
                                           int n) {
    if (n > 0) {
        Comparable<? super T> key = (Comparable<? super T>)x;
        int half = n >>> 1;           // loop while a non-leaf
        while (k < half) {
            int child = (k << 1) + 1; // assume left child is least
            Object c = array[child];
            int right = child + 1;
            if (right < n &&
                ((Comparable<? super T>) c).compareTo((T) array[right]) > 0)
                c = array[child = right];
            if (key.compareTo((T) c) <= 0)
                break;
            array[k] = c;
            k = child;
        }
        array[k] = key;
    }
}

这个方法，从最大的节点（已空）开始从左右儿子选择大者放入空节点，并尝试在下一个循环中，把空了的儿子节点用其儿子节点较大者填上。

 

这个队列通过大根堆的方式高效完成线程池等待队列当中线程任务触发先后的判断。

这也是PrioritizedEsThreadPoolExecutor优先级线程池的实现。