From 6ca9a0983209a0e3fd470556c3bd9d8b5fda2ee7 Mon Sep 17 00:00:00 2001 From: ffccites <99155080+PDGGK@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:28:44 +1000 Subject: [PATCH] Release HBaseIO resources even when an earlier close() throws HBaseIO released its resources as consecutive, unguarded statements in three teardowns, so a throwing earlier close() skipped everything after it: HBaseReader.close() scanner.close() -> connection.close() HBaseWriterFn.tearDown() mutator.close() -> connection.close() WriteRowMutationsFn.tearDown() table.close() -> HBaseSharedConnection.close() The writer case is not a corner case: BufferedMutator.close() is documented to perform a flush, so any bundle whose final buffered write fails already leaked the Connection it created in @Setup. The row-mutation case is worse than an ordinary leak. The skipped call is a reference-count decrement -- HBaseSharedConnection keeps a static pool, increments on getOrCreate and only closes the underlying Connection once the count reaches zero. Missing the decrement strands that entry, and its ZooKeeper session, for the lifetime of the JVM, and every later getOrCreate hands back the same unreleasable connection. Each teardown now runs every step and keeps the first failure, attaching later ones as suppressed. A plain nested try/finally would guarantee the calls happen but would silently swap which exception the caller sees, so it is only half a fix. The collection and rethrow logic is shared by the three sites as two package-private helpers on HBaseIO. HBaseReader, HBaseWriterFn and WriteRowMutationsFn drop `private` so the new test can construct them; they stay nested and are not part of any public API. Reverting each teardown individually fails exactly and only its own test. Fixes #39710 --- sdks/java/io/hbase/build.gradle | 1 + .../org/apache/beam/sdk/io/hbase/HBaseIO.java | 91 ++++++++-- .../beam/sdk/io/hbase/HBaseIOCloseTest.java | 160 ++++++++++++++++++ 3 files changed, 241 insertions(+), 11 deletions(-) create mode 100644 sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOCloseTest.java diff --git a/sdks/java/io/hbase/build.gradle b/sdks/java/io/hbase/build.gradle index c41e7edb3177..7a06f6cc3cd4 100644 --- a/sdks/java/io/hbase/build.gradle +++ b/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") diff --git a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java b/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java index bc575b50af54..b1efb4114812 100644 --- a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java +++ b/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 getOutputCoder() { } } - private static class HBaseReader extends BoundedSource.BoundedReader { + static class HBaseReader extends BoundedSource.BoundedReader { 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 { + class HBaseWriterFn extends DoFn { 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, Integer> { + class WriteRowMutationsFn extends DoFn, 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 diff --git a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOCloseTest.java b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOCloseTest.java new file mode 100644 index 000000000000..d4aee9a7c7d0 --- /dev/null +++ b/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. + * + *

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); + } +}