perf(restarts): Optimized RestartStageHandler. (#4621)

Co-authored-by: Valerii Sitkevich <[email protected]> Co-authored-by: Jason <[email protected]>
spinnaker · Jan 29, 2024 · 8c0c189 · 8c0c189
1 parent a587515
commit 8c0c189
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 6 deletions.
diff --git a/orca-core/src/main/java/com/netflix/spinnaker/orca/pipeline/model/StageExecutionImpl.java b/orca-core/src/main/java/com/netflix/spinnaker/orca/pipeline/model/StageExecutionImpl.java
@@ -550,7 +550,7 @@ public List<StageExecution> allDownstreamStages() {
     List<StageExecution> children = new ArrayList<>();
 
     if (execution != null) {
-      HashSet<String> visited = new HashSet<>();
+      List<StageExecution> notVisited = new ArrayList<>(getExecution().getStages());
       LinkedList<StageExecution> queue = new LinkedList<>();
 
       queue.push(this);
@@ -563,11 +563,12 @@ public List<StageExecution> allDownstreamStages() {
         }
 
         first = false;
-        visited.add(stage.getRefId());
+        notVisited.remove(stage);
 
-        List<StageExecution> childStages = stage.downstreamStages();
-
-        childStages.stream().filter(s -> !visited.contains(s.getRefId())).forEach(queue::add);
+        notVisited.stream()
+            .filter(
+                s -> s.getRequisiteStageRefIds().contains(stage.getRefId()) && !queue.contains(s))
+            .forEach(queue::add);
       }
     }
 

diff --git a/orca-core/src/test/groovy/com/netflix/spinnaker/orca/pipeline/model/StageSpec.groovy b/orca-core/src/test/groovy/com/netflix/spinnaker/orca/pipeline/model/StageSpec.groovy
@@ -206,6 +206,57 @@ class StageSpec extends Specification {
     descendants.find {it.refId == "7"} != null
   }
 
+  def "descendents do not multiply"() {
+    given:
+    def pipeline = pipeline {
+      stage {
+        refId = "0"
+      }
+      stage {
+        refId = "1"
+        requisiteStageRefIds = ["0"]
+      }
+      stage {
+        refId = "2"
+        requisiteStageRefIds = ["0"]
+      }
+      stage {
+        refId = "3"
+        requisiteStageRefIds = ["0"]
+      }
+      stage {
+        refId = "4"
+        requisiteStageRefIds = ["1", "2", "3"]
+      }
+      stage {
+        refId = "5"
+        requisiteStageRefIds = ["1", "2", "3"]
+      }
+      stage {
+        refId = "6"
+        requisiteStageRefIds = ["1", "2", "3"]
+      }
+      stage {
+        refId = "7"
+        requisiteStageRefIds = ["4", "5", "6"]
+      }
+    }
+
+    def stage = pipeline.stageByRef("0")
+
+    when:
+    def descendants = stage.allDownstreamStages()
+
+    then:
+    descendants.size() == 7
+    descendants.find {it.refId == "2"} != null
+    descendants.find {it.refId == "3"} != null
+    descendants.find {it.refId == "4"} != null
+    descendants.find {it.refId == "5"} != null
+    descendants.find {it.refId == "6"} != null
+    descendants.find {it.refId == "7"} != null
+  }
+
   def "should not fail on no descendents"() {
     given:
     def pipeline = pipeline {

diff --git a/orca-queue/src/main/kotlin/com/netflix/spinnaker/orca/q/handler/RestartStageHandler.kt b/orca-queue/src/main/kotlin/com/netflix/spinnaker/orca/q/handler/RestartStageHandler.kt
@@ -52,6 +52,7 @@ class RestartStageHandler(
       if (topStage.status.isComplete || topStage.status == NOT_STARTED) {
         topStage.addRestartDetails(message.user)
         topStage.reset()
+        topStage.resetChildren()
         if (stage.execution.shouldQueue()) {
           // this pipeline is already running and has limitConcurrent = true
           if (topStage.execution.status == NOT_STARTED) {
@@ -126,8 +127,10 @@ class RestartStageHandler(
 
       removeSynthetics()
     }
+  }
 
-    downstreamStages().forEach { it.reset() }
+  private fun StageExecution.resetChildren() {
+    allDownstreamStages().forEach { it.reset() }
   }
 
   private fun StageExecution.removeSynthetics() {