Add option for additional debug verification of task writes

Author: ssalinas

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

@@ -458,6 +458,7 @@ public class SingularityConfiguration extends Configuration {
   private boolean skipPersistingTooLongTaskIds = false;
 
   private boolean allowEmptyRequestInstances = false;
+  private boolean verifyTaskDataWrites = false;
 
   public long getAskDriverToKillTasksAgainAfterMillis() {
     return askDriverToKillTasksAgainAfterMillis;
@@ -2151,4 +2152,12 @@ public boolean allowEmptyRequestInstances() {
   public void setAllowEmptyRequestInstances(boolean allowEmptyRequestInstances) {
     this.allowEmptyRequestInstances = allowEmptyRequestInstances;
   }
+
+  public boolean isVerifyTaskDataWrites() {
+    return verifyTaskDataWrites;
+  }
+
+  public void setVerifyTaskDataWrites(boolean verifyTaskDataWrites) {
+    this.verifyTaskDataWrites = verifyTaskDataWrites;
+  }
 }

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

@@ -1402,32 +1402,22 @@ private void createTaskAndDeletePendingTaskPrivate(SingularityTask task)
     AtomicBoolean hasErr = new AtomicBoolean(false);
     try {
       final String path = getTaskPath(task.getTaskId());
-      curator
-        .transaction()
-        .forOperations(
-          curator.transactionOp().create().forPath(path, taskTranscoder.toBytes(task)),
-          curator
-            .transactionOp()
-            .create()
-            .forPath(
-              getLastActiveTaskStatusPath(task.getTaskId()),
-              taskStatusTranscoder.toBytes(taskStatusHolder)
-            )
-        )
-        .forEach(
-          r -> {
-            if (r.getError() > 0) {
-              LOG.error("Error committing new task {} to zk {}", task.getTaskId(), r);
-              hasErr.set(true);
-            }
-          }
-        );
-
+      saveTaskDeletePendingInTransaction(hasErr, path, task, taskStatusHolder);
       // Not checking isActive here, already called within offer check flow
       leaderCache.putActiveTask(task.getTaskId());
       taskCache.set(path, task);
+      if (configuration.isVerifyTaskDataWrites()) {
+        Optional<SingularityTask> maybeTask = getTaskCheckCache(task.getTaskId(), true);
+        if (!maybeTask.isPresent()) {
+          LOG.error("Found empty task after write for {}", task.getTaskId());
+          saveTaskDeletePendingInTransaction(hasErr, path, task, taskStatusHolder);
+        }
+      }
     } catch (KeeperException.NodeExistsException nee) {
       LOG.error("Task or active path already existed for {}", task.getTaskId());
+    } catch (Exception e) {
+      LOG.error("Could not save task data for {}", task.getTaskId(), e);
+      throw new RuntimeException(e);
     } finally {
       if (hasErr.get()) {
         // Rare case, but in case the transaction failed to write task data, try again from the cache data
@@ -1436,6 +1426,39 @@ private void createTaskAndDeletePendingTaskPrivate(SingularityTask task)
     }
   }
 
+  private void saveTaskDeletePendingInTransaction(
+    AtomicBoolean hasErr,
+    String taskPath,
+    SingularityTask task,
+    SingularityTaskStatusHolder taskStatusHolder
+  )
+    throws Exception {
+    byte[] taskBytes = taskTranscoder.toBytes(task);
+    if (taskBytes == null || taskBytes.length == 0) {
+      LOG.error("Encountered null or empty task bytes for {}", task.getTaskId());
+    }
+    curator
+      .transaction()
+      .forOperations(
+        curator.transactionOp().create().forPath(taskPath, taskBytes),
+        curator
+          .transactionOp()
+          .create()
+          .forPath(
+            getLastActiveTaskStatusPath(task.getTaskId()),
+            taskStatusTranscoder.toBytes(taskStatusHolder)
+          )
+      )
+      .forEach(
+        r -> {
+          if (r.getError() > 0) {
+            LOG.error("Error committing new task {} to zk {}", task.getTaskId(), r);
+            hasErr.set(true);
+          }
+        }
+      );
+  }
+
   public List<SingularityTaskId> getLBCleanupTasks() {
     return getChildrenAsIds(LB_CLEANUP_PATH_ROOT, taskIdTranscoder);
   }