Java: ExecutorService that blocks on submission after a certain queue size

前端 未结 7 1391
Happy的楠姐
Happy的楠姐 2020-11-30 17:38

I am trying to code a solution in which a single thread produces I/O-intensive tasks that can be performed in parallel. Each task have significant in-memory data. So I want

7条回答
  •  一生所求
    2020-11-30 18:04

    I had the similar problem and I implemented that by using beforeExecute/afterExecute hooks from ThreadPoolExecutor:

    import java.util.concurrent.BlockingQueue;
    import java.util.concurrent.ThreadPoolExecutor;
    import java.util.concurrent.TimeUnit;
    import java.util.concurrent.locks.Condition;
    import java.util.concurrent.locks.ReentrantLock;
    
    /**
     * Blocks current task execution if there is not enough resources for it.
     * Maximum task count usage controlled by maxTaskCount property.
     */
    public class BlockingThreadPoolExecutor extends ThreadPoolExecutor {
    
        private final ReentrantLock taskLock = new ReentrantLock();
        private final Condition unpaused = taskLock.newCondition();
        private final int maxTaskCount;
    
        private volatile int currentTaskCount;
    
        public BlockingThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
                long keepAliveTime, TimeUnit unit,
                BlockingQueue workQueue, int maxTaskCount) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
            this.maxTaskCount = maxTaskCount;
        }
    
        /**
         * Executes task if there is enough system resources for it. Otherwise
         * waits.
         */
        @Override
        protected void beforeExecute(Thread t, Runnable r) {
            super.beforeExecute(t, r);
            taskLock.lock();
            try {
                // Spin while we will not have enough capacity for this job
                while (maxTaskCount < currentTaskCount) {
                    try {
                        unpaused.await();
                    } catch (InterruptedException e) {
                        t.interrupt();
                    }
                }
                currentTaskCount++;
            } finally {
                taskLock.unlock();
            }
        }
    
        /**
         * Signalling that one more task is welcome
         */
        @Override
        protected void afterExecute(Runnable r, Throwable t) {
            super.afterExecute(r, t);
            taskLock.lock();
            try {
                currentTaskCount--;
                unpaused.signalAll();
            } finally {
                taskLock.unlock();
            }
        }
    }
    

    This should be good enough for you. Btw, original implementation was task size based because one task could be larger 100 time than another and submitting two huge tasks was killing the box, but running one big and plenty of small was Okay. If your I/O-intensive tasks are roughly the same size you could use this class, otherwise just let me know and I'll post size based implementation.

    P.S. You would want to check ThreadPoolExecutor javadoc. It's really nice user guide from Doug Lea about how it could be easily customized.

提交回复
热议问题