ParallelCombiner: Fix buffer leak on exception in "combine". (#5630)

Once a buffer is acquired, we need to make sure to release it if an
exception is thrown before the closeable iterator is created.
This commit is contained in:
Gian Merlino 2018-04-11 20:39:39 -04:00 committed by GitHub
parent e6efd75a3d
commit 72d6dcda4f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 51 additions and 40 deletions

View File

@ -136,7 +136,11 @@ public class ParallelCombiner<KeyType>
) )
{ {
// CombineBuffer is initialized when this method is called and closed after the result iterator is done // CombineBuffer is initialized when this method is called and closed after the result iterator is done
final Closer closer = Closer.create();
final ResourceHolder<ByteBuffer> combineBufferHolder = combineBufferSupplier.get(); final ResourceHolder<ByteBuffer> combineBufferHolder = combineBufferSupplier.get();
closer.register(combineBufferHolder);
try {
final ByteBuffer combineBuffer = combineBufferHolder.get(); final ByteBuffer combineBuffer = combineBufferHolder.get();
final int minimumRequiredBufferCapacity = StreamingMergeSortedGrouper.requiredBufferCapacity( final int minimumRequiredBufferCapacity = StreamingMergeSortedGrouper.requiredBufferCapacity(
combineKeySerdeFactory.factorizeWithDictionary(mergedDictionary), combineKeySerdeFactory.factorizeWithDictionary(mergedDictionary),
@ -168,13 +172,20 @@ public class ParallelCombiner<KeyType>
final CloseableIterator<Entry<KeyType>> combineIterator = Iterables.getOnlyElement(combineIteratorAndFutures.lhs); final CloseableIterator<Entry<KeyType>> combineIterator = Iterables.getOnlyElement(combineIteratorAndFutures.lhs);
final List<Future> combineFutures = combineIteratorAndFutures.rhs; final List<Future> combineFutures = combineIteratorAndFutures.rhs;
final Closer closer = Closer.create();
closer.register(combineBufferHolder);
closer.register(() -> checkCombineFutures(combineFutures)); closer.register(() -> checkCombineFutures(combineFutures));
return CloseableIterators.wrap(combineIterator, closer); return CloseableIterators.wrap(combineIterator, closer);
} }
catch (Throwable t) {
try {
closer.close();
}
catch (Throwable t2) {
t.addSuppressed(t2);
}
throw t;
}
}
private static void checkCombineFutures(List<Future> combineFutures) private static void checkCombineFutures(List<Future> combineFutures)
{ {