diff --git a/api/src/org/labkey/api/data/ConnectionWrapper.java b/api/src/org/labkey/api/data/ConnectionWrapper.java index 0c15cb5edb1..6e67fdb90ca 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,38 @@ public static Set getSPIDsForThread(Thread t) return result; } + /** + * 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> getConnectionsByScopeForThread(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); + } + } + } + return result; + } + + /** + * @return whether this wrapper is still open, and so whether its SPID still refers to the session it was handed out + * for. Removal happens before the underlying connection returns to the pool, so a true answer can't be stale in the + * dangerous direction. + */ + public boolean isAllocated() + { + return _openConnections.contains(this); + } + public Integer getSPID() { return _spid; 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 e8a6c5e6b7d..1af10dcbedc 100644 --- a/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java +++ b/api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java @@ -17,6 +17,7 @@ package org.labkey.api.data.dialect; import org.apache.commons.lang3.StringUtils; +import org.apache.logging.log4j.Level; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import org.labkey.api.collections.CaseInsensitiveMapWrapper; @@ -30,6 +31,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; @@ -37,6 +39,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 +61,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 +142,53 @@ public SQLFragment getDatabaseSizeSql(String databaseName) return new SQLFragment("SELECT pg_database_size(?)", databaseName); } + @Override + public boolean cancelQueries(DbScope scope, Collection connections, boolean terminate) + { + // Run the cancel on our own connection; the target connection belongs to the thread we're interrupting. + String function = terminate ? "pg_terminate_backend" : "pg_cancel_backend"; + + try (Connection conn = scope.getPooledConnection()) + { + // Spring's translator would hand back a DataAccessException, which the per-connection catch below can't narrow on + SqlExecutor executor = new SqlExecutor(scope, conn).setExceptionFramework(ExceptionFramework.JDBC); + for (ConnectionWrapper connection : connections) + { + Integer spid = connection.getSPID(); + + // Re-check as late as possible: if the thread let go while we were getting the connection above, the pool + // may have handed that physical connection, and this SPID, straight back out to someone else + if (!connection.isAllocated()) + { + LOG.debug("Skipping {}({}); the thread released that connection first", function, spid); + continue; + } + + try + { + // Reports an already-exited backend by returning false, not by throwing + boolean signalled = executor.executeWithResults(new SQLFragment("SELECT " + function + "(?)", spid), (rs, c) -> rs.next() && rs.getBoolean(1)); + + if (!signalled) + { + // A cancel losing the race to the query finishing is routine; a terminate finding nothing means we're out of options + LOG.log(terminate ? Level.WARN : Level.DEBUG, "{}({}) found no such backend", function, spid); + } + } + catch (RuntimeSQLException e) + { + LOG.warn("{}({}) 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..60b33cb9db9 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 sessions behind the given connections, which typically belong to a thread that's + * blocked in JDBC and can't respond to a request to stop. Cancelling throws an exception so the thread 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 connections, 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..44c1fe1c351 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; @@ -104,12 +106,15 @@ 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 { 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 +292,70 @@ public void cancelForJob(String jobGuid) killProcessesForThread(jobThread); - if (PipelineSchema.getInstance().getSqlDialect().isPostgreSQL()) + // Issue 50131: the job thread may be blocked in a slow query, 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. + Map> connectionsByScope = ConnectionWrapper.getConnectionsByScopeForThread(jobThread); + + connectionsByScope.forEach((scope, connections) -> cancelQueries(scope, connections, false, jobGuid)); + + // Escalate for anything that ignored the cancel, such as a task that swallows the exception and retries + waitForConnectionsToClose(jobThread, connectionsByScope).forEach((scope, connections) -> { + LOG.warn("Job thread {} still holds SPIDs {} after cancel; terminating those sessions", jobThread.getName(), getSPIDs(connections)); + cancelQueries(scope, connections, true, jobGuid); + }); + } + + /** Isolated per scope so that one failure, such as waiting out maxWaitMillis for a connection, doesn't skip the remaining scopes or the escalation */ + private static void cancelQueries(DbScope scope, Set connections, boolean terminate, String jobGuid) + { + try { - // 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. + // A dialect that can't do this on the cancel pass can't do it on the terminate pass either, so only warn once + if (!scope.getSqlDialect().cancelQueries(scope, connections, terminate) && !terminate) + { + LOG.warn("{} can't cancel queries out of band; job {} may run until its current query returns", scope.getSqlDialect(), jobGuid); + } + } + catch (RuntimeException e) + { + LOG.error("Failed to {} queries on {} for job {}", terminate ? "terminate" : "cancel", scope.getDisplayName(), jobGuid, e); + } + } - // 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. + private static Set getSPIDs(Collection connections) + { + return connections.stream().map(ConnectionWrapper::getSPID).collect(Collectors.toSet()); + } - // Piggyback on the job thread so we can shut down open connections on its behalf - try (DbScope.ConnectionSharingCloseable ignored = DbScope.shareConnections(jobThread, Thread.currentThread())) + /** + * 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.getConnectionsByScopeForThread(jobThread); + current.forEach((scope, connections) -> connections.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; } } }