<commit_before>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();
                
                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);
                        }

                        // 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()) {
                    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);<commit_msg>fix bug where number of assigned workers can be greater than number of configured workers<commit_after>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);
{'fraction_non_alphanumeric': 0.09793611147627186, 'fraction_numerical': 0.0010516629420270803, 'mean_word_length': 1.9216589861751152, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3861616', 'n_tokens_mistral': 1991, 'n_tokens_neox': 1920, 'n_words': 416}