[refactor] Split XQueryUtil into compile and execute phases

Author: adamretter

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

@@ -1389,7 +1389,7 @@ protected void search(final DBBroker broker, final Txn transaction, final String
             try (final XQueryUtil.QueryResult 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.QueryResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
+                response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
 
                 if (LOG.isDebugEnabled()) {
                     LOG.debug("Found {} in {}ms.", queryResult.result.getItemCount(), queryResult.executionTime);
@@ -1604,7 +1604,7 @@ private void executeXQuery(final DBBroker broker, final Txn transaction, final D
         try (final XQueryUtil.QueryResult queryResult = XQueryUtil.query(broker, source, true, null, outputProperties, setupXqueryContextPreCompilation, setupXqueryContextPreExecution, setupXqueryContextPostExecution)) {
 
             // Special header to indicate whether the compiled query is returned from the cache
-            response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.QueryResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
+            response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
 
             final boolean wrap = "yes".equals(outputProperties.getProperty("_wrap"));
             writeResults(response, broker, transaction, queryResult, -1, 1, false, outputProperties, wrap);
@@ -1663,7 +1663,7 @@ private void executeXProc(final DBBroker broker, final Txn transaction, final Do
         try (final XQueryUtil.QueryResult queryResult = XQueryUtil.query(broker, source, true, null, null, setupXqueryContextPreCompilation, setupXqueryContextPreExecution, setupXqueryContextPostExecution)) {
 
             // special header to indicate that the query is not returned from cache
-            response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.QueryResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
+            response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
 
             writeResults(response, broker, transaction, queryResult, -1, 1, false, outputProperties, false);
         } catch (final IOException e) {

exist-core/src/main/java/org/exist/http/servlets/RedirectorServlet.java

@@ -302,7 +302,7 @@ private XQueryUtil.QueryResult executeQuery(final Source source, final RequestWr
             final XQueryUtil.QueryResult queryResult = XQueryUtil.query(broker, source, true, null, outputProperties, setupXqueryContextPreCompilation, null, null);
 
             // special header to indicate that the query is not returned from cache
-            response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.QueryResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
+            response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
 
             return queryResult;
         }

exist-core/src/main/java/org/exist/http/servlets/XQueryServlet.java

@@ -508,7 +508,7 @@ protected void process(HttpServletRequest request, HttpServletResponse response)
             try (final XQueryUtil.QueryResult queryResult = XQueryUtil.query(broker, source, true, null, outputProperties, null, setupXqueryContextPreExecution, setupXqueryContextPostExecution)) {
 
                 // special header to indicate that the query is not returned from cache
-                response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.QueryResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
+                response.setHeader(XQUERY_CACHED_RESPONSE_HEADER, queryResult.compilationTime == XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY ? "true" : "false");
 
                 final String mediaType = outputProperties.getProperty(OutputKeys.MEDIA_TYPE);
                 if (mediaType != null) {

exist-core/src/main/java/org/exist/xquery/XQueryUtil.java

@@ -22,7 +22,6 @@
 
 import com.evolvedbinary.j8fu.function.BiConsumerE;
 import com.evolvedbinary.j8fu.function.ConsumerE;
-import com.evolvedbinary.j8fu.function.Function2E;
 import org.exist.security.PermissionDeniedException;
 import org.exist.source.Source;
 import org.exist.storage.BrokerPool;
@@ -34,10 +33,17 @@
 import java.io.IOException;
 import java.util.Properties;
 
+/**
+ * Utility to make compiling and executing XQuery simpler
+ * and safer by ensuring cleanup when the returned object is closed.
+ *
+ * @author <a href="mailto:adam@evolvedbinary.com">Adam Retter</a>
+ */
 public class XQueryUtil {
 
     /**
      * Execute an XPath or XQuery.
+     * Performs both the compile and execute steps.
      *
      * @param broker the database broker.
      * @param source the source of the query.
@@ -60,6 +66,7 @@ public static XQueryUtil.QueryResult query(final DBBroker broker, final Source s
 
     /**
      * Execute an XPath or XQuery.
+     * Performs both the compile and execute steps.
      *
      * @param broker the database broker.
      * @param source the source of the query.
@@ -78,6 +85,26 @@ public static XQueryUtil.QueryResult query(final DBBroker broker, final Source s
      * @throws XPathException if the query raises an error.
      */
     public static XQueryUtil.QueryResult query(final DBBroker broker, final Source source, final boolean isXPointer, final boolean cacheQuery, @Nullable final Sequence contextSequence, @Nullable final Properties outputProperties, @Nullable final ConsumerE<XQueryContext, XPathException> preCompilationContext, @Nullable final ConsumerE<XQueryContext, XPathException> preExecutionContext, @Nullable final BiConsumerE<XQueryContext, QueryResult, XPathException> postExecutionContext) throws XPathException, PermissionDeniedException, IOException {
+        final CompilationResult compilationResult = compile(broker, source, isXPointer, cacheQuery, preCompilationContext);
+        return execute(broker, compilationResult, contextSequence, outputProperties, preExecutionContext, postExecutionContext);
+    }
+
+    /**
+     * Compile an XPath or XQuery.
+     *
+     * @param broker the database broker.
+     * @param source the source of the query.
+     * @param isXPointer true if the query is an XPointer, false otherwise.
+     * @param cacheQuery true if the compiled query should be cached, false otherwise.
+     * @param preCompilationContext access to the XQuery Context pre-compilation, or null if access is not required.
+     *
+     * @return the result of the compilation and the compilation time.
+     *
+     * @throws IOException if the source of the query cannot be read.
+     * @throws PermissionDeniedException if the calling user does not have permission to execute the query.
+     * @throws XPathException if the query raises an error.
+     */
+    public static CompilationResult compile(final DBBroker broker, final Source source, final boolean isXPointer, final boolean cacheQuery, @Nullable final ConsumerE<XQueryContext, XPathException> preCompilationContext) throws XPathException, PermissionDeniedException, IOException {
         final BrokerPool brokerPool = broker.getBrokerPool();
         final XQuery xquery = brokerPool.getXQueryService();
 
@@ -104,7 +131,7 @@ public static XQueryUtil.QueryResult query(final DBBroker broker, final Source s
                 compiledXquery.getContext().updateContext(xqueryContext);
                 xqueryContext.getWatchDog().reset();
 
-                compilationTime = XQueryUtil.QueryResult.RETRIEVED_CACHED_COMPILED_QUERY;
+                compilationTime = XQueryUtil.CompilationResult.RETRIEVED_CACHED_COMPILED_QUERY;
 
             } else {
                 xqueryContext = new XQueryContext(brokerPool);
@@ -118,146 +145,162 @@ public static XQueryUtil.QueryResult query(final DBBroker broker, final Source s
                 compilationTime = System.currentTimeMillis() - compilationStart;
             }
 
+            return new CompilationResult(source, xqueryPool, compiledXquery, xqueryContext, compilationTime);
+
+        } catch (final XPathException | PermissionDeniedException | IOException e) {
+            // Make sure to clean-up in case of an exception!
+            if (xqueryContext != null) {
+                xqueryContext.runCleanupTasks();
+            }
+
+            if (xqueryPool != null && compiledXquery != null) {
+                xqueryPool.returnCompiledXQuery(source, compiledXquery);
+            }
+
+            throw e;
+        }
+    }
+
+    /**
+     * Execute a compiled XPath or XQuery.
+     *
+     * @param broker the database broker.
+     * @param compilationResult the compiled query.
+     * @param contextSequence the context sequence to use when executing the query, or null if there is no context sequence.
+     * @param outputProperties any output properties, or null.
+     * @param preExecutionContext access to the XQuery Context pre-execution, or null if access is not required.
+     * @param postExecutionContext access to the XQuery Context post-execution, or null if access is not required.
+     *
+     * @return the result of the query and compilation and execution times.
+     *
+     * @throws PermissionDeniedException if the calling user does not have permission to execute the query.
+     * @throws XPathException if the query raises an error.
+     */
+    public static QueryResult execute(final DBBroker broker, final CompilationResult compilationResult, @Nullable final Sequence contextSequence, @Nullable final Properties outputProperties, @Nullable final ConsumerE<XQueryContext, XPathException> preExecutionContext, @Nullable final BiConsumerE<XQueryContext, QueryResult, XPathException> postExecutionContext) throws XPathException, PermissionDeniedException {
+        final BrokerPool brokerPool = broker.getBrokerPool();
+        final XQuery xquery = brokerPool.getXQueryService();
+
+        try {
             if (preExecutionContext != null) {
-                preExecutionContext.accept(xqueryContext);
+                preExecutionContext.accept(compilationResult.xqueryContext);
             }
 
             final long executionStart = System.currentTimeMillis();
-            final Sequence result = xquery.execute(broker, compiledXquery, null, contextSequence, outputProperties, true);
+            final Sequence result = xquery.execute(broker, compilationResult.compiledXquery, null, contextSequence, outputProperties, true);
             final long executionTime = System.currentTimeMillis() - executionStart;
 
-            final QueryResult queryResult = new XQueryUtil.QueryResult(source, xqueryPool, compiledXquery, xqueryContext, compilationTime, executionTime, result);
+            final QueryResult queryResult = new XQueryUtil.QueryResult(compilationResult, executionTime, result);
 
             if (postExecutionContext != null) {
-                postExecutionContext.accept(xqueryContext, queryResult);
+                postExecutionContext.accept(compilationResult.xqueryContext, queryResult);
             }
 
             return queryResult;
 
-        } catch (final XPathException | PermissionDeniedException | IOException e) {
+        } catch (final XPathException | PermissionDeniedException e) {
             // Make sure to clean-up in case of an exception!
-            if (xqueryContext != null) {
-                xqueryContext.runCleanupTasks();
-            }
+            compilationResult.xqueryContext.runCleanupTasks();
 
-            if (xqueryPool != null && compiledXquery != null) {
-                xqueryPool.returnCompiledXQuery(source, compiledXquery);
+            if (compilationResult.xqueryPool != null) {
+                compilationResult.xqueryPool.returnCompiledXQuery(compilationResult.source, compilationResult.compiledXquery);
             }
 
             throw e;
         }
     }
 
-    public static class QueryResult implements AutoCloseable {
+    public static class CompilationResult implements AutoCloseable {
         /**
          * Indicates that the query did not need to be compiled as a cached version was available.
          */
         public static final int RETRIEVED_CACHED_COMPILED_QUERY = -1;
 
         private final Source source;
         private @Nullable final XQueryPool xqueryPool;
-        private @Nullable final CompiledXQuery compiledXquery;
+        private final CompiledXQuery compiledXquery;
         private final XQueryContext xqueryContext;
+        private boolean closed = false;
 
         public final long compilationTime;
-        public final long executionTime;
-        public final Sequence result;
 
-        public QueryResult(final Source source, @Nullable final XQueryPool xqueryPool, @Nullable final CompiledXQuery compiledXquery, final XQueryContext xqueryContext, final long compilationTime, final long executionTime, final Sequence result) {
+        public CompilationResult(final Source source, @Nullable final XQueryPool xqueryPool, final CompiledXQuery compiledXquery, final XQueryContext xqueryContext, final long compilationTime) {
             this.source = source;
             this.xqueryPool = xqueryPool;
             this.compiledXquery = compiledXquery;
             this.xqueryContext = xqueryContext;
-
             this.compilationTime = compilationTime;
-            this.executionTime = executionTime;
-            this.result = result;
         }
 
         @Override
         public void close() {
+            if (closed) {
+                return;
+            }
+
             // NOTE(AR) Only when the user has finished with the Query Result i.e. {@link #result}, can we then clean-up any associated resources.
             xqueryContext.runCleanupTasks();
 
             // Once we have cleaned-up if the query should be cached we return it to the query pool.
             if (xqueryPool != null && compiledXquery != null) {
                 xqueryPool.returnCompiledXQuery(source, compiledXquery);
             }
+
+            closed = true;
         }
+
     }
 
-    /**
-     * Compile a query and perform an operation with the compiled query.
-     *
-     * @param <T> the type of the result of the operation.
-     *
-     * @param broker the database broker.
-     * @param source the source of the query.
-     *
-     * @param op the operation to perform with the compiled query.
-     *
-     * @return the result of the operation.
-     */
-    public static <T> T withCompiledQuery(final DBBroker broker, final Source source, final Function2E<CompiledXQuery, T, XPathException, PermissionDeniedException> op) throws XPathException, PermissionDeniedException, IOException {
-        final BrokerPool pool = broker.getBrokerPool();
-        final XQuery xqueryService = pool.getXQueryService();
-        final XQueryPool xqueryPool = pool.getXQueryPool();
-        final CompiledXQuery compiledQuery = compileQuery(broker, xqueryService, xqueryPool, source);
-        try {
-            return op.apply(compiledQuery);
-        } finally {
-            if (compiledQuery != null) {
-                if (compiledQuery.getContext() != null) {
-                    compiledQuery.getContext().runCleanupTasks();
-                }
-                xqueryPool.returnCompiledXQuery(source, compiledQuery);
-            }
+    public static class QueryResult implements AutoCloseable {
+        private final Source source;
+        private @Nullable final XQueryPool xqueryPool;
+        private @Nullable final CompiledXQuery compiledXquery;
+        private final XQueryContext xqueryContext;
+        private boolean closed = false;
+
+        public final long compilationTime;
+        public final long executionTime;
+        public final Sequence result;
+        
+        public QueryResult(final Source source, @Nullable final XQueryPool xqueryPool, @Nullable final CompiledXQuery compiledXquery, final XQueryContext xqueryContext, final long compilationTime, final long executionTime, final Sequence result) {
+            this.source = source;
+            this.xqueryPool = xqueryPool;
+            this.compiledXquery = compiledXquery;
+            this.xqueryContext = xqueryContext;
+
+            this.compilationTime = compilationTime;
+            this.executionTime = executionTime;
+            this.result = result;
         }
-    }
 
-    private static CompiledXQuery compileQuery(final DBBroker broker, final XQuery xqueryService, final XQueryPool xqueryPool, final Source query) throws PermissionDeniedException, XPathException, IOException {
-        @Nullable CompiledXQuery compiled = null;
-        @Nullable XQueryContext context = null;
-        try {
-            compiled = xqueryPool.borrowCompiledXQuery(broker, query);
-            if (compiled == null) {
-                context = new XQueryContext(broker.getBrokerPool());
-            } else {
-                context = compiled.getContext();
-                context.prepareForReuse();
-            }
+        private QueryResult(final CompilationResult compilationResult, final long executionTime, final Sequence result) {
+            this.source = compilationResult.source;
+            this.xqueryPool = compilationResult.xqueryPool;
+            this.compiledXquery = compilationResult.compiledXquery;
+            this.xqueryContext = compilationResult.xqueryContext;
 
-            if (compiled == null) {
-                compiled = xqueryService.compile(context, query);
-            } else {
-                compiled.getContext().updateContext(context);
-                context.getWatchDog().reset();
-            }
+            this.compilationTime = compilationResult.compilationTime;
+            this.executionTime = executionTime;
+            this.result = result;
 
-            return compiled;
+            // transfer ownership of the resources owned by compilationResult to this class
+            compilationResult.closed = true;
+        }
 
-        } catch (final PermissionDeniedException | XPathException | IOException e) {
-            if (context != null) {
-                context.runCleanupTasks();
+        @Override
+        public void close() {
+            if (closed) {
+                return;
             }
-            if (compiled != null) {
-                xqueryPool.returnCompiledXQuery(query, compiled);
+
+            // NOTE(AR) Only when the user has finished with the Query Result i.e. {@link #result}, can we then clean-up any associated resources.
+            xqueryContext.runCleanupTasks();
+
+            // Once we have cleaned-up if the query should be cached we return it to the query pool.
+            if (xqueryPool != null && compiledXquery != null) {
+                xqueryPool.returnCompiledXQuery(source, compiledXquery);
             }
-            throw e;
-        }
-    }
 
-    /**
-     * Execute a compiled query.
-     *
-     * @param broker the database broker.
-     * @param compiledXQuery the compiled query.
-     *
-     * @return the result sequence.
-     */
-    public static Sequence executeQuery(final DBBroker broker, final CompiledXQuery compiledXQuery) throws PermissionDeniedException, XPathException {
-        final BrokerPool pool = broker.getBrokerPool();
-        final XQuery xqueryService = pool.getXQueryService();
-        return xqueryService.execute(broker, compiledXQuery, null, null, null, true);
+            closed = true;
+        }
     }
 }