[bugfix] Make sure that the REST Server correctly caches and frees query results in its Session

Author: adamretter

exist-core/src/main/java/org/exist/http/RESTServer.java

@@ -1387,7 +1387,9 @@ protected void search(final DBBroker broker, final Txn transaction, final String
             @Nullable final Item contextItem = extractContextItem(contextItemParam);
             final Sequence contextSequence = contextItem != null ? new ValueSequence(contextItem) : null;
 
-            try (final XQueryUtil.QueryResult queryResult = XQueryUtil.query(broker, source, true, contextSequence, outputProperties, setupXqueryContextPreCompilation, setupXqueryContextPreExecution, setupXqueryContextPostExecution)) {
+            @Nullable XQueryUtil.QueryResult queryResult = null;
+            try {
+                queryResult = XQueryUtil.query(broker, source, true, contextSequence, outputProperties, setupXqueryContextPreCompilation, setupXqueryContextPreExecution, setupXqueryContextPostExecution);
 
                 // special header to indicate that the query is not returned from cache
                 response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
@@ -1405,6 +1407,11 @@ protected void search(final DBBroker broker, final Txn transaction, final String
                 }
 
                 writeResults(response, broker, transaction, queryResult, howmany, start, typed, outputProperties, wrap);
+            } finally {
+                if (!cache && queryResult != null) {
+                    // NOTE(AR) we can only close the query result if we are not caching it for reuse in the sessionManager, otherwise it has to be closed when it is later removed from the sessionManager
+                    queryResult.close();
+                }
             }
 
         } catch (final IOException e) {

exist-core/src/main/java/org/exist/http/SessionManager.java

@@ -47,11 +47,13 @@
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.RemovalListener;
 import net.jcip.annotations.ThreadSafe;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.exist.xquery.XQueryUtil;
 
+import javax.annotation.Nullable;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -64,6 +66,15 @@ public class SessionManager {
     private static final Logger LOG = LogManager.getLogger(SessionManager.class);
     private static final long TIMEOUT = 120_000;  // ms (e.g. 2 minutes)
 
+    private static final RemovalListener<Integer, QueryAndResult> REMOVAL_LISTENER = (sessionId, queryAndResult, removalCause) -> {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Removing cached query result for session: {}", sessionId);
+        }
+
+        // NOTE(AR) make sure to release any resources still held by the query result and potentially send it back to the query pool for reuse
+        queryAndResult.result.close();
+    };
+
     private final AtomicInteger sessionIdCounter = new AtomicInteger();
     private final Cache<Integer, QueryAndResult> cache;
 
@@ -78,12 +89,10 @@ private QueryAndResult(final String query, final XQueryUtil.QueryResult result)
     }
 
     public SessionManager() {
-        final Caffeine<Object, Object> cacheBuilder = Caffeine.newBuilder()
-                .expireAfterAccess(TIMEOUT, TimeUnit.MILLISECONDS);
-        if(LOG.isDebugEnabled()) {
-            cacheBuilder.removalListener((key, value, cause) -> LOG.debug("Removing cached query result for session: {}", key));
-        }
-        cache = cacheBuilder.build();
+        this.cache = Caffeine.newBuilder()
+                .expireAfterAccess(TIMEOUT, TimeUnit.MILLISECONDS)
+                .removalListener(REMOVAL_LISTENER)
+                .build();
     }
 
     public int add(final String query, final XQueryUtil.QueryResult result) {
@@ -92,7 +101,7 @@ public int add(final String query, final XQueryUtil.QueryResult result) {
         return sessionId;
     }
 
-    public XQueryUtil.QueryResult get(final String query, final int sessionId) {
+    public @Nullable XQueryUtil.QueryResult get(final String query, final int sessionId) {
         if (sessionId < 0 || sessionId >= sessionIdCounter.get()) {
             return null; // out of scope
         }