Adding eis connection pool to diagnostics

Author: jonathan-buttner

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -366,6 +366,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SIMULATE_INGEST_EFFECTIVE_MAPPING = def(9_140_0_00);
     public static final TransportVersion RESOLVE_INDEX_MODE_ADDED = def(9_141_0_00);
     public static final TransportVersion DATA_STREAM_WRITE_INDEX_ONLY_SETTINGS = def(9_142_0_00);
+    public static final TransportVersion INFERENCE_DIAGNOSTICS_FOR_EIS = def(9_143_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/GetInferenceDiagnosticsAction.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.core.inference.action;
 
 import org.apache.http.pool.PoolStats;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.action.support.nodes.BaseNodeResponse;
@@ -116,29 +117,44 @@ public int hashCode() {
 
     public static class NodeResponse extends BaseNodeResponse implements ToXContentFragment {
         static final String CONNECTION_POOL_STATS_FIELD_NAME = "connection_pool_stats";
+        static final String EIS_CONNECTION_POOL_STATS_FIELD_NAME = "eis_connection_pool_stats";
 
-        private final ConnectionPoolStats connectionPoolStats;
+        private final ConnectionPoolStats externalConnectionPoolStats;
+        private final ConnectionPoolStats eisConnectionPoolStats;
 
-        public NodeResponse(DiscoveryNode node, PoolStats poolStats) {
+        public NodeResponse(DiscoveryNode node, PoolStats externalPoolStats, PoolStats eisPoolStats) {
             super(node);
-            connectionPoolStats = ConnectionPoolStats.of(poolStats);
+            externalConnectionPoolStats = ConnectionPoolStats.of(externalPoolStats);
+            eisConnectionPoolStats = ConnectionPoolStats.of(eisPoolStats);
         }
 
         public NodeResponse(StreamInput in) throws IOException {
             super(in);
 
-            connectionPoolStats = new ConnectionPoolStats(in);
+            externalConnectionPoolStats = new ConnectionPoolStats(in);
+            if (in.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_DIAGNOSTICS_FOR_EIS)) {
+                eisConnectionPoolStats = new ConnectionPoolStats(in);
+            } else {
+                eisConnectionPoolStats = null; // EIS stats are not available in older versions
+            }
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
-            connectionPoolStats.writeTo(out);
+            externalConnectionPoolStats.writeTo(out);
+
+            if (out.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_DIAGNOSTICS_FOR_EIS)) {
+                eisConnectionPoolStats.writeTo(out);
+            }
         }
 
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-            builder.field(CONNECTION_POOL_STATS_FIELD_NAME, connectionPoolStats, params);
+            builder.field(CONNECTION_POOL_STATS_FIELD_NAME, externalConnectionPoolStats, params);
+            if (eisConnectionPoolStats != null) {
+                builder.field(EIS_CONNECTION_POOL_STATS_FIELD_NAME, eisConnectionPoolStats, params);
+            }
             return builder;
         }
 
@@ -147,16 +163,21 @@ public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
             NodeResponse response = (NodeResponse) o;
-            return Objects.equals(connectionPoolStats, response.connectionPoolStats);
+            return Objects.equals(externalConnectionPoolStats, response.externalConnectionPoolStats)
+                && Objects.equals(eisConnectionPoolStats, response.eisConnectionPoolStats);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(connectionPoolStats);
+            return Objects.hash(externalConnectionPoolStats, eisConnectionPoolStats);
+        }
+
+        ConnectionPoolStats getExternalConnectionPoolStats() {
+            return externalConnectionPoolStats;
         }
 
-        ConnectionPoolStats getConnectionPoolStats() {
-            return connectionPoolStats;
+        ConnectionPoolStats getEisConnectionPoolStats() {
+            return eisConnectionPoolStats;
         }
 
         static class ConnectionPoolStats implements ToXContentObject, Writeable {

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/GetInferenceDiagnosticsActionNodeResponseTests.java

@@ -40,7 +40,7 @@ protected GetInferenceDiagnosticsAction.NodeResponse createTestInstance() {
     protected GetInferenceDiagnosticsAction.NodeResponse mutateInstance(GetInferenceDiagnosticsAction.NodeResponse instance)
         throws IOException {
         var select = randomIntBetween(0, 3);
-        var connPoolStats = instance.getConnectionPoolStats();
+        var connPoolStats = instance.getExternalConnectionPoolStats();
 
         return switch (select) {
             case 0 -> new GetInferenceDiagnosticsAction.NodeResponse(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -219,6 +219,8 @@ public class InferencePlugin extends Plugin
     private final SetOnce<ModelRegistry> modelRegistry = new SetOnce<>();
     private List<InferenceServiceExtension> inferenceServiceExtensions;
 
+    public record ClientManagers(HttpClientManager externalManager, HttpClientManager eisManager) {}
+
     public InferencePlugin(Settings settings) {
         this.settings = settings;
     }
@@ -302,6 +304,8 @@ public Collection<?> createComponents(PluginServices services) {
             inferenceServiceSettings.getConnectionTtl()
         );
 
+        var clientManagers = new ClientManagers(httpClientManager, elasticInferenceServiceHttpClientManager);
+
         var elasticInferenceServiceRequestSenderFactory = new HttpRequestSender.Factory(
             serviceComponents.get(),
             elasticInferenceServiceHttpClientManager,
@@ -373,7 +377,7 @@ public Collection<?> createComponents(PluginServices services) {
 
         components.add(serviceRegistry);
         components.add(modelRegistry.get());
-        components.add(httpClientManager);
+        components.add(clientManagers);
         components.add(inferenceStatsBinding);
 
         // Only add InferenceServiceNodeLocalRateLimitCalculator (which is a ClusterStateListener) for cluster aware rate limiting,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportGetInferenceDiagnosticsAction.java

@@ -18,7 +18,7 @@
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceDiagnosticsAction;
-import org.elasticsearch.xpack.inference.external.http.HttpClientManager;
+import org.elasticsearch.xpack.inference.InferencePlugin;
 
 import java.io.IOException;
 import java.util.List;
@@ -31,15 +31,15 @@ public class TransportGetInferenceDiagnosticsAction extends TransportNodesAction
     GetInferenceDiagnosticsAction.NodeResponse,
     Void> {
 
-    private final HttpClientManager httpClientManager;
+    private final InferencePlugin.ClientManagers httpClientManagers;
 
     @Inject
     public TransportGetInferenceDiagnosticsAction(
         ThreadPool threadPool,
         ClusterService clusterService,
         TransportService transportService,
         ActionFilters actionFilters,
-        HttpClientManager httpClientManager
+        InferencePlugin.ClientManagers managers
     ) {
         super(
             GetInferenceDiagnosticsAction.NAME,
@@ -50,7 +50,7 @@ public TransportGetInferenceDiagnosticsAction(
             threadPool.executor(ThreadPool.Names.MANAGEMENT)
         );
 
-        this.httpClientManager = Objects.requireNonNull(httpClientManager);
+        this.httpClientManagers = Objects.requireNonNull(managers);
     }
 
     @Override
@@ -74,6 +74,10 @@ protected GetInferenceDiagnosticsAction.NodeResponse newNodeResponse(StreamInput
 
     @Override
     protected GetInferenceDiagnosticsAction.NodeResponse nodeOperation(GetInferenceDiagnosticsAction.NodeRequest request, Task task) {
-        return new GetInferenceDiagnosticsAction.NodeResponse(transportService.getLocalNode(), httpClientManager.getPoolStats());
+        return new GetInferenceDiagnosticsAction.NodeResponse(
+            transportService.getLocalNode(),
+            httpClientManagers.externalManager().getPoolStats(),
+            httpClientManagers.eisManager().getPoolStats()
+        );
     }
 }