Api#2880: Close the underlying iterator in ClosingIterator in hasNext() call (#2881)

Author: Saurabh Agarwal

api/src/main/java/org/apache/iceberg/io/ClosingIterator.java

@@ -23,35 +23,38 @@
 import java.io.UncheckedIOException;
 import java.util.Iterator;
 
+/**
+ * A convenience wrapper around {@link CloseableIterator}, providing auto-close
+ * functionality when all of the elements in the iterator are consumed.
+ */
 public class ClosingIterator<T> implements Iterator<T> {
   private final CloseableIterator<T> iterator;
-  private boolean shouldClose = false;
+  private boolean isClosed;
 
   public ClosingIterator(CloseableIterator<T> iterator) {
     this.iterator = iterator;
-
   }
 
   @Override
   public boolean hasNext() {
     boolean hasNext = iterator.hasNext();
-    this.shouldClose = !hasNext;
+    if (!hasNext && !isClosed) {
+      close();
+    }
     return hasNext;
   }
 
   @Override
   public T next() {
-    T next = iterator.next();
+    return iterator.next();
+  }
 
-    if (shouldClose) {
-      // this will only be called once because iterator.next would throw NoSuchElementException
-      try {
-        iterator.close();
-      } catch (IOException e) {
-        throw new UncheckedIOException(e);
-      }
+  private void close() {
+    try {
+      iterator.close();
+      isClosed = true;
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
     }
-
-    return next;
   }
 }

api/src/test/java/org/apache/iceberg/io/TestClosingIterator.java

@@ -0,0 +1,71 @@
+/*
+ * 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.iceberg.io;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class TestClosingIterator {
+  @Test
+  public void testEmptyIterator() {
+    CloseableIterator<String> underlying = mock(CloseableIterator.class);
+    ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+    assertFalse(closingIterator.hasNext());
+  }
+
+  @Test
+  public void testHasNextAndNext() {
+    CloseableIterator<String> underlying = mock(CloseableIterator.class);
+    when(underlying.hasNext()).thenReturn(true);
+    when(underlying.next()).thenReturn("hello");
+    ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+    assertTrue(closingIterator.hasNext());
+    assertEquals("hello", closingIterator.next());
+  }
+
+  @Test
+  public void testUnderlyingIteratorCloseWhenElementsAreExhausted() throws Exception {
+    CloseableIterator<String> underlying = mock(CloseableIterator.class);
+    when(underlying.hasNext()).thenReturn(true).thenReturn(false);
+    when(underlying.next()).thenReturn("hello");
+    ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+    assertTrue(closingIterator.hasNext());
+    assertEquals("hello", closingIterator.next());
+
+    assertFalse(closingIterator.hasNext());
+    verify(underlying, times(1)).close();
+  }
+
+  @Test
+  public void testCloseCalledOnceForMultipleHasNextCalls() throws Exception {
+    CloseableIterator<String> underlying = mock(CloseableIterator.class);
+    ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+    assertFalse(closingIterator.hasNext());
+    assertFalse(closingIterator.hasNext());
+    verify(underlying, times(1)).close();
+  }
+}