支持优先队列操作的TreeSet

目的

由于项目需求，需要一个去重的优先级队列，参考了一些资料，jdk没有找到合适的结构，考虑重写一个。

思路

大体实现的思路有2个
1.维护一个优先级队列和HashSet，元素入队列出队列时同步到HashSet，且入队列时确保HashSet不包含当前元素，否则直接返回成功而不入队列
个人觉得同步两个容器的状态不够elegant

2.参考PriorityBlockingQueue实现一个去重的优先级队列
PriorityBlockingQueue用最小(大)堆实现，查找操作需要遍历整个堆，时间复杂度O(n)不太符合要求

3.用排序结构实现(如TreeSet，ConcurrentSkipListSet)，出队列操作选取排序第一个元素，元素入队列则直接插入

这里选用第三种方法。
由于TreeSet没有提供最大长度限制，且取元素操作是非阻塞的，不符合项目需求，在这里进行简单的修改。
1.添加capacity属性，限制长度
2.添加notEmpty Condition，TreeSet为空时阻塞take操作，添加元素时signalAll

代码

import java.util.Comparator;
import java.util.TreeSet;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 支持并发的有限队列，且保证元素唯一
 * @Author peng.liu
 * @Date 17-1-16
 */
public class ConcurrentPriorityUniqueQueue<T> {

    private TreeSet<T> treeSet = null;
    private int capacity;
    private ReentrantLock lock;

    private final Condition notEmpty;

    /**
     * 构造器
     * @param capacity 容量
     * @param comparator 比较器
     */
    public ConcurrentPriorityUniqueQueue(int capacity, Comparator<T> comparator){
        if(capacity <= 0){
            throw new IllegalArgumentException("queue capacity <= 0");
        }
        treeSet = new TreeSet<T>(comparator);
        this.capacity = capacity;
        this.lock = new ReentrantLock();
        this.notEmpty = lock.newCondition();
    }

    /**
     * 添加元素,即时返回
     * @param element
     * @return
     */
    public boolean offer(T element) {
        lock.lock();
        try {
            if (treeSet.size() == capacity) {
                return false;
            } else {
                treeSet.add(element);
                notEmpty.signalAll();
                return true;
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 获取优先队列队头,阻塞
     * @return
     * @throws InterruptedException
     */
    public T take() throws InterruptedException{
        lock.lockInterruptibly();
        try{
            while(treeSet.size() == 0){
                notEmpty.await();
            }
            return treeSet.pollFirst();
        }finally {
            lock.unlock();
        }
    }

    public int size(){
        return treeSet.size();
    }

}

注意

TreeSet根据传入的Comparator判断元素是否相等，对于一些优先级相等但实际不相同的元素在compare时注意下操作