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();

        }



    }



}