Switch to PG-side cancel/terminate of long-running queries

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,28 @@ public static Set<Integer> 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<DbScope, Set<ConnectionWrapper>> getConnectionsByScopeForThread(Thread t)
+    {
+        Map<DbScope, Set<ConnectionWrapper>> 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;
+    }
+
     public Integer getSPID()
     {
         return _spid;

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
     {

api/src/org/labkey/api/data/dialect/BasePostgreSqlDialect.java

@@ -30,13 +30,15 @@
 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;
 import org.labkey.api.data.RuntimeSQLException;
 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 +60,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 +141,38 @@ public SQLFragment getDatabaseSizeSql(String databaseName)
         return new SQLFragment("SELECT pg_database_size(?)", databaseName);
     }
 
+    @Override
+    public boolean cancelQueries(DbScope scope, Collection<Integer> 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())
+        {
+            // 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
+                {
+                    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)
     {

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<Integer> spids, boolean terminate)
+    {
+        return false;
+    }
+
     public boolean updateStatistics(TableInfo table)
     {
         SQLFragment sql = getAnalyzeCommandForTable(table.getSelectName());

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,60 @@ 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<DbScope, Set<ConnectionWrapper>> connectionsByScope = ConnectionWrapper.getConnectionsByScopeForThread(jobThread);
+
+        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
+        waitForConnectionsToClose(jobThread, connectionsByScope).forEach((scope, connections) -> {
+            Set<Integer> spids = getSPIDs(connections);
+            LOG.warn("Job thread {} still holds SPIDs {} after cancel; terminating those sessions", jobThread.getName(), spids);
+            scope.getSqlDialect().cancelQueries(scope, spids, true);
+        });
+    }
+
+    private static Set<Integer> getSPIDs(Collection<ConnectionWrapper> connections)
+    {
+        return connections.stream().map(ConnectionWrapper::getSPID).collect(Collectors.toSet());
+    }
 
-            // 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.
+    /**
+     * 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<DbScope, Set<ConnectionWrapper>> waitForConnectionsToClose(Thread jobThread, Map<DbScope, Set<ConnectionWrapper>> original)
+    {
+        long deadline = System.currentTimeMillis() + CANCEL_GRACE_PERIOD.toMillis();
+
+        while (true)
+        {
+            Map<DbScope, Set<ConnectionWrapper>> current = ConnectionWrapper.getConnectionsByScopeForThread(jobThread);
+            current.forEach((scope, connections) -> connections.retainAll(original.getOrDefault(scope, Set.of())));
+            current.values().removeIf(Set::isEmpty);
 
-            // Piggyback on the job thread so we can shut down open connections on its behalf
-            try (DbScope.ConnectionSharingCloseable ignored = DbScope.shareConnections(jobThread, Thread.currentThread()))
+            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;
             }
         }
     }