From 6f18dd7465bd8226c25d7f8e2f0d0acbf85956b9 Mon Sep 17 00:00:00 2001
From: Simon Brandhof <simon.brandhof@sonarsource.com>
Date: Thu, 3 Sep 2015 11:07:04 +0200
Subject: [PATCH] Add CloseableIterator#wrap(CloseableIterator,
 AutoCloseable...)

---
 .../sonar/core/util/CloseableIterator.java    | 47 ++++++++++++++++++-
 .../core/util/CloseableIteratorTest.java      | 32 +++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/sonar-core/src/main/java/org/sonar/core/util/CloseableIterator.java b/sonar-core/src/main/java/org/sonar/core/util/CloseableIterator.java
index ba8e6dca5fc..ec4050a8efc 100644
--- a/sonar-core/src/main/java/org/sonar/core/util/CloseableIterator.java
+++ b/sonar-core/src/main/java/org/sonar/core/util/CloseableIterator.java
@@ -20,9 +20,12 @@
 package org.sonar.core.util;
 
 import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableList;
 import java.util.Iterator;
+import java.util.List;
 import java.util.NoSuchElementException;
 import javax.annotation.CheckForNull;
+import org.apache.commons.lang.ArrayUtils;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
@@ -56,13 +59,23 @@ public abstract class CloseableIterator<O> implements Iterator<O>, AutoCloseable
    *
    * @throws IllegalArgumentException if the specified {@link Iterator} is a CloseableIterator
    */
-  public static <T> CloseableIterator<T> from(final Iterator<T> iterator) {
+  public static <T> CloseableIterator<T> from(Iterator<T> iterator) {
     // early fail
     requireNonNull(iterator);
     checkArgument(!(iterator instanceof AutoCloseable), "This method does not support creating a CloseableIterator from an Iterator which is Closeable");
     return new RegularIteratorWrapper<>(iterator);
   }
 
+  /**
+   * Wraps a {@code CloseableIterator} and optionally other instances of {@code AutoCloseable} that must be closed
+   * at the same time. The wrapped iterator is closed first then the other {@code AutoCloseable} in the defined order.
+   * 
+   * @throws IllegalArgumentException if the parameter {@code otherCloseables} contains the wrapped iterator
+   */
+  public static <T> CloseableIterator<T> wrap(CloseableIterator<T> iterator, AutoCloseable... otherCloseables) {
+    return new CloseablesIteratorWrapper<>(iterator, otherCloseables);
+  }
+
   private O nextObject = null;
   boolean isClosed = false;
 
@@ -90,7 +103,7 @@ public abstract class CloseableIterator<O> implements Iterator<O>, AutoCloseable
   }
 
   /**
-   * Reads next item and returns null if no more items.
+   * Reads next item and returns {@code null} if no more items.
    */
   @CheckForNull
   protected abstract O doNext();
@@ -165,4 +178,34 @@ public abstract class CloseableIterator<O> implements Iterator<O>, AutoCloseable
       // do nothing
     }
   }
+
+  private static class CloseablesIteratorWrapper<T> extends CloseableIterator<T> {
+    private final CloseableIterator<T> iterator;
+    private final List<AutoCloseable> otherCloseables;
+
+    private CloseablesIteratorWrapper(CloseableIterator<T> iterator, AutoCloseable... otherCloseables) {
+      requireNonNull(iterator);
+      checkArgument(!ArrayUtils.contains(otherCloseables, iterator));
+      this.iterator = iterator;
+      // the advantage of using ImmutableList is that it does not accept null elements, so it fails fast, during
+      // construction of the wrapper, but not in close()
+      this.otherCloseables = ImmutableList.copyOf(otherCloseables);
+    }
+
+    @Override
+    protected T doNext() {
+      return iterator.hasNext() ? iterator.next() : null;
+    }
+
+    @Override
+    protected void doClose() throws Exception {
+      // iterator can be already closed by doNext(), but closing here ensures
+      // that iterator is closed when it is not fully traversed.
+      iterator.close();
+
+      for (AutoCloseable otherCloseable : otherCloseables) {
+        otherCloseable.close();
+      }
+    }
+  }
 }
diff --git a/sonar-core/src/test/java/org/sonar/core/util/CloseableIteratorTest.java b/sonar-core/src/test/java/org/sonar/core/util/CloseableIteratorTest.java
index e2ed7b82a28..dd7feb31cf8 100644
--- a/sonar-core/src/test/java/org/sonar/core/util/CloseableIteratorTest.java
+++ b/sonar-core/src/test/java/org/sonar/core/util/CloseableIteratorTest.java
@@ -24,9 +24,12 @@ import java.util.Collections;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 import org.junit.Test;
+import org.mockito.InOrder;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
 
 public class CloseableIteratorTest {
 
@@ -138,6 +141,35 @@ public class CloseableIteratorTest {
     CloseableIterator.from(new CloseableIt());
   }
 
+  @Test
+  public void wrap_closeables() throws Exception {
+    AutoCloseable closeable1 = mock(AutoCloseable.class);
+    AutoCloseable closeable2 = mock(AutoCloseable.class);
+    CloseableIterator iterator = new SimpleCloseableIterator();
+
+    CloseableIterator wrapper = CloseableIterator.wrap(iterator, closeable1, closeable2);
+
+    assertThat(wrapper.next()).isEqualTo(1);
+    assertThat(wrapper.next()).isEqualTo(2);
+    assertThat(wrapper.hasNext()).isFalse();
+    assertThat(wrapper.isClosed).isTrue();
+    assertThat(iterator.isClosed).isTrue();
+    InOrder order = inOrder(closeable1, closeable2);
+    order.verify(closeable1).close();
+    order.verify(closeable2).close();
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void wrap_fails_if_iterator_declared_in_other_closeables() throws Exception {
+    CloseableIterator iterator = new SimpleCloseableIterator();
+    CloseableIterator.wrap(iterator, iterator);
+  }
+
+  @Test(expected = NullPointerException.class)
+  public void wrap_fails_if_null_closeable() throws Exception {
+    CloseableIterator.wrap(new SimpleCloseableIterator(), null);
+  }
+
   private static class CloseableIt implements Iterator<String>, AutoCloseable {
     private final Iterator<String> delegate = Collections.<String>emptyList().iterator();
 
-- 
2.39.5