[ML] fix bug where initial scale from 0->1 could scale too high (#84244) (#84248)

A bug introduced in #72423 could cause initial scaling from 0->1 to be much too high for the given tier.

This only occurs if the autoscaling capacity API was called numerous times in a row while there were no machine learning nodes, but autoscaling was enabled.

Author: benwtrent

docs/changelog/84244.yaml

@@ -0,0 +1,5 @@
+pr: 84244
+summary: Fix bug where initial scale from 0->1 could scale too high
+area: Machine Learning
+type: bug
+issues: []

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlAutoscalingDeciderService.java

@@ -702,14 +702,14 @@ AutoscalingDeciderResult scaleUpFromZero(
         // If we still have calculated zero, this means the ml memory tracker does not have the required info.
         // So, request a scale for the default. This is only for the 0 -> N scaling case.
         if (updatedCapacity.getNodeMlNativeMemoryRequirement() == 0L) {
-            updatedCapacity.merge(
+            updatedCapacity = updatedCapacity.merge(
                 new NativeMemoryCapacity(
                     ByteSizeValue.ofMb(AnalysisLimits.DEFAULT_MODEL_MEMORY_LIMIT_MB).getBytes(),
                     ByteSizeValue.ofMb(AnalysisLimits.DEFAULT_MODEL_MEMORY_LIMIT_MB).getBytes()
                 )
             );
         }
-        updatedCapacity.merge(
+        updatedCapacity = updatedCapacity.merge(
             new NativeMemoryCapacity(
                 MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
                 MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes()

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/NativeMemoryCapacity.java

@@ -26,9 +26,9 @@ static NativeMemoryCapacity from(NativeMemoryCapacity capacity) {
         return new NativeMemoryCapacity(capacity.tierMlNativeMemoryRequirement, capacity.nodeMlNativeMemoryRequirement, capacity.jvmSize);
     }
 
-    private long tierMlNativeMemoryRequirement;
-    private long nodeMlNativeMemoryRequirement;
-    private Long jvmSize;
+    private final long tierMlNativeMemoryRequirement;
+    private final long nodeMlNativeMemoryRequirement;
+    private final Long jvmSize;
 
     public NativeMemoryCapacity(long tierMlNativeMemoryRequirement, long nodeMlNativeMemoryRequirement, Long jvmSize) {
         this.tierMlNativeMemoryRequirement = tierMlNativeMemoryRequirement;
@@ -39,17 +39,22 @@ public NativeMemoryCapacity(long tierMlNativeMemoryRequirement, long nodeMlNativ
     NativeMemoryCapacity(long tierMlNativeMemoryRequirement, long nodeMlNativeMemoryRequirement) {
         this.tierMlNativeMemoryRequirement = tierMlNativeMemoryRequirement;
         this.nodeMlNativeMemoryRequirement = nodeMlNativeMemoryRequirement;
+        this.jvmSize = null;
     }
 
-    NativeMemoryCapacity merge(NativeMemoryCapacity nativeMemoryCapacity) {
-        this.tierMlNativeMemoryRequirement += nativeMemoryCapacity.tierMlNativeMemoryRequirement;
-        if (nativeMemoryCapacity.nodeMlNativeMemoryRequirement > this.nodeMlNativeMemoryRequirement) {
-            this.nodeMlNativeMemoryRequirement = nativeMemoryCapacity.nodeMlNativeMemoryRequirement;
-            // If the new node size is bigger, we have no way of knowing if the JVM size would stay the same
-            // So null out
-            this.jvmSize = null;
-        }
-        return this;
+    /**
+     * Merges the passed capacity with the current one. A new instance is created and returned
+     * @param nativeMemoryCapacity the capacity to merge with
+     * @return a new instance with the merged capacity values
+     */
+    NativeMemoryCapacity merge(final NativeMemoryCapacity nativeMemoryCapacity) {
+        if (this == nativeMemoryCapacity) return this;
+        long tier = this.tierMlNativeMemoryRequirement + nativeMemoryCapacity.tierMlNativeMemoryRequirement;
+        long node = Math.max(nativeMemoryCapacity.nodeMlNativeMemoryRequirement, this.nodeMlNativeMemoryRequirement);
+        // If the new node size is bigger, we have no way of knowing if the JVM size would stay the same
+        // So null out
+        Long jvmSize = nativeMemoryCapacity.nodeMlNativeMemoryRequirement > this.nodeMlNativeMemoryRequirement ? null : this.jvmSize;
+        return new NativeMemoryCapacity(tier, node, jvmSize);
     }
 
     public AutoscalingCapacity autoscalingCapacity(int maxMemoryPercent, boolean useAuto) {

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/autoscaling/NativeMemoryCapacityTests.java

@@ -30,12 +30,12 @@ public void testMerge() {
             ByteSizeValue.ofMb(200).getBytes(),
             ByteSizeValue.ofMb(50).getBytes()
         );
-        capacity.merge(new NativeMemoryCapacity(ByteSizeValue.ofGb(1).getBytes(), ByteSizeValue.ofMb(100).getBytes()));
+        capacity = capacity.merge(new NativeMemoryCapacity(ByteSizeValue.ofGb(1).getBytes(), ByteSizeValue.ofMb(100).getBytes()));
         assertThat(capacity.getTierMlNativeMemoryRequirement(), equalTo(ByteSizeValue.ofGb(1).getBytes() * 2L));
         assertThat(capacity.getNodeMlNativeMemoryRequirement(), equalTo(ByteSizeValue.ofMb(200).getBytes()));
         assertThat(capacity.getJvmSize(), equalTo(ByteSizeValue.ofMb(50).getBytes()));
 
-        capacity.merge(new NativeMemoryCapacity(ByteSizeValue.ofGb(1).getBytes(), ByteSizeValue.ofMb(300).getBytes()));
+        capacity = capacity.merge(new NativeMemoryCapacity(ByteSizeValue.ofGb(1).getBytes(), ByteSizeValue.ofMb(300).getBytes()));
 
         assertThat(capacity.getTierMlNativeMemoryRequirement(), equalTo(ByteSizeValue.ofGb(1).getBytes() * 3L));
         assertThat(capacity.getNodeMlNativeMemoryRequirement(), equalTo(ByteSizeValue.ofMb(300).getBytes()));