Discard excess tracking if the limit is reduced

Author: nicktindall

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/UndesiredAllocationsTracker.java

@@ -22,6 +22,11 @@
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
 /**
  * Keeps track of a limited number of shards that are currently in undesired allocations. If the
  * shards remain in undesired allocations for longer than a configurable threshold, it will log
@@ -114,6 +119,7 @@ public void cleanup(RoutingNodes routingNodes) {
                 return true;
             }
         });
+        shrinkIfOversized();
     }
 
     /**
@@ -187,4 +193,22 @@ private void logUndesiredShardDetails(
             allocation.setDebugMode(originalDebugMode);
         }
     }
+
+    /**
+     * If the maximum to track was reduced, and we are tracking more than the new maximum, purge the most recent entries
+     * to bring us under the new limit
+     */
+    private void shrinkIfOversized() {
+        if (undesiredAllocations.size() > maxUndesiredAllocationsToTrack) {
+            final var newestExcessValues = StreamSupport.stream(
+                Spliterators.spliteratorUnknownSize(undesiredAllocations.iterator(), Spliterator.ORDERED),
+                false
+            )
+                .sorted((a, b) -> Long.compare(b.value, a.value))
+                .limit(undesiredAllocations.size() - maxUndesiredAllocationsToTrack)
+                .map(olc -> olc.key)
+                .collect(Collectors.toSet());
+            undesiredAllocations.removeAll(newestExcessValues::contains);
+        }
+    }
 }