diff --git a/api/src/org/labkey/api/data/DbScope.java b/api/src/org/labkey/api/data/DbScope.java index d84f495a64d..e8bfbf7b492 100644 --- a/api/src/org/labkey/api/data/DbScope.java +++ b/api/src/org/labkey/api/data/DbScope.java @@ -1202,6 +1202,11 @@ public synchronized boolean release(Connection conn) return 0 == _refCount; } + + public synchronized int getRefCount() + { + return _refCount; + } } /** @@ -1259,6 +1264,20 @@ private Connection getCurrentConnection(@Nullable Logger log) throws SQLExceptio return getConnectionHolder().get(log); } + /** + * @return true if this thread already holds the shared, ref-counted thread connection — i.e., some code up the stack + * called {@link #getConnection()} and hasn't released it. Reports the existing {@link ConnectionHolder} ref count + * that governs {@link ConnectionType#Thread} sharing. + *

+ * A caller that borrows the thread connection and temporarily changes its state (e.g. disabling JDBC caching for a + * streaming read) must do so only when this returns false, so it is the outermost borrower and can restore the + * original state via runOnClose, which {@link ConnectionType#Thread} fires when the last holder releases it. + */ + public boolean isThreadConnectionActive() + { + return getConnectionHolder().getRefCount() > 0; + } + /** * Get a fresh connection directly from the pool... not part of the current transaction, not shared with the thread, etc. */ diff --git a/api/src/org/labkey/api/data/SqlExecutingSelector.java b/api/src/org/labkey/api/data/SqlExecutingSelector.java index 5349c107320..6351f6c52a6 100644 --- a/api/src/org/labkey/api/data/SqlExecutingSelector.java +++ b/api/src/org/labkey/api/data/SqlExecutingSelector.java @@ -19,6 +19,8 @@ import org.apache.logging.log4j.Logger; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import org.labkey.api.cache.Cache; +import org.labkey.api.cache.CacheManager; import org.labkey.api.data.dialect.SqlDialect; import org.labkey.api.data.dialect.SqlDialect.ExecutionPlanType; import org.labkey.api.data.dialect.StatementWrapper; @@ -33,9 +35,12 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import static org.labkey.api.util.ExceptionUtil.CALCULATED_COLUMN_SQL_TAG; @@ -46,10 +51,18 @@ public abstract class SqlExecutingSelector LARGE_RESULT_WARNING_THROTTLE = CacheManager.getCache(1000, CacheManager.DAY, "SqlSelector large result warnings"); + int _maxRows = Table.ALL_ROWS; protected long _offset = Table.NO_OFFSET; @Nullable Map _namedParameters = null; - private ConnectionFactory _connectionFactory = super::getConnection; + private @Nullable ConnectionFactory _connectionFactory = null; // null means "no explicit choice"; see getEffectiveConnectionFactory() + private boolean _jdbcCachingExplicitlySet = false; private Integer _fetchSize = null; // By default, use the standard fetch size private @Nullable AsyncQueryRequest _asyncRequest = null; @@ -79,21 +92,60 @@ public interface ConnectionFactory @Override public Connection getConnection() throws SQLException { - return _connectionFactory.get(); + // Public callers may hold the returned Connection beyond this call, so treat it as "escaping": don't borrow the + // thread's shared connection. Internal callers that fully consume and close the ResultSet within a single + // selector call use getConnection(true); see getEffectiveConnectionFactory(). + return getConnection(false); + } + + Connection getConnection(boolean selfContained) throws SQLException + { + return getEffectiveConnectionFactory(selfContained).get(); + } + + /** + * Determines which {@link ConnectionFactory} to use for this query, resolved lazily so the transaction check + * reflects execution-time state. An explicit {@link #setJdbcCaching(boolean)} call or a Connection supplied at + * construction is honored; otherwise JDBC caching is disabled by default (via the dialect) so the driver won't + * buffer the whole ResultSet in memory. The dialect returns null — use the shared Connection with default caching — + * when that default is unnecessary or unsafe: in a transaction, non-PostgreSQL, or not a SELECT. + *

+ * {@code selfContained} is passed through to {@link SqlDialect#getConnectionFactory}, which documents how it governs + * whether the thread's shared connection may be borrowed. + */ + private ConnectionFactory getEffectiveConnectionFactory(boolean selfContained) + { + // Honor an explicit setJdbcCaching() call (which populated _connectionFactory)... + if (_jdbcCachingExplicitlySet) + return _connectionFactory; + + // ...or a Connection supplied at construction time (super::getConnection returns the stashed _conn) + if (null != _conn) + return super::getConnection; + + ConnectionFactory factory = getScope().getSqlDialect().getConnectionFactory(false, selfContained, getScope(), + new SQLFragment("SELECT FakeColumn FROM FakeTable") /* SqlExecutingSelector always generates SELECT statements */); + + return null != factory ? factory : super::getConnection; } /** *

Calling this method with cache=false ensures that the JDBC driver will not cache the produced ResultSet in * memory, which is useful when potentially working with very large (e.g., > 100MB) ResultSets. Calling it with - * cache=true (the default setting) ensures the JDBC driver's default caching behavior.

+ * cache=true ensures the JDBC driver's default caching behavior.

* *

By default, the PostgreSQL JDBC driver caches every ResultSet in its entirety. This can lead to * OutOfMemoryErrors when working with very large ResultSets. When the underlying database is PostgreSQL, calling * this method with false instructs this SqlExecutingSelector to use an unshared Connection and configure it with * special settings that disable the driver caching. The trade-off is that the underlying database query will not * use the shared Connection that other code on the thread (up or down the call stack) may be using, making - * Connection exhaustion more likely; that's why JDBC caching is on by default. Calling this method is not - * compatible with passing in an explicit Connection to the constructor.

+ * Connection exhaustion more likely. Calling this method is not compatible with passing in an explicit Connection to + * the constructor.

+ * + *

When neither this method nor an explicit Connection is supplied, JDBC caching is disabled by default whenever + * it's safe (PostgreSQL, no active transaction, SELECT) — see {@link #getEffectiveConnectionFactory(boolean)}. Callers that + * require the driver's default caching (e.g. to share the thread's Connection) must opt in by calling this with + * cache=true.

* *

When the underlying database is not PostgreSQL, calling this method has no effect, other than validating that * the stashed Connection is null.

@@ -106,13 +158,50 @@ public SELECTOR setJdbcCaching(boolean cache) if (null != _conn) throw new IllegalStateException("Calling setJdbcCaching() is not valid when a Connection has already been provided"); - ConnectionFactory factory = getScope().getSqlDialect().getConnectionFactory(cache, getScope(), + // Explicitly disabling caching is documented to use an unshared Connection, so this path is never self-contained + ConnectionFactory factory = getScope().getSqlDialect().getConnectionFactory(cache, false, getScope(), new SQLFragment("SELECT FakeColumn FROM FakeTable") /* SqlExecutingSelector always generates SELECT statements */); _connectionFactory = null != factory ? factory : super::getConnection; + _jdbcCachingExplicitlySet = true; return getThis(); } + /** + * Overridden to warn when a large number of rows is pulled into a Java collection. Loading many rows into memory + * (here plus, potentially, in the JDBC driver's buffer) is a common source of OutOfMemoryErrors; callers should + * generally prefer a streaming method — {@link #forEach(Class, Selector.ForEachBlock)}, {@link #forEachBatch}, or {@link #uncachedStream} — that + * processes rows without materializing them all at once. {@code getArray}, {@code getCollection}, + * {@code getMapArray}, and {@code getMapCollection} all delegate here, so they're covered as well. + */ + @Override + public @NotNull ArrayList getArrayList(Class clazz) + { + ArrayList result = super.getArrayList(clazz); + + if (result.size() >= LARGE_RESULT_THRESHOLD) + { + Throwable stackTrace = new Throwable("Stack trace for large collection load"); + String stackKey = getStackKey(stackTrace); + + // Warn at most once per day (tolerating a race condition) per unique call stack to avoid flooding the log. + if (null == LARGE_RESULT_WARNING_THROTTLE.get(stackKey)) + { + LARGE_RESULT_WARNING_THROTTLE.put(stackKey, Boolean.TRUE); + LOGGER.warn("{} rows loaded into a collection via {}. Consider switching to forEach(), forEachBatch(), or uncachedStream() to reduce memory usage. SQL: {}", + result.size(), getClass().getSimpleName(), getSqlFactory(false).getSql(), stackTrace); + } + } + + return result; + } + + // Builds a stable key from a Throwable's stack trace so identical call stacks map to the same throttle entry + private static String getStackKey(Throwable t) + { + return Arrays.stream(t.getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("\n")); + } + /** * Set a ResultSet fetch size that differs from the default value (1,000 rows on PostgreSQL). This is normally a * fine fetch size, but not when dealing with rows containing large TEXT or BYTEA columns. @@ -396,7 +485,9 @@ public T handleResultSet(ResultSetHandler handler) if (null != _sql) { DbScope scope = getScope(); - conn = getConnection(); + // _closeResultSet means this factory fully consumes and closes the ResultSet within this call, so + // the connection can be borrowed from the thread and restored rather than dedicated to the caller + conn = getConnection(_closeResultSet); try { diff --git a/api/src/org/labkey/api/data/SqlSelectorTestCase.java b/api/src/org/labkey/api/data/SqlSelectorTestCase.java index e77de43d433..8bbe42e0631 100644 --- a/api/src/org/labkey/api/data/SqlSelectorTestCase.java +++ b/api/src/org/labkey/api/data/SqlSelectorTestCase.java @@ -186,13 +186,23 @@ public void testJdbcUncached() throws SQLException DbScope scope = CoreSchema.getInstance().getScope(); try (Connection conn = scope.getConnection()) { - // Default setting is to cache and share the connection + // Default (no explicit setJdbcCaching() call) now auto-disables JDBC caching when it's safe: a separate, + // uncached Connection on PostgreSQL (outside a transaction), but still the shared Connection on SQL Server. try (Connection conn2 = new SqlSelector(scope, "SELECT RowId, Body FROM comm.Announcements").getConnection()) { - assertEquals(conn, conn2); + if (scope.getSqlDialect().isPostgreSQL()) + { + assertNotEquals(conn, conn2); + assertEquals(TRANSACTION_READ_UNCOMMITTED, conn2.getTransactionIsolation()); + assertFalse(conn2.getAutoCommit()); + } + else + { + assertEquals(conn, conn2); + } } - // Same as the default setting + // Explicitly requesting caching shares the connection, even on PostgreSQL try (Connection conn2 = new SqlSelector(scope, "SELECT RowId, Body FROM comm.Announcements").setJdbcCaching(true).getConnection()) { assertEquals(conn, conn2); @@ -221,6 +231,54 @@ public void testJdbcUncached() throws SQLException } } } + + // A "self-contained" read (getArrayList(), forEach(), getRowCount(), etc., which fully consume and close the + // ResultSet within the call) borrows the thread's shared connection rather than a dedicated one, so nested + // queries reuse it and connection-local state stays visible. On PostgreSQL the outermost borrower puts it into + // no-caching mode and restores it on release; on SQL Server it's simply the shared connection. + Connection borrowed = new SqlSelector(scope, "SELECT RowId, Body FROM comm.Announcements").getConnection(true); + try + { + // A plain thread-connection acquisition returns the very same object (it was borrowed, not dedicated) + try (Connection threadConn = scope.getConnection()) + { + assertEquals(borrowed, threadConn); + } + + // A nested self-contained read reuses the same connection rather than grabbing another one + try (Connection nested = new SqlSelector(scope, "SELECT RowId, Body FROM comm.Announcements").getConnection(true)) + { + assertEquals(borrowed, nested); + } + + if (scope.getSqlDialect().isPostgreSQL()) + { + assertEquals(TRANSACTION_READ_UNCOMMITTED, borrowed.getTransactionIsolation()); + assertFalse(borrowed.getAutoCommit()); + } + } + finally + { + borrowed.close(); + } + + // Once the outermost borrower releases it, the thread connection is restored to normal caching mode + try (Connection restored = scope.getConnection()) + { + assertTrue(restored.getAutoCommit()); + assertEquals(TRANSACTION_READ_COMMITTED, restored.getTransactionIsolation()); + } + + // Inside a transaction, the default must NOT grab a separate Connection, even on PostgreSQL: the caller may be + // relying on reading its own uncommitted writes, so we fall back to the shared, transactional Connection. + try (DbScope.Transaction tx = scope.ensureTransaction()) + { + try (Connection conn2 = new SqlSelector(scope, "SELECT RowId, Body FROM comm.Announcements").getConnection()) + { + assertEquals(scope.getConnection(), conn2); + } + tx.commit(); + } } // Passing in a Connection and calling setJdbcCaching() should throw diff --git a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java index fa53d3cf3c9..749691e16e8 100644 --- a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java @@ -599,7 +599,7 @@ private void initializeUserDefinedTypes(DbScope scope) String maxLength = rs.getString("character_maximum_length"); // VARCHAR with no specific size has null maxLength... but character_octet_length seems okay - scale = Integer.valueOf(null != maxLength ? maxLength : rs.getString("character_octet_length")); + scale = Integer.parseInt(null != maxLength ? maxLength : rs.getString("character_octet_length")); } else { @@ -903,10 +903,10 @@ public int readOutputParameters(DbScope scope, CallableStatement stmt, Map { + boolean alreadyHeld = scope.isThreadConnectionActive(); + Connection conn = scope.getConnection(); + + if (!alreadyHeld && conn instanceof ConnectionWrapper cw && cw.getAutoCommit()) + cw.setRunOnClose(configureToDisableJdbcCaching(cw, scope)); + + return conn; + }; + } else { - // Factory that gets a fresh, read-only connection directly from the pool (not shared with the thread) and - // configures it to not cache ResultSet data in the JDBC driver, making it suitable for streaming very large - // ResultSets. See #39753 and #39888. + // The connection escapes the selector call (a live, streaming ResultSet/Stream is handed back to the caller) + // or the caller explicitly disabled caching, so use a fresh, read-only connection directly from the pool + // (not shared with the thread) whose lifetime the caller controls. See #39753 and #39888. return () -> { ConnectionWrapper conn = scope.getPooledConnection(DbScope.ConnectionType.Pooled, null); Closer closer = configureToDisableJdbcCaching(conn, scope); diff --git a/api/src/org/labkey/api/data/dialect/SqlDialect.java b/api/src/org/labkey/api/data/dialect/SqlDialect.java index 73cdd4cadb4..c34bfd6b0c0 100644 --- a/api/src/org/labkey/api/data/dialect/SqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/SqlDialect.java @@ -365,8 +365,11 @@ public String getSqlCastTypeName(JdbcType type) return null; } - // Return a ConnectionFactory only if the default behavior needs to be overridden - public @Nullable ConnectionFactory getConnectionFactory(boolean useJdbcCaching, DbScope scope, SQLFragment sql) + // Return a ConnectionFactory only if the default behavior needs to be overridden. selfContained indicates that the + // ResultSet will be fully consumed and closed within a single selector call (so the shared, ref-counted thread + // connection can be borrowed and its state restored on release); when false, the connection escapes to the caller + // as a live ResultSet/Stream and must therefore be an unshared connection whose lifetime the caller controls. + public @Nullable ConnectionFactory getConnectionFactory(boolean useJdbcCaching, boolean selfContained, DbScope scope, SQLFragment sql) { return null; }