<commit_before>package spark.scheduler.cluster

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

import akka.actor.{Props, Actor, ActorRef, ActorSystem}
import akka.util.duration._
import akka.pattern.ask

import spark.{SparkException, Logging, TaskState}
import akka.dispatch.Await
import java.util.concurrent.atomic.AtomicInteger

/**
 * A standalone scheduler backend, which waits for standalone executors to connect to it through
 * Akka. These may be executed in a variety of ways, such as Mesos tasks for the coarse-grained
 * Mesos mode or standalone processes for Spark's standalone deploy mode (spark.deploy.*).
 */
class StandaloneSchedulerBackend(scheduler: ClusterScheduler, actorSystem: ActorSystem)
  extends SchedulerBackend with Logging {

  // Use an atomic variable to track total number of cores in the cluster for simplicity and speed
  var totalCoreCount = new AtomicInteger(0)

  class MasterActor(sparkProperties: Seq[(String, String)]) extends Actor {
    val slaveActor = new HashMap[String, ActorRef]
    val slaveHost = new HashMap[String, String]
    val freeCores = new HashMap[String, Int]

    def receive = {
      case RegisterSlave(slaveId, host, cores) =>
        if (slaveActor.contains(slaveId)) {
          sender ! RegisterSlaveFailed("Duplicate slave ID: " + slaveId)
        } else {
          logInfo("Registered slave: " + sender + " with ID " + slaveId)
          sender ! RegisteredSlave(sparkProperties)
          slaveActor(slaveId) = sender
          slaveHost(slaveId) = host
          freeCores(slaveId) = cores
          totalCoreCount.addAndGet(cores)
          makeOffers()
        }

      case StatusUpdate(slaveId, taskId, state, data) =>
        scheduler.statusUpdate(taskId, state, data.value)
        if (TaskState.isFinished(state)) {
          freeCores(slaveId) += 1
          makeOffers(slaveId)
        }

      case ReviveOffers =>
        makeOffers()

      case StopMaster =>
        sender ! true
        context.stop(self)

      // TODO: Deal with nodes disconnecting too! (Including decreasing totalCoreCount)
    }

    // Make fake resource offers on all slaves
    def makeOffers() {
      launchTasks(scheduler.resourceOffers(
        slaveHost.toArray.map {case (id, host) => new WorkerOffer(id, host, freeCores(id))}))
    }

    // Make fake resource offers on just one slave
    def makeOffers(slaveId: String) {
      launchTasks(scheduler.resourceOffers(
        Seq(new WorkerOffer(slaveId, slaveHost(slaveId), freeCores(slaveId)))))
    }

    // Launch tasks returned by a set of resource offers
    def launchTasks(tasks: Seq[Seq[TaskDescription]]) {
      for (task <- tasks.flatten) {
        freeCores(task.slaveId) -= 1
        slaveActor(task.slaveId) ! LaunchTask(task)
      }
    }
  }

  var masterActor: ActorRef = null
  val taskIdsOnSlave = new HashMap[String, HashSet[String]]

  def start() {
    val properties = new ArrayBuffer[(String, String)]
    val iterator = System.getProperties.entrySet.iterator
    while (iterator.hasNext) {
      val entry = iterator.next
      val (key, value) = (entry.getKey.toString, entry.getValue.toString)
      if (key.startsWith("spark.")) {
        properties += ((key, value))
      }
    }
    masterActor = actorSystem.actorOf(
      Props(new MasterActor(properties)), name = StandaloneSchedulerBackend.ACTOR_NAME)
  }

  def stop() {
    try {
      if (masterActor != null) {
        val timeout = 5.seconds
        val future = masterActor.ask(StopMaster)(timeout)
        Await.result(future, timeout)
      }
    } catch {
      case e: Exception =><commit_msg>Detect disconnected slaves in StandaloneScheduler<commit_after>package spark.scheduler.cluster

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

import akka.actor._
import akka.util.duration._
import akka.pattern.ask

import spark.{SparkException, Logging, TaskState}
import akka.dispatch.Await
import java.util.concurrent.atomic.AtomicInteger
import akka.remote.{RemoteClientShutdown, RemoteClientDisconnected, RemoteClientLifeCycleEvent}

/**
 * A standalone scheduler backend, which waits for standalone executors to connect to it through
 * Akka. These may be executed in a variety of ways, such as Mesos tasks for the coarse-grained
 * Mesos mode or standalone processes for Spark's standalone deploy mode (spark.deploy.*).
 */
class StandaloneSchedulerBackend(scheduler: ClusterScheduler, actorSystem: ActorSystem)
  extends SchedulerBackend with Logging {

  // Use an atomic variable to track total number of cores in the cluster for simplicity and speed
  var totalCoreCount = new AtomicInteger(0)

  class MasterActor(sparkProperties: Seq[(String, String)]) extends Actor {
    val slaveActor = new HashMap[String, ActorRef]
    val slaveAddress = new HashMap[String, Address]
    val slaveHost = new HashMap[String, String]
    val freeCores = new HashMap[String, Int]
    val actorToSlaveId = new HashMap[ActorRef, String]
    val addressToSlaveId = new HashMap[Address, String]

    override def preStart() {
      // Listen for remote client disconnection events, since they don't go through Akka's watch()
      context.system.eventStream.subscribe(self, classOf[RemoteClientLifeCycleEvent])
    }

    def receive = {
      case RegisterSlave(slaveId, host, cores) =>
        if (slaveActor.contains(slaveId)) {
          sender ! RegisterSlaveFailed("Duplicate slave ID: " + slaveId)
        } else {
          logInfo("Registered slave: " + sender + " with ID " + slaveId)
          sender ! RegisteredSlave(sparkProperties)
          context.watch(sender)
          slaveActor(slaveId) = sender
          slaveHost(slaveId) = host
          freeCores(slaveId) = cores
          slaveAddress(slaveId) = sender.path.address
          actorToSlaveId(sender) = slaveId
          addressToSlaveId(sender.path.address) = slaveId
          totalCoreCount.addAndGet(cores)
          makeOffers()
        }

      case StatusUpdate(slaveId, taskId, state, data) =>
        scheduler.statusUpdate(taskId, state, data.value)
        if (TaskState.isFinished(state)) {
          freeCores(slaveId) += 1
          makeOffers(slaveId)
        }

      case ReviveOffers =>
        makeOffers()

      case StopMaster =>
        sender ! true
        context.stop(self)

      case Terminated(actor) =>
        actorToSlaveId.get(actor).foreach(removeSlave)

      case RemoteClientDisconnected(transport, address) =>
        addressToSlaveId.get(address).foreach(removeSlave)

      case RemoteClientShutdown(transport, address) =>
        addressToSlaveId.get(address).foreach(removeSlave)
    }

    // Make fake resource offers on all slaves
    def makeOffers() {
      launchTasks(scheduler.resourceOffers(
        slaveHost.toArray.map {case (id, host) => new WorkerOffer(id, host, freeCores(id))}))
    }

    // Make fake resource offers on just one slave
    def makeOffers(slaveId: String) {
      launchTasks(scheduler.resourceOffers(
        Seq(new WorkerOffer(slaveId, slaveHost(slaveId), freeCores(slaveId)))))
    }

    // Launch tasks returned by a set of resource offers
    def launchTasks(tasks: Seq[Seq[TaskDescription]]) {
      for (task <- tasks.flatten) {
        freeCores(task.slaveId) -= 1
        slaveActor(task.slaveId) ! LaunchTask(task)
      }
    }

    // Remove a disconnected slave from the cluster
    def removeSlave(slaveId: String) {
      logInfo("Slave " + slaveId + " disconnected, so removing it")
      val numCores = freeCores(slaveId)
      actorToSlaveId -= slaveActor(slaveId)
      addressToSlaveId -= slaveAddress(slaveId)
      slaveActor -= slaveId
      slaveHost -= slaveId
      freeCores -= slaveId
      slaveHost -= slaveId
      totalCoreCount.addAndGet(-numCores)
      scheduler.slaveLost(slaveId)
    }
  }

  var masterActor: ActorRef = null
  val taskIdsOnSlave = new HashMap[String, HashSet[String]]

  def start() {
    val properties = new ArrayBuffer[(String, String)]
    val iterator = System.getProperties.entrySet.iterator
    while (iterator.hasNext) {
      val entry = iterator.next
      val (key, value) = (entry.getKey.toString, entry.getValue.toString)
      if (key.startsWith("spark.")) {
        properties += ((key, value))
      }
    }
    masterActor = actorSystem.actorOf(
      Props(new MasterActor(properties)), name = StandaloneSchedulerBackend.ACTOR_NAME)
  }

  def stop() {
    try {
      if (masterActor != null) {
        val timeout = 5.seconds
        val future = masterActor.ask(StopMaster)(timeout)
        Await.result(future, timeout)
      }
    } catch {
      case e: Exception =>
{'fraction_non_alphanumeric': 0.09514990160898253, 'fraction_numerical': 0.0009260331056835282, 'mean_word_length': 3.088973024136299, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 18, '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': '6818799', 'n_tokens_mistral': 2533, 'n_tokens_neox': 2443, 'n_words': 654}