Report bytes written for failed queries. (#18842)

This patch updates things so the query/bytes metric, and the field
in logged requests, are both included for failed requests. This helps
see how many bytes were written before a query failed.

Author: gianm

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

@@ -156,9 +156,9 @@ public void recordSuccess(long numBytes)
         }
 
         @Override
-        public void recordFailure(Exception e)
+        public void recordFailure(Exception e, long bytesWritten)
         {
-          queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), -1);
+          queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), bytesWritten);
         }
 
         @Override

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

@@ -254,7 +254,8 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio
     incrementQueryCounterForException(e);
 
     if (resultsWriter != null) {
-      resultsWriter.recordFailure(e);
+      final long bytesWritten = accumulator != null ? accumulator.getNumBytesSent() : 0;
+      resultsWriter.recordFailure(e, bytesWritten);
 
       if (accumulator != null && accumulator.isInitialized()) {
         // We already started sending a response when we got the error message.  In this case we write the exception
@@ -361,7 +362,7 @@ public interface ResultsWriter extends Closeable
 
     void recordSuccess(long numBytes);
 
-    void recordFailure(Exception e);
+    void recordFailure(Exception e, long bytesWritten);
   }
 
   public interface Writer extends Closeable

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

@@ -89,7 +89,7 @@ public void recordSuccess(long numBytes)
       }
 
       @Override
-      public void recordFailure(Exception e)
+      public void recordFailure(Exception e, long bytesWritten)
       {
         assertTrue(Throwables.getStackTraceAsString(e).contains(embeddedExceptionMessage));
         recordFailureInvoked.set(true);

sql/src/main/java/org/apache/druid/sql/SqlExecutionReporter.java

@@ -67,9 +67,22 @@ public SqlExecutionReporter(
     this.startNs = System.nanoTime();
   }
 
+  /**
+   * Report a query failure with an unknown number of byte written. The {@code sqlQuery/bytes} metric will
+   * not be emitted.
+   */
   public void failed(Throwable e)
+  {
+    failed(e, -1);
+  }
+
+  /**
+   * Report a query failure with a known number of byte written. It will be emitted as {@code sqlQuery/bytes}.
+   */
+  public void failed(Throwable e, long bytesWritten)
   {
     this.e = e;
+    this.bytesWritten = bytesWritten;
   }
 
   public void succeeded(final long bytesWritten)

sql/src/main/java/org/apache/druid/sql/http/SqlResourceQueryResultPusher.java

@@ -161,14 +161,14 @@ public void recordSuccess(long numBytes)
       }
 
       @Override
-      public void recordFailure(Exception e)
+      public void recordFailure(Exception e, long bytesWritten)
       {
         if (QueryLifecycle.shouldLogStackTrace(e, sqlQuery.queryContext())) {
           log.warn(e, "Exception while processing sqlQueryId[%s]", sqlQueryId);
         } else {
           log.noStackTrace().warn(e, "Exception while processing sqlQueryId[%s]", sqlQueryId);
         }
-        stmt.reporter().failed(e);
+        stmt.reporter().failed(e, bytesWritten);
       }
 
       @Override