disable TEs on scaleDown (#563)

Netflix · Oct 9, 2023 · d1bef3c · d1bef3c
1 parent d94a7b0
commit d1bef3c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/...e-server/src/main/java/io/mantisrx/master/resourcecluster/ResourceClusterScalerActor.java b/...e-server/src/main/java/io/mantisrx/master/resourcecluster/ResourceClusterScalerActor.java
@@ -44,6 +44,7 @@
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
@@ -247,6 +248,16 @@ private void onGetClusterIdleInstancesResponse(GetClusterIdleInstancesResponse r
                 .idleInstances(response.getInstanceIds())
                 .build(),
             self());
+
+        // also disable the scale down targets to avoid them being used during the scale down process.
+        response.getInstanceIds().forEach(id ->
+            this.resourceClusterActor.tell(new DisableTaskExecutorsRequest(
+                Collections.emptyMap(),
+                this.clusterId,
+                Instant.now().plus(Duration.ofHours(24)),
+                Optional.of(id)),
+                self()
+        ));
     }
 
     private void onTriggerClusterUsageRequest(TriggerClusterUsageRequest req) {

diff --git a/...ver/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterScalerActorTests.java b/...ver/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterScalerActorTests.java
@@ -47,6 +47,7 @@
 import io.mantisrx.server.master.resourcecluster.ContainerSkuID;
 import io.mantisrx.server.master.resourcecluster.TaskExecutorID;
 import io.mantisrx.shaded.com.google.common.collect.ImmutableList;
+import io.mantisrx.shaded.com.google.common.collect.ImmutableSet;
 import java.io.IOException;
 import java.time.Clock;
 import java.time.Duration;
@@ -172,7 +173,9 @@ public void testScaler() {
                 .build()));
 
         // Test callback from fetch idle list.
-        ImmutableList<TaskExecutorID> idleInstances = ImmutableList.of(TaskExecutorID.of("agent1"));
+        ImmutableList<TaskExecutorID> idleInstances = ImmutableList.of(
+            TaskExecutorID.of("agent1"),
+            TaskExecutorID.of("agent2"));
         scalerActor.tell(
             GetClusterIdleInstancesResponse.builder()
                 .clusterId(CLUSTER_ID)
@@ -192,6 +195,17 @@ public void testScaler() {
                 .build(),
             hostActorProbe.expectMsgClass(ScaleResourceRequest.class));
 
+        // validate the idle intances are disabled
+        io.mantisrx.master.resourcecluster.DisableTaskExecutorsRequest disableTEReq =
+            clusterActorProbe.expectMsgClass(io.mantisrx.master.resourcecluster.DisableTaskExecutorsRequest.class);
+        io.mantisrx.master.resourcecluster.DisableTaskExecutorsRequest disableTEReq2 =
+            clusterActorProbe.expectMsgClass(io.mantisrx.master.resourcecluster.DisableTaskExecutorsRequest.class);
+        assertTrue(disableTEReq.getTaskExecutorID().isPresent());
+        assertTrue(disableTEReq2.getTaskExecutorID().isPresent());
+        assertEquals(
+            ImmutableSet.of(disableTEReq2.getTaskExecutorID().get(), disableTEReq.getTaskExecutorID().get()),
+            ImmutableSet.copyOf(idleInstances));
+
         // Test trigger again
         GetClusterUsageRequest req2 = clusterActorProbe.expectMsgClass(GetClusterUsageRequest.class);
         assertEquals(CLUSTER_ID, req2.getClusterID());