Fix query failed metric double count bug (#17454)

apache · Nov 9, 2024 · f906d0d · f906d0d
1 parent 0dcc2bc
commit f906d0d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 17 deletions.
diff --git a/server/src/main/java/org/apache/druid/server/QueryResultPusher.java b/server/src/main/java/org/apache/druid/server/QueryResultPusher.java
@@ -229,24 +229,8 @@ private Response handleQueryException(ResultsWriter resultsWriter, QueryExceptio
     return handleDruidException(resultsWriter, DruidException.fromFailure(new QueryExceptionCompat(e)));
   }
 
-  private Response handleDruidException(ResultsWriter resultsWriter, DruidException e)
+  private void incrementQueryCounterForException(final DruidException e)
   {
-    if (resultsWriter != null) {
-      resultsWriter.recordFailure(e);
-      counter.incrementFailed();
-
-      if (accumulator != null && accumulator.isInitialized()) {
-        // We already started sending a response when we got the error message.  In this case we just give up
-        // and hope that the partial stream generates a meaningful failure message for our client.  We could consider
-        // also throwing the exception body into the response to make it easier for the client to choke if it manages
-        // to parse a meaningful object out, but that's potentially an API change so we leave that as an exercise for
-        // the future.
-        trailerFields.put(QueryResource.ERROR_MESSAGE_TRAILER_HEADER, e.getMessage());
-        trailerFields.put(QueryResource.RESPONSE_COMPLETE_TRAILER_HEADER, "false");
-        return null;
-      }
-    }
-
     switch (e.getCategory()) {
       case INVALID_INPUT:
       case UNAUTHORIZED:
@@ -264,6 +248,26 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio
         counter.incrementTimedOut();
         break;
     }
+  }
+
+  private Response handleDruidException(ResultsWriter resultsWriter, DruidException e)
+  {
+    incrementQueryCounterForException(e);
+
+    if (resultsWriter != null) {
+      resultsWriter.recordFailure(e);
+
+      if (accumulator != null && accumulator.isInitialized()) {
+        // We already started sending a response when we got the error message.  In this case we just give up
+        // and hope that the partial stream generates a meaningful failure message for our client.  We could consider
+        // also throwing the exception body into the response to make it easier for the client to choke if it manages
+        // to parse a meaningful object out, but that's potentially an API change so we leave that as an exercise for
+        // the future.
+        trailerFields.put(QueryResource.ERROR_MESSAGE_TRAILER_HEADER, e.getMessage());
+        trailerFields.put(QueryResource.RESPONSE_COMPLETE_TRAILER_HEADER, "false");
+        return null;
+      }
+    }
 
     if (response == null) {
       final Response.ResponseBuilder bob = Response

diff --git a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java
@@ -1245,6 +1245,8 @@ public void testTooManyQuery() throws InterruptedException, ExecutionException
     for (Future<Boolean> theFuture : back2) {
       Assert.assertTrue(theFuture.get());
     }
+    Assert.assertEquals(2, queryResource.getSuccessfulQueryCount());
+    Assert.assertEquals(1, queryResource.getFailedQueryCount());
   }
 
   @Test(timeout = 10_000L)