Lazily initialize TransportActionStats in RequestHandlerRegistry (#117435) (#118381)

Same as #114107 but for the transport side. We are consuming more than
0.5M for these stat instances and a large fraction if not the majority
will go unused in many use-cases, given how expensive stats are in terms
of volatile operations the overhead of an added `getAcquire` is trivial
in the absence of writes. Outside of reducing the prod footprint this is
also kind of nice in x-pack internal cluster tests in particular where
it saves quite a bit of heap.

Co-authored-by: Dimitris Rempapis <[email protected]>

Author: original-brownbear

server/src/main/java/org/elasticsearch/bootstrap/Elasticsearch.java

@@ -50,6 +50,8 @@
 import org.elasticsearch.node.NodeValidationException;
 import org.elasticsearch.plugins.PluginBundle;
 import org.elasticsearch.plugins.PluginsLoader;
+import org.elasticsearch.rest.MethodHandlers;
+import org.elasticsearch.transport.RequestHandlerRegistry;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -209,7 +211,11 @@ private static void initPhase2(Bootstrap bootstrap) throws IOException {
             SubscribableListener.class,
             RunOnce.class,
             // We eagerly initialize to work around log4j permissions & JDK-8309727
-            VectorUtil.class
+            VectorUtil.class,
+            // RequestHandlerRegistry and MethodHandlers classes do nontrivial static initialization which should always succeed but load
+            // it now (before SM) to be sure
+            RequestHandlerRegistry.class,
+            MethodHandlers.class
         );
 
         // load the plugin Java modules and layers now for use in entitlements

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

@@ -22,13 +22,13 @@
 /**
  * Encapsulate multiple handlers for the same path, allowing different handlers for different HTTP verbs and versions.
  */
-final class MethodHandlers {
+public final class MethodHandlers {
 
     private final String path;
     private final Map<RestRequest.Method, Map<RestApiVersion, RestHandler>> methodHandlers;
 
     @SuppressWarnings("unused") // only accessed via #STATS_TRACKER_HANDLE, lazy initialized because instances consume non-trivial heap
-    private volatile HttpRouteStatsTracker statsTracker;
+    private HttpRouteStatsTracker statsTracker;
 
     private static final VarHandle STATS_TRACKER_HANDLE;
 

server/src/main/java/org/elasticsearch/transport/RequestHandlerRegistry.java

@@ -18,6 +18,8 @@
 import org.elasticsearch.tasks.TaskManager;
 
 import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 import java.util.concurrent.Executor;
 
 import static org.elasticsearch.core.Releasables.assertOnce;
@@ -31,7 +33,19 @@ public class RequestHandlerRegistry<Request extends TransportRequest> implements
     private final Executor executor;
     private final TaskManager taskManager;
     private final Writeable.Reader<Request> requestReader;
-    private final TransportActionStatsTracker statsTracker = new TransportActionStatsTracker();
+    @SuppressWarnings("unused") // only accessed via #STATS_TRACKER_HANDLE, lazy initialized because instances consume non-trivial heap
+    private TransportActionStatsTracker statsTracker;
+
+    private static final VarHandle STATS_TRACKER_HANDLE;
+
+    static {
+        try {
+            STATS_TRACKER_HANDLE = MethodHandles.lookup()
+                .findVarHandle(RequestHandlerRegistry.class, "statsTracker", TransportActionStatsTracker.class);
+        } catch (Exception e) {
+            throw new ExceptionInInitializerError(e);
+        }
+    }
 
     public RequestHandlerRegistry(
         String action,
@@ -113,15 +127,34 @@ public static <R extends TransportRequest> RequestHandlerRegistry<R> replaceHand
     }
 
     public void addRequestStats(int messageSize) {
-        statsTracker.addRequestStats(messageSize);
+        statsTracker().addRequestStats(messageSize);
     }
 
     @Override
     public void addResponseStats(int messageSize) {
-        statsTracker.addResponseStats(messageSize);
+        statsTracker().addResponseStats(messageSize);
     }
 
     public TransportActionStats getStats() {
+        var statsTracker = existingStatsTracker();
+        if (statsTracker == null) {
+            return TransportActionStats.EMPTY;
+        }
         return statsTracker.getStats();
     }
+
+    private TransportActionStatsTracker statsTracker() {
+        var tracker = existingStatsTracker();
+        if (tracker == null) {
+            var newTracker = new TransportActionStatsTracker();
+            if ((tracker = (TransportActionStatsTracker) STATS_TRACKER_HANDLE.compareAndExchange(this, null, newTracker)) == null) {
+                tracker = newTracker;
+            }
+        }
+        return tracker;
+    }
+
+    private TransportActionStatsTracker existingStatsTracker() {
+        return (TransportActionStatsTracker) STATS_TRACKER_HANDLE.getAcquire(this);
+    }
 }

server/src/main/java/org/elasticsearch/transport/TransportActionStats.java

@@ -27,6 +27,8 @@ public record TransportActionStats(
     long[] responseSizeHistogram
 ) implements Writeable, ToXContentObject {
 
+    public static final TransportActionStats EMPTY = new TransportActionStats(0, 0, new long[0], 0, 0, new long[0]);
+
     public TransportActionStats(StreamInput in) throws IOException {
         this(in.readVLong(), in.readVLong(), in.readVLongArray(), in.readVLong(), in.readVLong(), in.readVLongArray());
     }