This is an automated email from the ASF dual-hosted git repository.
yuluo-yx pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hertzbeat.git
The following commit(s) were added to refs/heads/master by this push:
new 90f4c28b86 [collector]bugfix: close jdbc connection on statement
creation failure (#4184)
90f4c28b86 is described below
commit 90f4c28b86471a3849789b6953f0f011dd8d7757
Author: hutiefang76 <[email protected]>
AuthorDate: Sat Jul 11 16:08:34 2026 +0800
[collector]bugfix: close jdbc connection on statement creation failure
(#4184)
Co-authored-by: hutiefang <[email protected]>
Co-authored-by: Duansg <[email protected]>
Signed-off-by: hutiefang <[email protected]>
---
.../collect/database/JdbcCommonCollect.java | 66 ++++++++--
.../collect/database/JdbcCommonCollectTest.java | 133 +++++++++++++++++++++
2 files changed, 188 insertions(+), 11 deletions(-)
diff --git
a/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollect.java
b/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollect.java
index 76e9421a6a..7be1a44ce9 100644
---
a/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollect.java
+++
b/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollect.java
@@ -63,6 +63,8 @@ public class JdbcCommonCollect extends AbstractCollect {
private static final String QUERY_TYPE_MULTI_ROW = "multiRow";
private static final String QUERY_TYPE_COLUMNS = "columns";
private static final String RUN_SCRIPT = "runScript";
+ private static final int CONNECTION_LOCK_STRIPES = 64;
+ private static final Object[] CONNECTION_LOCKS = createConnectionLocks();
private static final String[] VULNERABLE_KEYWORDS =
{"allowLoadLocalInfile", "allowLoadLocalInfileInPath", "useLocalInfile"};
@@ -317,6 +319,17 @@ public class JdbcCommonCollect extends AbstractCollect {
CacheIdentifier identifier = CacheIdentifier.builder()
.ip(url)
.username(username).password(password).build();
+ if (!reuseConnection) {
+ return getConnection(username, password, url, timeout, false,
identifier);
+ }
+ Object lock = CONNECTION_LOCKS[Math.floorMod(identifier.hashCode(),
CONNECTION_LOCKS.length)];
+ synchronized (lock) {
+ return getConnection(username, password, url, timeout, true,
identifier);
+ }
+ }
+
+ private Statement getConnection(String username, String password, String
url, Integer timeout,
+ boolean reuseConnection, CacheIdentifier
identifier) throws Exception {
Statement statement = null;
if (reuseConnection) {
Optional<AbstractConnection<?>> cacheOption =
connectionCommonCache.getCache(identifier, true);
@@ -348,21 +361,52 @@ public class JdbcCommonCollect extends AbstractCollect {
return statement;
}
}
- // renew connection when failed
- Connection connection = DriverManager.getConnection(url, username,
password);
- connection.setReadOnly(true);
- statement = connection.createStatement();
- int timeoutSecond = timeout / 1000;
- timeoutSecond = timeoutSecond <= 0 ? 1 : timeoutSecond;
- statement.setQueryTimeout(timeoutSecond);
- statement.setMaxRows(1000);
- if (reuseConnection) {
- JdbcConnect jdbcConnect = new JdbcConnect(connection);
- connectionCommonCache.addCache(identifier, jdbcConnect);
+ Connection connection = null;
+ try {
+ // renew connection when failed
+ connection = openConnection(url, username, password);
+ connection.setReadOnly(true);
+ statement = connection.createStatement();
+ int timeoutSecond = timeout / 1000;
+ timeoutSecond = timeoutSecond <= 0 ? 1 : timeoutSecond;
+ statement.setQueryTimeout(timeoutSecond);
+ statement.setMaxRows(1000);
+ if (reuseConnection) {
+ JdbcConnect jdbcConnect = new JdbcConnect(connection);
+ connectionCommonCache.addCache(identifier, jdbcConnect);
+ }
+ } catch (Exception exception) {
+ try {
+ if (statement != null) {
+ statement.close();
+ }
+ } catch (Exception closeException) {
+ log.error("Jdbc close statement error: {}",
closeException.getMessage());
+ }
+ try {
+ if (connection != null) {
+ connection.close();
+ }
+ } catch (Exception closeException) {
+ log.error("Jdbc close connection error: {}",
closeException.getMessage());
+ }
+ throw exception;
}
return statement;
}
+ private static Object[] createConnectionLocks() {
+ Object[] locks = new Object[CONNECTION_LOCK_STRIPES];
+ for (int i = 0; i < locks.length; i++) {
+ locks[i] = new Object();
+ }
+ return locks;
+ }
+
+ Connection openConnection(String url, String username, String password)
throws SQLException {
+ return DriverManager.getConnection(url, username, password);
+ }
+
/**
* query one row record, response metrics header and one value row
* eg:
diff --git
a/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollectTest.java
b/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollectTest.java
index 88cc7adcb6..3c10a31765 100644
---
a/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollectTest.java
+++
b/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/database/JdbcCommonCollectTest.java
@@ -17,21 +17,41 @@
package org.apache.hertzbeat.collector.collect.database;
+import org.apache.hertzbeat.collector.collect.common.cache.CacheIdentifier;
+import
org.apache.hertzbeat.collector.collect.common.cache.GlobalConnectionCache;
import org.apache.hertzbeat.collector.dispatch.DispatchConstants;
import org.apache.hertzbeat.common.entity.job.Metrics;
import org.apache.hertzbeat.common.entity.job.protocol.JdbcProtocol;
import org.apache.hertzbeat.common.entity.message.CollectRep;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.sql.Statement;
import java.util.LinkedHashMap;
import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
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 static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
/**
* Test case for {@link JdbcCommonCollect}
@@ -165,6 +185,65 @@ class JdbcCommonCollectTest {
assertEquals("Not support database platform: invalid",
exception.getMessage());
}
+ @Test
+ void testCloseConnectionWhenCreateStatementFails() throws Exception {
+ String url = "jdbc:postgresql://localhost:5432/hertzbeat";
+ String username = "root";
+ String password = "root";
+ Connection connection = mock(Connection.class);
+ when(connection.createStatement()).thenThrow(new SQLException("create
statement failed"));
+
+ try (MockedStatic<DriverManager> driverManager =
mockStatic(DriverManager.class)) {
+ driverManager.when(() -> DriverManager.getConnection(url,
username, password)).thenReturn(connection);
+
+ SQLException exception = assertThrows(SQLException.class,
+ () -> getConnection(jdbcCommonCollect, username, password,
url, 1000, false));
+
+ assertEquals("create statement failed", exception.getMessage());
+ verify(connection).close();
+ }
+ }
+
+ @Test
+ void testReuseConnectionSerializesCreationForSameIdentifier() throws
Exception {
+ String url = "jdbc:hertzbeat-test:concurrent";
+ String username = "root";
+ String password = "root";
+ Connection firstConnection = mock(Connection.class);
+ Connection secondConnection = mock(Connection.class);
+ Statement firstStatement = mock(Statement.class);
+ Statement reusedStatement = mock(Statement.class);
+ when(firstConnection.createStatement()).thenReturn(firstStatement,
reusedStatement);
+
when(secondConnection.createStatement()).thenReturn(mock(Statement.class));
+
+ BlockingJdbcCommonCollect collect = new
BlockingJdbcCommonCollect(firstConnection, secondConnection);
+ CacheIdentifier identifier = CacheIdentifier.builder()
+ .ip(url).username(username).password(password).build();
+ GlobalConnectionCache.getInstance().removeCache(identifier);
+
+ try (ExecutorService executor = Executors.newFixedThreadPool(2)) {
+ Future<Statement> first = executor.submit(() ->
getConnection(collect, username, password, url, 1000, true));
+ assertTrue(collect.firstOpenStarted.await(5, TimeUnit.SECONDS));
+
+ CountDownLatch secondTaskStarted = new CountDownLatch(1);
+ Future<Statement> second = executor.submit(() -> {
+ secondTaskStarted.countDown();
+ return getConnection(collect, username, password, url, 1000,
true);
+ });
+ assertTrue(secondTaskStarted.await(5, TimeUnit.SECONDS));
+ boolean secondPhysicalConnectionOpened =
collect.secondOpenStarted.await(1, TimeUnit.SECONDS);
+ collect.allowFirstOpen.countDown();
+
+ assertSame(firstStatement, first.get(5, TimeUnit.SECONDS));
+ assertSame(reusedStatement, second.get(5, TimeUnit.SECONDS));
+ assertFalse(secondPhysicalConnectionOpened);
+ assertEquals(1, collect.openCount.get());
+ } finally {
+ collect.allowFirstOpen.countDown();
+ GlobalConnectionCache.getInstance().removeCache(identifier);
+ }
+ }
+
@Test
void testConstructDatabaseUrlSecurityInterception() {
JdbcCommonCollect jdbcCollect = new JdbcCommonCollect();
@@ -221,4 +300,58 @@ class JdbcCommonCollectTest {
}
}
+ private Statement getConnection(JdbcCommonCollect jdbcCollect, String
username, String password, String url,
+ Integer timeout, boolean reuseConnection)
throws Exception {
+ try {
+ Method getConnectionMethod = JdbcCommonCollect.class
+ .getDeclaredMethod("getConnection", String.class,
String.class, String.class,
+ Integer.class, boolean.class);
+ getConnectionMethod.setAccessible(true);
+ return (Statement) getConnectionMethod.invoke(jdbcCollect,
username, password, url, timeout,
+ reuseConnection);
+ } catch (InvocationTargetException e) {
+ Throwable cause = e.getCause();
+ if (cause instanceof Exception exception) {
+ throw exception;
+ }
+ if (cause instanceof Error error) {
+ throw error;
+ }
+ throw new RuntimeException(cause);
+ }
+ }
+
+ private static final class BlockingJdbcCommonCollect extends
JdbcCommonCollect {
+
+ private final Connection firstConnection;
+ private final Connection secondConnection;
+ private final CountDownLatch firstOpenStarted = new CountDownLatch(1);
+ private final CountDownLatch secondOpenStarted = new CountDownLatch(1);
+ private final CountDownLatch allowFirstOpen = new CountDownLatch(1);
+ private final AtomicInteger openCount = new AtomicInteger();
+
+ private BlockingJdbcCommonCollect(Connection firstConnection,
Connection secondConnection) {
+ this.firstConnection = firstConnection;
+ this.secondConnection = secondConnection;
+ }
+
+ @Override
+ Connection openConnection(String url, String username, String
password) throws SQLException {
+ if (openCount.incrementAndGet() == 1) {
+ firstOpenStarted.countDown();
+ try {
+ if (!allowFirstOpen.await(5, TimeUnit.SECONDS)) {
+ throw new SQLException("Timed out waiting to open the
first connection");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new SQLException("Interrupted while opening the
first connection", e);
+ }
+ return firstConnection;
+ }
+ secondOpenStarted.countDown();
+ return secondConnection;
+ }
+ }
+
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
