fix(titus): custom job timeout should match task timeout (#3680)

re-add back changes from https://github.com/spinnaker/orca/pull/3624/files
which were accidentally reverted as part of https://github.com/spinnaker/orca/pull/3644/files

orca-clouddriver/src/main/groovy/com/netflix/spinnaker/orca/clouddriver/tasks/job/WaitOnJobCompletion.groovy

@@ -45,7 +45,6 @@ public class WaitOnJobCompletion extends AbstractCloudProviderAwareTask implemen
   final long backoffPeriod = TimeUnit.SECONDS.toMillis(10)
   final long timeout = TimeUnit.MINUTES.toMillis(120)
 
-
   @Autowired
   KatoRestService katoRestService
 
@@ -60,12 +59,27 @@ public class WaitOnJobCompletion extends AbstractCloudProviderAwareTask implemen
 
   static final String REFRESH_TYPE = "Job"
   /**
-   * Extra minute to pad the timing supplied by the job provider.
-   * E.g. if TitusJobRunner says this task is limitted to 50minutes we will wait 50m + 5m(padding),
+   * Extra time to pad the timing supplied by the job provider.
+   * E.g. if TitusJobRunner says this task is limited to 50minutes we will wait 50m + 5m(padding),
    * we should wait a bit longer to allow for any inaccuracies of the clock across the systems
    */
   static final Duration PROVIDER_PADDING = Duration.ofMinutes(5)
 
+  @Override
+  long getDynamicTimeout(@Nonnull StageExecution stage) {
+    String jobTimeoutFromProvider = (stage.context.get("jobRuntimeLimit") as String)
+
+    if (jobTimeoutFromProvider != null) {
+      try {
+        return Duration.parse(jobTimeoutFromProvider).plus(PROVIDER_PADDING).toMillis()
+      } catch (DateTimeParseException e) {
+        log.warn("Failed to parse job timeout specified by provider: '${jobTimeoutFromProvider}', using default", e)
+      }
+    }
+
+    return getTimeout()
+  }
+
   @Override @Nullable
   TaskResult onTimeout(@Nonnull StageExecution stage) {
     jobUtils.cancelWait(stage)

orca-clouddriver/src/test/java/com/netflix/spinnaker/orca/clouddriver/tasks/job/WaitOnJobCompletionTest.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2020 Netflix, 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.clouddriver.tasks.job;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.google.common.collect.ImmutableMap;
+import com.netflix.spinnaker.orca.api.pipeline.models.ExecutionType;
+import com.netflix.spinnaker.orca.pipeline.model.PipelineExecutionImpl;
+import com.netflix.spinnaker.orca.pipeline.model.StageExecutionImpl;
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+@RunWith(JUnitPlatform.class)
+public final class WaitOnJobCompletionTest {
+  @Test
+  void jobTimeoutSpecifiedByRunJobTask() {
+    Duration duration = Duration.ofMinutes(10);
+
+    WaitOnJobCompletion task = new WaitOnJobCompletion();
+    StageExecutionImpl myStage =
+        createStageWithContext(ImmutableMap.of("jobRuntimeLimit", duration.toString()));
+    assertThat(task.getDynamicTimeout(myStage))
+        .isEqualTo((duration.plus(WaitOnJobCompletion.getPROVIDER_PADDING())).toMillis());
+
+    StageExecutionImpl myStageInvalid =
+        createStageWithContext(ImmutableMap.of("jobRuntimeLimit", "garbage"));
+    assertThat(task.getDynamicTimeout(myStageInvalid)).isEqualTo(task.getTimeout());
+  }
+
+  private StageExecutionImpl createStageWithContext(Map<String, ?> context) {
+    return new StageExecutionImpl(
+        new PipelineExecutionImpl(ExecutionType.PIPELINE, "test"), "test", new HashMap<>(context));
+  }
+}