fix issue where CloseableIterator.flatMap does not close inner CloseableIterator

core/src/main/java/org/apache/druid/java/util/common/parsers/CloseableIterator.java

@@ -75,6 +75,7 @@ private CloseableIterator<R> findNextIeteratorIfNecessary()
           if (iterator != null) {
             try {
               iterator.close();
+              iterator = null;
             }
             catch (IOException e) {
               throw new UncheckedIOException(e);
@@ -112,6 +113,10 @@ public R next()
       public void close() throws IOException
       {
         delegate.close();
+        if (iterator != null) {
+          iterator.close();
+          iterator = null;
+        }
       }
     };
   }

core/src/test/java/org/apache/druid/java/util/common/parsers/CloseableIteratorTest.java

@@ -23,6 +23,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -54,10 +55,18 @@ public void testMap()
   }
 
   @Test
-  public void testFlatMap()
+  public void testFlatMap() throws IOException
   {
-    final CloseableIterator<Integer> actual = generateTestIterator(8)
-        .flatMap(list -> CloseableIterators.withEmptyBaggage(list.iterator()));
+    List<CloseTrackingCloseableIterator<Integer>> innerIterators = new ArrayList<>();
+    final CloseTrackingCloseableIterator<Integer> actual = new CloseTrackingCloseableIterator<>(
+        generateTestIterator(8)
+            .flatMap(list -> {
+              CloseTrackingCloseableIterator<Integer> inner =
+                  new CloseTrackingCloseableIterator<>(CloseableIterators.withEmptyBaggage(list.iterator()));
+              innerIterators.add(inner);
+              return inner;
+            })
+    );
     final Iterator<Integer> expected = IntStream
         .range(0, 8)
         .flatMap(i -> IntStream.range(0, i))
@@ -67,6 +76,48 @@ public void testFlatMap()
     }
     Assert.assertFalse(actual.hasNext());
     Assert.assertFalse(expected.hasNext());
+    actual.close();
+    Assert.assertEquals(1, actual.closeCount);
+    for (CloseTrackingCloseableIterator iter : innerIterators) {
+      Assert.assertEquals(1, iter.closeCount);
+    }
+  }
+
+  @Test
+  public void testFlatMapClosedEarly() throws IOException
+  {
+    final int numIterations = 8;
+    List<CloseTrackingCloseableIterator<Integer>> innerIterators = new ArrayList<>();
+    final CloseTrackingCloseableIterator<Integer> actual = new CloseTrackingCloseableIterator<>(
+        generateTestIterator(numIterations)
+            .flatMap(list -> {
+              CloseTrackingCloseableIterator<Integer> inner =
+                  new CloseTrackingCloseableIterator<>(CloseableIterators.withEmptyBaggage(list.iterator()));
+              innerIterators.add(inner);
+              return inner;
+            })
+    );
+    final Iterator<Integer> expected = IntStream
+        .range(0, numIterations)
+        .flatMap(i -> IntStream.range(0, i))
+        .iterator();
+
+    // burn through the first few iterators
+    int cnt = 0;
+    int numFlatIterations = 5;
+    while (expected.hasNext() && actual.hasNext() && cnt++ < numFlatIterations) {
+      Assert.assertEquals(expected.next(), actual.next());
+    }
+    // but stop while we still have an open current inner iterator and a few remaining inner iterators
+    Assert.assertTrue(actual.hasNext());
+    Assert.assertTrue(expected.hasNext());
+    Assert.assertEquals(4, innerIterators.size());
+    Assert.assertTrue(innerIterators.get(innerIterators.size() - 1).hasNext());
+    actual.close();
+    Assert.assertEquals(1, actual.closeCount);
+    for (CloseTrackingCloseableIterator iter : innerIterators) {
+      Assert.assertEquals(1, iter.closeCount);
+    }
   }
 
   private static CloseableIterator<List<Integer>> generateTestIterator(int numIterates)
@@ -99,4 +150,36 @@ public void close()
       }
     };
   }
+
+  static class CloseTrackingCloseableIterator<T> implements CloseableIterator<T>
+  {
+    CloseableIterator<T> inner;
+    int closeCount;
+
+    public CloseTrackingCloseableIterator(CloseableIterator<T> toTrack)
+    {
+      this.inner = toTrack;
+      this.closeCount = 0;
+    }
+
+
+    @Override
+    public void close() throws IOException
+    {
+      inner.close();
+      closeCount++;
+    }
+
+    @Override
+    public boolean hasNext()
+    {
+      return inner.hasNext();
+    }
+
+    @Override
+    public T next()
+    {
+      return inner.next();
+    }
+  }
 }