<commit_before>package mesosphere.marathon.tasks

import java.io._
import javax.inject.Inject

import mesosphere.marathon.Protos._
import mesosphere.marathon.state.{ PathId, StateMetrics, Timestamp }
import mesosphere.marathon.{ Main, MarathonConf }
import com.codahale.metrics.MetricRegistry
import org.apache.log4j.Logger
import org.apache.mesos.Protos.TaskStatus
import org.apache.mesos.state.{ State, Variable }

import scala.collection.JavaConverters._
import scala.collection._
import scala.collection.concurrent.TrieMap
import scala.collection.immutable.Set
import scala.concurrent.Future

class TaskTracker @Inject() (
  state: State,
  config: MarathonConf,
  val registry: MetricRegistry)
    extends StateMetrics {

  import mesosphere.marathon.tasks.TaskTracker._
  import mesosphere.util.BackToTheFuture.futureToFuture
  import mesosphere.util.ThreadPoolContext.context

  implicit val timeout = config.zkFutureTimeout

  private[this] val log = Logger.getLogger(getClass.getName)

  val PREFIX = "task:"
  val ID_DELIMITER = ":"

  private[this] val apps = TrieMap[PathId, InternalApp]()

  private[tasks] def fetchFromState(id: String): Variable = timedRead {
    state.fetch(id).get(timeout.duration.length, timeout.duration.unit)
  }

  private[tasks] def getKey(appId: PathId, taskId: String): String = {
    PREFIX + appId.safePath + ID_DELIMITER + taskId
  }

  def get(appId: PathId): Set[MarathonTask] =
    getInternal(appId).values.toSet

  def getVersion(appId: PathId, taskId: String): Option[Timestamp] =
    get(appId).collectFirst {
      case mt: MarathonTask if mt.getId == taskId =>
        Timestamp(mt.getVersion)
    }

  private def getInternal(appId: PathId): TrieMap[String, MarathonTask] =
    apps.getOrElseUpdate(appId, fetchApp(appId)).tasks

  def list: Map[PathId, App] = apps.mapValues(_.toApp).toMap

  def count(appId: PathId): Int = getInternal(appId).size

  def contains(appId: PathId): Boolean = apps.contains(appId)

  def take(appId: PathId, n: Int): Set[MarathonTask] = get(appId).take(n)

  def created(appId: PathId, task: MarathonTask): Unit = {
    // Keep this here so running() can pick it up
    getInternal(appId) += (task.getId -> task)
  }

  def running(appId: PathId, status: TaskStatus): Future[MarathonTask] = {
    val taskId = status.getTaskId.getValue
    val task = get(appId).find(_.getId == taskId) match {
      case Some(stagedTask) =>
        stagedTask.toBuilder
          .setStartedAt(System.currentTimeMillis)
          .setStatus(status)
          .build

      case _ =>
        log.warn(s"No staged task for ID $taskId")
        // We lost track of the host and port of this task, but still need to keep track of it
        MarathonTask.newBuilder
          .setId(taskId)
          .setStagedAt(System.currentTimeMillis)
          .setStartedAt(System.currentTimeMillis)
          .setStatus(status)
          .build
    }
    getInternal(appId) += (task.getId -> task)
    store(appId, task).map(_ => task)
  }

  def terminated(appId: PathId, status: TaskStatus): Future[Option[MarathonTask]] = {
    val appTasks = getInternal(appId)
    val app = apps(appId)
    val taskId = status.getTaskId.getValue

    appTasks.get(taskId) match {
      case Some(task) =>
        app.tasks.remove(task.getId)

        val variable = fetchFromState(getKey(appId, taskId))
        timedWrite { state.expunge(variable) }

        log.info(s"Task $taskId expunged and removed from TaskTracker")

        if (app.shutdown && app.tasks.isEmpty) {
          // Are we shutting down this app? If so, remove it
          remove(appId)
        }

        Future.successful(Some(task))
      case None =><commit_msg>Don't create unknown tasks with incomplete data in TaskTracker.running<commit_after>package mesosphere.marathon.tasks

import java.io._
import javax.inject.Inject

import mesosphere.marathon.Protos._
import mesosphere.marathon.state.{ PathId, StateMetrics, Timestamp }
import mesosphere.marathon.{ Main, MarathonConf }
import com.codahale.metrics.MetricRegistry
import org.apache.log4j.Logger
import org.apache.mesos.Protos.TaskStatus
import org.apache.mesos.state.{ State, Variable }

import scala.collection.JavaConverters._
import scala.collection._
import scala.collection.concurrent.TrieMap
import scala.collection.immutable.Set
import scala.concurrent.Future

class TaskTracker @Inject() (
  state: State,
  config: MarathonConf,
  val registry: MetricRegistry)
    extends StateMetrics {

  import mesosphere.marathon.tasks.TaskTracker._
  import mesosphere.util.BackToTheFuture.futureToFuture
  import mesosphere.util.ThreadPoolContext.context

  implicit val timeout = config.zkFutureTimeout

  private[this] val log = Logger.getLogger(getClass.getName)

  val PREFIX = "task:"
  val ID_DELIMITER = ":"

  private[this] val apps = TrieMap[PathId, InternalApp]()

  private[tasks] def fetchFromState(id: String): Variable = timedRead {
    state.fetch(id).get(timeout.duration.length, timeout.duration.unit)
  }

  private[tasks] def getKey(appId: PathId, taskId: String): String = {
    PREFIX + appId.safePath + ID_DELIMITER + taskId
  }

  def get(appId: PathId): Set[MarathonTask] =
    getInternal(appId).values.toSet

  def getVersion(appId: PathId, taskId: String): Option[Timestamp] =
    get(appId).collectFirst {
      case mt: MarathonTask if mt.getId == taskId =>
        Timestamp(mt.getVersion)
    }

  private def getInternal(appId: PathId): TrieMap[String, MarathonTask] =
    apps.getOrElseUpdate(appId, fetchApp(appId)).tasks

  def list: Map[PathId, App] = apps.mapValues(_.toApp).toMap

  def count(appId: PathId): Int = getInternal(appId).size

  def contains(appId: PathId): Boolean = apps.contains(appId)

  def take(appId: PathId, n: Int): Set[MarathonTask] = get(appId).take(n)

  def created(appId: PathId, task: MarathonTask): Unit = {
    // Keep this here so running() can pick it up
    getInternal(appId) += (task.getId -> task)
  }

  def running(appId: PathId, status: TaskStatus): Future[MarathonTask] = {
    val taskId = status.getTaskId.getValue
    get(appId).find(_.getId == taskId) match {
      case Some(stagedTask) =>
        val task = stagedTask.toBuilder
          .setStartedAt(System.currentTimeMillis)
          .setStatus(status)
          .build
        getInternal(appId) += (task.getId -> task)
        store(appId, task).map(_ => task)
      case _ =>
        val msg = s"No staged task for ID $taskId, ignoring"
        log.warn(msg)
        Future.failed(new Exception(msg))
    }
  }

  def terminated(appId: PathId, status: TaskStatus): Future[Option[MarathonTask]] = {
    val appTasks = getInternal(appId)
    val app = apps(appId)
    val taskId = status.getTaskId.getValue

    appTasks.get(taskId) match {
      case Some(task) =>
        app.tasks.remove(task.getId)

        val variable = fetchFromState(getKey(appId, taskId))
        timedWrite { state.expunge(variable) }

        log.info(s"Task $taskId expunged and removed from TaskTracker")

        if (app.shutdown && app.tasks.isEmpty) {
          // Are we shutting down this app? If so, remove it
          remove(appId)
        }

        Future.successful(Some(task))
      case None =>
{'fraction_non_alphanumeric': 0.10989315942833357, 'fraction_numerical': 0.0002775079783543777, 'mean_word_length': 3.524795982423101, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 19, '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': '3017120', 'n_tokens_mistral': 2318, 'n_tokens_neox': 2253, 'n_words': 523}