more refactoring, more tests with canceled/failed deploys

SingularityBase/src/main/java/com/hubspot/singularity/TaskCleanupType.java

@@ -3,7 +3,7 @@
 public enum TaskCleanupType {
 
   USER_REQUESTED(true, true), USER_REQUESTED_TASK_BOUNCE(false, false), DECOMISSIONING(false, false), SCALING_DOWN(true, false), BOUNCING(false, false), INCREMENTAL_BOUNCE(false, false),
-  DEPLOY_FAILED(true, true), NEW_DEPLOY_SUCCEEDED(true, false), DEPLOY_STEP_FINISHED(true, false), DEPLOY_CANCELED(true, true), UNHEALTHY_NEW_TASK(true, true), OVERDUE_NEW_TASK(true, true);
+  DEPLOY_FAILED(false, true), NEW_DEPLOY_SUCCEEDED(true, false), DEPLOY_STEP_FINISHED(true, false), DEPLOY_CANCELED(false, true), UNHEALTHY_NEW_TASK(true, true), OVERDUE_NEW_TASK(true, true);
 
   private final boolean killLongRunningTaskInstantly;
   private final boolean killNonLongRunningTaskInstantly;

SingularityService/src/main/java/com/hubspot/singularity/scheduler/SingularityCleaner.java

@@ -134,13 +134,8 @@ private boolean shouldKillTask(SingularityTaskCleanup taskCleanup, List<Singular
 
     final String activeDeployId = deployState.get().getActiveDeploy().get().getDeployId();
 
-    if (!taskCleanup.getTaskId().getDeployId().equals(activeDeployId)) {
-      LOG.debug("Killing a task {} immediately because it is not part of the active deploy {}", taskCleanup, deployState.get().getActiveDeploy().get());
-      return true;
-    }
-
     // check to see if there are enough active tasks out there that have been active for long enough that we can safely shut this task down.
-    final List<SingularityTaskId> matchingTasks = SingularityTaskId.matchingAndNotIn(activeTaskIds, taskCleanup.getTaskId().getRequestId(), taskCleanup.getTaskId().getDeployId(), cleaningTasks);
+    final List<SingularityTaskId> matchingTasks = SingularityTaskId.matchingAndNotIn(activeTaskIds, taskCleanup.getTaskId().getRequestId(), activeDeployId, cleaningTasks);
 
     // For an incremental bounce, shut down old tasks as new ones are started
     final SingularityDeployKey key = SingularityDeployKey.fromTaskId(taskCleanup.getTaskId());

SingularityService/src/main/java/com/hubspot/singularity/scheduler/SingularityDeployChecker.java

@@ -276,21 +276,15 @@ private long getAllowedMillis(SingularityDeploy deploy) {
   private boolean isDeployOverdue(SingularityPendingDeploy pendingDeploy, Optional<SingularityDeploy> deploy) {
     if (!deploy.isPresent()) {
       LOG.warn("Can't determine if deploy {} is overdue because it was missing", pendingDeploy);
-
       return false;
     }
 
-    final long startTime;
-    if (pendingDeploy.getDeployProgress().isPresent()) {
-      if (pendingDeploy.getDeployProgress().get().isStepComplete()) {
-        return false;
-      } else {
-        startTime = pendingDeploy.getDeployProgress().get().getTimestamp();
-      }
-    } else {
-      startTime = pendingDeploy.getDeployMarker().getTimestamp();
+    if (pendingDeploy.getDeployProgress().isPresent() && pendingDeploy.getDeployProgress().get().isStepComplete()) {
+      return false;
     }
 
+    final long startTime = getStartTime(pendingDeploy);
+
     final long deployDuration = System.currentTimeMillis() - startTime;
 
     final long allowedTime = getAllowedMillis(deploy.get());
@@ -306,6 +300,14 @@ private boolean isDeployOverdue(SingularityPendingDeploy pendingDeploy, Optional
     }
   }
 
+  private long getStartTime(SingularityPendingDeploy pendingDeploy) {
+    if (pendingDeploy.getDeployProgress().isPresent()) {
+      return pendingDeploy.getDeployProgress().get().getTimestamp();
+    } else {
+      return pendingDeploy.getDeployMarker().getTimestamp();
+    }
+  }
+
   private List<SingularityTask> getTasks(Collection<SingularityTaskId> taskIds, Map<SingularityTaskId, SingularityTask> taskIdToTask) {
     final List<SingularityTask> tasks = Lists.newArrayListWithCapacity(taskIds.size());
 

SingularityService/src/test/java/com/hubspot/singularity/scheduler/SingularitySchedulerTest.java

@@ -538,6 +538,65 @@ public void testLbUpdatesAfterEachDeployStep() {
     Assert.assertEquals(DeployState.SUCCEEDED, deployManager.getDeployResult(requestId, secondDeployId).get().getDeployState());
   }
 
+  @Test
+  public void testCanceledDeployTasksStayActiveUntilReplaced() {
+    initRequest();
+
+    SingularityRequest request = requestResource.getRequest(requestId).getRequest();
+
+    requestResource.postRequest(request.toBuilder().setInstances(Optional.of(2)).build());
+
+    initFirstDeploy();
+
+    SingularityTask firstTask = launchTask(request, firstDeploy, 1, TaskState.TASK_RUNNING);
+    SingularityTask secondTask = launchTask(request, firstDeploy, 2, TaskState.TASK_RUNNING);
+
+    deploy(secondDeployId, Optional.<Boolean> absent(), Optional.of(1), Optional.<Boolean> absent(), false);
+    deployChecker.checkDeploys();
+    scheduler.drainPendingQueue(stateCacheProvider.get());
+    Assert.assertEquals(1, taskManager.getPendingTaskIds().size());
+
+    resourceOffers();
+    Assert.assertEquals(1, taskManager.getActiveTaskIdsForDeploy(requestId, secondDeployId).size());
+
+    SingularityTaskId firstNewTaskId = taskManager.getActiveTaskIdsForDeploy(requestId, secondDeployId).get(0);
+    statusUpdate(taskManager.getTask(firstNewTaskId).get(), TaskState.TASK_RUNNING);
+    deployChecker.checkDeploys();
+
+    Assert.assertEquals(1, taskManager.getCleanupTaskIds().size());
+    Assert.assertTrue(taskManager.getCleanupTaskIds().contains(firstTask.getTaskId()));
+    SingularityDeployProgress deployProgressStepOne = deployManager.getPendingDeploys().get(0).getDeployProgress().get();
+    Assert.assertTrue(deployProgressStepOne.isStepComplete());
+    Assert.assertEquals(1, deployProgressStepOne.getTargetActiveInstances());
+
+    cleaner.drainCleanupQueue();
+    statusUpdate(firstTask, TaskState.TASK_KILLED);
+
+    deployChecker.checkDeploys();
+    deployResource.cancelDeploy(requestId, secondDeployId);
+    deployChecker.checkDeploys();
+
+    scheduler.drainPendingQueue(stateCacheProvider.get());
+    List<SingularityPendingTaskId> pendingTaskIds = taskManager.getPendingTaskIds();
+    Assert.assertEquals(1, pendingTaskIds.size());
+    Assert.assertEquals(firstDeployId, pendingTaskIds.get(0).getDeployId());
+
+    cleaner.drainCleanupQueue();
+    List<SingularityTaskId> cleanupTaskIds = taskManager.getCleanupTaskIds();
+    Assert.assertEquals(1, cleanupTaskIds.size());
+    Assert.assertEquals(secondDeployId, cleanupTaskIds.get(0).getDeployId());
+
+    resourceOffers();
+    for (SingularityTaskId taskId : taskManager.getActiveTaskIdsForDeploy(requestId, firstDeployId)) {
+      statusUpdate(taskManager.getTask(taskId).get(), TaskState.TASK_RUNNING);
+    }
+
+    cleaner.drainCleanupQueue();
+
+    Assert.assertEquals(0, taskManager.getCleanupTaskIds().size());
+    Assert.assertEquals(2, taskManager.getActiveTaskIdsForDeploy(requestId, firstDeployId).size());
+  }
+
   @Test
   public void testAfterDeployWaitsForScheduledTaskToFinish() {
     initScheduledRequest();