Reliably stop abandoned read-only requests and their queries

api/src/org/labkey/api/ApiModule.java

@@ -54,6 +54,7 @@
 import org.labkey.api.data.AbstractForeignKey;
 import org.labkey.api.data.AccumulatingCommitTask;
 import org.labkey.api.data.Aggregate;
+import org.labkey.api.data.AsyncQueryRequest;
 import org.labkey.api.data.AtomicDatabaseInteger;
 import org.labkey.api.data.BindingTestCase;
 import org.labkey.api.data.BooleanFormat;
@@ -78,6 +79,7 @@
 import org.labkey.api.data.NameGenerator;
 import org.labkey.api.data.PropertyManager;
 import org.labkey.api.data.RecordFactory;
+import org.labkey.api.data.RequestCancellation;
 import org.labkey.api.data.RepeatedCommitTaskMonitor;
 import org.labkey.api.data.ResultSetSelectorTestCase;
 import org.labkey.api.data.RowTrackingResultSetWrapper;
@@ -92,6 +94,7 @@
 import org.labkey.api.data.Table;
 import org.labkey.api.data.TableSelectorTestCase;
 import org.labkey.api.data.TempTableInClauseGenerator;
+import org.labkey.api.data.TransactionFilter;
 import org.labkey.api.data.WorkbookContainerType;
 import org.labkey.api.data.dialect.JdbcMetaDataTest;
 import org.labkey.api.data.dialect.ParameterSubstitutionTest;
@@ -528,6 +531,7 @@ public void registerServlets(ServletContext servletCtx)
             ApiKeyManager.TestCase.class,
             ApiQueryResponse.TestCase.class,
             AppPropsTestCase.class,
+            AsyncQueryRequest.TestCase.class,
             AtomicDatabaseInteger.TestCase.class,
             BindingTestCase.class,
             BlockingCache.BlockingCacheTest.class,
@@ -573,6 +577,7 @@ public void registerServlets(ServletContext servletCtx)
             QueryDataIteratorBuilder.TestCase.class,
             QueryView.TestCase.class,
             RecordFactory.TestCase.class,
+            RequestCancellation.TestCase.class,
             RepeatedCommitTaskMonitor.IntegrationTestCase.class,
             ResultSetDataIterator.TestCase.class,
             ResultSetSelectorTestCase.class,
@@ -591,6 +596,7 @@ public void registerServlets(ServletContext servletCtx)
             Table.TestCase.class,
             TableSelectorTestCase.CoreTableSelectorTest.class,
             TempTableInClauseGenerator.TestCase.class,
+            TransactionFilter.TestCase.class,
             URLHelper.TestCase.class,
             UserKeySimpleCache.TestCase.class,
             UserManager.TestCase.class,

api/src/org/labkey/api/action/SpringActionController.java

@@ -584,6 +584,12 @@ public ModelAndView handleRequest(HttpServletRequest request, @NotNull HttpServl
                 QueryService.get().setEnvironment(QueryService.Environment.ACTION, actionAnnotation.value());
             }
 
+            if (ExportAction.class.isAssignableFrom(actionClass) ||
+                (actionClass.isAnnotationPresent(Action.class) && actionClass.getAnnotation(Action.class).value().equals(ActionType.Export.class)))
+            {
+                request.setAttribute(TransactionFilter.EXPORT_ATTRIBUTE_NAME, true);
+            }
+
             // After the permission check so unauthorized requests can't consume @ConcurrencyLimit permits, and held
             // through rendering because the action's memory is live for that whole time.
             try (ConcurrencyLimiter.Permit ignored = ConcurrencyLimiter.acquire(actionClass, context))

api/src/org/labkey/api/data/AsyncQueryRequest.java

@@ -28,6 +28,8 @@
 import org.apache.logging.log4j.ThreadContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.junit.Assert;
+import org.junit.Test;
 import org.labkey.api.miniprofiler.MiniProfiler;
 import org.labkey.api.miniprofiler.RequestInfo;
 import org.labkey.api.query.QueryService;
@@ -36,6 +38,7 @@
 import org.labkey.api.util.UnexpectedException;
 import org.labkey.api.util.logging.LogHelper;
 import org.labkey.api.view.MockHttpResponseWithRealPassthrough;
+import org.springframework.mock.web.MockHttpServletResponse;
 
 import jakarta.servlet.http.HttpServletResponse;
 import jakarta.servlet.http.HttpServletResponseWrapper;
@@ -45,10 +48,13 @@
 import java.sql.Statement;
 import java.util.Map;
 import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class AsyncQueryRequest<T>
 {
     private static final Logger _log = LogHelper.getLogger(AsyncQueryRequest.class, "Runs DB queries on a background thread");
+    private static final long WORKER_EXIT_TIMEOUT_MS = 10_000;
 
     private static class CancelledException extends RuntimeException
     {
@@ -63,6 +69,7 @@ private static class CancelledException extends RuntimeException
     private final @NotNull Map<String, String> _spanTags;
 
     boolean _cancelled;
+    boolean _finished;
     @Nullable Statement _statement;
     @Nullable Span _span;
 
@@ -192,7 +199,8 @@ synchronized public T waitForResult(final Callable<T> callable) throws SQLExcept
                     }
                     catch (InterruptedException ie)
                     {
-                        cancel();
+                        cancel(RequestCancellation.getSource(Thread.currentThread()));
+                        awaitWorkerAfterInterrupt(thread);
                         throw new AbortedRequestException("Interrupted before DB query completed");
                     }
                 }
@@ -241,6 +249,32 @@ synchronized public T waitForResult(final Callable<T> callable) throws SQLExcept
         }
     }
 
+    /**
+     * Keeps sharing connections until the worker releases the one it borrowed. Otherwise its release can't find the
+     * request thread's holder, which then stays referenced on an aborted connection that the thread's next request reuses.
+     */
+    synchronized private void awaitWorkerAfterInterrupt(Thread worker)
+    {
+        long deadline = System.currentTimeMillis() + WORKER_EXIT_TIMEOUT_MS;
+        while (!_finished)
+        {
+            long remaining = deadline - System.currentTimeMillis();
+            if (remaining <= 0)
+            {
+                _log.warn("Async query thread '{}' still running {} ms after its request was interrupted", worker.getName(), WORKER_EXIT_TIMEOUT_MS);
+                return;
+            }
+            try
+            {
+                wait(remaining);
+            }
+            catch (InterruptedException ignored)
+            {
+                // Already cancelling
+            }
+        }
+    }
+
     synchronized public void setResult(T result)
     {
         if (_cancelled && result instanceof AutoCloseable ac)
@@ -255,12 +289,14 @@ synchronized public void setResult(T result)
             }
         }
         _result = result;
+        _finished = true;
         notify();
     }
 
     synchronized public void setException(Throwable exception)
     {
         _exception = exception;
+        _finished = true;
         notify();
     }
 
@@ -269,11 +305,16 @@ synchronized private boolean isCancelled()
         return _cancelled;
     }
 
-    synchronized private void cancel()
+    synchronized private void cancel(@Nullable RequestCancellation.Source source)
     {
         _cancelled = true;
         if (_span != null)
+        {
             _span.setTag("labkey.async_query.cancelled", true);
+            if (source != null)
+                _span.setTag("labkey.cancel.source", source.toString());
+        }
+        RequestCancellation.statementCancelled(source, _statement != null);
         if (_statement != null)
         {
             try
@@ -300,7 +341,8 @@ private IOException checkCancelled()
         catch (IOException ioe)
         {
             _log.debug("ping client: failure: {}", ioe.getMessage());
-            cancel();
+            RequestCancellation.requestCancelled(Thread.currentThread(), TransactionFilter.getRequestSummary(Thread.currentThread()), RequestCancellation.Source.CLIENT_DISCONNECT);
+            cancel(RequestCancellation.Source.CLIENT_DISCONNECT);
             return ioe;
         }
         return null;
@@ -311,4 +353,68 @@ public StackTraceElement[] getCreationStackTrace()
     {
         return _creationStackTrace;
     }
+
+    public static class TestCase extends Assert
+    {
+        /** An interrupt must leave the request thread's connection holder released and its next query working */
+        @Test
+        public void testInterruptReleasesConnection() throws Exception
+        {
+            DbScope scope = DbScope.getLabKeyScope();
+
+            for (int i = 0; i < 5; i++)
+            {
+                AtomicReference<Throwable> failure = new AtomicReference<>();
+                AtomicBoolean holderReleased = new AtomicBoolean();
+                AtomicReference<Throwable> nextQueryFailure = new AtomicReference<>();
+
+                Thread requestThread = new Thread(() -> {
+                    try
+                    {
+                        AsyncQueryRequest<Object> request = new AsyncQueryRequest<>(new MockHttpServletResponse(), "testInterruptReleasesConnection", Map.of());
+                        SqlSelector selector = new SqlSelector(scope, "SELECT pg_sleep(30)");
+                        selector.setAsyncRequest(request);
+                        request.waitForResult(() -> selector.getObject(Object.class));
+                    }
+                    catch (Throwable t)
+                    {
+                        failure.set(t);
+                    }
+
+                    try
+                    {
+                        // Checked before finishedWithThread(), whose cleanup would hide a leak
+                        holderReleased.set(!scope.isThreadConnectionActive());
+                        new SqlSelector(scope, "SELECT 1").getObject(Integer.class);
+                    }
+                    catch (Throwable t)
+                    {
+                        nextQueryFailure.set(t);
+                    }
+                    finally
+                    {
+                        DbScope.finishedWithThread();
+                    }
+                }, "AsyncQueryRequest interrupt test " + i);
+
+                long start = System.currentTimeMillis();
+                requestThread.start();
+                while (ConnectionWrapper.getConnectionsByScopeForEffectiveThread(requestThread).isEmpty() && System.currentTimeMillis() - start < 10_000)
+                {
+                    //noinspection BusyWait
+                    Thread.sleep(50);
+                }
+                // Let the statement reach the server
+                Thread.sleep(300);
+
+                requestThread.interrupt();
+                requestThread.join(30_000);
+
+                assertFalse("Request thread didn't exit", requestThread.isAlive());
+                assertTrue("Expected an aborted request, got " + failure.get(), failure.get() instanceof AbortedRequestException);
+                assertTrue("Request thread's connection holder is still referenced", holderReleased.get());
+                assertNull("Next query on the request thread failed", nextQueryFailure.get());
+            }
+        }
+    }
 }

api/src/org/labkey/api/data/ConnectionWrapper.java

@@ -68,6 +68,7 @@
 import java.util.TreeSet;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Predicate;
 
 /**
  * Wrapper over another JDBC Connection object that provides logging for errors and some tracking of usage.
@@ -153,6 +154,9 @@ private void realClose()
 
     private volatile @NotNull Closer _runOnClose = () -> {};
 
+    // Set for ConnectionType.Thread; release must go back here even if the closing thread no longer shares the lender's connections
+    private volatile @Nullable DbScope.ConnectionHolder _threadConnectionHolder;
+
     /** Remember the first SQLException that happened on this connection, as it may have left the connection in a bad state */
     private SQLException _originalSqlException;
 
@@ -167,13 +171,12 @@ private void realClose()
     static
     {
         // Issue 51483: DB query can be left running after shutting down server
-        ContextListener.addShutdownListener(ShutdownListener.of("DB connection closer", null, () -> {
-
-            // Only kill connections that have been open for more than two seconds to avoid
-            // interrupting work related to the shutdown
-            Calendar cutoffCal = Calendar.getInstance();
-            cutoffCal.add(Calendar.SECOND, -2);
-            Date cutoffDate = new Date(cutoffCal.getTimeInMillis());
+        ContextListener.addShutdownListener(ShutdownListener.of("DB connection closer", () -> {
+            // Closing doesn't stop a busy backend, which notices its client is gone only once it sends results
+            Date cutoffDate = getShutdownCutoff();
+            RequestCancellation.cancelQueries(getConnectionsByScope(c -> c._allocationTime.before(cutoffDate)), RequestCancellation.Source.SHUTDOWN, true);
+        }, () -> {
+            Date cutoffDate = getShutdownCutoff();
 
             // Build up the list and release the lock on _openConnections
             List<ConnectionWrapper> toClose = new ArrayList<>();
@@ -201,6 +204,14 @@ private void realClose()
         }));
     }
 
+    // Only kill connections that have been open for more than two seconds to avoid interrupting work related to the shutdown
+    private static Date getShutdownCutoff()
+    {
+        Calendar cutoffCal = Calendar.getInstance();
+        cutoffCal.add(Calendar.SECOND, -2);
+        return new Date(cutoffCal.getTimeInMillis());
+    }
+
     private static boolean initializeExplicitLogger()
     {
         var loggerContext = (LoggerContext) LogManager.getContext(false);
@@ -370,14 +381,25 @@ public static Set<Integer> getSPIDsForThread(Thread t)
      * the pool re-hands the same physical connection back out under its cached SPID.
      */
     public static Map<DbScope, Set<ConnectionWrapper>> getConnectionsByScopeForThread(Thread t)
+    {
+        return getConnectionsByScope(c -> c._allocatingThread == t);
+    }
+
+    /** Like {@link #getConnectionsByScopeForThread}, but also includes connections allocated by threads currently sharing t's connections, such as AsyncQueryRequest workers */
+    public static Map<DbScope, Set<ConnectionWrapper>> getConnectionsByScopeForEffectiveThread(Thread t)
+    {
+        return getConnectionsByScope(c -> c._allocatingThread == t || DbScope.getEffectiveThread(c._allocatingThread) == t);
+    }
+
+    private static Map<DbScope, Set<ConnectionWrapper>> getConnectionsByScope(Predicate<ConnectionWrapper> filter)
     {
         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)
+                if (c._spid != null && c._spid >= 0 && filter.test(c))
                 {
                     result.computeIfAbsent(c._scope, k -> new HashSet<>()).add(c);
                 }
@@ -386,6 +408,22 @@ public static Map<DbScope, Set<ConnectionWrapper>> getConnectionsByScopeForThrea
         return result;
     }
 
+    /** The driver's own connection behind the pool's wrapper, for driver-specific operations like an out-of-band cancel */
+    public Connection getPhysicalConnection()
+    {
+        return DbScope.getDelegate(_connection);
+    }
+
+    @Nullable DbScope.ConnectionHolder getThreadConnectionHolder()
+    {
+        return _threadConnectionHolder;
+    }
+
+    void setThreadConnectionHolder(@Nullable DbScope.ConnectionHolder holder)
+    {
+        _threadConnectionHolder = holder;
+    }
+
     /**
      * @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