Release HBaseIO resources even when an earlier close() throws

sdks/java/io/hbase/build.gradle

@@ -46,6 +46,7 @@ dependencies {
   testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
   testImplementation library.java.junit
   testImplementation library.java.hamcrest
+  testImplementation library.java.mockito_core
   // shaded-testing-utils has shaded all Hadoop/HBase dependencies
   testImplementation("org.apache.hbase:hbase-shaded-testing-util:$hbase_version")
   testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")

sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java

@@ -48,6 +48,7 @@
 import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.TableName;
@@ -187,6 +188,33 @@ public class HBaseIO {
   /** Disallow construction of utility class. */
   private HBaseIO() {}
 
+  /**
+   * Collects a teardown failure. The first one is the one the caller ends up seeing; later ones are
+   * attached to it as suppressed, so nothing is lost and nothing replaces the original.
+   */
+  @VisibleForTesting
+  static Throwable appendSuppressed(@Nullable Throwable existingFailure, Throwable newFailure) {
+    if (existingFailure == null) {
+      return newFailure;
+    }
+    existingFailure.addSuppressed(newFailure);
+    return existingFailure;
+  }
+
+  /**
+   * Rethrows a failure collected by {@link #appendSuppressed}, preserving its type where it can.
+   */
+  @VisibleForTesting
+  static void rethrowCloseFailure(Throwable failure) throws IOException {
+    if (failure instanceof IOException) {
+      throw (IOException) failure;
+    }
+    if (failure instanceof RuntimeException) {
+      throw (RuntimeException) failure;
+    }
+    throw new IOException(failure);
+  }
+
   /**
    * Creates an uninitialized {@link HBaseIO.Read}. Before use, the {@code Read} must be initialized
    * with a {@link HBaseIO.Read#withConfiguration(Configuration)} that specifies the HBase instance,
@@ -492,7 +520,7 @@ public Coder<Result> getOutputCoder() {
     }
   }
 
-  private static class HBaseReader extends BoundedSource.BoundedReader<Result> {
+  static class HBaseReader extends BoundedSource.BoundedReader<Result> {
     private HBaseSource source;
     private Connection connection;
     private ResultScanner scanner;
@@ -549,14 +577,28 @@ public boolean advance() {
     @Override
     public void close() throws IOException {
       LOG.debug("Closing reader after reading {} records.", recordsReturned);
+      // Release everything even if an earlier step throws, and keep the first failure: a
+      // connection left behind here outlives the reader.
+      Throwable failure = null;
       if (scanner != null) {
-        scanner.close();
+        try {
+          scanner.close();
+        } catch (Exception e) {
+          failure = appendSuppressed(failure, e);
+        }
         scanner = null;
       }
       if (connection != null) {
-        connection.close();
+        try {
+          connection.close();
+        } catch (Exception e) {
+          failure = appendSuppressed(failure, e);
+        }
         connection = null;
       }
+      if (failure != null) {
+        rethrowCloseFailure(failure);
+      }
     }
 
     @Override
@@ -732,7 +774,7 @@ Object readResolve() {
 
     private final String tableId;
 
-    private class HBaseWriterFn extends DoFn<Mutation, Void> {
+    class HBaseWriterFn extends DoFn<Mutation, Void> {
 
       HBaseWriterFn(Write write) {
         checkNotNull(write.tableId, "tableId");
@@ -765,14 +807,28 @@ public void finishBundle() throws Exception {
 
       @Teardown
       public void tearDown() throws Exception {
+        // BufferedMutator.close() performs a flush, so a failed final batch is an expected way
+        // for this to throw. Release the connection anyway, and keep the flush failure.
+        Throwable failure = null;
         if (mutator != null) {
-          mutator.close();
+          try {
+            mutator.close();
+          } catch (Exception e) {
+            failure = appendSuppressed(failure, e);
+          }
           mutator = null;
         }
         if (connection != null) {
-          connection.close();
+          try {
+            connection.close();
+          } catch (Exception e) {
+            failure = appendSuppressed(failure, e);
+          }
           connection = null;
         }
+        if (failure != null) {
+          rethrowCloseFailure(failure);
+        }
       }
 
       @Override
@@ -900,7 +956,7 @@ Object readResolve() {
     private final String tableId;
 
     /** Function to write row mutations to a hbase table. */
-    private class WriteRowMutationsFn extends DoFn<KV<byte[], RowMutations>, Integer> {
+    class WriteRowMutationsFn extends DoFn<KV<byte[], RowMutations>, Integer> {
 
       public WriteRowMutationsFn(WriteRowMutations writeRowMutations) {
         checkNotNull(writeRowMutations.tableId, "tableId");
@@ -930,13 +986,26 @@ public void finishBundle() throws Exception {
 
       @Teardown
       public void tearDown() throws Exception {
-
+        // HBaseSharedConnection.close() is a reference-count decrement, not an ordinary close.
+        // Skipping it strands the entry in the static pool for the lifetime of the JVM, so it has
+        // to run even when the table fails to close.
+        Throwable failure = null;
         if (table != null) {
-          table.close();
+          try {
+            table.close();
+          } catch (Exception e) {
+            failure = appendSuppressed(failure, e);
+          }
           table = null;
         }
-
-        HBaseSharedConnection.close(configuration);
+        try {
+          HBaseSharedConnection.close(configuration);
+        } catch (Exception e) {
+          failure = appendSuppressed(failure, e);
+        }
+        if (failure != null) {
+          rethrowCloseFailure(failure);
+        }
       }
 
       @ProcessElement

sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOCloseTest.java

@@ -0,0 +1,160 @@
+/*
+ * 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.beam.sdk.io.hbase;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import java.io.IOException;
+import java.lang.reflect.Field;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.HBaseConfiguration;
+import org.apache.hadoop.hbase.client.BufferedMutator;
+import org.apache.hadoop.hbase.client.Connection;
+import org.apache.hadoop.hbase.client.ResultScanner;
+import org.apache.hadoop.hbase.client.Table;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * {@link HBaseIO} used to release its resources as consecutive, unguarded statements, so a throwing
+ * earlier {@code close()} skipped everything after it.
+ *
+ * <p>The sharpest case is {@link HBaseIO.WriteRowMutations.WriteRowMutationsFn#tearDown()}: the
+ * skipped call there is {@link HBaseSharedConnection#close(Configuration)}, which is a
+ * reference-count decrement rather than an ordinary close. Missing it strands the entry in a {@code
+ * static} pool for the lifetime of the JVM, so that test asserts the real count rather than a mock
+ * interaction.
+ */
+@RunWith(JUnit4.class)
+public class HBaseIOCloseTest {
+
+  private final Configuration configuration = HBaseConfiguration.create();
+
+  @After
+  public void resetConnectionPool() throws IOException {
+    HBaseSharedConnection.closeAll();
+  }
+
+  // ---------------------------------------------------------------- failure collection
+
+  @Test
+  public void appendSuppressedKeepsTheFirstFailure() {
+    IOException first = new IOException("first");
+    IOException second = new IOException("second");
+
+    assertSame(first, HBaseIO.appendSuppressed(null, first));
+
+    Throwable collected = HBaseIO.appendSuppressed(first, second);
+    assertSame(first, collected);
+    assertArrayEquals(new Throwable[] {second}, collected.getSuppressed());
+  }
+
+  @Test
+  public void rethrowCloseFailurePreservesTheType() {
+    IOException io = new IOException("io");
+    assertSame(io, assertThrows(IOException.class, () -> HBaseIO.rethrowCloseFailure(io)));
+
+    IllegalStateException unchecked = new IllegalStateException("unchecked");
+    assertSame(
+        unchecked,
+        assertThrows(IllegalStateException.class, () -> HBaseIO.rethrowCloseFailure(unchecked)));
+
+    // Anything else has to be wrapped, because the teardowns only declare IOException.
+    Throwable checked = new Exception("checked");
+    IOException wrapped =
+        assertThrows(IOException.class, () -> HBaseIO.rethrowCloseFailure(checked));
+    assertSame(checked, wrapped.getCause());
+  }
+
+  // ---------------------------------------------------------------- reader
+
+  @Test
+  public void readerClosesTheConnectionWhenTheScannerFailsToClose() throws Exception {
+    HBaseIO.Read read = HBaseIO.read().withConfiguration(configuration).withTableId("some_table");
+    HBaseIO.HBaseReader reader =
+        new HBaseIO.HBaseReader(new HBaseIO.HBaseSource(read, null /* estimatedSizeBytes */));
+
+    ResultScanner scanner = mock(ResultScanner.class);
+    Connection connection = mock(Connection.class);
+    IOException scannerFailure = new IOException("scanner close failed");
+    doThrow(scannerFailure).when(scanner).close();
+    // The reader only acquires these in start(), which would need a live cluster.
+    set(reader, "scanner", scanner);
+    set(reader, "connection", connection);
+
+    assertSame(scannerFailure, assertThrows(IOException.class, reader::close));
+    verify(connection).close();
+  }
+
+  // ---------------------------------------------------------------- mutation writer
+
+  @Test
+  public void writerClosesTheConnectionWhenTheFinalFlushFails() throws Exception {
+    HBaseIO.Write write =
+        HBaseIO.write().withConfiguration(configuration).withTableId("some_table");
+    HBaseIO.Write.HBaseWriterFn fn = write.new HBaseWriterFn(write);
+
+    BufferedMutator mutator = mock(BufferedMutator.class);
+    Connection connection = mock(Connection.class);
+    // BufferedMutator.close() performs a flush, so this is the expected failure mode.
+    IOException flushFailure = new IOException("flush on close failed");
+    doThrow(flushFailure).when(mutator).close();
+    set(fn, "mutator", mutator);
+    set(fn, "connection", connection);
+
+    assertSame(flushFailure, assertThrows(IOException.class, fn::tearDown));
+    verify(connection).close();
+  }
+
+  // ---------------------------------------------------------------- row-mutation writer
+
+  @Test
+  public void rowMutationWriterReleasesTheSharedConnectionWhenTheTableFailsToClose()
+      throws Exception {
+    HBaseSharedConnection.getOrCreate(configuration);
+    assertEquals(1, HBaseSharedConnection.getConnectionCount(configuration));
+
+    HBaseIO.WriteRowMutations write =
+        HBaseIO.writeRowMutations().withConfiguration(configuration).withTableId("some_table");
+    HBaseIO.WriteRowMutations.WriteRowMutationsFn fn = write.new WriteRowMutationsFn(write);
+
+    Table table = mock(Table.class);
+    IOException tableFailure = new IOException("table close failed");
+    doThrow(tableFailure).when(table).close();
+    set(fn, "table", table);
+
+    assertSame(tableFailure, assertThrows(IOException.class, fn::tearDown));
+    // The point of the fix: the reference count still went back down, so the pooled connection is
+    // releasable instead of being stranded for the lifetime of the JVM.
+    assertEquals(0, HBaseSharedConnection.getConnectionCount(configuration));
+  }
+
+  private static void set(Object target, String fieldName, Object value) throws Exception {
+    Field field = target.getClass().getDeclaredField(fieldName);
+    field.setAccessible(true);
+    field.set(target, value);
+  }
+}