fix(plugins): allow more than one SimpleStage impl (#3574)

* fix(plugins): allow more than one SimpleStage impl

Multiple SimpleStages keeps Orca from starting because SimpleTask had a Component annotation and SimpleStage in the constructor. Orca would try to construct all the Tasks for taskResolver and when it got to SimpleTask it wouldn't know which SimpleStage to inject. So SimpleTask obviously shouldn't be a component. But once the Component annotation was removed it would not execute SimpleStages.
The TaskResolver now needs to load all the SimpleStages and wrap each on as a SimpleTask. They need to be keyed by the SimpleStage impl class name, they can't all be keyed as SimpleTask. But you could only use classes of type Task, even though it ultimately just gets turned into a String.
On the other side the RunTaskHandler looked up tasks based on a list of Task components and SimpleTask is no longer in that list. Now RunTaskHandler looks up tasks using the TaskResolver. Doesn't it make more sense to use the TaskResolver to resolve tasks anyway? I also uses the implementingClass to look up the task rather than the RunTask.taskType because it would only be SimpleTask when someone is using SimpleStages. This seems cleaner but we still use RunTask.taskType and resolve the task based on type as a backup mainly for existing tests.

* add docs to SimpleTaskDefinition

orca-api/src/main/java/com/netflix/spinnaker/orca/api/pipeline/graph/TaskNode.java

@@ -109,6 +109,17 @@ public Builder withTask(String name, Class<? extends Task> implementingClass) {
       return this;
     }
 
+    /**
+     * Adds a task to the current graph.
+     *
+     * @param task the task node to add
+     * @return this builder with the new task appended.
+     */
+    public Builder withTask(TaskNode task) {
+      graph.add(task);
+      return this;
+    }
+
     /**
      * Adds a sub-graph of tasks that may loop if any of them return {@link
      * ExecutionStatus#REDIRECT}. The sub-graph will run after any previously added tasks and before
@@ -162,8 +173,23 @@ public GraphType getType() {
     }
   }
 
+  /**
+   * This is an abstraction above TaskDefinition that allows more flexibility for the implementing
+   * class name.
+   */
+  interface DefinedTask {
+
+    /** @return name of the task */
+    @Nonnull
+    String getName();
+
+    /** @return name of the class implementing the stage */
+    @Nonnull
+    String getImplementingClassName();
+  }
+
   /** An individual task. */
-  class TaskDefinition implements TaskNode {
+  class TaskDefinition implements TaskNode, DefinedTask {
     private final String name;
     private final Class<? extends Task> implementingClass;
 
@@ -179,5 +205,10 @@ public TaskDefinition(@Nonnull String name, @Nonnull Class<? extends Task> imple
     public @Nonnull Class<? extends Task> getImplementingClass() {
       return implementingClass;
     }
+
+    @Override
+    public @Nonnull String getImplementingClassName() {
+      return getImplementingClass().getCanonicalName();
+    }
   }
 }

orca-api/src/main/java/com/netflix/spinnaker/orca/api/simplestage/SimpleTaskDefinition.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2020 Armory, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.netflix.spinnaker.orca.api.simplestage;
+
+import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode;
+import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode.DefinedTask;
+import javax.annotation.Nonnull;
+
+/**
+ * Provides a TaskDefinition for a SimpleStage which will be wrapped in a SimpleTask. SimpleTask
+ * alone is not enough information to disambiguate the implementing class.
+ */
+public class SimpleTaskDefinition implements TaskNode, DefinedTask {
+  private final String name;
+  private final Class<? extends SimpleStage> implementingClass;
+
+  /**
+   * @param name The name of the SimpleStage that gets executed
+   * @param implementingClass The SimpleStage class type that will be referenced by the task
+   *     resolver
+   */
+  public SimpleTaskDefinition(
+      @Nonnull String name, @Nonnull Class<? extends SimpleStage> implementingClass) {
+    this.name = name;
+    this.implementingClass = implementingClass;
+  }
+
+  @Override
+  public @Nonnull String getName() {
+    return name;
+  }
+
+  /** @return the SimpleStage class that will be executed by the task */
+  public @Nonnull Class<? extends SimpleStage> getImplementingClass() {
+    return implementingClass;
+  }
+
+  @Override
+  public @Nonnull String getImplementingClassName() {
+    return getImplementingClass().getCanonicalName();
+  }
+}

orca-core/src/main/java/com/netflix/spinnaker/orca/TaskResolver.java

@@ -18,9 +18,9 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.netflix.spinnaker.orca.api.pipeline.Task;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
+import com.netflix.spinnaker.orca.api.simplestage.SimpleStage;
+import com.netflix.spinnaker.orca.pipeline.SimpleTask;
+import java.util.*;
 import javax.annotation.Nonnull;
 
 /**
@@ -44,6 +44,10 @@ public TaskResolver(Collection<Task> tasks) {
    *     alias
    */
   public TaskResolver(Collection<Task> tasks, boolean allowFallback) {
+    this(tasks, Collections.emptyList(), true);
+  }
+
+  public TaskResolver(Collection<Task> tasks, List<SimpleStage> stages, boolean allowFallback) {
     for (Task task : tasks) {
       taskByAlias.put(task.getClass().getCanonicalName(), task);
       for (String alias : task.aliases()) {
@@ -59,6 +63,9 @@ public TaskResolver(Collection<Task> tasks, boolean allowFallback) {
         taskByAlias.put(alias, task);
       }
     }
+    for (SimpleStage stage : stages) {
+      taskByAlias.put(stage.getClass().getCanonicalName(), new SimpleTask(stage));
+    }
 
     this.allowFallback = allowFallback;
   }
@@ -81,6 +88,23 @@ public Task getTask(@Nonnull String taskTypeIdentifier) {
     return task;
   }
 
+  /**
+   * Fetch a {@code Task} by {@code Class type}.
+   *
+   * @param taskType Task type (class of task)
+   * @return the Task matching {@code taskType}
+   * @throws NoSuchTaskException if Task does not exist
+   */
+  @Nonnull
+  public Task getTask(@Nonnull Class<? extends Task> taskType) {
+    Task matchingTask =
+        taskByAlias.values().stream()
+            .filter((Task task) -> taskType.isAssignableFrom(task.getClass()))
+            .findFirst()
+            .orElseThrow(() -> new NoSuchTaskException(taskType.getCanonicalName()));
+    return matchingTask;
+  }
+
   /**
    * @param taskTypeIdentifier Task identifier (class name or alias)
    * @return Task Class
@@ -103,13 +127,13 @@ public Class<? extends Task> getTaskClass(@Nonnull String taskTypeIdentifier) {
     }
   }
 
-  class DuplicateTaskAliasException extends IllegalStateException {
+  public class DuplicateTaskAliasException extends IllegalStateException {
     DuplicateTaskAliasException(String message) {
       super(message);
     }
   }
 
-  class NoSuchTaskException extends IllegalArgumentException {
+  public class NoSuchTaskException extends IllegalArgumentException {
     NoSuchTaskException(String taskTypeIdentifier) {
       super("No task found for '" + taskTypeIdentifier + "'");
     }

orca-core/src/main/java/com/netflix/spinnaker/orca/config/OrcaConfiguration.java

@@ -208,8 +208,10 @@ public TaskScheduler taskScheduler() {
   }
 
   @Bean
-  public TaskResolver taskResolver(Collection<Task> tasks) {
-    return new TaskResolver(tasks, true);
+  public TaskResolver taskResolver(
+      Collection<Task> tasks, Optional<List<SimpleStage>> simpleStages) {
+    List<SimpleStage> stages = simpleStages.orElseGet(ArrayList::new);
+    return new TaskResolver(tasks, stages, true);
   }
 
   @Bean

orca-core/src/main/java/com/netflix/spinnaker/orca/pipeline/SimpleStageDefinitionBuilder.java

@@ -20,6 +20,7 @@
 import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode;
 import com.netflix.spinnaker.orca.api.pipeline.models.StageExecution;
 import com.netflix.spinnaker.orca.api.simplestage.SimpleStage;
+import com.netflix.spinnaker.orca.api.simplestage.SimpleTaskDefinition;
 import javax.annotation.Nonnull;
 
 public class SimpleStageDefinitionBuilder implements StageDefinitionBuilder {
@@ -30,7 +31,6 @@ public SimpleStageDefinitionBuilder(SimpleStage simpleStage) {
   }
 
   public void taskGraph(@Nonnull StageExecution stage, @Nonnull TaskNode.Builder builder) {
-    SimpleTask task = new SimpleTask(simpleStage);
-    builder.withTask(simpleStage.getName(), task.getClass());
+    builder.withTask(new SimpleTaskDefinition(simpleStage.getName(), simpleStage.getClass()));
   }
 }

orca-core/src/main/java/com/netflix/spinnaker/orca/pipeline/SimpleTask.java

@@ -36,14 +36,12 @@
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.core.GenericTypeResolver;
 import org.springframework.core.ResolvableType;
-import org.springframework.stereotype.Component;
 
 @Slf4j
-@Component
 public class SimpleTask implements Task {
   private SimpleStage simpleStage;
 
-  SimpleTask(@Nullable SimpleStage simpleStage) {
+  public SimpleTask(@Nullable SimpleStage simpleStage) {
     this.simpleStage = simpleStage;
   }
 

orca-queue/src/main/kotlin/com/netflix/spinnaker/orca/q/StageDefinitionBuilders.kt

@@ -21,7 +21,7 @@ import com.netflix.spinnaker.orca.api.pipeline.models.StageExecution
 import com.netflix.spinnaker.orca.pipeline.RestrictExecutionDuringTimeWindow
 import com.netflix.spinnaker.orca.api.pipeline.graph.StageDefinitionBuilder
 import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode
-import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode.TaskDefinition
+import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode.DefinedTask
 import com.netflix.spinnaker.orca.api.pipeline.graph.TaskNode.TaskGraph
 import com.netflix.spinnaker.orca.pipeline.graph.StageGraphBuilderImpl
 import com.netflix.spinnaker.orca.api.pipeline.SyntheticStageOwner
@@ -52,11 +52,11 @@ private fun processTaskNode(
 ) {
   element.apply {
     when (value) {
-      is TaskDefinition -> {
+      is DefinedTask -> {
         val task = TaskExecutionImpl()
         task.id = (stage.tasks.size + 1).toString()
         task.name = value.name
-        task.implementingClass = value.implementingClass.name
+        task.implementingClass = value.implementingClassName
         if (isSubGraph) {
           task.isLoopStart = isFirst
           task.isLoopEnd = isLast

orca-queue/src/main/kotlin/com/netflix/spinnaker/orca/q/handler/RunTaskHandler.kt

@@ -33,6 +33,7 @@ import com.netflix.spinnaker.orca.api.pipeline.OverridableTimeoutRetryableTask
 import com.netflix.spinnaker.orca.api.pipeline.RetryableTask
 import com.netflix.spinnaker.orca.api.pipeline.Task
 import com.netflix.spinnaker.orca.TaskExecutionInterceptor
+import com.netflix.spinnaker.orca.TaskResolver
 import com.netflix.spinnaker.orca.api.pipeline.TaskResult
 import com.netflix.spinnaker.orca.clouddriver.utils.CloudProviderAware
 import com.netflix.spinnaker.orca.exceptions.ExceptionHandler
@@ -76,7 +77,7 @@ class RunTaskHandler(
   override val stageNavigator: StageNavigator,
   override val stageDefinitionBuilderFactory: StageDefinitionBuilderFactory,
   override val contextParameterProcessor: ContextParameterProcessor,
-  private val tasks: Collection<Task>,
+  private val taskResolver: TaskResolver,
   private val clock: Clock,
   private val exceptionHandlers: List<ExceptionHandler>,
   private val taskExecutionInterceptors: List<TaskExecutionInterceptor>,
@@ -200,15 +201,18 @@ class RunTaskHandler(
 
   private fun RunTask.withTask(block: (StageExecution, TaskExecution, Task) -> Unit) =
     withTask { stage, taskModel ->
-      tasks
-        .find { taskType.isAssignableFrom(it.javaClass) }
-        .let { task ->
-          if (task == null) {
-            queue.push(InvalidTaskType(this, taskType.name))
-          } else {
-            block.invoke(stage, taskModel, task)
-          }
+      try {
+        taskResolver.getTask(taskModel.implementingClass)
+      } catch (e: TaskResolver.NoSuchTaskException) {
+        try {
+          taskResolver.getTask(taskType)
+        } catch (e: TaskResolver.NoSuchTaskException) {
+          queue.push(InvalidTaskType(this, taskType.name))
+          null
         }
+      }?.let {
+        block.invoke(stage, taskModel, it)
+      }
     }
 
   private fun Task.backoffPeriod(taskModel: TaskExecution, stage: StageExecution): TemporalAmount =