feat: interruptible VIATRA engine

Reduce server load by introducing a timeout for semantics analysis.

5 files changed, 37 insertions, 5 deletions

diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java
index d31325f1..6b3be115 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java
@@ -5,14 +5,15 @@
  */
 package tools.refinery.store.query.viatra;
 
-import tools.refinery.viatra.runtime.api.ViatraQueryEngineOptions;
-import tools.refinery.viatra.runtime.matchers.backend.IQueryBackendFactory;
-import tools.refinery.viatra.runtime.matchers.backend.QueryEvaluationHint;
 import tools.refinery.store.model.ModelStore;
 import tools.refinery.store.query.ModelQueryBuilder;
 import tools.refinery.store.query.dnf.AnyQuery;
 import tools.refinery.store.query.dnf.Dnf;
 import tools.refinery.store.query.rewriter.DnfRewriter;
+import tools.refinery.viatra.runtime.CancellationToken;
+import tools.refinery.viatra.runtime.api.ViatraQueryEngineOptions;
+import tools.refinery.viatra.runtime.matchers.backend.IQueryBackendFactory;
+import tools.refinery.viatra.runtime.matchers.backend.QueryEvaluationHint;
 
 import java.util.Collection;
 import java.util.function.Function;
@@ -29,6 +30,8 @@ public interface ViatraModelQueryBuilder extends ModelQueryBuilder {
 
         ViatraModelQueryBuilder searchBackend(IQueryBackendFactory queryBackendFactory);
 
+        ViatraModelQueryBuilder cancellationToken(CancellationToken cancellationToken);
+
         @Override
         default ViatraModelQueryBuilder queries(AnyQuery... queries) {
                 ModelQueryBuilder.super.queries(queries);
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java
index f1209f69..ad754988 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java
@@ -18,6 +18,7 @@ import tools.refinery.store.query.viatra.ViatraModelQueryAdapter;
 import tools.refinery.store.query.viatra.internal.matcher.FunctionalViatraMatcher;
 import tools.refinery.store.query.viatra.internal.matcher.RawPatternMatcher;
 import tools.refinery.store.query.viatra.internal.matcher.RelationalViatraMatcher;
+import tools.refinery.viatra.runtime.CancellationToken;
 import tools.refinery.viatra.runtime.api.AdvancedViatraQueryEngine;
 import tools.refinery.viatra.runtime.api.GenericQueryGroup;
 import tools.refinery.viatra.runtime.api.IQuerySpecification;
@@ -81,6 +82,10 @@ public class ViatraModelQueryAdapterImpl implements ViatraModelQueryAdapter, Mod
                 return storeAdapter;
         }
 
+        public CancellationToken getCancellationToken() {
+                return storeAdapter.getCancellationToken();
+        }
+
         @Override
         public <T> ResultSet<T> getResultSet(Query<T> query) {
                 var canonicalQuery = storeAdapter.getCanonicalQuery(query);
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java
index cfdc43ba..bb0630f3 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java
@@ -17,6 +17,7 @@ import tools.refinery.store.query.viatra.ViatraModelQueryBuilder;
 import tools.refinery.store.query.viatra.internal.localsearch.FlatCostFunction;
 import tools.refinery.store.query.viatra.internal.matcher.RawPatternMatcher;
 import tools.refinery.store.query.viatra.internal.pquery.Dnf2PQuery;
+import tools.refinery.viatra.runtime.CancellationToken;
 import tools.refinery.viatra.runtime.api.IQuerySpecification;
 import tools.refinery.viatra.runtime.api.ViatraQueryEngineOptions;
 import tools.refinery.viatra.runtime.localsearch.matcher.integration.LocalSearchGenericBackendFactory;
@@ -35,6 +36,7 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder<Via
                         // Use a cost function that ignores the initial (empty) model but allows higher arity input keys.
                         LocalSearchHintOptions.PLANNER_COST_FUNCTION, new FlatCostFunction()
         ), (IQueryBackendFactory) null);
+        private CancellationToken cancellationToken = CancellationToken.NONE;
         private final CompositeRewriter rewriter;
         private final Dnf2PQuery dnf2PQuery = new Dnf2PQuery();
         private final Set<AnyQuery> queries = new LinkedHashSet<>();
@@ -85,6 +87,12 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder<Via
         }
 
         @Override
+        public ViatraModelQueryBuilder cancellationToken(CancellationToken cancellationToken) {
+                this.cancellationToken = cancellationToken;
+                return this;
+        }
+
+        @Override
         public ViatraModelQueryBuilder queries(Collection<? extends AnyQuery> queries) {
                 checkNotConfigured();
                 this.queries.addAll(queries);
@@ -136,7 +144,7 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder<Via
                 validateSymbols(store);
                 return new ViatraModelQueryStoreAdapterImpl(store, buildEngineOptions(), dnf2PQuery.getSymbolViews(),
                                 Collections.unmodifiableMap(canonicalQueryMap), Collections.unmodifiableMap(querySpecifications),
-                                Collections.unmodifiableSet(vacuousQueries));
+                                Collections.unmodifiableSet(vacuousQueries), cancellationToken);
         }
 
         private ViatraQueryEngineOptions buildEngineOptions() {
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java
index 5407ad01..f32e1cc6 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java
@@ -5,6 +5,7 @@
  */
 package tools.refinery.store.query.viatra.internal;
 
+import tools.refinery.viatra.runtime.CancellationToken;
 import tools.refinery.viatra.runtime.api.IQuerySpecification;
 import tools.refinery.viatra.runtime.api.ViatraQueryEngineOptions;
 import tools.refinery.viatra.runtime.matchers.context.IInputKey;
@@ -26,18 +27,20 @@ public class ViatraModelQueryStoreAdapterImpl implements ViatraModelQueryStoreAd
         private final Map<AnyQuery, IQuerySpecification<RawPatternMatcher>> querySpecifications;
         private final Set<AnyQuery> vacuousQueries;
         private final Set<AnyQuery> allQueries;
+        private final CancellationToken cancellationToken;
 
         ViatraModelQueryStoreAdapterImpl(ModelStore store, ViatraQueryEngineOptions engineOptions,
                                                                          Map<AnySymbolView, IInputKey> inputKeys,
                                                                          Map<AnyQuery, AnyQuery> canonicalQueryMap,
                                                                          Map<AnyQuery, IQuerySpecification<RawPatternMatcher>> querySpecifications,
-                                                                         Set<AnyQuery> vacuousQueries) {
+                                                                         Set<AnyQuery> vacuousQueries, CancellationToken cancellationToken) {
                 this.store = store;
                 this.engineOptions = engineOptions;
                 this.inputKeys = inputKeys;
                 this.canonicalQueryMap = canonicalQueryMap;
                 this.querySpecifications = querySpecifications;
                 this.vacuousQueries = vacuousQueries;
+                this.cancellationToken = cancellationToken;
                 var mutableAllQueries = new LinkedHashSet<AnyQuery>(querySpecifications.size() + vacuousQueries.size());
                 mutableAllQueries.addAll(querySpecifications.keySet());
                 mutableAllQueries.addAll(vacuousQueries);
@@ -62,6 +65,10 @@ public class ViatraModelQueryStoreAdapterImpl implements ViatraModelQueryStoreAd
                 return allQueries;
         }
 
+        public CancellationToken getCancellationToken() {
+                return cancellationToken;
+        }
+
         @Override
         public <T> Query<T> getCanonicalQuery(Query<T> query) {
                 // We know that canonical forms of queries do not change output types.
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java
index d1fa5239..dadab5dd 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java
@@ -5,6 +5,7 @@
  */
 package tools.refinery.store.query.viatra.internal.context;
 
+import tools.refinery.viatra.runtime.CancellationToken;
 import tools.refinery.viatra.runtime.matchers.context.*;
 import tools.refinery.viatra.runtime.matchers.tuple.ITuple;
 import tools.refinery.viatra.runtime.matchers.tuple.Tuple;
@@ -32,10 +33,13 @@ public class RelationalRuntimeContext implements IQueryRuntimeContext {
 
         private final Model model;
 
+        private final CancellationToken cancellationToken;
+
         RelationalRuntimeContext(ViatraModelQueryAdapterImpl adapter) {
                 model = adapter.getModel();
                 metaContext = new RelationalQueryMetaContext(adapter.getStoreAdapter().getInputKeys());
                 modelUpdateListener = new ModelUpdateListener(adapter);
+                cancellationToken = adapter.getCancellationToken();
         }
 
         @Override
@@ -192,4 +196,9 @@ public class RelationalRuntimeContext implements IQueryRuntimeContext {
         public void executeAfterTraversal(Runnable runnable) {
                 runnable.run();
         }
+
+        @Override
+        public CancellationToken getCancellationToken() {
+                return cancellationToken;
+        }
 }