package com.hubspot.singularity.scheduler;

import java.util.Iterator;
import java.util.List;
import java.util.Set;

import org.apache.mesos.Protos.Status;
import org.apache.mesos.Protos.TaskStatus;
import org.apache.mesos.SchedulerDriver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Sets;
import com.google.inject.Inject;
import com.hubspot.mesos.JavaUtils;
import com.hubspot.singularity.SingularityTaskId;
import com.hubspot.singularity.data.TaskManager;

public class SingularityTaskReconciliation {

  private static final Logger LOG = LoggerFactory.getLogger(SingularityTaskReconciliation.class);

  private final TaskManager taskManager;

  @Inject
  public SingularityTaskReconciliation(TaskManager taskManager) {
    this.taskManager = taskManager;
  }

  public void reconcileTasks(SchedulerDriver driver) {
    final long start = System.currentTimeMillis();

    Set<SingularityTaskId> taskIds = Sets.newHashSet(taskManager.getActiveTaskIds());
    List<TaskStatus> taskStatuses = taskManager.getLastActiveTaskStatuses();

    for (Iterator<TaskStatus> taskStatusItr = taskStatuses.iterator(); taskStatusItr.hasNext();) {
      TaskStatus taskStatus = taskStatusItr.next();
      SingularityTaskId taskId = SingularityTaskId.fromString(taskStatus.getTaskId().getValue());

      if (!taskIds.contains(taskId)) {
        LOG.info("Task {} not found, deleting taskStatus", taskId);
        taskManager.deleteLastActiveTaskStatus(taskId);
        taskStatusItr.remove();
      }
    }

    LOG.info("Requesting reconciliation for {} taskStatuses", taskStatuses.size());

    Status status = driver.reconcileTasks(taskStatuses);

    LOG.info("Requested reconciliation of {} taskStatuses with driver status {} in {}", taskStatuses.size(), status, JavaUtils.duration(start));
  }

}

@@ -141,32 +141,113 @@
 .Protos.Status;%0A
+import org.apache.mesos.Protos.TaskID;%0Aimport org.apache.mesos.Protos.TaskState;%0A
 import org.apach
@@ -265,32 +265,32 @@
 tos.TaskStatus;%0A
-
 import org.apach
@@ -372,16 +372,56 @@
 ctory;%0A%0A
+import com.google.common.base.Optional;%0A
 import c
@@ -520,16 +520,206 @@
 aUtils;%0A
+import com.hubspot.singularity.SingularityTask;%0Aimport com.hubspot.singularity.SingularityTaskHistoryUpdate;%0Aimport com.hubspot.singularity.SingularityTaskHistoryUpdate.SimplifiedTaskState;%0A
 import c
@@ -1834,15 +1834,1167 @@
 ();%0A
-
       %7D
+ else %7B%0A        taskIds.remove(taskId);%0A      %7D%0A    %7D%0A%0A    // didn't find a taskStatus -- this is temporary for migration%0A    for (SingularityTaskId taskId : taskIds) %7B%0A      Optional%3CSingularityTask%3E task = taskManager.getActiveTask(taskId.getId());%0A%0A      if (!task.isPresent()) %7B%0A        LOG.info(%22Task %7B%7D didn't have an active task%22, taskId);%0A        continue;%0A      %7D%0A%0A      List%3CSingularityTaskHistoryUpdate%3E updates = taskManager.getTaskHistoryUpdates(taskId);%0A%0A      if (updates.isEmpty()) %7B%0A        LOG.info(%22Task %7B%7D didn't have updates%22, taskId);%0A        continue;%0A      %7D%0A%0A      SimplifiedTaskState simplifiedTaskState = SingularityTaskHistoryUpdate.getCurrentState(updates);%0A%0A      TaskState taskState = TaskState.TASK_RUNNING;%0A%0A      if (simplifiedTaskState == SimplifiedTaskState.UNKNOWN) %7B%0A        taskState = TaskState.TASK_STARTING;%0A      %7D%0A%0A      LOG.info(%22Adding a task status %7B%7D for %7B%7D%22, taskState, taskId);%0A%0A      taskStatuses.add(TaskStatus.newBuilder()%0A          .setSlaveId(task.get().getOffer().getSlaveId())%0A          .setState(taskState)%0A          .setTaskId(TaskID.newBuilder().setValue(taskId.getId()))%0A          .build());
 %0A   

add task status taskIds we don't have taskStatus