From 8c6100e6d016143a5a8729cea6ce650b07fe6729 Mon Sep 17 00:00:00 2001
From: Armin Braun <me@obrown.io>
Date: Thu, 5 May 2022 18:17:30 +0200
Subject: [PATCH] Make AllocationDecider canRemain loop tighter

No need to collect a multi decision here (which is surprisingly prominent in the profiling)
when we're not using it in the end anyway.
---
 .../decider/AllocationDeciders.java           | 53 ++++++++++++-------
 .../decider/AllocationDecidersTests.java      |  7 ++-
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDeciders.java b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDeciders.java
index 1f675b6a0ec6c..02546ccbd9b90 100644
--- a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDeciders.java
+++ b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDeciders.java
@@ -86,29 +86,44 @@ public Decision canRemain(ShardRouting shardRouting, RoutingNode node, RoutingAl
             }
             return Decision.NO;
         }
-        Decision.Multi ret = new Decision.Multi();
-        for (AllocationDecider allocationDecider : allocations) {
-            Decision decision = allocationDecider.canRemain(shardRouting, node, allocation);
-            // short track if a NO is returned.
-            if (decision.type() == Decision.Type.NO) {
-                if (logger.isTraceEnabled()) {
-                    logger.trace(
-                        "Shard [{}] can not remain on node [{}] due to [{}]",
-                        shardRouting,
-                        node.nodeId(),
-                        allocationDecider.getClass().getSimpleName()
-                    );
-                }
-                if (allocation.debugDecision() == false) {
-                    return Decision.NO;
-                } else {
+        if (allocation.debugDecision()) {
+            Decision.Multi ret = new Decision.Multi();
+            for (AllocationDecider allocationDecider : allocations) {
+                Decision decision = allocationDecider.canRemain(shardRouting, node, allocation);
+                // short track if a NO is returned.
+                if (decision.type() == Decision.Type.NO) {
+                    maybeTraceLogNoDecision(shardRouting, node, allocationDecider);
                     ret.add(decision);
+                } else {
+                    addDecision(ret, decision, allocation);
                 }
-            } else {
-                addDecision(ret, decision, allocation);
             }
+            return ret;
+        } else {
+            // tighter loop if debug information is not collected: don't collect yes decisions + break out right away on NO
+            Decision ret = Decision.YES;
+            for (AllocationDecider allocationDecider : allocations) {
+                switch (allocationDecider.canRemain(shardRouting, node, allocation).type()) {
+                    case NO -> {
+                        maybeTraceLogNoDecision(shardRouting, node, allocationDecider);
+                        return Decision.NO;
+                    }
+                    case THROTTLE -> ret = Decision.THROTTLE;
+                }
+            }
+            return ret;
+        }
+    }
+
+    private void maybeTraceLogNoDecision(ShardRouting shardRouting, RoutingNode node, AllocationDecider allocationDecider) {
+        if (logger.isTraceEnabled()) {
+            logger.trace(
+                "Shard [{}] can not remain on node [{}] due to [{}]",
+                shardRouting,
+                node.nodeId(),
+                allocationDecider.getClass().getSimpleName()
+            );
         }
-        return ret;
     }
 
     public Decision canAllocate(IndexMetadata indexMetadata, RoutingNode node, RoutingAllocation allocation) {
diff --git a/server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDecidersTests.java b/server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDecidersTests.java
index 1c071a21605bc..a970b0fb55d7d 100644
--- a/server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDecidersTests.java
+++ b/server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/AllocationDecidersTests.java
@@ -97,7 +97,12 @@ public Decision canRebalance(RoutingAllocation allocation) {
         verify(deciders.canAllocate(shardRouting, allocation), matcher);
         verify(deciders.canRebalance(shardRouting, allocation), matcher);
         verify(deciders.canRebalance(allocation), matcher);
-        verify(deciders.canRemain(shardRouting, routingNode, allocation), matcher);
+        final Decision canRemainResult = deciders.canRemain(shardRouting, routingNode, allocation);
+        if (allocation.debugDecision()) {
+            verify(canRemainResult, matcher);
+        } else {
+            assertSame(canRemainResult, Decision.YES);
+        }
         verify(deciders.canForceAllocatePrimary(shardRouting, routingNode, allocation), matcher);
         verify(deciders.shouldAutoExpandToNode(idx, null, allocation), matcher);
     }