Record chunked response size before sending last chunk (elastic#102586)

We expect that the HTTP response stats should reflect the responses
received by the client, but today we record that a chunked-encoded
response was sent _after_ sending the last chunk, which might mean that
the client can receive the complete response body and then retrieve
stats that do not include that response. With this commit we record the
stats before sending the last chunk.

Closes elastic#102547

qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/HttpStatsIT.java

@@ -52,8 +52,6 @@ public void testNodeHttpStats() throws IOException {
         assertHttpStats(new XContentTestUtils.JsonMapView((Map<String, Object>) nodesMap.get(nodeId)));
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/102547")
-    @SuppressWarnings("unchecked")
     public void testClusterInfoHttpStats() throws IOException {
         internalCluster().ensureAtLeastNumDataNodes(3);
         performHttpRequests();

server/src/main/java/org/elasticsearch/rest/RestController.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.common.path.PathTrie;
 import org.elasticsearch.common.recycler.Recycler;
 import org.elasticsearch.common.util.Maps;
+import org.elasticsearch.common.util.concurrent.RunOnce;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.RestApiVersion;
@@ -820,12 +821,12 @@ private void close() {
     private static class EncodedLengthTrackingChunkedRestResponseBody implements ChunkedRestResponseBody {
 
         private final ChunkedRestResponseBody delegate;
-        private final MethodHandlers methodHandlers;
+        private final RunOnce onCompletion;
         private long encodedLength = 0;
 
         private EncodedLengthTrackingChunkedRestResponseBody(ChunkedRestResponseBody delegate, MethodHandlers methodHandlers) {
             this.delegate = delegate;
-            this.methodHandlers = methodHandlers;
+            this.onCompletion = new RunOnce(() -> methodHandlers.addResponseStats(encodedLength));
         }
 
         @Override
@@ -837,6 +838,9 @@ public boolean isDone() {
         public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) throws IOException {
             final ReleasableBytesReference bytesReference = delegate.encodeChunk(sizeHint, recycler);
             encodedLength += bytesReference.length();
+            if (isDone()) {
+                onCompletion.run();
+            }
             return bytesReference;
         }
 
@@ -848,7 +852,9 @@ public String getResponseContentTypeString() {
         @Override
         public void close() {
             delegate.close();
-            methodHandlers.addResponseStats(encodedLength);
+            // the client might close the connection before we send the last chunk, in which case we won't have recorded the response in the
+            // stats yet, so we do it now:
+            onCompletion.run();
         }
     }