Self Service Shuffle Opt Out

SingularityService/src/main/java/com/hubspot/singularity/config/ApiPaths.java

@@ -15,6 +15,7 @@ public class ApiPaths {
   public static final String REQUEST_RESOURCE_PATH = API_BASE_PATH + "/requests";
   public static final String S3_LOG_RESOURCE_PATH = API_BASE_PATH + "/logs";
   public static final String SANDBOX_RESOURCE_PATH = API_BASE_PATH + "/sandbox";
+  public static final String SHUFFLE_RESOURCE_PATH = API_BASE_PATH + "/shuffle";
   public static final String SLAVE_RESOURCE_PATH = API_BASE_PATH + "/slaves";
   public static final String STATE_RESOURCE_PATH = API_BASE_PATH + "/state";
   public static final String TASK_RESOURCE_PATH = API_BASE_PATH + "/tasks";

SingularityService/src/main/java/com/hubspot/singularity/config/SingularityConfiguration.java

@@ -93,6 +93,7 @@ public class SingularityConfiguration extends Configuration {
 
   private int maxTasksToShufflePerHost = 2;
 
+  @Deprecated
   private List<String> doNotShuffleRequests = new ArrayList<>();
 
   private int minutesBeforeNewTaskEligibleForShuffle = 15;

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

@@ -0,0 +1,49 @@
+package com.hubspot.singularity.data;
+
+import com.codahale.metrics.MetricRegistry;
+import com.google.inject.Inject;
+import com.hubspot.singularity.SingularityCreateResult;
+import com.hubspot.singularity.SingularityDeleteResult;
+import com.hubspot.singularity.config.SingularityConfiguration;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.utils.ZKPaths;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+
+public class ShuffleConfigurationManager extends CuratorAsyncManager {
+  private static final Logger LOG = LoggerFactory.getLogger(RequestManager.class);
+
+  private static final String ZK_ROOT_PATH = "/shuffle";
+  private static final String BLACKLIST_PATH = ZK_ROOT_PATH + "/blacklist";
+
+  @Inject
+  public ShuffleConfigurationManager(
+      CuratorFramework curator,
+      SingularityConfiguration configuration,
+      MetricRegistry metricRegistry
+  ) {
+    super(curator, configuration, metricRegistry);
+  }
+
+  private String getShuffleBlacklistPath(String requestId) {
+    return ZKPaths.makePath(BLACKLIST_PATH, requestId);
+  }
+
+  public List<String> getShuffleBlacklist() {
+    return getChildren(BLACKLIST_PATH);
+  }
+
+  public SingularityCreateResult addToShuffleBlacklist(String requestId) {
+    return create(getShuffleBlacklistPath(requestId));
+  }
+
+  public SingularityDeleteResult removeFromShuffleBlacklist(String requestId) {
+    return delete(getShuffleBlacklistPath(requestId));
+  }
+
+  public boolean isOnShuffleBlacklist(String requestId) {
+    return exists(getShuffleBlacklistPath(requestId));
+  }
+}

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

@@ -36,6 +36,7 @@ protected void configure() {
     bind(SingularityValidator.class).in(Scopes.SINGLETON);
     bind(UserManager.class).in(Scopes.SINGLETON);
     bind(UsageManager.class).in(Scopes.SINGLETON);
+    bind(ShuffleConfigurationManager.class).in(Scopes.SINGLETON);
 
     bind(WebhookManager.class).in(Scopes.SINGLETON);
 

SingularityService/src/main/java/com/hubspot/singularity/data/zkmigrations/ShuffleBlacklistMigration.java

@@ -0,0 +1,27 @@
+package com.hubspot.singularity.data.zkmigrations;
+
+import com.google.inject.Inject;
+import com.google.inject.Singleton;
+import com.hubspot.singularity.config.SingularityConfiguration;
+import com.hubspot.singularity.data.ShuffleConfigurationManager;
+
+@Singleton
+public class ShuffleBlacklistMigration extends ZkDataMigration {
+  private final SingularityConfiguration configuration;
+  private final ShuffleConfigurationManager manager;
+
+
+  @Inject
+  public ShuffleBlacklistMigration(SingularityConfiguration configuration, ShuffleConfigurationManager manager) {
+    super(16);
+    this.configuration = configuration;
+    this.manager = manager;
+  }
+
+  @Override
+  public void applyMigration() {
+    for (String requestId : configuration.getDoNotShuffleRequests()) {
+      manager.addToShuffleBlacklist(requestId);
+    }
+  }
+}

SingularityService/src/main/java/com/hubspot/singularity/data/zkmigrations/SingularityZkMigrationsModule.java

@@ -32,6 +32,7 @@ public void configure(Binder binder) {
     dataMigrations.addBinding().to(NamespacePendingTasksMigration.class);
     dataMigrations.addBinding().to(NamespaceActiveTasksMigration.class);
     dataMigrations.addBinding().to(ClearUsagesMigration.class);
+    dataMigrations.addBinding().to(ShuffleBlacklistMigration.class);
   }
 
   @Provides

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

@@ -18,6 +18,7 @@
 import javax.ws.rs.Consumes;
 import javax.ws.rs.DELETE;
 import javax.ws.rs.GET;
+import javax.ws.rs.HEAD;
 import javax.ws.rs.POST;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
@@ -199,7 +200,6 @@ private void submitRequest(SingularityRequest request, Optional<SingularityReque
           }
         });
 
-
         if (oldRequest.get().getInstancesSafe() > rackManager.getNumActive()) {
           if (request.isRackSensitive() && configuration.isRebalanceRacksOnScaleDown()) {
             rebalancingHelper.rebalanceRacks(request, remainingActiveTasks, user.getEmail());

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

@@ -0,0 +1,96 @@
+package com.hubspot.singularity.resources;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Inject;
+import com.hubspot.singularity.SingularityAuthorizationScope;
+import com.hubspot.singularity.SingularityRequestWithState;
+import com.hubspot.singularity.SingularityUser;
+import com.hubspot.singularity.WebExceptions;
+import com.hubspot.singularity.auth.SingularityAuthorizationHelper;
+import com.hubspot.singularity.config.ApiPaths;
+import com.hubspot.singularity.data.RequestManager;
+import com.hubspot.singularity.data.ShuffleConfigurationManager;
+import com.ning.http.client.AsyncHttpClient;
+import io.dropwizard.auth.Auth;
+import io.swagger.v3.oas.annotations.Operation;
+import io.swagger.v3.oas.annotations.Parameter;
+import io.swagger.v3.oas.annotations.media.Schema;
+import io.swagger.v3.oas.annotations.tags.Tag;
+import io.swagger.v3.oas.annotations.tags.Tags;
+import org.apache.curator.framework.recipes.leader.LeaderLatch;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.ws.rs.DELETE;
+import javax.ws.rs.GET;
+import javax.ws.rs.POST;
+import javax.ws.rs.Path;
+import javax.ws.rs.PathParam;
+import javax.ws.rs.Produces;
+import javax.ws.rs.core.MediaType;
+import java.util.List;
+import java.util.Optional;
+
+@Path(ApiPaths.SHUFFLE_RESOURCE_PATH)
+@Produces({ MediaType.APPLICATION_JSON })
+@Schema(title = "Manages shuffle configuration of Singularity.")
+@Tags({@Tag(name = "Shuffle")})
+public class ShuffleConfigurationResource extends AbstractLeaderAwareResource {
+  private static final Logger LOG = LoggerFactory.getLogger(ShuffleConfigurationResource.class);
+
+  private final ShuffleConfigurationManager shuffleCfgManager;
+  private final SingularityAuthorizationHelper authorizationHelper;
+
+  @Inject
+  public ShuffleConfigurationResource(
+      AsyncHttpClient httpClient, LeaderLatch leaderLatch, ObjectMapper objectMapper,
+      ShuffleConfigurationManager shuffleCfgManager,
+      SingularityAuthorizationHelper authorizationHelper
+  ) {
+    super(httpClient, leaderLatch, objectMapper);
+
+    this.shuffleCfgManager = shuffleCfgManager;
+    this.authorizationHelper = authorizationHelper;
+  }
+
+  @GET
+  @Path("/blacklist")
+  @Operation(summary = "Retrieve the set of request IDs that should not be shuffled.")
+  public List<String> getShuffleBlacklist(@Parameter(hidden = true) @Auth SingularityUser user) {
+    authorizationHelper.checkReadAuthorization(user);
+    return shuffleCfgManager.getShuffleBlacklist();
+  }
+
+  @GET
+  @Path("/blacklist/{requestId}")
+  @Operation(summary = "Check if a request ID is on the shuffle blacklist.")
+  public boolean isOnShuffleBlacklist(
+      @Parameter(hidden = true) @Auth SingularityUser user,
+      @Parameter(required = true, description = "The Request ID to fetch crash loops for") @PathParam("requestId") String requestId
+  ) {
+    authorizationHelper.checkForAuthorizationByRequestId(requestId, user, SingularityAuthorizationScope.READ);
+    return shuffleCfgManager.isOnShuffleBlacklist(requestId);
+  }
+
+  @POST
+  @Path("/blacklist/{requestId}")
+  @Operation(summary = "Add a request ID to the shuffle blacklist. No effect if already present.")
+  public void addToShuffleBlacklist(
+      @Parameter(hidden = true) @Auth SingularityUser user,
+      @Parameter(required = true, description = "The Request ID to fetch crash loops for") @PathParam("requestId") String requestId
+  ) {
+    authorizationHelper.checkForAuthorizationByRequestId(requestId, user, SingularityAuthorizationScope.WRITE);
+    shuffleCfgManager.addToShuffleBlacklist(requestId);
+  }
+
+  @DELETE
+  @Path("/blacklist/{requestId}")
+  @Operation(summary = "Remove a request ID from the shuffle blacklist. No effect if not present.")
+  public void removeFromShuffleBlacklist(
+      @Parameter(hidden = true) @Auth SingularityUser user,
+      @Parameter(required = true, description = "The Request ID to fetch crash loops for") @PathParam("requestId") String requestId
+  ) {
+    authorizationHelper.checkForAuthorizationByRequestId(requestId, user, SingularityAuthorizationScope.WRITE);
+    shuffleCfgManager.removeFromShuffleBlacklist(requestId);
+  }
+}

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

@@ -40,7 +40,9 @@ protected void configure() {
     bind(UsageResource.class);
     bind(RequestGroupResource.class);
     bind(InactiveSlaveResource.class);
+    bind(ShuffleConfigurationResource.class);
     bind(TaskTrackerResource.class);
+    bind(NotificationsResource.class);
 
     bind(RebalancingHelper.class).in(Scopes.SINGLETON);
 

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

@@ -37,6 +37,7 @@ protected void configure() {
     bind(SingularityPriorityKillPoller.class).in(Scopes.SINGLETON);
     bind(SingularityUsageCleanerPoller.class).in(Scopes.SINGLETON);
     bind(SingularityUsagePoller.class).in(Scopes.SINGLETON);
+    bind(SingularityTaskShuffler.class).in(Scopes.SINGLETON);
     bind(SingularityMesosTaskPrioritizer.class).in(Scopes.SINGLETON);
     bind(SingularityMesosOfferScheduler.class).in(Scopes.SINGLETON);
     bind(SingularityLeaderCache.class).in(Scopes.SINGLETON);