package com.hubspot.singularity.data;

import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.inject.Inject;
import com.hubspot.singularity.SingularityDeploy;
import com.hubspot.singularity.SingularityDeployKey;
import com.hubspot.singularity.SingularityPendingTask;
import com.hubspot.singularity.SingularityRequestWithState;
import com.hubspot.singularity.SingularityTaskRequest;

public class TaskRequestManager {

  private final static Logger LOG = LoggerFactory.getLogger(TaskRequestManager.class);
  
  private final DeployManager deployManager;
  private final RequestManager requestManager;
  
  @Inject
  public TaskRequestManager(DeployManager deployManager, RequestManager requestManager) {
    this.deployManager = deployManager;
    this.requestManager = requestManager;
  }

  public List<SingularityTaskRequest> getTaskRequests(List<SingularityPendingTask> tasks) {
    final Map<String, SingularityPendingTask> requestIdToPendingTaskId = Maps.newHashMapWithExpectedSize(tasks.size());
    
    for (SingularityPendingTask task : tasks) {
      requestIdToPendingTaskId.put(task.getPendingTaskId().getRequestId(), task);
    }
    
    final List<SingularityRequestWithState> matchingRequests = requestManager.getRequests(requestIdToPendingTaskId.keySet());
    
    final Map<SingularityPendingTask, SingularityDeployKey> deployKeys = SingularityDeployKey.fromPendingTasks(requestIdToPendingTaskId.values());
    final Map<SingularityDeployKey, SingularityDeploy> matchingDeploys = deployManager.getDeploysForKeys(deployKeys.values());
    
    final List<SingularityTaskRequest> taskRequests = Lists.newArrayListWithCapacity(matchingRequests.size());
    
    for (SingularityRequestWithState request : matchingRequests) {
      SingularityPendingTask task = requestIdToPendingTaskId.get(request.getRequest().getId());
    
      SingularityDeploy foundDeploy = matchingDeploys.get(deployKeys.get(task));
      
      if (foundDeploy == null) {
        LOG.warn("Couldn't find a matching deploy for pending task {}", task);
        continue;
      }
      
      if (!request.getState().isRunnable()) {
        LOG.warn("Request was in state {} for pending task {}", request.getState(), task);
        continue;
      }
      
      taskRequests.add(new SingularityTaskRequest(request.getRequest(), foundDeploy, task));
    }
    
    return taskRequests;
  }
  
}

package com.hubspot.singularity.data;

import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;
import com.google.inject.Inject;
import com.hubspot.singularity.SingularityDeploy;
import com.hubspot.singularity.SingularityDeployKey;
import com.hubspot.singularity.SingularityPendingTask;
import com.hubspot.singularity.SingularityRequestWithState;
import com.hubspot.singularity.SingularityTaskRequest;

public class TaskRequestManager {

  private final static Logger LOG = LoggerFactory.getLogger(TaskRequestManager.class);
  
  private final DeployManager deployManager;
  private final RequestManager requestManager;
  
  @Inject
  public TaskRequestManager(DeployManager deployManager, RequestManager requestManager) {
    this.deployManager = deployManager;
    this.requestManager = requestManager;
  }

  public List<SingularityTaskRequest> getTaskRequests(List<SingularityPendingTask> tasks) {
    final Multimap<String, SingularityPendingTask> requestIdToPendingTaskId = ArrayListMultimap.create(tasks.size(), 1);
    
    for (SingularityPendingTask task : tasks) {
      requestIdToPendingTaskId.put(task.getPendingTaskId().getRequestId(), task);
    }
    
    final List<SingularityRequestWithState> matchingRequests = requestManager.getRequests(requestIdToPendingTaskId.keySet());
    
    final Map<SingularityPendingTask, SingularityDeployKey> deployKeys = SingularityDeployKey.fromPendingTasks(requestIdToPendingTaskId.values());
    final Map<SingularityDeployKey, SingularityDeploy> matchingDeploys = deployManager.getDeploysForKeys(Sets.newHashSet(deployKeys.values()));
    
    final List<SingularityTaskRequest> taskRequests = Lists.newArrayListWithCapacity(matchingRequests.size());
    
    for (SingularityRequestWithState request : matchingRequests) {
      for (SingularityPendingTask task : requestIdToPendingTaskId.get(request.getRequest().getId())) {
        SingularityDeploy foundDeploy = matchingDeploys.get(deployKeys.get(task));
        
        if (foundDeploy == null) {
          LOG.warn("Couldn't find a matching deploy for pending task {}", task);
          continue;
        }
        
        if (!request.getState().isRunnable()) {
          LOG.warn("Request was in state {} for pending task {}", request.getState(), task);
          continue;
        }
        
        taskRequests.add(new SingularityTaskRequest(request.getRequest(), foundDeploy, task));
      }
    }
    
    return taskRequests;
  }
  
}

fix bug with multiple instances