Fix various Yielder leaks.

core/src/main/java/org/apache/druid/common/guava/CombiningSequence.java

@@ -29,8 +29,6 @@
 import java.util.Comparator;
 import java.util.function.BinaryOperator;
 
-/**
- */
 public class CombiningSequence<T> implements Sequence<T>
 {
   public static <T> CombiningSequence<T> create(
@@ -76,9 +74,22 @@ public <OutType> Yielder<OutType> toYielder(OutType initValue, final YieldingAcc
         new CombiningYieldingAccumulator<>(ordering, mergeFn, accumulator);
 
     combiningAccumulator.setRetVal(initValue);
-    Yielder<T> baseYielder = baseSequence.toYielder(null, combiningAccumulator);
 
-    return makeYielder(baseYielder, combiningAccumulator, false);
+    final Yielder<T> baseYielder = baseSequence.toYielder(null, combiningAccumulator);
+
+    try {
+      return makeYielder(baseYielder, combiningAccumulator, false);
+    }
+    catch (Throwable t1) {
+      try {
+        baseYielder.close();
+      }
+      catch (Throwable t2) {
+        t1.addSuppressed(t2);
+      }
+
+      throw t1;
+    }
   }
 
   private <OutType> Yielder<OutType> makeYielder(

core/src/main/java/org/apache/druid/java/util/common/guava/ExplodingSequence.java

@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 org.apache.druid.java.util.common.guava;
+
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Wraps an underlying sequence and allows us to force it to explode at various points.
+ */
+public class ExplodingSequence<T> extends YieldingSequenceBase<T>
+{
+  private final Sequence<T> baseSequence;
+  private final boolean getThrowsException;
+  private final boolean closeThrowsException;
+  private final AtomicLong closed = new AtomicLong();
+
+  public ExplodingSequence(Sequence<T> baseSequence, boolean getThrowsException, boolean closeThrowsException)
+  {
+    this.baseSequence = baseSequence;
+    this.getThrowsException = getThrowsException;
+    this.closeThrowsException = closeThrowsException;
+  }
+
+  @Override
+  public <OutType> Yielder<OutType> toYielder(
+      final OutType initValue,
+      final YieldingAccumulator<OutType, T> accumulator
+  )
+  {
+    return wrapYielder(baseSequence.toYielder(initValue, accumulator));
+  }
+
+  public long getCloseCount()
+  {
+    return closed.get();
+  }
+
+  private <OutType> Yielder<OutType> wrapYielder(final Yielder<OutType> baseYielder)
+  {
+    return new Yielder<OutType>()
+    {
+      @Override
+      public OutType get()
+      {
+        if (getThrowsException) {
+          throw new RuntimeException("get");
+        } else {
+          return baseYielder.get();
+        }
+      }
+
+      @Override
+      public Yielder<OutType> next(OutType initValue)
+      {
+        return wrapYielder(baseYielder.next(initValue));
+      }
+
+      @Override
+      public boolean isDone()
+      {
+        return baseYielder.isDone();
+      }
+
+      @Override
+      public void close() throws IOException
+      {
+        closed.incrementAndGet();
+
+        if (closeThrowsException) {
+          throw new IOException("close");
+        } else {
+          baseYielder.close();
+        }
+      }
+    };
+  }
+}

core/src/main/java/org/apache/druid/java/util/common/guava/MergeSequence.java

@@ -58,38 +58,62 @@ public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulat
         )
     );
 
-    pQueue = baseSequences.accumulate(
-        pQueue,
-        (queue, in) -> {
-          final Yielder<T> yielder = in.toYielder(
-              null,
-              new YieldingAccumulator<T, T>()
-              {
-                @Override
-                public T accumulate(T accumulated, T in)
+    try {
+      pQueue = baseSequences.accumulate(
+          pQueue,
+          (queue, in) -> {
+            final Yielder<T> yielder = in.toYielder(
+                null,
+                new YieldingAccumulator<T, T>()
                 {
-                  yield();
-                  return in;
+                  @Override
+                  public T accumulate(T accumulated, T in)
+                  {
+                    yield();
+                    return in;
+                  }
                 }
+            );
+
+            if (!yielder.isDone()) {
+              try {
+                queue.add(yielder);
               }
-          );
+              catch (Throwable t1) {
+                try {
+                  yielder.close();
+                }
+                catch (Throwable t2) {
+                  t1.addSuppressed(t2);
+                }
 
-          if (!yielder.isDone()) {
-            queue.add(yielder);
-          } else {
-            try {
-              yielder.close();
-            }
-            catch (IOException e) {
-              throw new RuntimeException(e);
+                throw t1;
+              }
+            } else {
+              try {
+                yielder.close();
+              }
+              catch (IOException e) {
+                throw new RuntimeException(e);
+              }
             }
+
+            return queue;
           }
+      );
 
-          return queue;
-        }
-    );
+      return makeYielder(pQueue, initValue, accumulator);
+    }
+    catch (Throwable t1) {
+      try {
+        closeAll(pQueue);
+      }
+      catch (Throwable t2) {
+        t1.addSuppressed(t2);
+      }
 
-    return makeYielder(pQueue, initValue, accumulator);
+      throw t1;
+    }
   }
 
   private <OutType> Yielder<OutType> makeYielder(
@@ -101,8 +125,22 @@ private <OutType> Yielder<OutType> makeYielder(
     OutType retVal = initVal;
     while (!accumulator.yielded() && !pQueue.isEmpty()) {
       Yielder<T> yielder = pQueue.remove();
-      retVal = accumulator.accumulate(retVal, yielder.get());
-      yielder = yielder.next(null);
+
+      try {
+        retVal = accumulator.accumulate(retVal, yielder.get());
+        yielder = yielder.next(null);
+      }
+      catch (Throwable t1) {
+        try {
+          yielder.close();
+        }
+        catch (Throwable t2) {
+          t1.addSuppressed(t2);
+        }
+
+        throw t1;
+      }
+
       if (yielder.isDone()) {
         try {
           yielder.close();
@@ -144,12 +182,21 @@ public boolean isDone()
       @Override
       public void close() throws IOException
       {
-        Closer closer = Closer.create();
-        while (!pQueue.isEmpty()) {
-          closer.register(pQueue.remove());
-        }
-        closer.close();
+        closeAll(pQueue);
       }
     };
   }
+
+  private static <T> void closeAll(final PriorityQueue<Yielder<T>> pQueue) throws IOException
+  {
+    Closer closer = Closer.create();
+    while (!pQueue.isEmpty()) {
+      final Yielder<T> yielder = pQueue.poll();
+      if (yielder != null) {
+        // Note: yielder can be null if our comparator threw an exception during queue.add.
+        closer.register(yielder);
+      }
+    }
+    closer.close();
+  }
 }

core/src/main/java/org/apache/druid/java/util/common/guava/YieldingSequenceBase.java

@@ -19,6 +19,8 @@
 
 package org.apache.druid.java.util.common.guava;
 
+import java.io.IOException;
+
 /**
  * A Sequence that is based entirely on the Yielder implementation.
  * <p/>
@@ -29,16 +31,33 @@ public abstract class YieldingSequenceBase<T> implements Sequence<T>
   @Override
   public <OutType> OutType accumulate(OutType initValue, Accumulator<OutType, T> accumulator)
   {
+    final OutType retVal;
     Yielder<OutType> yielder = toYielder(initValue, YieldingAccumulators.fromAccumulator(accumulator));
 
     try {
       while (!yielder.isDone()) {
         yielder = yielder.next(yielder.get());
       }
-      return yielder.get();
+      retVal = yielder.get();
     }
-    finally {
-      CloseQuietly.close(yielder);
+    catch (Throwable t1) {
+      try {
+        yielder.close();
+      }
+      catch (Throwable t2) {
+        t1.addSuppressed(t2);
+      }
+
+      throw t1;
     }
+
+    try {
+      yielder.close();
+    }
+    catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+
+    return retVal;
   }
 }

core/src/test/java/org/apache/druid/common/guava/CombiningSequenceTest.java

@@ -20,17 +20,21 @@
 package org.apache.druid.common.guava;
 
 import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
 import org.apache.druid.java.util.common.Pair;
+import org.apache.druid.java.util.common.guava.ExplodingSequence;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.java.util.common.guava.Yielder;
 import org.apache.druid.java.util.common.guava.YieldingAccumulator;
+import org.hamcrest.CoreMatchers;
 import org.junit.Assert;
 import org.junit.Test;
+import org.junit.internal.matchers.ThrowableMessageMatcher;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 
@@ -39,6 +43,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
@@ -194,6 +199,40 @@ public void testNothing() throws Exception
     testCombining(Collections.emptyList(), Collections.emptyList());
   }
 
+  @Test
+  public void testExplodingSequence()
+  {
+    final ExplodingSequence<Integer> bomb =
+        new ExplodingSequence<>(Sequences.simple(ImmutableList.of(1, 2, 2)), false, true);
+
+    final CombiningSequence<Integer> combiningSequence =
+        CombiningSequence.create(bomb, Comparator.naturalOrder(), (a, b) -> a);
+
+    try {
+      combiningSequence.toYielder(
+          null,
+          new YieldingAccumulator<Integer, Integer>()
+          {
+            @Override
+            public Integer accumulate(Integer accumulated, Integer in)
+            {
+              if (in > 1) {
+                throw new RuntimeException("boom");
+              }
+
+              return in;
+            }
+          }
+      );
+      Assert.fail("Expected exception");
+    }
+    catch (Exception e) {
+      Assert.assertThat(e, ThrowableMessageMatcher.hasMessage(CoreMatchers.equalTo("boom")));
+    }
+
+    Assert.assertEquals("Closes resources", 1, bomb.getCloseCount());
+  }
+
   private void testCombining(List<Pair<Integer, Integer>> pairs, List<Pair<Integer, Integer>> expected)
       throws Exception
   {