Use the last good NodeUsageStatsForThreadPools when a node returns an error

Author: nicktindall

server/src/main/java/org/elasticsearch/cluster/NodeUsageStatsForThreadPoolsCollector.java

@@ -11,16 +11,22 @@
 
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.action.admin.cluster.node.usage.NodeUsageStatsForThreadPoolsAction;
 import org.elasticsearch.action.admin.cluster.node.usage.TransportNodeUsageStatsForThreadPoolsAction;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.cluster.node.DiscoveryNode;
 
+import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Collects the thread pool usage stats for each node in the cluster.
  * <p>
- * Results are returned as a map of node ID to node usage stats.
+ * Results are returned as a map of node ID to node usage stats. Keeps track of the most recent
+ * usage stats for each node, which will be returned in the event of a failure response from that node.
  */
 public class NodeUsageStatsForThreadPoolsCollector {
     public static final NodeUsageStatsForThreadPoolsCollector EMPTY = new NodeUsageStatsForThreadPoolsCollector() {
@@ -37,6 +43,8 @@ public void collectUsageStats(
         "transport_node_usage_stats_for_thread_pools_action"
     );
 
+    private final Map<String, NodeUsageStatsForThreadPools> lastNodeUsageStatsForThreadPools = new ConcurrentHashMap<>();
+
     /**
      * Collects the thread pool usage stats ({@link NodeUsageStatsForThreadPools}) for each node in the cluster.
      *
@@ -47,15 +55,39 @@ public void collectUsageStats(
         ClusterState clusterState,
         ActionListener<Map<String, NodeUsageStatsForThreadPools>> listener
     ) {
-        var dataNodeIds = clusterState.nodes().getDataNodes().values().stream().map(node -> node.getId()).toArray(String[]::new);
+        var dataNodeIds = clusterState.nodes().getDataNodes().values().stream().map(DiscoveryNode::getId).toArray(String[]::new);
+        // Discard last-seen values for any nodes no longer present in the cluster state
+        lastNodeUsageStatsForThreadPools.keySet().retainAll(Arrays.asList(dataNodeIds));
         if (clusterState.getMinTransportVersion().supports(TRANSPORT_NODE_USAGE_STATS_FOR_THREAD_POOLS_ACTION)) {
             client.execute(
                 TransportNodeUsageStatsForThreadPoolsAction.TYPE,
                 new NodeUsageStatsForThreadPoolsAction.Request(dataNodeIds),
-                listener.map(response -> response.getAllNodeUsageStatsForThreadPools())
+                listener.map(this::replaceFailuresWithLastSeenValues)
             );
         } else {
             listener.onResponse(Map.of());
         }
     }
+
+    private Map<String, NodeUsageStatsForThreadPools> replaceFailuresWithLastSeenValues(
+        NodeUsageStatsForThreadPoolsAction.Response response
+    ) {
+        final Map<String, NodeUsageStatsForThreadPools> returnedUsageStats = response.getAllNodeUsageStatsForThreadPools();
+        // Update the last-seen usage stats
+        this.lastNodeUsageStatsForThreadPools.putAll(returnedUsageStats);
+
+        if (response.hasFailures() == false) {
+            return returnedUsageStats;
+        }
+
+        // Add in the last-seen usage stats for any nodes that failed to respond
+        final Map<String, NodeUsageStatsForThreadPools> cachedValuesForFailed = new HashMap<>(returnedUsageStats);
+        for (FailedNodeException failedNodeException : response.failures()) {
+            NodeUsageStatsForThreadPools nodeUsageStatsForThreadPools = lastNodeUsageStatsForThreadPools.get(failedNodeException.nodeId());
+            if (nodeUsageStatsForThreadPools != null) {
+                cachedValuesForFailed.put(failedNodeException.nodeId(), nodeUsageStatsForThreadPools);
+            }
+        }
+        return cachedValuesForFailed;
+    }
 }