java 可伸缩阻塞队列实现

最近一年多写的最虐心的代码。必须好好复习java并发了。搞了一晚上终于测试都跑通过了，特此纪念，以资鼓励!

 

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 实现可调整大小的阻塞队列，支持数据迁移平衡reader,writer读取吸入速度，达到最大吞吐
 * @author hanjie
 *
 */
public class RecordBuffer {

    public static final Record CLOSE_RECORD = new Record() {

        @Override
        public Object getColumnValue(String columnName) {
            // TODO Auto-generated method stub
            return null;
        }
    };
    
    public static final Record SWITCH_QUEUE_RECORD = new Record() {

        @Override
        public Object getColumnValue(String columnName) {
            // TODO Auto-generated method stub
            return null;
        }
    };
    
    public Lock switchingQueueLock = new ReentrantLock();
    public Condition readerSwitched = switchingQueueLock.newCondition();
    public Condition writerSwitched = switchingQueueLock.newCondition();
    public Condition switchFinished = switchingQueueLock.newCondition();
    
    public volatile boolean readerSwitchSuccess = true;
    public volatile boolean writerSwitchSuccess = true;
    public volatile boolean switchingQueue = false;
    public volatile boolean closed = false;
    private volatile ArrayBlockingQueue<Record> queue;
    private TaskCounter taskCounter;
    

    public RecordBuffer(TaskCounter taskCounter, int size) {
        this.queue = new ArrayBlockingQueue<Record>(size);
        this.taskCounter = taskCounter;
    }



    public void resize(int newSize) {
        try {
            
            if(closed){
                return;
            }
    
            switchingQueueLock.lock();
            try {
                //double check下，要不可能writer收到CLOSED_record已经 退出了。writerSwitched.await() 会hang住
                if(closed){
                    return;
                }
                this.switchingQueue = true;
        
                ArrayBlockingQueue<Record> oldQueue = queue;
                queue = new ArrayBlockingQueue<Record>(newSize);
                this.readerSwitchSuccess = false;
                this.writerSwitchSuccess = false;
                
                //先拯救下writer，可能writer刚好阻塞到take上,失败也没关系，说明老队列不空，writer不会阻塞到take
                oldQueue.offer(SWITCH_QUEUE_RECORD);

                while (!writerSwitchSuccess) {
                    writerSwitched.await();
                }
                //writer先切换队列，然后reader可能阻塞在最后一个put上,清空下老队列拯救reader，让它顺利醒来
                transferOldQueueRecordsToNewQueue(oldQueue);
                
                
                while (!readerSwitchSuccess) {
                    readerSwitched.await();
                }
                //前面的清空，刚好碰到reader要put最后一个，非阻塞式清空动作就有残留最后一个put
                transferOldQueueRecordsToNewQueue(oldQueue);
                
                this.switchingQueue = false;
                this.switchFinished.signalAll();

            } finally {
                switchingQueueLock.unlock();
            }

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(e);
        }
    }

    private void transferOldQueueRecordsToNewQueue(ArrayBlockingQueue<Record> oldQueue)
            throws InterruptedException {
        List<Record> oldRecords = new ArrayList<Record>(oldQueue.size());
        Record record = null;
        while ((record = oldQueue.poll()) != null) {
            oldRecords.add(record);
        }
        // 转移老队列剩下的记录到新队列
        for (int i = 0; i < oldRecords.size(); i++) {
            queue.put(oldRecords.get(i));
        }
    }

    public void close() {
        this.closed = true;
        switchingQueueLock.lock();
        try {
            //如果正在切换队列， 等切换做完才能，发送最后一个CLOSE
            while (switchingQueue) {
                switchFinished.await();
            }
            
            this.queue.put(CLOSE_RECORD);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(e);
        }
        finally{
            switchingQueueLock.unlock();
        }
    }

    public void put(Record record) {
        try {
            
            if (!queue.offer(record)) {
                taskCounter.incrBufferFullCount();
                if (!readerSwitchSuccess) {
                    notifyReaderSwitchSuccess();
                }
                queue.put(record);
            }
            taskCounter.incrReadCount();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(e);
        }
    }

    private void notifyReaderSwitchSuccess() {
        System.out.println("reader switch");
        switchingQueueLock.lock();
        try {
            readerSwitchSuccess = true;
            readerSwitched.signalAll();
        } finally {
            switchingQueueLock.unlock();
        }
    }

    public Record take() {
        try {
            
            Record record = queue.poll();
            //如果拿到了切换记录，则切换队列重试
            if(record == SWITCH_QUEUE_RECORD){
                if (!writerSwitchSuccess) {
                    notifyWriterSwitchSuccess();
                }
                record = queue.poll();
            }
            
            if (record == null) {
                taskCounter.incrBufferEmptyCount();
                
                //调用take先检查是否正在切换,保证拿到新的队列
                if (!writerSwitchSuccess) {
                    notifyWriterSwitchSuccess();
                }
                record = queue.take();
                //如果很不幸刚好在take阻塞时候，切换，只能发送一个切换记录将其唤醒
                if(record == SWITCH_QUEUE_RECORD){
                    if (!writerSwitchSuccess) {
                        notifyWriterSwitchSuccess();
                    }
                    record = queue.take();
                }
            }
            if (record == CLOSE_RECORD) {
                if (!writerSwitchSuccess) {
                    notifyWriterSwitchSuccess();
                }
                return null;
            }
            taskCounter.incrWriteCount();
            return record;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(e);
        }
    }

    private void notifyWriterSwitchSuccess() {

        System.out.println("writer switch");
        switchingQueueLock.lock();
        try {
            writerSwitchSuccess = true;
            writerSwitched.signalAll();
        } finally {
            switchingQueueLock.unlock();
        }

    }

}