diff --git a/api/src/org/labkey/api/data/DbScope.java b/api/src/org/labkey/api/data/DbScope.java index d84f495a64d..4ec508a649c 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,22 @@ 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 yet. This simply reports the existing {@link + * ConnectionHolder} reference count (the same count that already governs {@link ConnectionType#Thread} sharing); it + * is NOT a separate/new counter maintained for this purpose. + *

+ * Callers that borrow the thread connection and temporarily modify its state (e.g., disabling JDBC caching for a + * streaming read) must do so only when this returns false, so that they are the outermost borrower and can safely + * restore the original state — via the connection's runOnClose, which {@link ConnectionType#Thread} fires when the + * last holder releases it (ref count returns to 0). + */ + 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..ed1bae67286 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,69 @@ 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. When a caller has explicitly chosen a caching + * behavior via {@link #setJdbcCaching(boolean)} or supplied a Connection at construction time, that choice is + * honored. Otherwise, JDBC caching is disabled by default: we ask the dialect for a ConnectionFactory so the driver + * won't buffer the entire ResultSet in memory. + *

+ * The {@code selfContained} flag reflects how the ResultSet is consumed. When true (e.g. {@link #getArrayList}, + * {@link #forEach}, {@link #getRowCount}), the ResultSet is fully consumed and closed within this selector call, so + * the dialect may borrow the thread's shared, ref-counted connection — nested queries then reuse it (avoiding + * connection-pool exhaustion) and connection-local state (temp tables, search_path) stays visible — because its + * state can be restored before control returns to the caller. When false (e.g. {@code getResultSet(false)}, + * {@link #uncachedStream}), a live ResultSet/Stream is handed back to the caller, so the dialect uses a dedicated, + * unshared connection whose lifetime the caller controls. + *

+ * The dialect returns null (meaning "use the shared Connection with the driver's default caching") when a + * transaction is active, the dialect is not PostgreSQL, or the statement is not a SELECT, so this default is safe by + * construction. Resolving lazily here (rather than at construction) ensures the transaction check reflects the state + * at execution time. + */ + 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.

+ * + *

Note that when neither this method nor an explicit Connection is supplied, JDBC caching is disabled by default + * whenever it's safe to do so (PostgreSQL, no active transaction, SELECT statement) — see + * {@link #getEffectiveConnectionFactory()}. Callers that require the driver's default caching behavior (e.g., to + * share the thread's Connection) must therefore opt in explicitly by calling this method 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 +167,51 @@ 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}, {@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 per unique call stack to avoid flooding the log. A benign race (two threads + // logging the same stack at once) is acceptable for a throttle. + 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 +495,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..ad9fba5b8f8 100644 --- a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java @@ -1039,7 +1039,7 @@ public String getExtraInfo(SQLException e) } @Override - public ConnectionFactory getConnectionFactory(boolean useJdbcCaching, DbScope scope, SQLFragment sql) + public ConnectionFactory getConnectionFactory(boolean useJdbcCaching, boolean selfContained, DbScope scope, SQLFragment sql) { // Fiddle with the Connection settings only if asked to turn off JDBC caching, we're not inside a transaction, // and it's a read-only statement (a SELECT), so we won't mess up any state the caller is relying on. @@ -1047,11 +1047,31 @@ public ConnectionFactory getConnectionFactory(boolean useJdbcCaching, DbScope sc { return null; } + else if (selfContained) + { + // The ResultSet will be fully consumed and the connection released within a single selector call, so borrow + // the thread's shared, ref-counted connection (the same one scope.getConnection() hands out) instead of + // grabbing a separate one. Nested queries then reuse it (avoiding connection-pool exhaustion), and + // connection-local state (temp tables, search_path, session settings) stays visible. We piggyback on the + // existing thread-connection reference count rather than tracking our own: only the outermost borrower (when + // isThreadConnectionActive() is false) flips the connection into no-JDBC-caching mode and registers the + // restore via runOnClose, which ConnectionType.Thread fires when the last holder releases it (ref count + // returns to 0). Nested borrows find it already configured and reuse it as-is. + return () -> { + boolean alreadyHeld = scope.isThreadConnectionActive(); + Connection conn = scope.getConnection(); + + if (!alreadyHeld && conn instanceof ConnectionWrapper cw) + 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 946edfa95b2..c2fff7572bd 100644 --- a/api/src/org/labkey/api/data/dialect/SqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/SqlDialect.java @@ -361,8 +361,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; } @@ -863,25 +866,25 @@ public SQLFragment getNumericCast(SQLFragment expression) * @param arguments Arguments passed from the LK SQL * @return the dialect equivalent SQLFragrment */ - public SQLFragment getGreatestAndLeastSQL(String method, SQLFragment... arguments) - { - throw new UnsupportedOperationException(getClass().getSimpleName() + " does not implement"); - } - - public boolean supportsIsNumeric() - { - return false; - } - - public SQLFragment isNumericExpr(SQLFragment expression) - { - throw new UnsupportedOperationException(getClass().getSimpleName() + " does not implement"); - } - - public void handleCreateDatabaseException(SQLException e) throws ServletException - { - throw(new ServletException("Can't create database", e)); - } + public SQLFragment getGreatestAndLeastSQL(String method, SQLFragment... arguments) + { + throw new UnsupportedOperationException(getClass().getSimpleName() + " does not implement"); + } + + public boolean supportsIsNumeric() + { + return false; + } + + public SQLFragment isNumericExpr(SQLFragment expression) + { + throw new UnsupportedOperationException(getClass().getSimpleName() + " does not implement"); + } + + public void handleCreateDatabaseException(SQLException e) throws ServletException + { + throw(new ServletException("Can't create database", e)); + } /** * Wrap one or more INSERT statements to allow explicit specification