diff --git a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
index a39004a9d0..ba5a539a87 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
@@ -22,7 +22,9 @@
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
+import java.util.stream.Stream;
import java.util.stream.StreamSupport;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Utilities for AutoCloseable classes. */
public final class AutoCloseables {
@@ -33,7 +35,8 @@ private AutoCloseables() {}
* Returns a new {@link AutoCloseable} that calls {@link #close(Iterable)} on autoCloseables
* when close is called.
*/
- public static AutoCloseable all(final Collection extends AutoCloseable> autoCloseables) {
+ public static AutoCloseable all(
+ final @Nullable Collection extends @Nullable AutoCloseable> autoCloseables) {
return new AutoCloseable() {
@Override
public void close() throws Exception {
@@ -48,7 +51,10 @@ public void close() throws Exception {
* @param t the throwable to add suppressed exception to
* @param autoCloseables the closeables to close
*/
- public static void close(Throwable t, AutoCloseable... autoCloseables) {
+ public static void close(Throwable t, @Nullable AutoCloseable... autoCloseables) {
+ if (autoCloseables == null) {
+ return;
+ }
close(t, Arrays.asList(autoCloseables));
}
@@ -58,7 +64,8 @@ public static void close(Throwable t, AutoCloseable... autoCloseables) {
* @param t the throwable to add suppressed exception to
* @param autoCloseables the closeables to close
*/
- public static void close(Throwable t, Iterable extends AutoCloseable> autoCloseables) {
+ public static void close(
+ Throwable t, @Nullable Iterable extends @Nullable AutoCloseable> autoCloseables) {
try {
close(autoCloseables);
} catch (Exception e) {
@@ -71,7 +78,10 @@ public static void close(Throwable t, Iterable extends AutoCloseable> autoClos
*
* @param autoCloseables the closeables to close
*/
- public static void close(AutoCloseable... autoCloseables) throws Exception {
+ public static void close(@Nullable AutoCloseable... autoCloseables) throws Exception {
+ if (autoCloseables == null) {
+ return;
+ }
close(Arrays.asList(autoCloseables));
}
@@ -80,7 +90,8 @@ public static void close(AutoCloseable... autoCloseables) throws Exception {
*
* @param ac the closeables to close
*/
- public static void close(Iterable extends AutoCloseable> ac) throws Exception {
+ public static void close(@Nullable Iterable extends @Nullable AutoCloseable> ac)
+ throws Exception {
// this method can be called on a single object if it implements Iterable
// like for example VectorContainer make sure we handle that properly
if (ac == null) {
@@ -111,12 +122,17 @@ public static void close(Iterable extends AutoCloseable> ac) throws Exception
/** Calls {@link #close(Iterable)} on the flattened list of closeables. */
@SafeVarargs
- public static void close(Iterable extends AutoCloseable>... closeables) throws Exception {
+ public static void close(@Nullable Iterable extends @Nullable AutoCloseable>... closeables)
+ throws Exception {
+ if (closeables == null) {
+ return;
+ }
close(flatten(closeables));
}
@SafeVarargs
- private static Iterable flatten(Iterable extends AutoCloseable>... closeables) {
+ private static Iterable flatten(
+ Iterable extends @Nullable AutoCloseable>... closeables) {
return new Iterable() {
// Cast from Iterable extends AutoCloseable> to Iterable is safe in this
// context
@@ -127,16 +143,18 @@ public Iterator iterator() {
return Arrays.stream(closeables)
.flatMap(
(Iterable extends AutoCloseable> i) ->
- StreamSupport.stream(
- ((Iterable) i).spliterator(), /* parallel= */ false))
+ i == null
+ ? Stream.empty()
+ : StreamSupport.stream(
+ ((Iterable) i).spliterator(), /* parallel= */ false))
.iterator();
}
};
}
/** Converts ac to a {@link Iterable} filtering out any null values. */
- public static Iterable iter(AutoCloseable... ac) {
- if (ac.length == 0) {
+ public static Iterable iter(@Nullable AutoCloseable... ac) {
+ if (ac == null || ac.length == 0) {
return Collections.emptyList();
} else {
final List nonNullAc = new ArrayList<>();
@@ -153,10 +171,11 @@ public static Iterable iter(AutoCloseable... ac) {
public static class RollbackCloseable implements AutoCloseable {
private boolean commit = false;
- private List closeables;
+ private final List closeables;
- public RollbackCloseable(AutoCloseable... closeables) {
- this.closeables = new ArrayList<>(Arrays.asList(closeables));
+ public RollbackCloseable(@Nullable AutoCloseable... closeables) {
+ this.closeables =
+ closeables == null ? new ArrayList<>() : new ArrayList<>(Arrays.asList(closeables));
}
public T add(T t) {
@@ -165,12 +184,18 @@ public T add(T t) {
}
/** Add all of list to the rollback list. */
- public void addAll(AutoCloseable... list) {
+ public void addAll(@Nullable AutoCloseable... list) {
+ if (list == null) {
+ return;
+ }
closeables.addAll(Arrays.asList(list));
}
/** Add all of list to the rollback list. */
- public void addAll(Iterable extends AutoCloseable> list) {
+ public void addAll(@Nullable Iterable extends @Nullable AutoCloseable> list) {
+ if (list == null) {
+ return;
+ }
for (AutoCloseable ac : list) {
closeables.add(ac);
}
@@ -189,7 +214,7 @@ public void close() throws Exception {
}
/** Creates an {@link RollbackCloseable} from the given closeables. */
- public static RollbackCloseable rollbackable(AutoCloseable... closeables) {
+ public static RollbackCloseable rollbackable(@Nullable AutoCloseable... closeables) {
return new RollbackCloseable(closeables);
}
@@ -203,7 +228,7 @@ public static RollbackCloseable rollbackable(AutoCloseable... closeables) {
* @throws RuntimeException if an Exception occurs; the Exception is wrapped by the
* RuntimeException
*/
- public static void closeNoChecked(final AutoCloseable autoCloseable) {
+ public static void closeNoChecked(final @Nullable AutoCloseable autoCloseable) {
if (autoCloseable != null) {
try {
autoCloseable.close();
diff --git a/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java
new file mode 100644
index 0000000000..ba5b78178a
--- /dev/null
+++ b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java
@@ -0,0 +1,268 @@
+/*
+ * 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.arrow.util;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+public class TestAutoCloseables {
+
+ /** Closeable that records that it was closed and can optionally throw. */
+ private static final class TrackCloseable implements AutoCloseable {
+ private boolean closed;
+ private final Exception toThrow;
+
+ TrackCloseable() {
+ this.toThrow = null;
+ }
+
+ TrackCloseable(Exception toThrow) {
+ this.toThrow = toThrow;
+ }
+
+ @Override
+ public void close() throws Exception {
+ closed = true;
+ if (toThrow != null) {
+ throw toThrow;
+ }
+ }
+
+ boolean isClosed() {
+ return closed;
+ }
+ }
+
+ @Test
+ public void testCloseVarargsIgnoresNulls() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.close(a, null, b);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testCloseVarargsThrowsFirstExceptionAndSuppressesRest() throws Exception {
+ Exception e1 = new Exception("first");
+ Exception e2 = new Exception("second");
+ TrackCloseable c1 = new TrackCloseable(e1);
+ TrackCloseable c2 = new TrackCloseable(e2);
+ Exception thrown = assertThrows(Exception.class, () -> AutoCloseables.close(c1, c2));
+ assertEquals("first", thrown.getMessage());
+ assertTrue(Arrays.asList(thrown.getSuppressed()).contains(e2));
+ }
+
+ @Test
+ public void testCloseIterableNullIterableReturns() throws Exception {
+ AutoCloseables.close((List) null); // no exception
+ }
+
+ @Test
+ public void testCloseIterableIgnoresNullElements() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ List list = Arrays.asList(a, null, b);
+ AutoCloseables.close(list);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testCloseIterableWhenIterableIsAlsoAutoCloseable() throws Exception {
+ TrackCloseable iter = new TrackCloseable();
+ TrackCloseable inner = new TrackCloseable();
+ // When the Iterable itself implements AutoCloseable (e.g. VectorContainer),
+ // close(Iterable) calls close() on it and does not iterate over elements
+ class IterableCloseable implements Iterable, AutoCloseable {
+ @Override
+ @SuppressWarnings("unchecked")
+ public Iterator iterator() {
+ return (Iterator) Collections.singletonList(inner);
+ }
+
+ @Override
+ public void close() throws Exception {
+ iter.close();
+ }
+ }
+ AutoCloseables.close(new IterableCloseable());
+ assertTrue(iter.isClosed());
+ assertFalse(inner.isClosed());
+ }
+
+ @Test
+ public void testCloseIterableVarargsWithNullIterables() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ TrackCloseable c = new TrackCloseable();
+ List list1 = Arrays.asList(null, a, b);
+ List list2 = Collections.singletonList(c);
+ AutoCloseables.close(list1, null, list2);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ assertTrue(c.isClosed());
+ }
+
+ @Test
+ public void testCloseThrowableSuppressesException() {
+ Exception e = new Exception("from close");
+ TrackCloseable c = new TrackCloseable(e);
+ Exception main = new Exception("main");
+ AutoCloseables.close(main, c);
+ assertTrue(c.isClosed());
+ assertEquals(1, main.getSuppressed().length);
+ assertEquals(e, main.getSuppressed()[0]);
+ }
+
+ @Test
+ public void testCloseThrowableWithNullCloseables() {
+ Exception main = new Exception("main");
+ AutoCloseables.close(main, (AutoCloseable) null);
+ assertEquals(0, main.getSuppressed().length);
+
+ AutoCloseables.close(main, (AutoCloseable[]) null); // no exception
+ }
+
+ @Test
+ public void testIterFiltersNulls() {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ Iterable it = AutoCloseables.iter(a, null, b);
+ List list = new ArrayList<>();
+ it.forEach(list::add);
+ assertEquals(2, list.size());
+ assertTrue(list.contains(a));
+ assertTrue(list.contains(b));
+ }
+
+ @Test
+ public void testIterEmptyVarargs() {
+ Iterable it = AutoCloseables.iter();
+ List list = new ArrayList<>();
+ it.forEach(list::add);
+ assertTrue(list.isEmpty());
+ }
+
+ @Test
+ public void testIterWithNull() {
+ AutoCloseables.iter((AutoCloseable) null); // no exception
+ }
+
+ @Test
+ public void testCloseNoCheckedWithNull() {
+ AutoCloseables.closeNoChecked(null); // no exception
+ }
+
+ @Test
+ public void testCloseNoCheckedWrapsException() {
+ Exception e = new Exception("close failed");
+ TrackCloseable c = new TrackCloseable(e);
+ RuntimeException re =
+ assertThrows(RuntimeException.class, () -> AutoCloseables.closeNoChecked(c));
+ assertSame(re.getCause(), e);
+ assertTrue(re.getMessage().contains("close failed"));
+ }
+
+ @Test
+ public void testNoop() throws Exception {
+ AutoCloseable noop = AutoCloseables.noop();
+ assertSame(noop, AutoCloseables.noop());
+ noop.close(); // no exception
+ }
+
+ @Test
+ public void testAllClosesCollectionOnClose() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ List list = Arrays.asList(a, b);
+ AutoCloseable all = AutoCloseables.all(list);
+ assertFalse(a.isClosed());
+ assertFalse(b.isClosed());
+ all.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testAllWithNullCollection() throws Exception {
+ AutoCloseable all = AutoCloseables.all(null);
+ all.close(); // no exception
+ }
+
+ @Test
+ public void testRollbackCloseableClosesWhenNotCommitted() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b);
+ rb.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableDoesNotCloseWhenCommitted() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b);
+ rb.commit();
+ rb.close();
+ assertFalse(a.isClosed());
+ assertFalse(b.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableAddAndAddAll() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ TrackCloseable c = new TrackCloseable();
+ TrackCloseable d = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a);
+ rb.add(b);
+ rb.addAll(c, d);
+ rb.addAll((AutoCloseable[]) null); // null varargs shouldn't fail
+ rb.addAll((List) null); // null Iterable shouldn't fail
+ rb.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ assertTrue(c.isClosed());
+ assertTrue(d.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableWithNull() throws Exception {
+ AutoCloseables.rollbackable((AutoCloseable) null); // no exception
+ }
+
+ @Test
+ public void testRollbackCloseableWithNulls() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, null);
+ rb.close();
+ assertTrue(a.isClosed());
+ }
+}