This is an automated email from the ASF dual-hosted git repository.
jerryshao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new f3c4396c05 [#13449] fix(core): close the JDBC DataSource in
JdbcPartitionStatisticStorage.close (#13450)
f3c4396c05 is described below
commit f3c4396c0593dd19c1e0a16bd54fe7b3fd12a132
Author: YangJie <[email protected]>
AuthorDate: Wed Sep 23 23:22:23 2026 -0400
[#13449] fix(core): close the JDBC DataSource in
JdbcPartitionStatisticStorage.close (#13450)
### What changes were proposed in this pull request?
`JdbcPartitionStatisticStorage.close()` now closes an `AutoCloseable`
`DataSource` (such as DBCP2's `BasicDataSource`) it owns, matching the
Lance storage that releases its resources on close. A
non-`AutoCloseable` (externally managed or injected) `DataSource` keeps
its existing lifecycle and is left untouched.
### Why are the changes needed?
The factory that was assumed to own the `DataSource` is a
constructor-local in `StatisticManager` and unreachable at shutdown, so
the no-op `close()` leaked the DBCP2 pool and its sockets on every
create/close cycle.
Fix: #13449
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added
`TestJdbcPartitionStatisticStorage.testCloseClosesOwnedDataSource`,
which pins that an owned `AutoCloseable` `DataSource` is closed (and
idempotently), while an injected non-`AutoCloseable` one is left open.
It fails against the pre-fix code.
---------
Co-authored-by: Qi Yu <[email protected]>
---
.../storage/JdbcPartitionStatisticStorage.java | 17 +++++++++-
.../JdbcPartitionStatisticStorageFactory.java | 36 ++++++++--------------
.../storage/TestJdbcPartitionStatisticStorage.java | 19 ++++++++++++
3 files changed, 47 insertions(+), 25 deletions(-)
diff --git
a/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorage.java
b/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorage.java
index a976489302..a2788071ff 100644
---
a/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorage.java
+++
b/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorage.java
@@ -30,6 +30,7 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
import javax.sql.DataSource;
import org.apache.gravitino.Entity;
@@ -63,6 +64,7 @@ public class JdbcPartitionStatisticStorage implements
PartitionStatisticStorage
private final DataSource dataSource;
private final EntityStore entityStore;
private final DatabaseType databaseType;
+ private final AtomicBoolean closed = new AtomicBoolean(false);
/** Supported database types. */
private enum DatabaseType {
@@ -356,8 +358,21 @@ public class JdbcPartitionStatisticStorage implements
PartitionStatisticStorage
@Override
public void close() throws IOException {
- // DataSource lifecycle is managed externally by the factory
+ if (!closed.compareAndSet(false, true)) {
+ return;
+ }
LOG.debug("Closing JdbcPartitionStatisticStorage");
+ // This storage is the only reachable owner of the pooled DataSource: the
+ // factory that created it is discarded by the manager, so close must
+ // release the pool. DataSources that do not implement AutoCloseable keep
+ // their externally-managed lifecycle.
+ if (dataSource instanceof AutoCloseable) {
+ try {
+ ((AutoCloseable) dataSource).close();
+ } catch (Exception e) {
+ throw new IOException("Failed to close JDBC DataSource", e);
+ }
+ }
}
/**
diff --git
a/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorageFactory.java
b/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorageFactory.java
index b7a5d66fa9..49ba834f99 100644
---
a/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorageFactory.java
+++
b/core/src/main/java/org/apache/gravitino/stats/storage/JdbcPartitionStatisticStorageFactory.java
@@ -68,9 +68,6 @@ public class JdbcPartitionStatisticStorageFactory implements
PartitionStatisticS
// Default values
private static final String DEFAULT_JDBC_DRIVER = "com.mysql.cj.jdbc.Driver";
- // Keep reference to DataSource for cleanup
- private BasicDataSource dataSource;
-
@Override
public PartitionStatisticStorage create(Map<String, String> properties) {
LOG.info(
@@ -79,16 +76,20 @@ public class JdbcPartitionStatisticStorageFactory
implements PartitionStatisticS
validateRequiredProperties(properties);
+ // The created storage owns the DataSource and closes it on its own
close().
+ BasicDataSource localDataSource;
+ try {
+ localDataSource = createDataSource(properties);
+ } catch (Exception e) {
+ throw new GravitinoRuntimeException(e, "Failed to create
JdbcPartitionStatisticStorage");
+ }
try {
- dataSource = createDataSource(properties);
- return new JdbcPartitionStatisticStorage(dataSource);
+ return new JdbcPartitionStatisticStorage(localDataSource);
} catch (Exception e) {
- if (dataSource != null) {
- try {
- dataSource.close();
- } catch (SQLException closeException) {
- LOG.error("Failed to close data source after creation error",
closeException);
- }
+ try {
+ localDataSource.close();
+ } catch (SQLException closeException) {
+ LOG.error("Failed to close data source after creation error",
closeException);
}
throw new GravitinoRuntimeException(e, "Failed to create
JdbcPartitionStatisticStorage");
}
@@ -184,17 +185,4 @@ public class JdbcPartitionStatisticStorageFactory
implements PartitionStatisticS
}
return masked;
}
-
- /**
- * Closes the data source if it was created by this factory.
- *
- * @throws SQLException if closing fails
- */
- public void close() throws SQLException {
- if (dataSource != null) {
- LOG.info("Closing JDBC DataSource");
- dataSource.close();
- dataSource = null;
- }
- }
}
diff --git
a/core/src/test/java/org/apache/gravitino/stats/storage/TestJdbcPartitionStatisticStorage.java
b/core/src/test/java/org/apache/gravitino/stats/storage/TestJdbcPartitionStatisticStorage.java
index b5242eb967..a9a3bd1ad0 100644
---
a/core/src/test/java/org/apache/gravitino/stats/storage/TestJdbcPartitionStatisticStorage.java
+++
b/core/src/test/java/org/apache/gravitino/stats/storage/TestJdbcPartitionStatisticStorage.java
@@ -18,6 +18,7 @@
*/
package org.apache.gravitino.stats.storage;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -40,6 +41,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.sql.DataSource;
+import org.apache.commons.dbcp2.BasicDataSource;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.apache.gravitino.EntityStore;
import org.apache.gravitino.GravitinoEnv;
@@ -101,6 +103,23 @@ public class TestJdbcPartitionStatisticStorage {
storage = new JdbcPartitionStatisticStorage(mockDataSource);
}
+ @Test
+ public void testCloseClosesOwnedDataSource() throws Exception {
+ // The storage is the only reachable owner of the pooled DataSource (the
+ // factory is discarded by the manager), so close() must release it.
+ BasicDataSource dataSource = new BasicDataSource();
+ dataSource.setDriverClassName("org.h2.Driver");
+ dataSource.setUrl("jdbc:h2:mem:stats_close_test;DB_CLOSE_DELAY=-1");
+
+ JdbcPartitionStatisticStorage owned = new
JdbcPartitionStatisticStorage(dataSource);
+ owned.close();
+ assertTrue(dataSource.isClosed());
+
+ // close() is idempotent: a second call must not throw and must leave the
pool closed.
+ assertDoesNotThrow(owned::close);
+ assertTrue(dataSource.isClosed());
+ }
+
@AfterEach
public void tearDown() throws Exception {
if (storage != null) {
