<commit_before>package mesosphere.marathon.tasks

import com.twitter.util.NonFatal
import mesosphere.marathon.state.DiskSource
import mesosphere.mesos.protos
import org.apache.mesos.Protos.Resource.DiskInfo.Source
import org.apache.mesos.Protos.Resource.{ DiskInfo, ReservationInfo }
import org.apache.mesos.{ Protos => MesosProtos }
import org.slf4j.LoggerFactory
import scala.collection.JavaConverters._

object ResourceUtil {
  implicit class RichResource(resource: MesosProtos.Resource) {
    def getDiskSourceOption: Option[Source] =
      if (resource.hasDisk && resource.getDisk.hasSource)
        Some(resource.getDisk.getSource)
      else
        None

    def getStringification: String = {
      require(resource.getName == protos.Resource.DISK)
      val diskSource = DiskSource.fromMesos(getDiskSourceOption)
      /* TODO - make this match mesos stringification */
      (List(
        resource.getName,
        diskSource.diskType.toString,
        resource.getScalar.getValue.toString) ++
        diskSource.path).mkString(":")
    }
  }

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

  /**
    * The resources in launched tasks, should
    * be consumed from resources in the offer with the same [[ResourceMatchKey]].
    */
  private[this] case class ResourceMatchKey(
    role: String, name: String,
    reservation: Option[ReservationInfo], disk: Option[DiskInfo])

  private[this] object ResourceMatchKey {
    def apply(resource: MesosProtos.Resource): ResourceMatchKey = {
      val reservation = if (resource.hasReservation) Some(resource.getReservation) else None
      val disk = if (resource.hasDisk) Some(resource.getDisk) else None
      ResourceMatchKey(resource.getRole, resource.getName, reservation, disk)
    }
  }

  /**
    * Decrements the scalar resource by amount
    *
    */
  def subtractScalarValue(resource: MesosProtos.Resource, amount: Double): Option[MesosProtos.Resource] = {
    require(resource.getType == MesosProtos.Value.Type.SCALAR)
    val isMountDiskResource =
      resource.hasDisk && resource.getDisk.hasSource &&
        (resource.getDisk.getSource.getType == Source.Type.MOUNT)

    val leftOver: Double = resource.getScalar.getValue - amount
    if (leftOver <= 0 || isMountDiskResource) {
      None
    } else {
      Some(resource
        .toBuilder
        .setScalar(
          MesosProtos.Value.Scalar
            .newBuilder().setValue(leftOver))
        .build())
    }
  }

  def subtractSetValue(resource: MesosProtos.Resource, set: Set[String]): Option[MesosProtos.Resource] = {
    require(resource.getType == MesosProtos.Value.Type.SET)
    val baseSet: Set[String] = resource.getSet.getItemList.asScala.toSet
    require(set subsetOf baseSet, s"set must be subset of $baseSet")

    val resultSet: Set[String] = baseSet -- set

    if (resultSet.nonEmpty)
      Some(
        resource
          .toBuilder
          .setSet(MesosProtos.Value.Set.newBuilder().addAllItem(resultSet.asJava))
          .build()
      )
    else
      None
  }

  def deductRange(
    baseRange: MesosProtos.Value.Range,
    usedRange: MesosProtos.Value.Range): Seq[MesosProtos.Value.Range] = {
    if (baseRange.getEnd < usedRange.getBegin || baseRange.getBegin > usedRange.getEnd) {
      // baseRange completely before or after usedRange
      Seq(baseRange)
    } else {
      val rangeBefore: Option[MesosProtos.Value.Range] = if (baseRange.getBegin < usedRange.getBegin)
        Some(baseRange.toBuilder.setEnd(usedRange.getBegin - 1).build())
      else
        None

      val rangeAfter: Option[MesosProtos.Value.Range] = if (baseRange.getEnd > usedRange.getEnd)
        Some(baseRange.toBuilder.setBegin(usedRange.getEnd + 1).build())
      else
        None

      Seq(rangeBefore, rangeAfter).flatten
    }
  }

  def subtractRangeValues(resource: MesosProtos.Resource, ranges: Seq[MesosProtos.Value.Range]): Option[MesosProtos.Resource] = {
    val baseRanges = resource.getRanges.getRangeList.asScala

    // FIXME: too expensive?
    val diminished = baseRanges.flatMap { baseRange =>
      ranges.foldLeft(Seq(baseRange)) {
        case (result, used) =>
          result.flatMap(deductRange(_, used))
      }
    }

    val rangesBuilder = MesosProtos.Value.Ranges.newBuilder()
    diminished.foreach(rangesBuilder.addRange)

    val result = resource
      .toBuilder
      .setRanges(rangesBuilder)
      .build()

    if (result.getRanges.getRangeCount > 0)
      Some(result)
    else
      None
  }

  /**
    * Deduct usedResource from resource. If nothing is left, None is returned.
    */
  def consumeResource(
    resource: MesosProtos.Resource,
    usedResource: MesosProtos.Resource): Option[MesosProtos.Resource] = {
    require(resource.getType == usedResource.getType)

    resource.getType match {
      case MesosProtos.Value.Type.SCALAR => subtractScalarValue(resource, usedResource.getScalar.getValue)
      case MesosProtos.Value.Type.RANGES => subtractRangeValues(resource, usedResource.getRanges.getRangeList.asScala)
      case MesosProtos.Value.Type.SET => subtractSetValue(resource, usedResource.getSet.getItemList.asScala.toSet)

      case unexpectedResourceType: MesosProtos.Value.Type =>
        log.warn("unexpected resourceType {} for resource {}", Seq(unexpectedResourceType, resource.getName): _*)
        // we don't know the resource, thus we consume it completely
        None
    }
  }

  /**
    * Deduct usedResources from resources by matching them by name and role.
    */
  def consumeResources(
    resources: Iterable[MesosProtos.Resource],
    usedResources: Iterable[MesosProtos.Resource]): Iterable[MesosProtos.Resource] = {
    val usedResourceMap: Map[ResourceMatchKey, Seq[MesosProtos.Resource]] =
      usedResources.groupBy(ResourceMatchKey(_)).mapValues(_.to[Seq])

    resources.flatMap { resource: MesosProtos.Resource =>
      usedResourceMap.get(ResourceMatchKey(resource)) match {
        case Some(usedResources: Iterable[MesosProtos.Resource]) =>
          usedResources.foldLeft(Some(resource): Option[MesosProtos.Resource]) {
            case (Some(resource), usedResource) =>
              if (resource.getType != usedResource.getType) {
                log.warn(
                  "Different resource types for resource {}: {} and {}",
                  resource.getName, resource.getType, usedResource.getType)
                None
              } else
                try ResourceUtil.consumeResource(resource, usedResource)
                catch {
                  case NonFatal(e) =>
                    log.warn("while consuming {} of type {}", resource.getName, resource.getType, e)
                    None
                }

            case (None, _) => None
          }
        case None => // if the resource isn't used, we keep it
          Some(resource)
      }
    }
  }

  /**
    * Deduct usedResources from resources in the offer.
    */<commit_msg>rename ResourceMatchKey to ResourceConsumeKey to reduce confusion<commit_after>package mesosphere.marathon.tasks

import com.twitter.util.NonFatal
import mesosphere.marathon.state.DiskSource
import mesosphere.mesos.protos
import org.apache.mesos.Protos.Resource.DiskInfo.Source
import org.apache.mesos.Protos.Resource.{ DiskInfo, ReservationInfo }
import org.apache.mesos.{ Protos => MesosProtos }
import org.slf4j.LoggerFactory
import scala.collection.JavaConverters._

object ResourceUtil {
  implicit class RichResource(resource: MesosProtos.Resource) {
    def getDiskSourceOption: Option[Source] =
      if (resource.hasDisk && resource.getDisk.hasSource)
        Some(resource.getDisk.getSource)
      else
        None

    def getStringification: String = {
      require(resource.getName == protos.Resource.DISK)
      val diskSource = DiskSource.fromMesos(getDiskSourceOption)
      /* TODO - make this match mesos stringification */
      (List(
        resource.getName,
        diskSource.diskType.toString,
        resource.getScalar.getValue.toString) ++
        diskSource.path).mkString(":")
    }
  }

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

  /**
    * The resources in launched tasks, should
    * be consumed from resources in the offer with the same [[ResourceMatchKey]].
    */
  private[this] case class ResourceConsumeKey(
    role: String, name: String,
    reservation: Option[ReservationInfo], disk: Option[DiskInfo])

  private[this] object ResourceConsumeKey {
    def apply(resource: MesosProtos.Resource): ResourceConsumeKey = {
      val reservation = if (resource.hasReservation) Some(resource.getReservation) else None
      val disk = if (resource.hasDisk) Some(resource.getDisk) else None
      ResourceConsumeKey(resource.getRole, resource.getName, reservation, disk)
    }
  }

  /**
    * Decrements the scalar resource by amount
    *
    */
  def subtractScalarValue(resource: MesosProtos.Resource, amount: Double): Option[MesosProtos.Resource] = {
    require(resource.getType == MesosProtos.Value.Type.SCALAR)
    val isMountDiskResource =
      resource.hasDisk && resource.getDisk.hasSource &&
        (resource.getDisk.getSource.getType == Source.Type.MOUNT)

    val leftOver: Double = resource.getScalar.getValue - amount
    if (leftOver <= 0 || isMountDiskResource) {
      None
    } else {
      Some(resource
        .toBuilder
        .setScalar(
          MesosProtos.Value.Scalar
            .newBuilder().setValue(leftOver))
        .build())
    }
  }

  def subtractSetValue(resource: MesosProtos.Resource, set: Set[String]): Option[MesosProtos.Resource] = {
    require(resource.getType == MesosProtos.Value.Type.SET)
    val baseSet: Set[String] = resource.getSet.getItemList.asScala.toSet
    require(set subsetOf baseSet, s"set must be subset of $baseSet")

    val resultSet: Set[String] = baseSet -- set

    if (resultSet.nonEmpty)
      Some(
        resource
          .toBuilder
          .setSet(MesosProtos.Value.Set.newBuilder().addAllItem(resultSet.asJava))
          .build()
      )
    else
      None
  }

  def deductRange(
    baseRange: MesosProtos.Value.Range,
    usedRange: MesosProtos.Value.Range): Seq[MesosProtos.Value.Range] = {
    if (baseRange.getEnd < usedRange.getBegin || baseRange.getBegin > usedRange.getEnd) {
      // baseRange completely before or after usedRange
      Seq(baseRange)
    } else {
      val rangeBefore: Option[MesosProtos.Value.Range] = if (baseRange.getBegin < usedRange.getBegin)
        Some(baseRange.toBuilder.setEnd(usedRange.getBegin - 1).build())
      else
        None

      val rangeAfter: Option[MesosProtos.Value.Range] = if (baseRange.getEnd > usedRange.getEnd)
        Some(baseRange.toBuilder.setBegin(usedRange.getEnd + 1).build())
      else
        None

      Seq(rangeBefore, rangeAfter).flatten
    }
  }

  def subtractRangeValues(resource: MesosProtos.Resource, ranges: Seq[MesosProtos.Value.Range]): Option[MesosProtos.Resource] = {
    val baseRanges = resource.getRanges.getRangeList.asScala

    // FIXME: too expensive?
    val diminished = baseRanges.flatMap { baseRange =>
      ranges.foldLeft(Seq(baseRange)) {
        case (result, used) =>
          result.flatMap(deductRange(_, used))
      }
    }

    val rangesBuilder = MesosProtos.Value.Ranges.newBuilder()
    diminished.foreach(rangesBuilder.addRange)

    val result = resource
      .toBuilder
      .setRanges(rangesBuilder)
      .build()

    if (result.getRanges.getRangeCount > 0)
      Some(result)
    else
      None
  }

  /**
    * Deduct usedResource from resource. If nothing is left, None is returned.
    */
  def consumeResource(
    resource: MesosProtos.Resource,
    usedResource: MesosProtos.Resource): Option[MesosProtos.Resource] = {
    require(resource.getType == usedResource.getType)

    resource.getType match {
      case MesosProtos.Value.Type.SCALAR => subtractScalarValue(resource, usedResource.getScalar.getValue)
      case MesosProtos.Value.Type.RANGES => subtractRangeValues(resource, usedResource.getRanges.getRangeList.asScala)
      case MesosProtos.Value.Type.SET => subtractSetValue(resource, usedResource.getSet.getItemList.asScala.toSet)

      case unexpectedResourceType: MesosProtos.Value.Type =>
        log.warn("unexpected resourceType {} for resource {}", Seq(unexpectedResourceType, resource.getName): _*)
        // we don't know the resource, thus we consume it completely
        None
    }
  }

  /**
    * Deduct usedResources from resources by matching them by name and role.
    */
  def consumeResources(
    resources: Iterable[MesosProtos.Resource],
    usedResources: Iterable[MesosProtos.Resource]): Iterable[MesosProtos.Resource] = {
    val usedResourceMap: Map[ResourceConsumeKey, Seq[MesosProtos.Resource]] =
      usedResources.groupBy(ResourceConsumeKey(_)).mapValues(_.to[Seq])

    resources.flatMap { resource: MesosProtos.Resource =>
      usedResourceMap.get(ResourceConsumeKey(resource)) match {
        case Some(usedResources: Iterable[MesosProtos.Resource]) =>
          usedResources.foldLeft(Some(resource): Option[MesosProtos.Resource]) {
            case (Some(resource), usedResource) =>
              if (resource.getType != usedResource.getType) {
                log.warn(
                  "Different resource types for resource {}: {} and {}",
                  resource.getName, resource.getType, usedResource.getType)
                None
              } else
                try ResourceUtil.consumeResource(resource, usedResource)
                catch {
                  case NonFatal(e) =>
                    log.warn("while consuming {} of type {}", resource.getName, resource.getType, e)
                    None
                }

            case (None, _) => None
          }
        case None => // if the resource isn't used, we keep it
          Some(resource)
      }
    }
  }

  /**
    * Deduct usedResources from resources in the offer.
    */
{'fraction_non_alphanumeric': 0.1022580181216741, 'fraction_numerical': 0.0007191140514885661, 'mean_word_length': 2.934087694483734, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 35, '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': '7148645', 'n_tokens_mistral': 4123, 'n_tokens_neox': 4031, 'n_words': 818}