<commit_before>package spark

import java.io._
import java.net._
import java.util.{BitSet, Random, Timer, TimerTask, UUID}
import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.{Executors, ThreadPoolExecutor, ThreadFactory}

import scala.collection.mutable.{ArrayBuffer, HashMap}

/**
 * A simple implementation of shuffle using local files served through HTTP.
 *
 * TODO: Add support for compression when spark.compress is set to true.
 */
@serializable
class LocalFileShuffle[K, V, C] extends Shuffle[K, V, C] with Logging {
  @transient var totalSplits = 0
  @transient var hasSplits = 0 
  @transient var hasSplitsBitVector: BitSet = null
  @transient var combiners: HashMap[K,C] = null

  override def compute(input: RDD[(K, V)],
                       numOutputSplits: Int,
                       createCombiner: V => C,
                       mergeValue: (C, V) => C,
                       mergeCombiners: (C, C) => C)
  : RDD[(K, C)] =
  {
    val sc = input.sparkContext
    val shuffleId = LocalFileShuffle.newShuffleId()
    logInfo("Shuffle ID: " + shuffleId)

    val splitRdd = new NumberedSplitRDD(input)
    val numInputSplits = splitRdd.splits.size

    // Run a parallel map and collect to write the intermediate data files,
    // returning a list of inputSplitId -> serverUri pairs
    val outputLocs = splitRdd.map((pair: (Int, Iterator[(K, V)])) => {
      val myIndex = pair._1
      val myIterator = pair._2
      val buckets = Array.tabulate(numOutputSplits)(_ => new HashMap[K, C])
      for ((k, v) <- myIterator) {
        var bucketId = k.hashCode % numOutputSplits
        if (bucketId < 0) { // Fix bucket ID if hash code was negative
          bucketId += numOutputSplits
        }
        val bucket = buckets(bucketId)
        bucket(k) = bucket.get(k) match {
          case Some(c) => mergeValue(c, v)
          case None => createCombiner(v)
        }
      }
      
      for (i <- 0 until numOutputSplits) {
        val file = LocalFileShuffle.getOutputFile(shuffleId, myIndex, i)
        val writeStartTime = System.currentTimeMillis
        logInfo ("BEGIN WRITE: " + file)
        val out = new ObjectOutputStream(new FileOutputStream(file))
        buckets(i).foreach(pair => out.writeObject(pair))
        out.close()
        logInfo ("END WRITE: " + file)
        val writeTime = (System.currentTimeMillis - writeStartTime)
        logInfo ("Writing " + file + " of size " + file.length + " bytes took " + writeTime + " millis.")
      }
      (myIndex, LocalFileShuffle.serverAddress, LocalFileShuffle.serverPort)
    }).collect()

    val splitsByUri = new ArrayBuffer[(String, Int, Int)]
    for ((inputId, serverAddress, serverPort) <- outputLocs) {
      splitsByUri += ((serverAddress, serverPort, inputId))
    }

    // TODO: Could broadcast splitsByUri

    // Return an RDD that does each of the merges for a given partition
    val indexes = sc.parallelize(0 until numOutputSplits, numOutputSplits)
    return indexes.flatMap((myId: Int) => {
      totalSplits = splitsByUri.size
      hasSplitsBitVector = new BitSet (totalSplits)
      combiners = new HashMap[K, C]
      
      while (hasSplits < totalSplits) {
        // Select a random split to pull
        val splitIndex = selectRandomSplit
        val (serverAddress, serverPort, inputId) = 
          splitsByUri (splitIndex)

        val requestPath = "%d/%d/%d".format(shuffleId, inputId, myId)
        
        val shuffleClient = 
          new ShuffleClient(serverAddress, serverPort, requestPath)
        val readStartTime = System.currentTimeMillis
        logInfo ("BEGIN READ: " + requestPath)
        shuffleClient.start
        shuffleClient.join

        val inputStream = new ObjectInputStream (
          new ByteArrayInputStream(shuffleClient.byteArray))
        try {
          while (true) {
            val (k, c) = inputStream.readObject().asInstanceOf[(K, C)]
            combiners(k) = combiners.get(k) match {
              case Some(oldC) => mergeCombiners(oldC, c)
              case None => c
            }
          }
        } catch {
          case e: EOFException => {}
        }
        inputStream.close
        
        hasSplits += 1
        hasSplitsBitVector.synchronized {
          hasSplitsBitVector.set (splitIndex)
        }
        
        logInfo ("END READ: " + requestPath)
        val readTime = (System.currentTimeMillis - readStartTime)
        logInfo ("Reading " + requestPath + " took " + readTime + " millis.")
      }
      combiners
    })
  }
  
  def selectRandomSplit: Int = {
    var requiredSplits = new ArrayBuffer[Int]
    
    hasSplitsBitVector.synchronized {
      for (i <- 0 until totalSplits) {
        if (!hasSplitsBitVector.get(i)) {
          requiredSplits += i
        }
      }
    }
    
    if (requiredSplits.size > 0) {
      requiredSplits(LocalFileShuffle.ranGen.nextInt (requiredSplits.size))
    } else {
      -1
    }
  }
  
  class ShuffleClient (hostAddress: String, listenPort: Int, requestPath: String)
  extends Thread with Logging {
    private var peerSocketToSource: Socket = null
    private var oosSource: ObjectOutputStream = null
    private var oisSource: ObjectInputStream = null
    
    var byteArray: Array[Byte] = null

    override def run: Unit = {
      // Setup the timeout mechanism
      var timeOutTask = new TimerTask {
        override def run: Unit = {
          cleanUpConnections
        }
      }
      
      var timeOutTimer = new Timer
      // TODO: Set wait timer
      // TODO: If its too small, things FAIL
      timeOutTimer.schedule (timeOutTask, 10000)
      
      logInfo ("ShuffleClient started... => %s:%d#%s".format(hostAddress, listenPort, requestPath))
      
      try {
        // Connect to the source
        peerSocketToSource = new Socket (hostAddress, listenPort)
        oosSource =
          new ObjectOutputStream (peerSocketToSource.getOutputStream)
        oosSource.flush
        var isSource = peerSocketToSource.getInputStream
        oisSource = new ObjectInputStream (isSource)
        
        // Send the request
        oosSource.writeObject(requestPath)
        
        // Receive the length of the requested file
        var requestedFileLen = oisSource.readObject.asInstanceOf[Int]
        logInfo ("Received requestedFileLen = " + requestedFileLen)

        // Turn the timer OFF, if the sender responds before timeout
        timeOutTimer.cancel
        
        // Receive the file
        if (requestedFileLen != -1) {
          byteArray = new Array[Byte] (requestedFileLen)
          var bytesRead = isSource.read (byteArray, 0, byteArray.length)
          var alreadyRead = bytesRead

          while (alreadyRead < requestedFileLen) {
            bytesRead = isSource.read(byteArray, alreadyRead, 
              (byteArray.length - alreadyRead))
            if(bytesRead > 0) {
              alreadyRead = alreadyRead + bytesRead
            }
          } 
        } else {
          throw new SparkException("ShuffleServer " + hostAddress + " does not have " + requestPath)
        }
      } catch {<commit_msg>combining is happening inside the thread. Its still synchronized though.<commit_after>package spark

import java.io._
import java.net._
import java.util.{BitSet, Random, Timer, TimerTask, UUID}
import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.{Executors, ThreadPoolExecutor, ThreadFactory}

import scala.collection.mutable.{ArrayBuffer, HashMap}

/**
 * A simple implementation of shuffle using local files served through HTTP.
 *
 * TODO: Add support for compression when spark.compress is set to true.
 */
@serializable
class LocalFileShuffle[K, V, C] extends Shuffle[K, V, C] with Logging {
  @transient var totalSplits = 0
  @transient var hasSplits = 0 
  @transient var hasSplitsBitVector: BitSet = null
  @transient var combiners: HashMap[K,C] = null

  override def compute(input: RDD[(K, V)],
                       numOutputSplits: Int,
                       createCombiner: V => C,
                       mergeValue: (C, V) => C,
                       mergeCombiners: (C, C) => C)
  : RDD[(K, C)] =
  {
    val sc = input.sparkContext
    val shuffleId = LocalFileShuffle.newShuffleId()
    logInfo("Shuffle ID: " + shuffleId)

    val splitRdd = new NumberedSplitRDD(input)
    val numInputSplits = splitRdd.splits.size

    // Run a parallel map and collect to write the intermediate data files,
    // returning a list of inputSplitId -> serverUri pairs
    val outputLocs = splitRdd.map((pair: (Int, Iterator[(K, V)])) => {
      val myIndex = pair._1
      val myIterator = pair._2
      val buckets = Array.tabulate(numOutputSplits)(_ => new HashMap[K, C])
      for ((k, v) <- myIterator) {
        var bucketId = k.hashCode % numOutputSplits
        if (bucketId < 0) { // Fix bucket ID if hash code was negative
          bucketId += numOutputSplits
        }
        val bucket = buckets(bucketId)
        bucket(k) = bucket.get(k) match {
          case Some(c) => mergeValue(c, v)
          case None => createCombiner(v)
        }
      }
      
      for (i <- 0 until numOutputSplits) {
        val file = LocalFileShuffle.getOutputFile(shuffleId, myIndex, i)
        val writeStartTime = System.currentTimeMillis
        logInfo ("BEGIN WRITE: " + file)
        val out = new ObjectOutputStream(new FileOutputStream(file))
        buckets(i).foreach(pair => out.writeObject(pair))
        out.close()
        logInfo ("END WRITE: " + file)
        val writeTime = (System.currentTimeMillis - writeStartTime)
        logInfo ("Writing " + file + " of size " + file.length + " bytes took " + writeTime + " millis.")
      }
      (myIndex, LocalFileShuffle.serverAddress, LocalFileShuffle.serverPort)
    }).collect()

    val splitsByUri = new ArrayBuffer[(String, Int, Int)]
    for ((inputId, serverAddress, serverPort) <- outputLocs) {
      splitsByUri += ((serverAddress, serverPort, inputId))
    }

    // TODO: Could broadcast splitsByUri

    // Return an RDD that does each of the merges for a given partition
    val indexes = sc.parallelize(0 until numOutputSplits, numOutputSplits)
    return indexes.flatMap((myId: Int) => {
      totalSplits = splitsByUri.size
      hasSplitsBitVector = new BitSet (totalSplits)
      combiners = new HashMap[K, C]
      
      while (hasSplits < totalSplits) {
        // Select a random split to pull
        val splitIndex = selectRandomSplit
        val (serverAddress, serverPort, inputId) = 
          splitsByUri (splitIndex)

        val requestPath = "%d/%d/%d".format(shuffleId, inputId, myId)
        
        val shuffleClient = new ShuffleClient(serverAddress, serverPort, 
            requestPath, mergeCombiners)
        val readStartTime = System.currentTimeMillis
        logInfo ("BEGIN READ: " + requestPath)
        shuffleClient.start
        shuffleClient.join

        hasSplits += 1
        hasSplitsBitVector.synchronized {
          hasSplitsBitVector.set (splitIndex)
        }
        
        logInfo ("END READ: " + requestPath)
        val readTime = (System.currentTimeMillis - readStartTime)
        logInfo ("Reading " + requestPath + " took " + readTime + " millis.")
      }
      combiners
    })
  }
  
  def selectRandomSplit: Int = {
    var requiredSplits = new ArrayBuffer[Int]
    
    hasSplitsBitVector.synchronized {
      for (i <- 0 until totalSplits) {
        if (!hasSplitsBitVector.get(i)) {
          requiredSplits += i
        }
      }
    }
    
    if (requiredSplits.size > 0) {
      requiredSplits(LocalFileShuffle.ranGen.nextInt (requiredSplits.size))
    } else {
      -1
    }
  }
  
  class ShuffleClient (hostAddress: String, listenPort: Int, 
    requestPath: String, mergeCombiners: (C, C) => C)
  extends Thread with Logging {
    private var peerSocketToSource: Socket = null
    private var oosSource: ObjectOutputStream = null
    private var oisSource: ObjectInputStream = null
    
    private var byteArray: Array[Byte] = null

    override def run: Unit = {
      // Setup the timeout mechanism
      var timeOutTask = new TimerTask {
        override def run: Unit = {
          cleanUpConnections
        }
      }
      
      var timeOutTimer = new Timer
      // TODO: Set wait timer
      // TODO: If its too small, things FAIL
      timeOutTimer.schedule (timeOutTask, 100000)
      
      logInfo ("ShuffleClient started... => %s:%d#%s".format(hostAddress, listenPort, requestPath))
      
      try {
        // Connect to the source
        peerSocketToSource = new Socket (hostAddress, listenPort)
        oosSource =
          new ObjectOutputStream (peerSocketToSource.getOutputStream)
        oosSource.flush
        var isSource = peerSocketToSource.getInputStream
        oisSource = new ObjectInputStream (isSource)
        
        // Send the request
        oosSource.writeObject(requestPath)
        
        // Receive the length of the requested file
        var requestedFileLen = oisSource.readObject.asInstanceOf[Int]
        logInfo ("Received requestedFileLen = " + requestedFileLen)

        // Turn the timer OFF, if the sender responds before timeout
        timeOutTimer.cancel
        
        // Receive the file
        if (requestedFileLen != -1) {
          byteArray = new Array[Byte] (requestedFileLen)
          var bytesRead = isSource.read (byteArray, 0, byteArray.length)
          var alreadyRead = bytesRead

          while (alreadyRead < requestedFileLen) {
            bytesRead = isSource.read(byteArray, alreadyRead, 
              (byteArray.length - alreadyRead))
            if(bytesRead > 0) {
              alreadyRead = alreadyRead + bytesRead
            }
          }

          // Now add this to combiners
          combiners.synchronized {
            val inputStream = new ObjectInputStream (
              new ByteArrayInputStream(byteArray))
            try {
              while (true) {
                val (k, c) = inputStream.readObject.asInstanceOf[(K, C)]
                combiners(k) = combiners.get(k) match {
                  case Some(oldC) => mergeCombiners(oldC, c)
                  case None => c
                }
              }
            } catch {
              case e: Exception => {
                logInfo ("Merging to combiners had a " + e)
              }
            }
            inputStream.close
          }          
        } else {
          throw new SparkException("ShuffleServer " + hostAddress + " does not have " + requestPath)
        }
      } catch {
{'fraction_non_alphanumeric': 0.09190898467037877, 'fraction_numerical': 0.002680965147453083, 'mean_word_length': 2.21858407079646, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 36, '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': '3938732', 'n_tokens_mistral': 4261, 'n_tokens_neox': 4074, 'n_words': 1134}