Update transaction and attempt to repair task data

Author: ssalinas

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

@@ -50,12 +50,15 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.api.transaction.CuratorTransactionFinal;
+import org.apache.curator.framework.api.transaction.CuratorTransactionResult;
 import org.apache.curator.utils.ZKPaths;
 import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.KeeperException.NodeExistsException;
+import org.apache.zookeeper.OpResult.ErrorResult;
 import org.apache.zookeeper.data.Stat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -1304,6 +1307,21 @@ public void createTaskAndDeletePendingTask(SingularityTask task) {
     }
   }
 
+  public Optional<SingularityTask> tryRepairTask(SingularityTaskId taskId) {
+    try {
+      Optional<SingularityTask> maybeTask = getTask(taskId); // checks zkCache for task data
+      String path = getTaskPath(taskId);
+      if (maybeTask.isPresent() && !exists(path)) {
+        LOG.info("Found info for task {} from cache not in zk node, rewriting", taskId);
+        save(path, maybeTask.map(taskTranscoder::toBytes));
+        leaderCache.putActiveTask(taskId);
+      }
+      return maybeTask;
+    } catch (Exception e) {
+      return Optional.empty();
+    }
+  }
+
   public Map<SingularityTaskId, SingularityTask> getTasks(
     Iterable<SingularityTaskId> taskIds
   ) {
@@ -1379,24 +1397,35 @@ private void createTaskAndDeletePendingTaskPrivate(SingularityTask task)
     try {
       final String path = getTaskPath(task.getTaskId());
 
-      CuratorTransactionFinal transaction = curator
-        .inTransaction()
-        .create()
-        .forPath(path, taskTranscoder.toBytes(task))
-        .and();
-
-      transaction
-        .create()
-        .forPath(
-          getLastActiveTaskStatusPath(task.getTaskId()),
-          taskStatusTranscoder.toBytes(taskStatusHolder)
+      AtomicBoolean hasErr = new AtomicBoolean(false);
+      curator
+        .transaction()
+        .forOperations(
+          curator.transactionOp().create().forPath(path, taskTranscoder.toBytes(task)),
+          curator
+            .transactionOp()
+            .create()
+            .forPath(
+              getLastActiveTaskStatusPath(task.getTaskId()),
+              taskStatusTranscoder.toBytes(taskStatusHolder)
+            )
         )
-        .and()
-        .commit();
+        .forEach(
+          r -> {
+            if (r.getError() > 0) {
+              LOG.error("Error committing new task {} to zk {}", task.getTaskId(), r);
+              hasErr.set(true);
+            }
+          }
+        );
 
       // Not checking isActive here, already called within offer check flow
       leaderCache.putActiveTask(task.getTaskId());
       taskCache.set(path, task);
+      if (hasErr.get()) {
+        // Rare case, but in case the transaction failed to write task data, try again from the cache data
+        tryRepairTask(task.getTaskId());
+      }
     } catch (KeeperException.NodeExistsException nee) {
       LOG.error("Task or active path already existed for {}", task.getTaskId());
     }

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

@@ -515,7 +515,7 @@ private StatusUpdateResult unsafeProcessStatusUpdate(
       return StatusUpdateResult.IGNORED;
     }
 
-    final Optional<SingularityTask> task = taskManager.getTask(taskIdObj);
+    Optional<SingularityTask> task = taskManager.getTask(taskIdObj);
 
     if (status.getState() == TaskState.TASK_LOST) {
       boolean isMesosFailure =
@@ -542,6 +542,9 @@ private StatusUpdateResult unsafeProcessStatusUpdate(
     }
 
     if (!taskState.isDone()) {
+      if (!task.isPresent()) {
+        task = taskManager.tryRepairTask(taskIdObj);
+      }
       if (task.isPresent()) {
         final Optional<SingularityPendingDeploy> pendingDeploy = deployManager.getPendingDeploy(
           taskIdObj.getRequestId()

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

@@ -1495,9 +1495,33 @@ private SingularityDeployResult enqueueAndProcessLbRequest(
       toShutDown = otherActiveTasks;
     }
 
-    final Map<SingularityTaskId, SingularityTask> tasks = taskManager.getTasks(
-      Iterables.concat(deployActiveTasks, toShutDown)
-    );
+    Iterable<SingularityTaskId> allIds = Iterables.concat(deployActiveTasks, toShutDown);
+    final Map<SingularityTaskId, SingularityTask> tasks = taskManager.getTasks(allIds);
+    // Handle a case where occasionally a task's data cannot be found in zk, but the task is running
+    if (
+      !tasks.keySet().containsAll(deployActiveTasks) ||
+      !tasks.keySet().containsAll(toShutDown)
+    ) {
+      LOG.warn("Could not find task data for task to shut down, will check taskCache");
+      for (SingularityTaskId taskId : deployActiveTasks) {
+        if (!tasks.containsKey(taskId)) {
+          Optional<SingularityTask> maybeRepaired = taskManager.tryRepairTask(taskId);
+          if (maybeRepaired.isPresent()) {
+            LOG.info("Repaired and fetched definition for {}", taskId);
+            tasks.put(taskId, maybeRepaired.get());
+          }
+        }
+      }
+      for (SingularityTaskId taskId : toShutDown) {
+        if (!tasks.containsKey(taskId)) {
+          Optional<SingularityTask> maybeRepaired = taskManager.tryRepairTask(taskId);
+          if (maybeRepaired.isPresent()) {
+            LOG.info("Repaired and fetched definition for {}", taskId);
+            tasks.put(taskId, maybeRepaired.get());
+          }
+        }
+      }
+    }
     final LoadBalancerRequestId lbRequestId = SingularityDeployCheckHelper.getNewLoadBalancerRequestId(
       pendingDeploy
     );