From d95834d7d2a85a22904a3cc154b1a00e535ab615 Mon Sep 17 00:00:00 2001 From: labkey-jeckels Date: Fri, 14 Aug 2026 18:01:14 -0700 Subject: [PATCH 1/2] Switch to PG-side cancel/terminate of long-running queries --- .../labkey/api/data/ConnectionWrapper.java | 22 ++++++ .../data/dialect/BasePostgreSqlDialect.java | 33 ++++++++ .../labkey/api/data/dialect/SqlDialect.java | 11 +++ .../pipeline/api/PipelineJobServiceImpl.java | 76 +++++++++++-------- 4 files changed, 110 insertions(+), 32 deletions(-) diff --git a/api/src/org/labkey/api/data/ConnectionWrapper.java b/api/src/org/labkey/api/data/ConnectionWrapper.java index 0c15cb5edb1..8161e5e81a3 100644 --- a/api/src/org/labkey/api/data/ConnectionWrapper.java +++ b/api/src/org/labkey/api/data/ConnectionWrapper.java @@ -58,6 +58,7 @@ import java.util.Calendar; import java.util.Collections; import java.util.Date; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -353,6 +354,27 @@ public static Set getSPIDsForThread(Thread t) return result; } + /** + * SPIDs are only meaningful to the scope that handed them out, so group them for callers that need to issue + * commands against the right database. + */ + public static Map> getSPIDsByScopeForThread(Thread t) + { + Map> result = new HashMap<>(); + synchronized(_openConnections) + { + for (ConnectionWrapper c : _openConnections) + { + // A null or negative SPID means the dialect couldn't determine one, so there's nothing to cancel + if (c._allocatingThread == t && c._spid != null && c._spid >= 0) + { + result.computeIfAbsent(c._scope, k -> new HashSet<>()).add(c._spid); + } + } + } + return result; + } + public Integer getSPID() { return _spid; diff --git a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java index e8a6c5e6b7d..ba48ea84495 100644 --- a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java @@ -37,6 +37,7 @@ import org.labkey.api.data.SQLFragment; import org.labkey.api.data.Selector; import org.labkey.api.data.SqlExecutingSelector.ConnectionFactory; +import org.labkey.api.data.SqlExecutor; import org.labkey.api.data.SqlSelector; import org.labkey.api.data.Table; import org.labkey.api.data.TableInfo; @@ -58,6 +59,7 @@ import java.sql.Statement; import java.sql.Types; import java.util.Calendar; +import java.util.Collection; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.Map; @@ -138,6 +140,37 @@ public SQLFragment getDatabaseSizeSql(String databaseName) return new SQLFragment("SELECT pg_database_size(?)", databaseName); } + @Override + public boolean cancelQueries(DbScope scope, Collection spids, boolean terminate) + { + // Postgres delivers these on a side channel, so they land even when the target backend is mid-query. Run them + // on our own connection; the target's belongs to the thread we're interrupting. + String function = terminate ? "pg_terminate_backend" : "pg_cancel_backend"; + + try (Connection conn = scope.getPooledConnection()) + { + SqlExecutor executor = new SqlExecutor(scope, conn); + for (Integer spid : spids) + { + try + { + executor.execute(new SQLFragment("SELECT " + function + "(?)", spid)); + } + catch (RuntimeSQLException e) + { + // The backend may have exited between the SPID lookup and this call + LOG.debug("{}({}) failed", function, spid, e); + } + } + } + catch (SQLException e) + { + throw new RuntimeSQLException(e); + } + + return true; + } + @Override public StatementWrapper getStatementWrapper(ConnectionWrapper conn, Statement stmt, String sql) { diff --git a/api/src/org/labkey/api/data/dialect/SqlDialect.java b/api/src/org/labkey/api/data/dialect/SqlDialect.java index a7e14d12d95..4edceb34be8 100644 --- a/api/src/org/labkey/api/data/dialect/SqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/SqlDialect.java @@ -1508,6 +1508,17 @@ public Integer getSPID(Connection conn) throws SQLException } } + /** + * Stop work in progress on the given sessions, which typically belong to a thread that's blocked in JDBC and can't + * respond to a request to stop. Cancelling makes that thread's call throw so it unwinds and releases its own + * connection; terminating kills the session outright, leaving its pooled connection dead. + * @return false if this dialect can't stop queries out of band + */ + public boolean cancelQueries(DbScope scope, Collection spids, boolean terminate) + { + return false; + } + public boolean updateStatistics(TableInfo table) { SQLFragment sql = getAnalyzeCommandForTable(table.getSelectName()); diff --git a/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java b/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java index 554dd12ccf2..8d1f2814065 100644 --- a/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java +++ b/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java @@ -30,6 +30,7 @@ import org.labkey.api.collections.CaseInsensitiveHashMap; import org.labkey.api.collections.CaseInsensitiveHashSet; import org.labkey.api.collections.CopyOnWriteHashMap; +import org.labkey.api.data.ConnectionWrapper; import org.labkey.api.data.Container; import org.labkey.api.data.DbScope; import org.labkey.api.formSchema.CheckboxField; @@ -94,6 +95,7 @@ import java.io.FileNotFoundException; import java.io.IOException; import java.net.URI; +import java.time.Duration; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -110,6 +112,8 @@ public class PipelineJobServiceImpl implements PipelineJobService public static final String MODULE_PIPELINE_DIR = "pipeline"; public static final Logger LOG = LogManager.getLogger(PipelineJobServiceImpl.class); + /** How long to let a cancelled job wind down its queries before terminating its database sessions outright */ + private static final Duration CANCEL_GRACE_PERIOD = Duration.ofSeconds(2); private static final String PIPELINE_TOOLS_ERROR = "Failed to locate %s. Use the site pipeline tools settings to specify where it can be found. (Currently '%s')"; private static final String INSTALLED_PIPELINE_TOOL_ERROR = "Failed to locate %s. Check tool install location defined in pipelineConfig.xml. (Currently '%s')"; private static final String MODULE_TASKS_DIR = "tasks"; @@ -287,42 +291,50 @@ public void cancelForJob(String jobGuid) killProcessesForThread(jobThread); - if (PipelineSchema.getInstance().getSqlDialect().isPostgreSQL()) - { - // Issue 50131 - // Closing a Connection doesn't terminate all pending statements on SQL Server. Calling close - // and returning to the pool results in another thread getting a Connection that's still busy and - // blocks it. + // Issue 50131: the job thread may be blocked in a query that never returns, so we can't wait for it to notice + // the cancellation. Ask the database to abort its queries instead, which makes its JDBC calls throw so it + // unwinds and releases its own connections on its own thread. Closing them from this thread hands a connection + // that's still in use back to the pool, where another request picks it up mid-transaction. + Map> spidsByScope = ConnectionWrapper.getSPIDsByScopeForThread(jobThread); + + spidsByScope.forEach((scope, spids) -> { + if (!scope.getSqlDialect().cancelQueries(scope, spids, false)) + { + LOG.warn("{} can't cancel queries out of band; job {} may run until its current query returns", scope.getSqlDialect(), jobGuid); + } + }); - // If we need to support query cancellation on SQL Server, we'd have to start tracking the Statements - // the Connection hands out and kill them individually. + // Escalate for anything that ignored the cancel, such as a task that swallows the exception and retries + waitForSPIDsToClear(jobThread, spidsByScope).forEach((scope, spids) -> { + LOG.warn("Job thread {} still holds SPIDs {} after cancel; terminating those sessions", jobThread.getName(), spids); + scope.getSqlDialect().cancelQueries(scope, spids, true); + }); + } - // Piggyback on the job thread so we can shut down open connections on its behalf - try (DbScope.ConnectionSharingCloseable ignored = DbScope.shareConnections(jobThread, Thread.currentThread())) + /** @return the subset of the original SPIDs that the thread still holds after the grace period, grouped by scope */ + private Map> waitForSPIDsToClear(Thread jobThread, Map> original) + { + long deadline = System.currentTimeMillis() + CANCEL_GRACE_PERIOD.toMillis(); + + while (true) + { + Map> current = ConnectionWrapper.getSPIDsByScopeForThread(jobThread); + current.forEach((scope, spids) -> spids.retainAll(original.getOrDefault(scope, Set.of()))); + current.values().removeIf(Set::isEmpty); + + if (current.isEmpty() || System.currentTimeMillis() >= deadline) { - DbScope.DifferentConnectionException lastException; - int retry = 0; - do - { - try - { - DbScope.closeConnectionsForCurrentThreadWithoutReleasingLocks(); - lastException = null; - } - catch (DbScope.DifferentConnectionException e) - { - // The connection we tried to close has already been closed and the thread has already - // started using a different connection. Try again - lastException = e; - retry++; - } - } - while (lastException != null && retry < 3); + return current; + } - if (lastException != null) - { - throw lastException; - } + try + { + Thread.sleep(100); + } + catch (InterruptedException e) + { + Thread.currentThread().interrupt(); + return current; } } } From c77cc0758796f03c535adaa4dc93e3af8bb947ea Mon Sep 17 00:00:00 2001 From: labkey-jeckels Date: Fri, 14 Aug 2026 18:18:57 -0700 Subject: [PATCH 2/2] Claude review --- .../labkey/api/data/ConnectionWrapper.java | 11 ++++---- api/src/org/labkey/api/data/DbScope.java | 7 +++-- .../data/dialect/BasePostgreSqlDialect.java | 4 ++- .../pipeline/api/PipelineJobServiceImpl.java | 27 +++++++++++++------ 4 files changed, 31 insertions(+), 18 deletions(-) diff --git a/api/src/org/labkey/api/data/ConnectionWrapper.java b/api/src/org/labkey/api/data/ConnectionWrapper.java index 8161e5e81a3..89ffd94403a 100644 --- a/api/src/org/labkey/api/data/ConnectionWrapper.java +++ b/api/src/org/labkey/api/data/ConnectionWrapper.java @@ -355,12 +355,13 @@ public static Set getSPIDsForThread(Thread t) } /** - * SPIDs are only meaningful to the scope that handed them out, so group them for callers that need to issue - * commands against the right database. + * Connections whose SPID the dialect could determine, grouped by scope because a SPID is only meaningful to the + * scope that handed it out. Compare these wrappers, not their SPIDs, to tell whether a connection is still in use: + * the pool re-hands the same physical connection back out under its cached SPID. */ - public static Map> getSPIDsByScopeForThread(Thread t) + public static Map> getConnectionsByScopeForThread(Thread t) { - Map> result = new HashMap<>(); + Map> result = new HashMap<>(); synchronized(_openConnections) { for (ConnectionWrapper c : _openConnections) @@ -368,7 +369,7 @@ public static Map> getSPIDsByScopeForThread(Thread t) // A null or negative SPID means the dialect couldn't determine one, so there's nothing to cancel if (c._allocatingThread == t && c._spid != null && c._spid >= 0) { - result.computeIfAbsent(c._scope, k -> new HashSet<>()).add(c._spid); + result.computeIfAbsent(c._scope, k -> new HashSet<>()).add(c); } } } diff --git a/api/src/org/labkey/api/data/DbScope.java b/api/src/org/labkey/api/data/DbScope.java index f020ec50d63..feda3eb4467 100644 --- a/api/src/org/labkey/api/data/DbScope.java +++ b/api/src/org/labkey/api/data/DbScope.java @@ -1222,10 +1222,9 @@ public ConnectionAlreadyReleasedException(String s) } /** - * Special case for when the connection that we expected to close has already been closed and another connection - - * is in use instead. - * Makes it easier to retry in cases that are prone to race conditions, like killing pipeline jobs. + * Special case for when the connection that we expected to close has already been closed and another connection is + * in use instead. Expected rather than a bug because TransactionFilter's read-only request timeout closes a + * thread's connections from a different thread. */ public static class DifferentConnectionException extends IllegalStateException implements SkipMothershipLogging { diff --git a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java index ba48ea84495..6b33dd79739 100644 --- a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java @@ -30,6 +30,7 @@ import org.labkey.api.data.DbSchema; import org.labkey.api.data.DbScope; import org.labkey.api.data.DbScope.LabKeyDataSource; +import org.labkey.api.data.ExceptionFramework; import org.labkey.api.data.JdbcType; import org.labkey.api.data.MetadataSqlSelector; import org.labkey.api.data.PropertyStorageSpec; @@ -149,7 +150,8 @@ public boolean cancelQueries(DbScope scope, Collection spids, boolean t try (Connection conn = scope.getPooledConnection()) { - SqlExecutor executor = new SqlExecutor(scope, conn); + // Spring's translator would hand back a DataAccessException, which the per-SPID catch below can't narrow on + SqlExecutor executor = new SqlExecutor(scope, conn).setExceptionFramework(ExceptionFramework.JDBC); for (Integer spid : spids) { try diff --git a/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java b/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java index 8d1f2814065..ec696c666df 100644 --- a/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java +++ b/pipeline/src/org/labkey/pipeline/api/PipelineJobServiceImpl.java @@ -106,6 +106,7 @@ import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.TimeUnit; import java.util.regex.Matcher; +import java.util.stream.Collectors; public class PipelineJobServiceImpl implements PipelineJobService { @@ -295,31 +296,41 @@ public void cancelForJob(String jobGuid) // the cancellation. Ask the database to abort its queries instead, which makes its JDBC calls throw so it // unwinds and releases its own connections on its own thread. Closing them from this thread hands a connection // that's still in use back to the pool, where another request picks it up mid-transaction. - Map> spidsByScope = ConnectionWrapper.getSPIDsByScopeForThread(jobThread); + Map> connectionsByScope = ConnectionWrapper.getConnectionsByScopeForThread(jobThread); - spidsByScope.forEach((scope, spids) -> { - if (!scope.getSqlDialect().cancelQueries(scope, spids, false)) + connectionsByScope.forEach((scope, connections) -> { + if (!scope.getSqlDialect().cancelQueries(scope, getSPIDs(connections), false)) { LOG.warn("{} can't cancel queries out of band; job {} may run until its current query returns", scope.getSqlDialect(), jobGuid); } }); // Escalate for anything that ignored the cancel, such as a task that swallows the exception and retries - waitForSPIDsToClear(jobThread, spidsByScope).forEach((scope, spids) -> { + waitForConnectionsToClose(jobThread, connectionsByScope).forEach((scope, connections) -> { + Set spids = getSPIDs(connections); LOG.warn("Job thread {} still holds SPIDs {} after cancel; terminating those sessions", jobThread.getName(), spids); scope.getSqlDialect().cancelQueries(scope, spids, true); }); } - /** @return the subset of the original SPIDs that the thread still holds after the grace period, grouped by scope */ - private Map> waitForSPIDsToClear(Thread jobThread, Map> original) + private static Set getSPIDs(Collection connections) + { + return connections.stream().map(ConnectionWrapper::getSPID).collect(Collectors.toSet()); + } + + /** + * Tracks wrappers rather than SPIDs so a thread that caught the cancel and borrowed a fresh connection doesn't look + * like one that never let go; the pool hands the same physical connection, and its SPID, straight back out. + * @return the subset of the original connections the thread still holds after the grace period, grouped by scope + */ + private Map> waitForConnectionsToClose(Thread jobThread, Map> original) { long deadline = System.currentTimeMillis() + CANCEL_GRACE_PERIOD.toMillis(); while (true) { - Map> current = ConnectionWrapper.getSPIDsByScopeForThread(jobThread); - current.forEach((scope, spids) -> spids.retainAll(original.getOrDefault(scope, Set.of()))); + Map> current = ConnectionWrapper.getConnectionsByScopeForThread(jobThread); + current.forEach((scope, connections) -> connections.retainAll(original.getOrDefault(scope, Set.of()))); current.values().removeIf(Set::isEmpty); if (current.isEmpty() || System.currentTimeMillis() >= deadline)