<commit_before>package org.radargun;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.radargun.config.MasterConfig;
import org.radargun.state.MasterState;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * This is the master that will coordinate the {@link Slave}s in order to run the benchmark.
 *
 * @author Mircea.<EMAIL>
 */
public class Master {

   private static Log log = LogFactory.getLog(Master.class);
   public static final int DEFAULT_PORT = 2103;

   MasterConfig masterConfig;

   private ServerSocketChannel serverSocketChannel;
   private SocketChannel[] slaves;

   private Map<SocketChannel, ByteBuffer> writeBufferMap = new HashMap<SocketChannel, ByteBuffer>();
   private Map<SocketChannel, ByteBuffer> readBufferMap = new HashMap<SocketChannel, ByteBuffer>();
   private List<DistStageAck> responses = new ArrayList<DistStageAck>();
   private Selector communicationSelector;
   private Selector discoverySelector;
   private Map<SocketChannel, Integer> slave2Index = new HashMap<SocketChannel, Integer>();
   private MasterState state;
   int processedSlaves = 0;
   private static final int DEFAULT_READ_BUFF_CAPACITY = 1024;

   public Master(MasterConfig masterConfig) {
      this.masterConfig = masterConfig;
      state = new MasterState(masterConfig);
      slaves = new SocketChannel[masterConfig.getSlaveCount()];
      try {
         communicationSelector = Selector.open();
      } catch (IOException e) {
         throw new IllegalStateException(e);
      }
      Runtime.getRuntime().addShutdownHook(new ShutDownHook("Master process"));
   }

   public void start() throws Exception {
      try {
         startServerSocket();
         runDiscovery();
         prepareNextStage();
         startCommunicationWithSlaves();
      } catch (Exception e) {
         log.error(e.toString());
      } finally {
         releaseResources();
      }
   }

   private void prepareNextStage() throws Exception {
      DistStage toExecute = state.getNextDistStageToProcess();
      if (toExecute == null) {
         releaseResourcesAndExit();
      } else {
         runDistStage(toExecute, toExecute.getActiveSlaveCount());
      }
   }

   private void runDistStage(DistStage currentStage, int noSlaves) throws Exception {
      writeBufferMap.clear();
      DistStage toSerialize;
      for (int i = 0; i < noSlaves; i++) {
         SocketChannel slave = slaves[i];
         slave.configureBlocking(false);
         slave.register(communicationSelector, SelectionKey.OP_WRITE);
         toSerialize = currentStage.clone();
         toSerialize.initOnMaster(state, i);
         if (i == 0) {//only log this once
            if (log.isDebugEnabled())
               log.debug("Starting '" + toSerialize.getClass().getSimpleName() + "' on " + toSerialize.getActiveSlaveCount() + " slave nodes. Details: " + toSerialize);
            else
               log.info("Starting '" + toSerialize.getClass().getSimpleName() + "' on " + toSerialize.getActiveSlaveCount() + " slave nodes.");
         }
         byte[] bytes = SerializationHelper.prepareForSerialization(toSerialize);
         writeBufferMap.put(slave, ByteBuffer.wrap(bytes));
      }
   }

   private void releaseResources() {
      try {
         discoverySelector.close();
      } catch (Throwable e) {
         log.warn(e);
      }
      try {
         communicationSelector.close();
      } catch (Throwable e) {
         log.warn(e);
      }
      for (SocketChannel sc : slaves) {
         try {
            sc.socket().close();
         } catch (Throwable e) {
            log.warn(e);
         }
      }

      try {
         if (serverSocketChannel != null) serverSocketChannel.socket().close();
      } catch (Throwable e) {
         log.warn(e);
      }
   }

   private void runDiscovery() throws IOException {
      discoverySelector = Selector.open();
      serverSocketChannel.register(discoverySelector, SelectionKey.OP_ACCEPT);
      int slaveCount = 0;
      while (slaveCount < masterConfig.getSlaveCount()) {
         log.info("Awaiting registration from " + (masterConfig.getSlaveCount() - slaveCount) + " slaves.");
         discoverySelector.select();
         Set<SelectionKey> keySet = discoverySelector.selectedKeys();
         Iterator<SelectionKey> it = keySet.iterator();
         while (it.hasNext()) {
            SelectionKey selectionKey = it.next();
            it.remove();
            if (!selectionKey.isValid()) {
               continue;
            }
            ServerSocketChannel srvSocketChannel = (ServerSocketChannel) selectionKey.channel();
            SocketChannel socketChannel = srvSocketChannel.accept();            
            
            int slaveIndex = readIndex(socketChannel);
            if (slaveIndex >= slaves.length || slaves[slaveIndex] != null) {
               throw new IllegalArgumentException("Slave requests invalid slaveIndex " + slaveIndex);
            } else if (slaveIndex < 0) {
               for (int i = 0; i < slaves.length; ++i) {
                  if (slaves[i] == null) {
                     slaveIndex = i;
                     break;
                  }
               }
            }
            slaves[slaveIndex] = socketChannel;
            slaveCount++;
            slave2Index.put(socketChannel, slaveIndex);
            this.readBufferMap.put(socketChannel, ByteBuffer.allocate(DEFAULT_READ_BUFF_CAPACITY));
            if (log.isTraceEnabled())
               log.trace("Added new slave connection from: " + socketChannel.socket().getInetAddress());
         }
      }
      log.info("Connection established from " + slaveCount + " slaves.");
   }

   private int readIndex(SocketChannel socketChannel) throws IOException {      
      ByteBuffer slaveIndexBuffer = ByteBuffer.allocate(4);
      while (slaveIndexBuffer.hasRemaining()) {
         socketChannel.read(slaveIndexBuffer);      
      }
      slaveIndexBuffer.flip();
      return slaveIndexBuffer.getInt();
   }

   private void startCommunicationWithSlaves() throws Exception {
      while (true) {
         communicationSelector.select();
         Set<SelectionKey> keys = communicationSelector.selectedKeys();
         if (keys.size() > 0) {
            Iterator<SelectionKey> keysIt = keys.iterator();
            while (keysIt.hasNext()) {
               SelectionKey key = keysIt.next();
               keysIt.remove();
               if (!key.isValid()) {<commit_msg>Dumb bug fix<commit_after>package org.radargun;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.radargun.config.MasterConfig;
import org.radargun.state.MasterState;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * This is the master that will coordinate the {@link Slave}s in order to run the benchmark.
 *
 * @author <EMAIL>
 */
public class Master {

   private static Log log = LogFactory.getLog(Master.class);
   public static final int DEFAULT_PORT = 2103;

   MasterConfig masterConfig;

   private ServerSocketChannel serverSocketChannel;
   private SocketChannel[] slaves;

   private Map<SocketChannel, ByteBuffer> writeBufferMap = new HashMap<SocketChannel, ByteBuffer>();
   private Map<SocketChannel, ByteBuffer> readBufferMap = new HashMap<SocketChannel, ByteBuffer>();
   private List<DistStageAck> responses = new ArrayList<DistStageAck>();
   private Selector communicationSelector;
   private Selector discoverySelector;
   private Map<SocketChannel, Integer> slave2Index = new HashMap<SocketChannel, Integer>();
   private MasterState state;
   int processedSlaves = 0;
   private static final int DEFAULT_READ_BUFF_CAPACITY = 1024;

   public Master(MasterConfig masterConfig) {
      this.masterConfig = masterConfig;
      state = new MasterState(masterConfig);
      slaves = new SocketChannel[masterConfig.getSlaveCount()];
      try {
         communicationSelector = Selector.open();
      } catch (IOException e) {
         throw new IllegalStateException(e);
      }
      Runtime.getRuntime().addShutdownHook(new ShutDownHook("Master process"));
   }

   public void start() throws Exception {
      try {
         startServerSocket();
         runDiscovery();
         prepareNextStage();
         startCommunicationWithSlaves();
      } catch (Exception e) {
         log.error(e.toString());
      } finally {
         releaseResources();
      }
   }

   private void prepareNextStage() throws Exception {
      DistStage toExecute = state.getNextDistStageToProcess();
      if (toExecute == null) {
         releaseResourcesAndExit();
      } else {
         runDistStage(toExecute, toExecute.getActiveSlaveCount());
      }
   }

   private void runDistStage(DistStage currentStage, int noSlaves) throws Exception {
      writeBufferMap.clear();
      DistStage toSerialize;
      for (int i = 0; i < noSlaves; i++) {
         SocketChannel slave = slaves[i];
         slave.configureBlocking(false);
         slave.register(communicationSelector, SelectionKey.OP_WRITE);
         toSerialize = currentStage.clone();
         toSerialize.initOnMaster(state, i);
         if (i == 0) {//only log this once
            if (log.isDebugEnabled())
               log.debug("Starting '" + toSerialize.getClass().getSimpleName() + "' on " + toSerialize.getActiveSlaveCount() + " slave nodes. Details: " + toSerialize);
            else
               log.info("Starting '" + toSerialize.getClass().getSimpleName() + "' on " + toSerialize.getActiveSlaveCount() + " slave nodes.");
         }
         byte[] bytes = SerializationHelper.prepareForSerialization(toSerialize);
         writeBufferMap.put(slave, ByteBuffer.wrap(bytes));
      }
   }

   private void releaseResources() {
      try {
         discoverySelector.close();
      } catch (Throwable e) {
         log.warn(e);
      }
      try {
         communicationSelector.close();
      } catch (Throwable e) {
         log.warn(e);
      }
      for (SocketChannel sc : slaves) {
         try {
            sc.socket().close();
         } catch (Throwable e) {
            log.warn(e);
         }
      }

      try {
         if (serverSocketChannel != null) serverSocketChannel.socket().close();
      } catch (Throwable e) {
         log.warn(e);
      }
   }

   private void runDiscovery() throws IOException {
      discoverySelector = Selector.open();
      serverSocketChannel.register(discoverySelector, SelectionKey.OP_ACCEPT);
      int slaveCount = 0;
      while (slaveCount < masterConfig.getSlaveCount()) {
         log.info("Awaiting registration from " + (masterConfig.getSlaveCount() - slaveCount) + " slaves.");
         discoverySelector.select();
         Set<SelectionKey> keySet = discoverySelector.selectedKeys();
         Iterator<SelectionKey> it = keySet.iterator();
         while (it.hasNext()) {
            SelectionKey selectionKey = it.next();
            it.remove();
            if (!selectionKey.isValid()) {
               continue;
            }
            ServerSocketChannel srvSocketChannel = (ServerSocketChannel) selectionKey.channel();
            SocketChannel socketChannel = srvSocketChannel.accept();            
            
            int slaveIndex = readIndex(socketChannel);
            if (slaveIndex < 0) {
               for (int i = 0; i < slaves.length; ++i) {
                  if (slaves[i] == null) {
                     slaveIndex = i;
                     break;
                  }
               }
            } else if (slaveIndex >= slaves.length || slaves[slaveIndex] != null) {
               throw new IllegalArgumentException("Slave requests invalid slaveIndex " + slaveIndex);            
            }
            slaves[slaveIndex] = socketChannel;
            slaveCount++;
            slave2Index.put(socketChannel, slaveIndex);
            this.readBufferMap.put(socketChannel, ByteBuffer.allocate(DEFAULT_READ_BUFF_CAPACITY));
            if (log.isTraceEnabled())
               log.trace("Added new slave connection from: " + socketChannel.socket().getInetAddress());
         }
      }
      log.info("Connection established from " + slaveCount + " slaves.");
   }

   private int readIndex(SocketChannel socketChannel) throws IOException {      
      ByteBuffer slaveIndexBuffer = ByteBuffer.allocate(4);
      while (slaveIndexBuffer.hasRemaining()) {
         socketChannel.read(slaveIndexBuffer);      
      }
      slaveIndexBuffer.flip();
      return slaveIndexBuffer.getInt();
   }

   private void startCommunicationWithSlaves() throws Exception {
      while (true) {
         communicationSelector.select();
         Set<SelectionKey> keys = communicationSelector.selectedKeys();
         if (keys.size() > 0) {
            Iterator<SelectionKey> keysIt = keys.iterator();
            while (keysIt.hasNext()) {
               SelectionKey key = keysIt.next();
               keysIt.remove();
               if (!key.isValid()) {
{'fraction_non_alphanumeric': 0.10618496262641067, 'fraction_numerical': 0.002638135717426352, 'mean_word_length': 2.5706436420722136, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 33, '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': '7479972', 'n_tokens_mistral': 3590, 'n_tokens_neox': 3533, 'n_words': 804}