@Override
public void processElement(final ProcessContext c) throws Exception {
try {
workTickets.acquire();
} catch (InterruptedException e) {
throw new RuntimeException("Interrupted while scheduling work", e);
}
if (failure.get() != null) {
throw Throwables.propagate(failure.get());
}
executor.submit(
new Runnable() {
@Override
public void run() {
try {
doFn.processElement(new WrappedContext(c));
} catch (Throwable t) {
failure.compareAndSet(null, t);
Throwables.propagateIfPossible(t);
throw new AssertionError("Unexpected checked exception: " + t);
} finally {
workTickets.release();
}
}
});
}
@Override
public void finishBundle(Context c) throws Exception {
// Acquire all the work tickets to guarantee that all the previous
// processElement calls have finished.
workTickets.acquire(maxParallelism);
if (failure.get() != null) {
throw Throwables.propagate(failure.get());
}
doFn.finishBundle(c);
}