Add ability to enforce concurrent query limits

core/base-rest-responses/src/main/java/datawave/webservice/query/exception/DatawaveErrorCode.java

@@ -169,6 +169,7 @@ public enum DatawaveErrorCode {
     QUERY_PLAN_ERROR(500, 161, "Error retrieving plan for query."),
     QUERY_PREDICTIONS_ERROR(500, 162, "Error retrieving predictions for query."),
     QUERY_VALIDATION_ERROR(500, 163, "Error validating query."),
+    CONCURRENT_QUERY_LIMIT_ERROR(500, 164, "Error checking concurrent query limits."),
     // 204 No Content
     NO_QUERIES_FOUND(204, 1, "No queries found for user."),
     RESULTS_NOT_SENT(204, 2, "Results not sent."),
@@ -275,7 +276,8 @@ public enum DatawaveErrorCode {
     CURRENT_AND_PREVIOUS_EVENT_ORDER_INVALID(412, 16, "Current event and previous event are not in chronological order"),
     CURRENT_AND_NEXT_EVENT_ORDER_INVALID(412, 17, "Current event and next event are not in chronological order"),
     FIELD_PHRASE_QUERY_NOT_INDEXED(412, 18, "Field cannot be queried as a phrase since it was not indexed as such."),
-    NO_QUERY_VALIDATION_RULES_CONFIGURED(412, 19, "No query validation rules configured for the query logic.");
+    NO_QUERY_VALIDATION_RULES_CONFIGURED(412, 19, "No query validation rules configured for the query logic."),
+    CONCURRENT_QUERY_LIMIT_EXCEEDED(412, 20, "Concurrent query limit exceeded.");
 
     private String message;
     private int httpCode;

core/utils/type-utils/src/main/java/datawave/query/parser/JavaRegexAnalyzer.java

@@ -23,7 +23,7 @@ public class JavaRegexAnalyzer {
 
     // Types as applied to portions of the regex. We are interested in portions that
     // are literals and those that contain regex constructs.
-    private enum RegexType {
+    public enum RegexType {
         LITERAL(true), // a literal value
         ESCAPED_LITERAL(true), // an escaped literal (e.g. \[ or \.)
         REGEX(false), // a regex
@@ -42,7 +42,7 @@ public boolean isLiteral() {
         }
     }
 
-    private static class RegexPart {
+    public static class RegexPart {
         // the regex is not-final to allow applyRegexCaseSensitivity
         public String regex;
         public RegexType type;
@@ -58,6 +58,18 @@ public RegexPart(String reg, RegexType typ, int nonCapt) {
             this(reg, typ, (nonCapt > 0));
         }
 
+        public String getRegex() {
+            return regex;
+        }
+
+        public RegexType getType() {
+            return type;
+        }
+
+        public boolean isNonCapturing() {
+            return nonCapturing;
+        }
+
         @Override
         public boolean equals(Object o) {
             if (!(o instanceof RegexPart)) {
@@ -153,6 +165,10 @@ public String getRegex() {
         return getRegex(regexParts);
     }
 
+    public RegexPart[] getRegexParts() {
+        return regexParts;
+    }
+
     public static String getRegex(RegexPart[] regexParts) {
         StringBuilder regex = new StringBuilder();
         for (RegexPart part : regexParts) {

pom.xml

@@ -72,8 +72,8 @@
         <version.commons-pool>1.6</version.commons-pool>
         <version.commons-pool2>2.6.1</version.commons-pool2>
         <version.commons-text>1.11.0</version.commons-text>
-        <version.curator>5.7.1</version.curator>
-        <version.curator.test>5.7.1</version.curator.test>
+        <version.curator>5.9.0</version.curator>
+        <version.curator.test>5.9.0</version.curator.test>
         <version.datawave.accumulo-api>4.0.0</version.datawave.accumulo-api>
         <version.datawave.audit-api>4.0.0</version.datawave.audit-api>
         <version.datawave.authorization-api>4.0.0</version.datawave.authorization-api>

web-services/deploy/configuration/src/main/resources/beanRefContext.xml

@@ -18,6 +18,7 @@
                 <value>classpath*:datawave/query/QueryExpiration.xml</value>
                 <value>classpath*:datawave/query/QueryMetricsWriter.xml</value>
                 <value>classpath*:datawave/query/*QueryLogicFactory.xml</value>
+                <value>classpath*:datawave/query/QueryLimiterFactory.xml</value>
                 <value>classpath*:datawave/query/CachedResults*.xml</value>
 				<value>classpath*:datawave/modification/ModificationServices.xml</value>
                 <value>classpath*:org/apache/accumulo/operations/LookupBean.xml</value>

web-services/deploy/configuration/src/main/resources/datawave/query/QueryLimiterFactory.xml

@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<beans xmlns="http://www.springframework.org/schema/beans"
+       xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+       xmlns:util="http://www.springframework.org/schema/util"
+       xmlns:context="http://www.springframework.org/schema/context"
+       xsi:schemaLocation="http://www.springframework.org/schema/beans
+    http://www.springframework.org/schema/beans/spring-beans-4.0.xsd
+    http://www.springframework.org/schema/context
+    http://www.springframework.org/schema/context/spring-context-4.0.xsd
+    http://www.springframework.org/schema/util
+    http://www.springframework.org/schema/util/spring-util-4.0.xsd">
+
+    <!-- Configuration for the QueryLimitProvider. -->
+    <bean id="queryLimitConfiguration" scope="prototype" class="datawave.webservice.query.limit.QueryLimitConfiguration">
+        <property name="defaultSystemQueryLimit" value="20000"/>
+        <property name="defaultUserQueryLimit" value="100"/>
+    </bean>
+
+    <!-- Configuration for the QueryLimiter. -->
+    <bean id="queryLimiter" scope="prototype" class="datawave.webservice.query.limit.QueryLimiter" >
+        <property name="zookeeperConfig" value="${zookeeper.hosts}"/>
+        <property name="configuration" ref="queryLimitConfiguration"/>
+    </bean>
+
+</beans>

web-services/deploy/spring-framework-integration/src/test/resources/springFrameworkBeanRefContext.xml

@@ -15,6 +15,7 @@
                 <value>classpath*:datawave/security/PrincipalFactory.xml</value>
                 <value>classpath*:datawave/query/QueryExpiration.xml</value>
                 <value>classpath*:datawave/query/*QueryLogicFactory.xml</value>
+                <value>classpath*:datawave/query/QueryLimiterFactory.xml</value>
                 <value>classpath*:datawave/query/CachedResults*.xml</value>
                 <value>classpath*:datawave/mapreduce/MapReduceJobs.xml</value>
             </list>

web-services/query/pom.xml

@@ -250,11 +250,26 @@
             <artifactId>datawave-in-memory-accumulo</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.curator</groupId>
+            <artifactId>curator-test</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.javassist</groupId>
             <artifactId>javassist</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.springframework</groupId>
+            <artifactId>spring-test</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <finalName>${project.artifactId}</finalName>

web-services/query/src/main/java/datawave/webservice/query/cache/QueryHeartbeatCache.java

@@ -0,0 +1,52 @@
+package datawave.webservice.query.cache;
+
+import java.io.IOException;
+
+import javax.annotation.PostConstruct;
+import javax.inject.Singleton;
+
+import org.apache.log4j.Logger;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+
+import datawave.webservice.query.limit.QueryHeartbeat;
+
+/**
+ * A cache for storing query heartbeats of running queries.
+ */
+@Singleton
+public class QueryHeartbeatCache {
+
+    private static final Logger log = Logger.getLogger(QueryHeartbeatCache.class);
+
+    private Cache<String,QueryHeartbeat> cache;
+
+    @PostConstruct
+    public void init() {
+        cache = Caffeine.newBuilder().build();
+    }
+
+    public void put(String queryId, QueryHeartbeat heartbeat) {
+        cache.put(queryId, heartbeat);
+    }
+
+    public QueryHeartbeat get(String queryId) {
+        return cache.getIfPresent(queryId);
+    }
+
+    public void stopAndRemoveHeartbeat(String queryId) {
+        QueryHeartbeat heartbeat = cache.asMap().remove(queryId);
+        if (heartbeat != null) {
+            try {
+                heartbeat.stop();
+            } catch (IOException e) {
+                log.error("Error stopping query heartbeat", e);
+            }
+        }
+    }
+
+    public void clear() {
+        cache.asMap().clear();
+    }
+}

web-services/query/src/main/java/datawave/webservice/query/limit/ActiveQueryException.java

@@ -0,0 +1,27 @@
+package datawave.webservice.query.limit;
+
+/**
+ * Represents an exception that occurred when interacting with the {@link ActiveQueryTracker}.
+ */
+public class ActiveQueryException extends RuntimeException {
+
+    public ActiveQueryException() {
+        super();
+    }
+
+    public ActiveQueryException(String message) {
+        super(message);
+    }
+
+    public ActiveQueryException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public ActiveQueryException(Throwable cause) {
+        super(cause);
+    }
+
+    protected ActiveQueryException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+}