package org.infinispan.executors;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
/**
* Exectues given tasks in provided executor.
*
* @author Radim Vansa <rvansa@redhat.com>
*/
public class ExecutorAllCompletionService implements CompletionService<Void> {
private ExecutorCompletionService executorService;
private AtomicReference<ExecutionException> firstException = new AtomicReference<ExecutionException>();
private AtomicLong scheduled = new AtomicLong();
private AtomicLong completed = new AtomicLong();
public ExecutorAllCompletionService(Executor executor) {
this.executorService = new ExecutorCompletionService(executor);
}
@Override
public Future<Void> submit(final Callable<Void> task) {
scheduled.incrementAndGet();
Future<Void> future = executorService.submit(task);
pollUntilEmpty();
return future;
}
@Override
public Future<Void> submit(final Runnable task, Void result) {
scheduled.incrementAndGet();
Future<Void> future = executorService.submit(task, result);
pollUntilEmpty();
return future;
}
private void pollUntilEmpty() {
Future<Void> completedFuture;
while ((completedFuture = executorService.poll()) != null) {
try {
completedFuture.get();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
} catch (ExecutionException e) {
if (firstException.get() == null) {
firstException.compareAndSet(null, e);
}
} finally {
completed.incrementAndGet();
}
}
}
/**
* @return True if all currently scheduled tasks have already been completed, false otherwise;
*/
public boolean isAllCompleted() {
pollUntilEmpty();
return completed.get() >= scheduled.get();
}
public long getScheduledTasks() {
return scheduled.get();
}
public long getCompletedTasks() {
return completed.get();
}
public void waitUntilAllCompleted() {
while (completed.get() < scheduled.get()) {
// Here is a race - if we poll the last scheduled entry elsewhere, we may wait
// another 100 ms until we realize that everything has already completed.
// Nevertheless, that's not so bad.
try {
Future<Void> future = poll(100, TimeUnit.MILLISECONDS);
if (future != null) {
future.get();
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
return;
} catch (ExecutionException e) {
if (firstException.get() == null) {
firstException.compareAndSet(null, e);
}
}
}
}
public boolean isExceptionThrown() {
return firstException.get() != null;
}
public ExecutionException getFirstException() {
return firstException.get();
}
@Override
public Future<Void> take() throws InterruptedException {
Future<Void> future = executorService.take();
completed.incrementAndGet();
return future;
}
@Override
public Future<Void> poll() {
Future<Void> future = executorService.poll();
if (future != null) {
completed.incrementAndGet();
}
return future;
}
@Override
public Future<Void> poll(long timeout, TimeUnit unit) throws InterruptedException {
Future<Void> future = executorService.poll(timeout, unit);
if (future != null) {
completed.incrementAndGet();
}
return future;
}
}