diff --git a/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/BSimPostgresDBConnectionManager.java b/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/BSimPostgresDBConnectionManager.java index 70aea9aa2b..bd3ac7176c 100644 --- a/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/BSimPostgresDBConnectionManager.java +++ b/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/BSimPostgresDBConnectionManager.java @@ -18,6 +18,7 @@ package ghidra.features.bsim.query; import java.net.URL; import java.sql.Connection; import java.sql.SQLException; +import java.util.Collection; import java.util.HashMap; import javax.security.auth.callback.NameCallback; @@ -40,12 +41,14 @@ public class BSimPostgresDBConnectionManager { private static final int CONN_POOL_MAX_IDLE = 2; private static HashMap dataSourceMap = new HashMap<>(); + private static boolean shutdownHookInstalled = false; public static synchronized BSimPostgresDataSource getDataSource( BSimServerInfo postgresServerInfo) { if (postgresServerInfo.getDBType() != DBType.postgres) { throw new IllegalArgumentException("expected postgres server info"); } + enableShutdownHook(); return dataSourceMap.computeIfAbsent(postgresServerInfo, info -> new BSimPostgresDataSource(info)); } @@ -75,6 +78,30 @@ public class BSimPostgresDBConnectionManager { dataSourceMap.remove(serverInfo); } + private static synchronized void enableShutdownHook() { + if (shutdownHookInstalled) { + return; + } + Runtime.getRuntime().addShutdownHook(new Thread() { + @Override + public void run() { + Collection dataSources = dataSourceMap.values(); + for (BSimPostgresDataSource ds : dataSources) { + int activeConnections = ds.getActiveConnections(); + if (activeConnections != 0) { + Msg.error(BSimPostgresDBConnectionManager.class, + activeConnections + + " BSim active Postgres connections were not properly closed: " + + ds.serverInfo); + } + ds.close(); + } + dataSourceMap.clear(); + } + }); + shutdownHookInstalled = true; + } + public static class BSimPostgresDataSource implements BSimJDBCDataSource { // NOTE: can be renamed private final BSimServerInfo serverInfo; diff --git a/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/client/PostgresFunctionDatabase.java b/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/client/PostgresFunctionDatabase.java index ce23733cb1..f8fd89c7ed 100755 --- a/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/client/PostgresFunctionDatabase.java +++ b/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/client/PostgresFunctionDatabase.java @@ -4,9 +4,9 @@ * Licensed 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. @@ -42,7 +42,7 @@ import ghidra.features.bsim.query.protocol.*; * */ public final class PostgresFunctionDatabase - extends AbstractSQLFunctionDatabase { + extends AbstractSQLFunctionDatabase { // NOTE: Previously named ColumnDatabase @@ -101,7 +101,7 @@ public final class PostgresFunctionDatabase } private void changePassword(Connection c, String username, char[] newPassword) - throws SQLException { + throws SQLException { StringBuilder buffer = new StringBuilder(); buffer.append("ALTER ROLE \""); buffer.append(username); @@ -158,7 +158,7 @@ public final class PostgresFunctionDatabase */ private void serverLoadWeights(Connection db) throws SQLException { try (Statement st = db.createStatement(); - ResultSet rs = st.executeQuery("SELECT lsh_load()")) { + ResultSet rs = st.executeQuery("SELECT lsh_load()")) { while (rs.next()) { // int val = rs.getInt(1); } @@ -171,13 +171,16 @@ public final class PostgresFunctionDatabase Connection db = initConnection(); serverLoadWeights(db); - if (asynchronous) { - try (Statement st = db.createStatement()) { - // Tell server to do asynchronous commits. This speeds up large - // ingests with a (slight) danger of - // losing the most recent commits if the server crashes (NOTE: - // database integrity should still be recoverable) - st.executeUpdate("SET LOCAL synchronous_commit TO OFF"); + try (Statement st = db.createStatement()) { + // Tell server to do asynchronous commits. This speeds up large + // ingests with a (slight) danger of + // losing the most recent commits if the server crashes (NOTE: + // database integrity should still be recoverable) + if (asynchronous) { + st.executeUpdate("SET SESSION synchronous_commit TO OFF"); + } + else { + st.executeUpdate("SET SESSION synchronous_commit to ON"); } } @@ -223,12 +226,15 @@ public final class PostgresFunctionDatabase serverLoadWeights(db); + // Tell server to do asynchronous commits. This speeds up large + // ingests with a (slight) danger of + // losing the most recent commits if the server crashes (NOTE: + // database integrity should still be recoverable) if (asynchronous) { - // Tell server to do asynchronous commits. This speeds up large - // ingests with a (slight) danger of - // losing the most recent commits if the server crashes (NOTE: - // database integrity should still be recoverable) - st.executeUpdate("SET LOCAL synchronous_commit TO OFF"); + st.executeUpdate("SET SESSION synchronous_commit TO OFF"); + } + else { + st.executeUpdate("SET SESSION synchronous_commit to ON"); } } } @@ -253,7 +259,7 @@ public final class PostgresFunctionDatabase */ private void rebuildIndex(Connection c) throws SQLException { try (Statement st = c.createStatement(); - ResultSet rs = st.executeQuery("SELECT lsh_reload()")) { + ResultSet rs = st.executeQuery("SELECT lsh_reload()")) { st.execute("SET maintenance_work_mem TO '2GB'"); st.execute( "CREATE INDEX vectable_vec_idx ON vectable USING gin (vec gin_lshvector_ops)"); @@ -275,7 +281,7 @@ public final class PostgresFunctionDatabase * @throws SQLException if there is a problem creating or executing the query */ private int preWarm(Connection c, int mainIndex, int secondaryIndex, int vectors) - throws SQLException { + throws SQLException { try (Statement st = c.createStatement()) { // Try to load the entire main index into the PostgreSQL cache int res = -1; @@ -386,8 +392,8 @@ public final class PostgresFunctionDatabase * @throws SQLException if there is a problem creating or executing the query */ @Override - protected int queryNearestVector(List resultset, LSHVector vec, - double simthresh, double sigthresh, int max) throws SQLException { + protected int queryNearestVector(List resultset, LSHVector vec, double simthresh, + double sigthresh, int max) throws SQLException { PreparedStatement s = selectNearestVectorStatement.prepareIfNeeded(() -> initConnection().prepareStatement( "WITH const(cvec) AS (VALUES( lshvector_in( CAST( ? AS cstring) ) ) )," + @@ -469,7 +475,7 @@ public final class PostgresFunctionDatabase @Override protected VectorResult queryVectorId(long id) throws SQLException { PreparedStatement s = selectVectorByRowIdStatement.prepareIfNeeded(() -> initConnection() - .prepareStatement("SELECT id,count,vec FROM vectable WHERE id = ?")); + .prepareStatement("SELECT id,count,vec FROM vectable WHERE id = ?")); s.setLong(1, id); try (ResultSet rs = s.executeQuery()) { if (!rs.next()) { @@ -506,7 +512,7 @@ public final class PostgresFunctionDatabase @Override public QueryResponseRecord doQuery(BSimQuery query, Connection c) - throws SQLException, LSHException, DatabaseNonFatalException { + throws SQLException, LSHException, DatabaseNonFatalException { if (query instanceof PrewarmRequest q) { fdbPrewarm(q, c); diff --git a/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/file/BSimH2FileDBConnectionManager.java b/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/file/BSimH2FileDBConnectionManager.java index b7ad1d0629..1e3b9ffaca 100644 --- a/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/file/BSimH2FileDBConnectionManager.java +++ b/Ghidra/Features/BSim/src/main/java/ghidra/features/bsim/query/file/BSimH2FileDBConnectionManager.java @@ -40,6 +40,7 @@ public class BSimH2FileDBConnectionManager { * Data source map keyed by absolute DB file path */ private static HashMap dataSourceMap = new HashMap<>(); + private static boolean shutdownHookInstalled = false; /** * Get all H2 File DB data sorces which exist in the JVM. @@ -62,6 +63,7 @@ public class BSimH2FileDBConnectionManager { if (fileServerInfo.getDBType() != DBType.file) { throw new IllegalArgumentException("expected file info"); } + enableShutdownHook(); return dataSourceMap.computeIfAbsent(fileServerInfo, info -> new BSimH2FileDataSource(info)); } @@ -102,6 +104,30 @@ public class BSimH2FileDBConnectionManager { return true; } + private static synchronized void enableShutdownHook() { + if (shutdownHookInstalled) { + return; + } + Runtime.getRuntime().addShutdownHook(new Thread() { + @Override + public void run() { + Collection dataSources = dataSourceMap.values(); + for (BSimH2FileDataSource ds : dataSources) { + int activeConnections = ds.getActiveConnections(); + if (activeConnections != 0) { + Msg.error(BSimH2FileDBConnectionManager.class, + activeConnections + + " BSim active H2 File connections were not properly closed: " + + ds.serverInfo); + } + ds.close(); + } + dataSourceMap.clear(); + } + }); + shutdownHookInstalled = true; + } + /** * {@link BSimH2FileDataSource} provides a pooled DB data source for a specific H2 File DB. */