/**
 * Copyright 2015 Palantir Technologies
 *
 * Licensed under the BSD-3 License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://opensource.org/licenses/BSD-3-Clause
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.palantir.common.concurrent;

import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import com.google.common.base.Throwables;

public class BlockingWorkerPool {
    private final CompletionService<Void> service;
    private final int concurrentTaskLimit;

    private int currentTaskCount = 0;

    public BlockingWorkerPool(ExecutorService executor, int concurrentTaskLimit) {
        this.service = new ExecutorCompletionService<Void>(executor);
        this.concurrentTaskLimit = concurrentTaskLimit;
    }

    /**
     * Submits a task to the pool. If the pool is currently executing the
     * maximum number of tasks allowed, this call will block until one of the
     * currently executing tasks has finished.
     *
     * @throws RuntimeException wrapping an ExecutionException if a previously
     *         submitted task threw an exception.
     */
    public synchronized void submitTask(Runnable task)
            throws InterruptedException {
        waitForAvailability();

        assert currentTaskCount < concurrentTaskLimit;
        service.submit(task, null);
        ++currentTaskCount;
    }

    private void waitForSingleTask() throws InterruptedException {
        if (currentTaskCount <= 0) { return; }

        Future<Void> f = service.take();
        --currentTaskCount;
        try {
            f.get();
        } catch (ExecutionException e) {
            throw Throwables.propagate(e.getCause());
        }
    }

    /**
     * Waits for all tasks that have been previously submitted to the pool to
     * finish.
     *
     * @throws RuntimeException wrapping an ExecutionException if a previously
     *         submitted task threw an exception.
     */
    public synchronized void waitForSubmittedTasks()
            throws InterruptedException {
        while (currentTaskCount > 0) {
            waitForSingleTask();
        }
    }

    /**
     * Waits until the number of tasks drops below the concurrent task limit.
     * @throws InterruptedException
     */
    public synchronized void waitForAvailability() throws InterruptedException {
        if (currentTaskCount >= concurrentTaskLimit) {
            waitForSingleTask();
        }
    }
}

@@ -853,16 +853,66 @@
 .Future;
+%0Aimport java.util.concurrent.atomic.AtomicInteger;
 %0A%0Aimport
@@ -1089,19 +1089,35 @@
 private 
+f
 in
-t
+al AtomicInteger
  current
@@ -1128,17 +1128,35 @@
 Count = 
-0
+new AtomicInteger()
 ;%0A%0A    p
@@ -1877,16 +1877,22 @@
 askCount
+.get()
  %3C concu
@@ -1901,24 +1901,80 @@
 entTaskLimit
+ : %22currentTaskCount must be less than currentTaskLimit%22
 ;%0A        se
@@ -2011,10 +2011,8 @@
     
-++
 curr
@@ -2019,24 +2019,42 @@
 entTaskCount
+.incrementAndGet()
 ;%0A    %7D%0A%0A   
@@ -2145,16 +2145,22 @@
 askCount
+.get()
  %3C= 0) %7B
@@ -2224,10 +2224,8 @@
     
---
 curr
@@ -2236,16 +2236,34 @@
 askCount
+.decrementAndGet()
 ;%0A      
@@ -2774,16 +2774,22 @@
 askCount
+.get()
  %3E 0) %7B%0A
@@ -3077,16 +3077,22 @@
 askCount
+.get()
  %3E= conc

Clean up executors