modify QueryCapacityExceededException to provide better messaging (#9547)

* modify QueryCapacityExceededException to provide better messaging

* style

integration-tests/src/test/java/org/apache/druid/tests/query/ITWikipediaQueryTest.java

@@ -21,7 +21,6 @@ package org.apache.druid.tests.query;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
-import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
 import org.apache.druid.query.Druids;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
@@ -110,7 +109,7 @@ public class ITWikipediaQueryTest
       StatusResponseHolder status = future.get();
       if (status.getStatus().getCode() == QueryCapacityExceededException.STATUS_CODE) {
         limited++;
-        Assert.assertTrue(status.getContent().contains(StringUtils.format(QueryCapacityExceededException.ERROR_MESSAGE_TEMPLATE, "one")));
+        Assert.assertTrue(status.getContent().contains(QueryCapacityExceededException.makeLaneErrorMessage("one", 1)));
       } else if (status.getStatus().getCode() == HttpResponseStatus.OK.getCode()) {
         success++;
       }

server/src/main/java/org/apache/druid/server/QueryCapacityExceededException.java

@@ -21,6 +21,7 @@ package org.apache.druid.server;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.annotations.VisibleForTesting;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.query.QueryException;
 
@@ -33,20 +34,22 @@ import org.apache.druid.query.QueryException;
  */
 public class QueryCapacityExceededException extends QueryException
 {
+  private static final String TOTAL_ERROR_MESSAGE_TEMPLATE =
+      "Too many concurrent queries, total query capacity of %s exceeded. Please try your query again later.";
+  private static final String LANE_ERROR_MESSAGE_TEMPLATE =
+      "Too many concurrent queries for lane '%s', query capacity of %s exceeded. Please try your query again later.";
   private static final String ERROR_CLASS = QueryCapacityExceededException.class.getName();
   public static final String ERROR_CODE = "Query capacity exceeded";
-  public static final String ERROR_MESSAGE = "Total query capacity exceeded";
-  public static final String ERROR_MESSAGE_TEMPLATE = "Query capacity exceeded for lane '%s'";
   public static final int STATUS_CODE = 429;
 
-  public QueryCapacityExceededException()
+  public QueryCapacityExceededException(int capacity)
   {
-    super(ERROR_CODE, ERROR_MESSAGE, ERROR_CLASS, null);
+    super(ERROR_CODE, makeTotalErrorMessage(capacity), ERROR_CLASS, null);
   }
 
-  public QueryCapacityExceededException(String lane)
+  public QueryCapacityExceededException(String lane, int capacity)
   {
-    super(ERROR_CODE, StringUtils.format(ERROR_MESSAGE_TEMPLATE, lane), ERROR_CLASS, null);
+    super(ERROR_CODE, makeLaneErrorMessage(lane, capacity), ERROR_CLASS, null);
   }
 
   @JsonCreator
@@ -57,4 +60,16 @@ public class QueryCapacityExceededException extends QueryException
   {
     super(errorCode, errorMessage, errorClass, null);
   }
+
+  @VisibleForTesting
+  public static String makeTotalErrorMessage(int capacity)
+  {
+    return StringUtils.format(TOTAL_ERROR_MESSAGE_TEMPLATE, capacity);
+  }
+
+  @VisibleForTesting
+  public static String makeLaneErrorMessage(String lane, int capacity)
+  {
+    return StringUtils.format(LANE_ERROR_MESSAGE_TEMPLATE, lane, capacity);
+  }
 }

server/src/main/java/org/apache/druid/server/QueryScheduler.java

@@ -202,7 +202,7 @@ public class QueryScheduler implements QueryWatcher
       laneConfig.ifPresent(config -> {
         Bulkhead laneLimiter = laneRegistry.bulkhead(lane, config);
         if (!laneLimiter.tryAcquirePermission()) {
-          throw new QueryCapacityExceededException(lane);
+          throw new QueryCapacityExceededException(lane, config.getMaxConcurrentCalls());
         }
         hallPasses.add(laneLimiter);
       });
@@ -214,7 +214,7 @@ public class QueryScheduler implements QueryWatcher
       totalConfig.ifPresent(config -> {
         Bulkhead totalLimiter = laneRegistry.bulkhead(TOTAL, config);
         if (!totalLimiter.tryAcquirePermission()) {
-          throw new QueryCapacityExceededException();
+          throw new QueryCapacityExceededException(config.getMaxConcurrentCalls());
         }
         hallPasses.add(totalLimiter);
       });

server/src/test/java/org/apache/druid/server/QueryResourceTest.java

@@ -28,7 +28,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
 import org.apache.druid.jackson.DefaultObjectMapper;
-import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.concurrent.Execs;
 import org.apache.druid.java.util.common.guava.LazySequence;
 import org.apache.druid.java.util.common.guava.Sequences;
@@ -86,7 +85,8 @@ public class QueryResourceTest
 {
   private static final QueryToolChestWarehouse WAREHOUSE = new MapQueryToolChestWarehouse(ImmutableMap.of());
   private static final ObjectMapper JSON_MAPPER = new DefaultObjectMapper();
-  private static final AuthenticationResult AUTHENTICATION_RESULT = new AuthenticationResult("druid", "druid", null, null);
+  private static final AuthenticationResult AUTHENTICATION_RESULT =
+      new AuthenticationResult("druid", "druid", null, null);
 
   private final HttpServletRequest testServletRequest = EasyMock.createMock(HttpServletRequest.class);
 
@@ -728,7 +728,7 @@ public class QueryResourceTest
           catch (IOException e) {
             throw new RuntimeException(e);
           }
-          Assert.assertEquals(QueryCapacityExceededException.ERROR_MESSAGE, ex.getMessage());
+          Assert.assertEquals(QueryCapacityExceededException.makeTotalErrorMessage(2), ex.getMessage());
           Assert.assertEquals(QueryCapacityExceededException.ERROR_CODE, ex.getErrorCode());
         }
     );
@@ -770,10 +770,7 @@ public class QueryResourceTest
             throw new RuntimeException(e);
           }
           Assert.assertEquals(
-              StringUtils.format(
+              QueryCapacityExceededException.makeLaneErrorMessage(HiLoQueryLaningStrategy.LOW, 1),
-                  QueryCapacityExceededException.ERROR_MESSAGE_TEMPLATE,
-                  HiLoQueryLaningStrategy.LOW
-              ),
               ex.getMessage()
           );
           Assert.assertEquals(QueryCapacityExceededException.ERROR_CODE, ex.getErrorCode());
@@ -825,10 +822,7 @@ public class QueryResourceTest
             throw new RuntimeException(e);
           }
           Assert.assertEquals(
-              StringUtils.format(
+              QueryCapacityExceededException.makeLaneErrorMessage(HiLoQueryLaningStrategy.LOW, 1),
-                  QueryCapacityExceededException.ERROR_MESSAGE_TEMPLATE,
-                  HiLoQueryLaningStrategy.LOW
-              ),
               ex.getMessage()
           );
           Assert.assertEquals(QueryCapacityExceededException.ERROR_CODE, ex.getErrorCode());

server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java

@@ -37,7 +37,6 @@ import org.apache.druid.guice.JsonConfigProvider;
 import org.apache.druid.guice.JsonConfigurator;
 import org.apache.druid.guice.annotations.Global;
 import org.apache.druid.guice.annotations.Json;
-import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.concurrent.Execs;
 import org.apache.druid.java.util.common.guava.BaseSequence;
 import org.apache.druid.java.util.common.guava.LazySequence;
@@ -212,7 +211,7 @@ public class QuerySchedulerTest
   public void testHiLoFailsWhenOutOfLaneCapacity()
   {
     expected.expectMessage(
-        StringUtils.format(QueryCapacityExceededException.ERROR_MESSAGE_TEMPLATE, HiLoQueryLaningStrategy.LOW)
+        QueryCapacityExceededException.makeLaneErrorMessage(HiLoQueryLaningStrategy.LOW, TEST_LO_CAPACITY)
     );
     expected.expect(QueryCapacityExceededException.class);
 
@@ -237,7 +236,7 @@ public class QuerySchedulerTest
   @Test
   public void testHiLoFailsWhenOutOfTotalCapacity()
   {
-    expected.expectMessage(QueryCapacityExceededException.ERROR_MESSAGE);
+    expected.expectMessage(QueryCapacityExceededException.makeTotalErrorMessage(TEST_HI_CAPACITY));
     expected.expect(QueryCapacityExceededException.class);
 
     Query<?> interactive1 = scheduler.prioritizeAndLaneQuery(QueryPlus.wrap(makeInteractiveQuery()), ImmutableSet.of());

sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java

@@ -756,7 +756,7 @@ public class SqlResourceTest extends CalciteTestBase
         QueryException interruped = result.lhs;
         Assert.assertEquals(QueryCapacityExceededException.ERROR_CODE, interruped.getErrorCode());
         Assert.assertEquals(
-            StringUtils.format(QueryCapacityExceededException.ERROR_MESSAGE_TEMPLATE, HiLoQueryLaningStrategy.LOW),
+            QueryCapacityExceededException.makeLaneErrorMessage(HiLoQueryLaningStrategy.LOW, 2),
             interruped.getMessage()
         );
         limited++;