Add skip lb removal flag to DeleteRequestRequest

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

@@ -9,6 +9,7 @@ public class SingularityRequestCleanup {
   private final Optional<String> user;
   private final RequestCleanupType cleanupType;
   private final Optional<Boolean> killTasks;
+  private final Optional<Boolean> removeFromLoadBalancer;
   private final Optional<Boolean> skipHealthchecks;
   private final Optional<String> deployId;
   private final long timestamp;
@@ -18,16 +19,24 @@ public class SingularityRequestCleanup {
   private final Optional<SingularityShellCommand> runShellCommandBeforeKill;
 
   @JsonCreator
-  public SingularityRequestCleanup(@JsonProperty("user") Optional<String> user, @JsonProperty("cleanupType") RequestCleanupType cleanupType, @JsonProperty("timestamp") long timestamp,
-      @JsonProperty("killTasks") Optional<Boolean> killTasks, @JsonProperty("requestId") String requestId, @JsonProperty("deployId") Optional<String> deployId,
-      @JsonProperty("skipHealthchecks") Optional<Boolean> skipHealthchecks, @JsonProperty("message") Optional<String> message, @JsonProperty("actionId") Optional<String> actionId,
-      @JsonProperty("runShellCommandBeforeKill") Optional<SingularityShellCommand> runShellCommandBeforeKill) {
+  public SingularityRequestCleanup(@JsonProperty("user") Optional<String> user,
+                                   @JsonProperty("cleanupType") RequestCleanupType cleanupType,
+                                   @JsonProperty("timestamp") long timestamp,
+                                   @JsonProperty("killTasks") Optional<Boolean> killTasks,
+                                   @JsonProperty("removeFromLoadBalancer") Optional<Boolean> removeFromLoadBalancer,
+                                   @JsonProperty("requestId") String requestId,
+                                   @JsonProperty("deployId") Optional<String> deployId,
+                                   @JsonProperty("skipHealthchecks") Optional<Boolean> skipHealthchecks,
+                                   @JsonProperty("message") Optional<String> message,
+                                   @JsonProperty("actionId") Optional<String> actionId,
+                                   @JsonProperty("runShellCommandBeforeKill") Optional<SingularityShellCommand> runShellCommandBeforeKill) {
     this.user = user;
     this.cleanupType = cleanupType;
     this.timestamp = timestamp;
     this.requestId = requestId;
     this.deployId = deployId;
     this.killTasks = killTasks;
+    this.removeFromLoadBalancer = removeFromLoadBalancer;
     this.skipHealthchecks = skipHealthchecks;
     this.actionId = actionId;
     this.message = message;
@@ -42,6 +51,10 @@ public Optional<Boolean> getKillTasks() {
     return killTasks;
   }
 
+  public Optional<Boolean> getRemoveFromLoadBalancer() {
+    return removeFromLoadBalancer;
+  }
+
   public String getRequestId() {
     return requestId;
   }
@@ -80,6 +93,7 @@ public String toString() {
         "user=" + user +
         ", cleanupType=" + cleanupType +
         ", killTasks=" + killTasks +
+        ", removeFromLoadBalancer=" + removeFromLoadBalancer +
         ", skipHealthchecks=" + skipHealthchecks +
         ", deployId=" + deployId +
         ", timestamp=" + timestamp +

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

@@ -13,18 +13,32 @@ public class SingularityTaskCleanup {
   private final Optional<String> message;
   private final Optional<String> actionId;
   private final Optional<SingularityTaskShellCommandRequestId> runBeforeKillId;
+  private final Optional<Boolean> removeFromLoadBalancer;
 
   @JsonCreator
   public SingularityTaskCleanup(@JsonProperty("user") Optional<String> user, @JsonProperty("cleanupType") TaskCleanupType cleanupType, @JsonProperty("timestamp") long timestamp,
       @JsonProperty("taskId") SingularityTaskId taskId, @JsonProperty("message") Optional<String> message, @JsonProperty("actionId") Optional<String> actionId,
-      @JsonProperty("runBeforeKillId") Optional<SingularityTaskShellCommandRequestId> runBeforeKillId) {
+      @JsonProperty("runBeforeKillId") Optional<SingularityTaskShellCommandRequestId> runBeforeKillId,
+                                @JsonProperty("removeFromLoadBalancer") Optional<Boolean> removeFromLoadBalancer
+  ) {
     this.user = user;
     this.cleanupType = cleanupType;
     this.timestamp = timestamp;
     this.taskId = taskId;
     this.message = message;
     this.actionId = actionId;
     this.runBeforeKillId = runBeforeKillId;
+    this.removeFromLoadBalancer = removeFromLoadBalancer;
+  }
+
+  public SingularityTaskCleanup(Optional<String> user,
+                                TaskCleanupType cleanupType,
+                                long timestamp,
+                                SingularityTaskId taskId,
+                                Optional<String> message,
+                                Optional<String> actionId,
+                                Optional<SingularityTaskShellCommandRequestId> runBeforeKillId) {
+    this(user, cleanupType, timestamp, taskId, message, actionId, runBeforeKillId, Optional.absent());
   }
 
   public Optional<String> getActionId() {
@@ -55,6 +69,10 @@ public Optional<SingularityTaskShellCommandRequestId> getRunBeforeKillId() {
     return runBeforeKillId;
   }
 
+  public Optional<Boolean> getRemoveFromLoadBalancer() {
+    return removeFromLoadBalancer;
+  }
+
   @Override
   public String toString() {
     return "SingularityTaskCleanup{" +

SingularityBase/src/main/java/com/hubspot/singularity/api/SingularityDeleteRequestRequest.java

@@ -9,12 +9,15 @@ public class SingularityDeleteRequestRequest {
 
   private final Optional<String> message;
   private final Optional<String> actionId;
+  private final Optional<Boolean> deleteFromLoadBalancer;
 
   @JsonCreator
   public SingularityDeleteRequestRequest(@JsonProperty("message") Optional<String> message,
-      @JsonProperty("actionId") Optional<String> actionId) {
+                                         @JsonProperty("actionId") Optional<String> actionId,
+                                         @JsonProperty("deleteFromLoadBalancer") Optional<Boolean> deleteFromLoadBalancer) {
     this.message = message;
     this.actionId = actionId;
+    this.deleteFromLoadBalancer = deleteFromLoadBalancer;
   }
 
   @ApiModelProperty(required=false, value="A message to show to users about why this action was taken")
@@ -27,11 +30,17 @@ public Optional<String> getActionId() {
     return actionId;
   }
 
+  @ApiModelProperty(required = false, value = "Should the service associated with the request be removed from the load balancer")
+  public Optional<Boolean> getDeleteFromLoadBalancer() {
+    return deleteFromLoadBalancer;
+  }
+
   @Override
   public String toString() {
     return "SingularityDeleteRequestRequest{" +
         "message=" + message +
         ", actionId=" + actionId +
+        ", deleteFromLoadBalancer=" + deleteFromLoadBalancer +
         '}';
   }
 }

SingularityService/src/main/java/com/hubspot/singularity/data/RequestManager.java

@@ -398,11 +398,11 @@ public Optional<SingularityRequestWithState> getRequest(String requestId, boolea
     return getData(getRequestPath(requestId), requestTranscoder);
   }
 
-  public void startDeletingRequest(SingularityRequest request, Optional<String> user, Optional<String> actionId, Optional<String> message) {
+  public void startDeletingRequest(SingularityRequest request, Optional<Boolean> removeFromLoadBalancer, Optional<String> user, Optional<String> actionId, Optional<String> message) {
     final long now = System.currentTimeMillis();
 
     // delete it no matter if the delete request already exists.
-    createCleanupRequest(new SingularityRequestCleanup(user, RequestCleanupType.DELETING, now, Optional.of(Boolean.TRUE), request.getId(), Optional.<String> absent(),
+    createCleanupRequest(new SingularityRequestCleanup(user, RequestCleanupType.DELETING, now, Optional.of(Boolean.TRUE), removeFromLoadBalancer, request.getId(), Optional.<String> absent(),
         Optional.<Boolean> absent(), message, actionId, Optional.<SingularityShellCommand>absent()));
 
     markDeleting(request, System.currentTimeMillis(), user, message);

SingularityService/src/main/java/com/hubspot/singularity/helpers/RequestHelper.java

@@ -97,9 +97,10 @@ private void checkReschedule(SingularityRequest newRequest, Optional<Singularity
       if (maybeDeployId.isPresent()) {
         if (maybeBounceRequest.isPresent()) {
           Optional<String> actionId = maybeBounceRequest.get().getActionId().or(Optional.of(UUID.randomUUID().toString()));
+          Optional<Boolean> removeFromLoadBalancer = Optional.absent();
           SingularityCreateResult createResult = requestManager.createCleanupRequest(
             new SingularityRequestCleanup(user, maybeBounceRequest.get().getIncremental().or(true) ? RequestCleanupType.INCREMENTAL_BOUNCE : RequestCleanupType.BOUNCE,
-              System.currentTimeMillis(), Optional.<Boolean> absent(), newRequest.getId(), Optional.of(maybeDeployId.get()), skipHealthchecks, message, actionId, maybeBounceRequest.get().getRunShellCommandBeforeKill()));
+              System.currentTimeMillis(), Optional.<Boolean> absent(), removeFromLoadBalancer, newRequest.getId(), Optional.of(maybeDeployId.get()), skipHealthchecks, message, actionId, maybeBounceRequest.get().getRunShellCommandBeforeKill()));
 
           if (createResult != SingularityCreateResult.EXISTED) {
             requestManager.bounce(newRequest, System.currentTimeMillis(), user, Optional.of("Bouncing due to bounce after scale"));

SingularityService/src/main/java/com/hubspot/singularity/resources/RequestResource.java

@@ -238,7 +238,7 @@ public SingularityRequestParent bounce(String requestId, Optional<SingularityBou
 
     requestManager.createCleanupRequest(
         new SingularityRequestCleanup(JavaUtils.getUserEmail(user), isIncrementalBounce ? RequestCleanupType.INCREMENTAL_BOUNCE : RequestCleanupType.BOUNCE,
-            System.currentTimeMillis(), Optional.<Boolean> absent(), requestId, Optional.of(deployId), skipHealthchecks, message, actionId, runBeforeKill));
+            System.currentTimeMillis(), Optional.<Boolean> absent(), Optional.absent(), requestId, Optional.of(deployId), skipHealthchecks, message, actionId, runBeforeKill));
 
     requestManager.bounce(requestWithState.getRequest(), System.currentTimeMillis(), JavaUtils.getUserEmail(user), message);
 
@@ -378,9 +378,10 @@ public SingularityRequestParent pause(String requestId, Optional<SingularityPaus
     }
 
     final long now = System.currentTimeMillis();
+    Optional<Boolean> removeFromLoadBalancer = Optional.absent();
 
     SingularityCreateResult result = requestManager.createCleanupRequest(new SingularityRequestCleanup(JavaUtils.getUserEmail(user),
-        RequestCleanupType.PAUSING, now, killTasks, requestId, Optional.<String> absent(), Optional.<Boolean> absent(), message, actionId, runBeforeKill));
+        RequestCleanupType.PAUSING, now, killTasks, removeFromLoadBalancer, requestId, Optional.<String> absent(), Optional.<Boolean> absent(), message, actionId, runBeforeKill));
 
     checkConflict(result == SingularityCreateResult.CREATED, "%s is already pausing - try again soon", requestId, result);
 
@@ -617,13 +618,15 @@ public SingularityRequest deleteRequest(String requestId, Optional<SingularityDe
 
     Optional<String> message = Optional.absent();
     Optional<String> actionId = Optional.absent();
+    Optional<Boolean> deleteFromLoadBalancer = Optional.absent();
 
     if (deleteRequest.isPresent()) {
       actionId = deleteRequest.get().getActionId();
       message = deleteRequest.get().getMessage();
+      deleteFromLoadBalancer = deleteRequest.get().getDeleteFromLoadBalancer();
     }
 
-    requestManager.startDeletingRequest(request, JavaUtils.getUserEmail(user), actionId, message);
+    requestManager.startDeletingRequest(request, deleteFromLoadBalancer, JavaUtils.getUserEmail(user), actionId, message);
 
     mailer.sendRequestRemovedMail(request, JavaUtils.getUserEmail(user), message);
 

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

@@ -365,7 +365,9 @@ private void drainRequestCleanupQueue() {
           } else {
             Optional<SingularityRequestHistory> maybeHistory = requestHistoryHelper.getLastHistory(requestId);
             if (maybeHistory.isPresent()) {
-              if (maybeHistory.get().getRequest().isLoadBalanced() && configuration.isDeleteRemovedRequestsFromLoadBalancer()) {
+              if (maybeHistory.get().getRequest().isLoadBalanced()
+                  && configuration.isDeleteRemovedRequestsFromLoadBalancer()
+                  && requestCleanup.getRemoveFromLoadBalancer().or(true)) {
                 createLbCleanupRequest(requestId, matchingActiveTaskIds);
               }
               requestManager.markDeleted(maybeHistory.get().getRequest(), start, requestCleanup.getUser(), requestCleanup.getMessage());
@@ -494,7 +496,7 @@ private void delete(SingularityRequestCleanup requestCleanup, Iterable<Singulari
         runBeforeKillId = Optional.of(shellRequest.getId());
       }
 
-      taskManager.createTaskCleanup(new SingularityTaskCleanup(requestCleanup.getUser(), TaskCleanupType.REQUEST_DELETING, start, taskId, requestCleanup.getMessage(), requestCleanup.getActionId(), runBeforeKillId));
+      taskManager.createTaskCleanup(new SingularityTaskCleanup(requestCleanup.getUser(), TaskCleanupType.REQUEST_DELETING, start, taskId, requestCleanup.getMessage(), requestCleanup.getActionId(), runBeforeKillId, requestCleanup.getRemoveFromLoadBalancer()));
     }
   }
 
@@ -657,7 +659,7 @@ private void cleanupRequestIfNoRemainingTasks(SingularityTaskCleanup cleanupTask
         requestManager.createCleanupRequest(
             new SingularityRequestCleanup(
                 cleanupTask.getUser(), RequestCleanupType.DELETING, System.currentTimeMillis(),
-                Optional.of(Boolean.TRUE), requestId, Optional.absent(),
+                Optional.of(Boolean.TRUE), cleanupTask.getRemoveFromLoadBalancer(), requestId, Optional.absent(),
                 Optional.absent(), cleanupTask.getMessage(), Optional.absent(), Optional.absent()));
       }
     }

SingularityService/src/test/java/com/hubspot/singularity/SingularityHistoryTest.java

@@ -491,7 +491,7 @@ public void testMessage() {
     }
 
     requestResource.scale(requestId, new SingularityScaleRequest(Optional.of(2), Optional.<Long> absent(), Optional.<Boolean> absent(), Optional.<String> absent(), Optional.of(msg), Optional.<Boolean>absent(), Optional.<Boolean>absent()));
-    requestResource.deleteRequest(requestId, Optional.of(new SingularityDeleteRequestRequest(Optional.of("a msg"), Optional.<String> absent())));
+    requestResource.deleteRequest(requestId, Optional.of(new SingularityDeleteRequestRequest(Optional.of("a msg"), Optional.<String> absent(), Optional.absent())));
 
     cleaner.drainCleanupQueue();
 

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

@@ -48,7 +48,7 @@ public void testRequestAgePurging() {
 
     Assert.assertTrue(!requestManager.getRequestHistory(requestId).isEmpty());
 
-    requestManager.startDeletingRequest(request, user, Optional.<String> absent(), Optional.<String> absent());
+    requestManager.startDeletingRequest(request, Optional.absent(), user, Optional.<String> absent(), Optional.<String> absent());
 
     requestManager.deleteHistoryParent(requestId);
 
@@ -79,17 +79,17 @@ public void testRequestCountPurging() {
     configuration.setMaxRequestsWithHistoryInZkWhenNoDatabase(Optional.of(2));
     configuration.setDeleteStaleRequestsFromZkWhenNoDatabaseAfterHours(7);
 
-    requestManager.startDeletingRequest(requestOne, user, Optional.<String>absent(), Optional.<String>absent());
+    requestManager.startDeletingRequest(requestOne, Optional.absent(), user, Optional.<String>absent(), Optional.<String>absent());
     requestManager.deleteHistoryParent(requestOne.getId());
     requestManager.activate(requestOne, RequestHistoryType.CREATED, System.currentTimeMillis() - TimeUnit.HOURS.toMillis(4), Optional.<String> absent(), Optional.<String> absent());
     requestManager.cooldown(requestOne, System.currentTimeMillis() - TimeUnit.HOURS.toMillis(3));
 
-    requestManager.startDeletingRequest(requestTwo, user, Optional.<String>absent(), Optional.<String>absent());
+    requestManager.startDeletingRequest(requestTwo, Optional.absent(), user, Optional.<String>absent(), Optional.<String>absent());
     requestManager.deleteHistoryParent(requestTwo.getId());
     requestManager.activate(requestTwo, RequestHistoryType.CREATED, System.currentTimeMillis() - TimeUnit.HOURS.toMillis(4), Optional.<String> absent(), Optional.<String> absent());
     requestManager.cooldown(requestTwo, System.currentTimeMillis() - TimeUnit.HOURS.toMillis(3));
 
-    requestManager.startDeletingRequest(requestThree, user, Optional.<String>absent(), Optional.<String>absent());
+    requestManager.startDeletingRequest(requestThree, Optional.absent(), user, Optional.<String>absent(), Optional.<String>absent());
     requestManager.deleteHistoryParent(requestThree.getId());
     requestManager.activate(requestThree, RequestHistoryType.CREATED, System.currentTimeMillis() - TimeUnit.HOURS.toMillis(4), Optional.<String> absent(), Optional.<String> absent());
     requestManager.cooldown(requestThree, System.currentTimeMillis() - TimeUnit.HOURS.toMillis(3));
@@ -198,7 +198,7 @@ public void testPurgingDoesntApplyIfDatabasePresent() {
     initRequest();
     initFirstDeploy();
 
-    requestManager.startDeletingRequest(request, user, Optional.<String> absent(), Optional.<String> absent());
+    requestManager.startDeletingRequest(request, Optional.absent(), user, Optional.<String> absent(), Optional.<String> absent());
 
     requestManager.deleteHistoryParent(requestId);