Fetch the list of task ids to persist behind lock

Author: ssalinas

SingularityService/src/main/java/com/hubspot/singularity/data/history/SingularityTaskHistoryPersister.java

@@ -14,6 +14,7 @@
 import com.hubspot.singularity.config.SingularityConfiguration;
 import com.hubspot.singularity.data.DeployManager;
 import com.hubspot.singularity.data.TaskManager;
+import com.hubspot.singularity.mesos.SingularitySchedulerLock;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
@@ -42,6 +43,7 @@ public class SingularityTaskHistoryPersister
   private final DeployManager deployManager;
   private final HistoryManager historyManager;
   private final int agentReregisterTimeoutSeconds;
+  private final SingularitySchedulerLock singularitySchedulerLock;
 
   @Inject
   public SingularityTaskHistoryPersister(
@@ -50,6 +52,7 @@ public SingularityTaskHistoryPersister(
     DeployManager deployManager,
     HistoryManager historyManager,
     SingularityManagedThreadPoolFactory managedThreadPoolFactory,
+    SingularitySchedulerLock singularitySchedulerLock,
     @Named(SingularityHistoryModule.PERSISTER_LOCK) ReentrantLock persisterLock,
     @Named(
       SingularityHistoryModule.LAST_TASK_PERSISTER_SUCCESS
@@ -61,6 +64,7 @@ public SingularityTaskHistoryPersister(
     this.deployManager = deployManager;
     this.agentReregisterTimeoutSeconds =
       configuration.getMesosConfiguration().getAgentReregisterTimeoutSeconds();
+    this.singularitySchedulerLock = singularitySchedulerLock;
   }
 
   @Override
@@ -88,22 +92,30 @@ public void runActionOnPoll() {
             () -> {
               try {
                 LOG.debug("Checking request {}", requestId);
-                List<SingularityTaskId> taskIds = taskManager.getTaskIdsForRequest(
-                  requestId
+                List<SingularityTaskId> taskIds = singularitySchedulerLock.runWithRequestLockAndReturn(
+                  () -> {
+                    List<SingularityTaskId> ids = taskManager.getTaskIdsForRequest(
+                      requestId
+                    );
+                    ids.removeAll(taskManager.getActiveTaskIdsForRequest(requestId));
+                    ids.removeAll(taskManager.getLBCleanupTasks());
+                    List<SingularityPendingDeploy> pendingDeploys = deployManager.getPendingDeploys();
+                    ids =
+                      ids
+                        .stream()
+                        .filter(
+                          taskId ->
+                            !isPartOfPendingDeploy(pendingDeploys, taskId) &&
+                            !couldReturnWithRecoveredAgent(taskId)
+                        )
+                        .sorted(SingularityTaskId.STARTED_AT_COMPARATOR_DESC)
+                        .collect(Collectors.toList());
+                    return ids;
+                  },
+                  requestId,
+                  "task history persister fetch"
                 );
-                taskIds.removeAll(taskManager.getActiveTaskIdsForRequest(requestId));
-                taskIds.removeAll(taskManager.getLBCleanupTasks());
-                List<SingularityPendingDeploy> pendingDeploys = deployManager.getPendingDeploys();
-                taskIds =
-                  taskIds
-                    .stream()
-                    .filter(
-                      taskId ->
-                        !isPartOfPendingDeploy(pendingDeploys, taskId) &&
-                        !couldReturnWithRecoveredAgent(taskId)
-                    )
-                    .sorted(SingularityTaskId.STARTED_AT_COMPARATOR_DESC)
-                    .collect(Collectors.toList());
+
                 int forRequest = 0;
                 int transferred = 0;
                 for (SingularityTaskId taskId : taskIds) {

SingularityService/src/main/java/com/hubspot/singularity/mesos/SingularityMesosStatusUpdateHandler.java

@@ -373,10 +373,16 @@ private StatusUpdateResult unsafeProcessStatusUpdate(
     // If a task is missing data in Singularity there is not much we can do to recover it
     Optional<SingularityTask> maybeTask = taskManager.getTask(taskIdObj);
     if (!maybeTask.isPresent()) {
-      maybeTask = tryFindMissingTaskData(taskIdObj, taskId);
+      maybeTask = tryFindMissingTaskData(taskIdObj, taskId, taskState);
     }
     if (!maybeTask.isPresent()) {
-      handledMissingTaskData(taskIdObj, taskId, newTaskStatusHolder, taskState, now);
+      return handledMissingTaskData(
+        taskIdObj,
+        taskId,
+        newTaskStatusHolder,
+        taskState,
+        now
+      );
     }
 
     SingularityTask task = maybeTask.get();
@@ -636,15 +642,38 @@ private StatusUpdateResult tryRecoverTask(
 
   private Optional<SingularityTask> tryFindMissingTaskData(
     SingularityTaskId taskIdObj,
-    String taskId
+    String taskId,
+    ExtendedTaskState taskState
   ) {
     LOG.warn("Missing task data for {}, trying to recover", taskId);
+    // If found in this first step, it was a bad zk write and everything should just work
     Optional<SingularityTask> maybeTask = taskManager.tryRepairTask(taskIdObj);
     if (!maybeTask.isPresent()) {
       // Ensure history manager calls cannot interrupt the status update path
       try {
-        maybeTask =
-          historyManager.getTaskHistory(taskId).map(SingularityTaskHistory::getTask);
+        Optional<SingularityTaskHistory> maybeTaskHistory = historyManager.getTaskHistory(
+          taskId
+        );
+        if (maybeTaskHistory.isPresent()) {
+          maybeTask = maybeTaskHistory.map(SingularityTaskHistory::getTask);
+          if (
+            maybeTaskHistory
+              .get()
+              .getLastTaskUpdate()
+              .map(SingularityTaskHistoryUpdate::getTaskState)
+              .orElse(taskState)
+              .isDone() &&
+            !taskState.isDone()
+          ) {
+            // Don't bother with LB state/etc recovery, let the task get killed and replaced as a cleaner replacement
+            LOG.info(
+              "Recovered task {} was previously marked as done. Will not reactivate fully",
+              taskId
+            );
+            taskManager.repairFoundTask(maybeTask.get());
+            return Optional.empty();
+          }
+        }
       } catch (Exception e) {
         LOG.error("Could not fetch {} from history", taskId, e);
       }