package storm.mesos;

import backtype.storm.scheduler.Cluster;
import backtype.storm.scheduler.ExecutorDetails;
import backtype.storm.scheduler.IScheduler;
import backtype.storm.scheduler.SchedulerAssignment;
import backtype.storm.scheduler.Topologies;
import backtype.storm.scheduler.TopologyDetails;
import backtype.storm.scheduler.WorkerSlot;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.log4j.Logger;

public class MesosNimbusScheduler implements IScheduler {
    public static final Logger LOG = Logger.getLogger(MesosNimbusScheduler.class);

    
    private Map<WorkerSlot, List<ExecutorDetails>> getSlotAllocations(SchedulerAssignment a) {
        if(a==null) return new HashMap();
        Map<ExecutorDetails, WorkerSlot> em =  a.getExecutorToSlot();
        Map<WorkerSlot, List<ExecutorDetails>> ret = new HashMap();
        for(ExecutorDetails e: em.keySet()) {
            WorkerSlot s = em.get(e);
            if(!ret.containsKey(s)) ret.put(s, new ArrayList());
            ret.get(s).add(e);
        }
        return ret;
    }
    
    @Override
    public void schedule(Topologies topologies, Cluster cluster) {
        for(TopologyDetails details: topologies.getTopologies()) {
            List<WorkerSlot> availableSlots = cluster.getAvailableSlots();
            
            
            List<WorkerSlot> myAvailable = new ArrayList<WorkerSlot>();
            for(WorkerSlot s: availableSlots) {
                if(MesosNimbus.getTopologyIdFromNodeId(s.getNodeId()).equals(details.getId())) {
                    myAvailable.add(s);
                }
            }

            if(!myAvailable.isEmpty()) {
                SchedulerAssignment a = cluster.getAssignmentById(details.getId());

                Map<WorkerSlot, List<ExecutorDetails>> slotAllocations = getSlotAllocations(a);
                
                
                int avg = details.getExecutors().size() / details.getNumWorkers();
                
                int keptWorkers = slotAllocations.size();
                if(slotAllocations.size() < details.getNumWorkers() && myAvailable.size() > slotAllocations.size()) {
                    for(WorkerSlot s: slotAllocations.keySet()) {
                        //TODO: would be more precise to check the distribution.. can end up with unbalanced topology
                        // this way if it shrinks and then grows to numWorkers
                        int numE = slotAllocations.get(s).size();
                        if(numE != avg && numE != avg + 1) {
                            cluster.freeSlot(s);
                            keptWorkers--;
                        }

                        // can't do this because it may not be a recognized slot by mesos anymore
                        // myAvailable.add(s);
                    }
                }
                
                
                //keep at most numWorkers
                while(myAvailable.size() > details.getNumWorkers() - keptWorkers) {
                    myAvailable.remove(myAvailable.size()-1);
                }
                Map<WorkerSlot, List<ExecutorDetails>> toAssign = new HashMap();
                for(WorkerSlot s: myAvailable) {
                    toAssign.put(s, new ArrayList());
                }
                int i=0;
                Map<String, List<ExecutorDetails>> toSchedule = cluster.getNeedsSchedulingComponentToExecutors(details);

                for(String c: toSchedule.keySet()) {
                    List<ExecutorDetails> executors = toSchedule.get(c);
                    for(ExecutorDetails e: executors) {
                        WorkerSlot s = myAvailable.get(i);
                        toAssign.get(s).add(e);
                        i = (i + 1) % myAvailable.size();
                    }
                }
                for(WorkerSlot s: toAssign.keySet()) {
                    List<ExecutorDetails> executors = toAssign.get(s);
                    if(executors.size()>0) {
                        cluster.assign(s, details.getId(), executors);                        
                    }
                }
            }
        }
    }
    
}

@@ -350,27 +350,34 @@
 ort 
-java.util.ArrayList
+backtype.storm.utils.Utils
 ;%0Aim
@@ -391,23 +391,25 @@
 va.util.
-HashMap
+ArrayList
 ;%0Aimport
@@ -423,19 +423,19 @@
 til.Hash
-Set
+Map
 ;%0Aimport
@@ -477,30 +477,8 @@
 ap;%0A
-import java.util.Set;%0A
 impo
@@ -1964,18 +1964,78 @@
     
-int avg = 
+Map%3CInteger, Integer%3E distribution = new HashMap(Utils.integerDivided(
 deta
@@ -2059,18 +2059,17 @@
 ).size()
- /
+,
  details
@@ -2084,17 +2084,35 @@
 orkers()
-;
+));                
 %0A       
@@ -2221,17 +2221,18 @@
 .size() 
-%3C
+!=
  details
@@ -2250,17 +2250,42 @@
 ers() &&
- 
+%0A                        (
 myAvaila
@@ -2319,16 +2319,91 @@
 s.size()
+%0A                         %7C%7C myAvailable.size() %3E= details.getNumWorkers())
 ) %7B%0A    
@@ -2492,305 +2492,254 @@
     
-//TODO: would be more precise to check the distribution.. can end up with unbalanced topology%0A                        // this way if it shrinks and then grows to numWorkers%0A                        int numE = slotAllocations.get(s).size();%0A                        if(numE != avg && numE != avg + 1)
+int numE = slotAllocations.get(s).size();%0A                        Integer amt = distribution.get(numE);%0A                        if(amt!=null && amt %3E 0) %7B%0A                            distribution.put(numE, amt - 1);%0A                        %7D else
  %7B%0A 
@@ -2828,16 +2828,44 @@
 rkers--;
+                            
 %0A       
@@ -3187,16 +3187,28 @@
 size() %3E
+ Math.max(1,
  details
@@ -3238,16 +3238,17 @@
 Workers)
+)
  %7B%0A     

fix how scheduler decides how many workers each topology gets