<commit_before>/**
 * Copyright (C) 2013 Carnegie Mellon University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package thomasdb.master

import akka.actor.{Actor, ActorLogging, ActorRef, Props}
import akka.pattern.{ask, pipe}
import java.io.File
import scala.collection.mutable.{Buffer, Map}
import scala.concurrent.{Await, Future, Promise}
import scala.util.{Failure, Success, Try}

import thomasdb.{Adjustable, ThomasDB}
import thomasdb.Constants._
import thomasdb.datastore.Datastore
import thomasdb.messages._
import thomasdb.list._
import thomasdb.stats.{Stats, WorkerInfo}
import thomasdb.worker.{Task, Worker}

object Master {
  def props(): Props = Props(classOf[Master])
}

class Master extends Actor with ActorLogging {
  import context.dispatcher

  Stats.registeredWorkers.clear()

  log.info("Master launched.")

  private val workers = Map[WorkerId, ActorRef]()

  // Maps workerIds to Datastores.
  private val datastoreRefs = Map[WorkerId, ActorRef]()

  // Maps mutatorIds to the Task the mutator's computation was launched on.
  private val tasks = Map[Int, ActorRef]()

  private var nextMutatorId = 0

  private var nextWorkerId: WorkerId = 0

  // The next Worker to schedule a Task on.
  private var nextWorker: WorkerId = 0

  def cycleWorkers() {
    nextWorker = (incrementWorkerId(nextWorker) % workers.size).toShort
  }

  def receive = {
    // Worker
    case RegisterWorkerMessage(
        workerRef: ActorRef, datastoreRef: ActorRef, webuiAddress: String) =>
      log.info("Registering worker at " + workerRef)


      val workerId = nextWorkerId
      workers(workerId) = workerRef
      datastoreRefs(workerId) = datastoreRef

      sender ! workerId
      nextWorkerId = incrementWorkerId(nextWorkerId)

      Stats.registeredWorkers += WorkerInfo(workerId, webuiAddress)

      for ((thatWorkerId, thatDatastoreRef) <- datastoreRefs) {
        thatDatastoreRef ! RegisterDatastoreMessage(workerId, datastoreRef)
        datastoreRef ! RegisterDatastoreMessage(thatWorkerId, thatDatastoreRef)
      }

    case ScheduleTaskMessage(parent: ActorRef, workerId: WorkerId) =>
      if (workerId == -1) {
        (workers(nextWorker) ? CreateTaskMessage(parent)) pipeTo sender
        cycleWorkers()
      } else {
        (workers(workerId) ? CreateTaskMessage(parent)) pipeTo sender
      }

    // Mutator
    case RegisterMutatorMessage =>
      log.info("Registering mutator " + nextMutatorId)
      sender ! nextMutatorId
      nextMutatorId += 1

    case RunMutatorMessage(adjust: Adjustable[_], mutatorId: Int) =>
      log.info("Starting initial run for mutator " + mutatorId)

      val taskRefFuture = workers(nextWorker) ?
        CreateTaskMessage(workers(nextWorker))
      cycleWorkers()
      val taskRef = Await.result(taskRefFuture.mapTo[ActorRef], DURATION)

      (taskRef ? RunTaskMessage(adjust)) pipeTo sender

      tasks(mutatorId) = taskRef

    case PropagateMutatorMessage(mutatorId: Int) =>
      log.info("Initiating change propagation for mutator " + mutatorId)

      (tasks(mutatorId) ? PropagateTaskMessage) pipeTo sender

    case GetMutatorDDGMessage(mutatorId: Int) =>
      (tasks(mutatorId) ? GetTaskDDGMessage) pipeTo sender

    case ShutdownMutatorMessage(mutatorId: Int) =>
      log.info("Shutting down mutator " + mutatorId)
      (tasks(mutatorId) ? ShutdownTaskMessage) pipeTo sender
      tasks -= mutatorId

      for ((workerId, datastoreRef) <- datastoreRefs) {
        datastoreRef ! ClearMessage()
      }

    // Datastore
    case CreateModMessage(value: Any) =>
      (workers(nextWorker) ? CreateModMessage(value)) pipeTo sender
      cycleWorkers()

    case CreateModMessage(null) =>
      (workers(nextWorker) ? CreateModMessage(null)) pipeTo sender
      cycleWorkers()

    case GetModMessage(modId: ModId, null) =>
      val workerId = getWorkerId(modId)
      (datastoreRefs(workerId) ? GetModMessage(modId, null)) pipeTo sender

    case UpdateModMessage(modId: ModId, value: Any, null) =>
      val workerId = getWorkerId(modId)
      (datastoreRefs(workerId) ? UpdateModMessage(modId, value, null)) pipeTo sender

    case UpdateModMessage(modId: ModId, null, null) =>
      val workerId = getWorkerId(modId)
      (datastoreRefs(workerId) ? UpdateModMessage(modId, null, null)) pipeTo sender

    case CreateListMessage(conf: ListConf) =>
      def makeList(f: (String, ActorRef) => ListInput[Any, Any]) = {
        val datastoreRef = datastoreRefs(nextWorker)
        val future = datastoreRef ? CreateListIdsMessage(conf, 1, 1, 0)
        val listIds = Await.result(future.mapTo[Buffer[String]], DURATION)
        cycleWorkers()

        f(listIds.head, datastoreRef)
      }

      def makePartitions(f: (String, ActorRef, WorkerId) => Unit) {
        val futures = Buffer[(Future[Buffer[String]], ActorRef, WorkerId)]()

        val partitionsPerWorker = conf.partitions / workers.size

        var index = 0
        for ((workerId, workerRef) <- workers) {
          val datastoreRef = datastoreRefs(workerId)

          val partitions =
            if (index == workers.size - 1 &&
                workers.size % partitionsPerWorker != 0) {
              workers.size % partitionsPerWorker
            } else {
              partitionsPerWorker
            }
          val future = datastoreRef ? CreateListIdsMessage(
            conf, partitions, workers.size, index)
          futures += ((future.mapTo[Buffer[String]], datastoreRef, workerId))

          index += 1
        }

        /*var index = 0
        for (i <- 1 to conf.partitions) {
          val datastoreRef = datastoreRefs(nextWorker)<commit_msg>Fixed a bug with the number of partitions we were creating.<commit_after>/**
 * Copyright (C) 2013 Carnegie Mellon University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package thomasdb.master

import akka.actor.{Actor, ActorLogging, ActorRef, Props}
import akka.pattern.{ask, pipe}
import java.io.File
import scala.collection.mutable.{Buffer, Map}
import scala.concurrent.{Await, Future, Promise}
import scala.util.{Failure, Success, Try}

import thomasdb.{Adjustable, ThomasDB}
import thomasdb.Constants._
import thomasdb.datastore.Datastore
import thomasdb.messages._
import thomasdb.list._
import thomasdb.stats.{Stats, WorkerInfo}
import thomasdb.worker.{Task, Worker}

object Master {
  def props(): Props = Props(classOf[Master])
}

class Master extends Actor with ActorLogging {
  import context.dispatcher

  Stats.registeredWorkers.clear()

  log.info("Master launched.")

  private val workers = Map[WorkerId, ActorRef]()

  // Maps workerIds to Datastores.
  private val datastoreRefs = Map[WorkerId, ActorRef]()

  // Maps mutatorIds to the Task the mutator's computation was launched on.
  private val tasks = Map[Int, ActorRef]()

  private var nextMutatorId = 0

  private var nextWorkerId: WorkerId = 0

  // The next Worker to schedule a Task on.
  private var nextWorker: WorkerId = 0

  def cycleWorkers() {
    nextWorker = (incrementWorkerId(nextWorker) % workers.size).toShort
  }

  def receive = {
    // Worker
    case RegisterWorkerMessage(
        workerRef: ActorRef, datastoreRef: ActorRef, webuiAddress: String) =>
      log.info("Registering worker at " + workerRef)


      val workerId = nextWorkerId
      workers(workerId) = workerRef
      datastoreRefs(workerId) = datastoreRef

      sender ! workerId
      nextWorkerId = incrementWorkerId(nextWorkerId)

      Stats.registeredWorkers += WorkerInfo(workerId, webuiAddress)

      for ((thatWorkerId, thatDatastoreRef) <- datastoreRefs) {
        thatDatastoreRef ! RegisterDatastoreMessage(workerId, datastoreRef)
        datastoreRef ! RegisterDatastoreMessage(thatWorkerId, thatDatastoreRef)
      }

    case ScheduleTaskMessage(parent: ActorRef, workerId: WorkerId) =>
      if (workerId == -1) {
        (workers(nextWorker) ? CreateTaskMessage(parent)) pipeTo sender
        cycleWorkers()
      } else {
        (workers(workerId) ? CreateTaskMessage(parent)) pipeTo sender
      }

    // Mutator
    case RegisterMutatorMessage =>
      log.info("Registering mutator " + nextMutatorId)
      sender ! nextMutatorId
      nextMutatorId += 1

    case RunMutatorMessage(adjust: Adjustable[_], mutatorId: Int) =>
      log.info("Starting initial run for mutator " + mutatorId)

      val taskRefFuture = workers(nextWorker) ?
        CreateTaskMessage(workers(nextWorker))
      cycleWorkers()
      val taskRef = Await.result(taskRefFuture.mapTo[ActorRef], DURATION)

      (taskRef ? RunTaskMessage(adjust)) pipeTo sender

      tasks(mutatorId) = taskRef

    case PropagateMutatorMessage(mutatorId: Int) =>
      log.info("Initiating change propagation for mutator " + mutatorId)

      (tasks(mutatorId) ? PropagateTaskMessage) pipeTo sender

    case GetMutatorDDGMessage(mutatorId: Int) =>
      (tasks(mutatorId) ? GetTaskDDGMessage) pipeTo sender

    case ShutdownMutatorMessage(mutatorId: Int) =>
      log.info("Shutting down mutator " + mutatorId)
      (tasks(mutatorId) ? ShutdownTaskMessage) pipeTo sender
      tasks -= mutatorId

      for ((workerId, datastoreRef) <- datastoreRefs) {
        datastoreRef ! ClearMessage()
      }

    // Datastore
    case CreateModMessage(value: Any) =>
      (workers(nextWorker) ? CreateModMessage(value)) pipeTo sender
      cycleWorkers()

    case CreateModMessage(null) =>
      (workers(nextWorker) ? CreateModMessage(null)) pipeTo sender
      cycleWorkers()

    case GetModMessage(modId: ModId, null) =>
      val workerId = getWorkerId(modId)
      (datastoreRefs(workerId) ? GetModMessage(modId, null)) pipeTo sender

    case UpdateModMessage(modId: ModId, value: Any, null) =>
      val workerId = getWorkerId(modId)
      (datastoreRefs(workerId) ? UpdateModMessage(modId, value, null)) pipeTo sender

    case UpdateModMessage(modId: ModId, null, null) =>
      val workerId = getWorkerId(modId)
      (datastoreRefs(workerId) ? UpdateModMessage(modId, null, null)) pipeTo sender

    case CreateListMessage(conf: ListConf) =>
      def makeList(f: (String, ActorRef) => ListInput[Any, Any]) = {
        val datastoreRef = datastoreRefs(nextWorker)
        val future = datastoreRef ? CreateListIdsMessage(conf, 1, 1, 0)
        val listIds = Await.result(future.mapTo[Buffer[String]], DURATION)
        cycleWorkers()

        f(listIds.head, datastoreRef)
      }

      def makePartitions(f: (String, ActorRef, WorkerId) => Unit) {
        val futures = Buffer[(Future[Buffer[String]], ActorRef, WorkerId)]()

        val partitionsPerWorker = conf.partitions / workers.size

        var index = 0
        for ((workerId, workerRef) <- workers) {
          val datastoreRef = datastoreRefs(workerId)
          val partitions =
            if (index == workers.size - 1 &&
                conf.partitions % workers.size != 0) {
              workers.size % partitionsPerWorker
            } else {
              partitionsPerWorker
            }
          val future = datastoreRef ? CreateListIdsMessage(
            conf, partitions, workers.size, index)
          futures += ((future.mapTo[Buffer[String]], datastoreRef, workerId))

          index += 1
        }

        /*var index = 0
        for (i <- 1 to conf.partitions) {
          val datastoreRef = datastoreRefs(nextWorker)
{'fraction_non_alphanumeric': 0.09560912263614967, 'fraction_numerical': 0.0035711387062738412, 'mean_word_length': 3.2814454482279363, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '1782693', 'n_tokens_mistral': 3777, 'n_tokens_neox': 3588, 'n_words': 963}