src/jmh/java/org/apache/aurora/benchmark/Offers.java

/**
 * 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 org.apache.aurora.benchmark;

import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;

import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Data;
import org.apache.aurora.scheduler.base.Numbers;
import org.apache.aurora.scheduler.offers.HostOffer;
import org.apache.aurora.scheduler.offers.OfferManager;
import org.apache.aurora.scheduler.resources.ResourceType;
import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
import org.apache.mesos.v1.Protos;

import static org.apache.aurora.scheduler.resources.ResourceType.CPUS;
import static org.apache.aurora.scheduler.resources.ResourceType.DISK_MB;
import static org.apache.aurora.scheduler.resources.ResourceType.PORTS;
import static org.apache.aurora.scheduler.resources.ResourceType.RAM_MB;

/**
 * Offer factory.
 */
final class Offers {
  private Offers() {
    // Utility class.
  }

  /**
   * Saves offers into the {@link OfferManager}.
   *
   * @param offerManager {@link OfferManager} to save into.
   * @param offers Offers to save.
   */
  static void addOffers(OfferManager offerManager, Iterable<HostOffer> offers) {
    for (HostOffer offer : offers) {
      offerManager.add(offer);
    }
  }

  /**
   * Builds offers for the specified configuration.
   */
  static final class Builder {
    private static final String OFFER_ID_FORMAT = "offer-%s";
    private static final String FRAMEWORK_ID = "framework_id";

    private double cpu = 8.0;
    private Amount<Long, Data> ram = Amount.of(16L, Data.GB);
    private Amount<Long, Data> disk = Amount.of(256L, Data.GB);
    private int ports = 1024;

    Builder setCpu(double newCpu) {
      cpu = newCpu;
      return this;
    }

    Builder setRam(Amount<Long, Data> newRam) {
      ram = newRam;
      return this;
    }

    Builder setDisk(Amount<Long, Data> newDisk) {
      disk = newDisk;
      return this;
    }

    Builder setPorts(int newPorts) {
      ports = newPorts;
      return this;
    }

    /**
     * Builds a set of {@link HostOffer} for the current configuration.
     *
     * @param hostAttributes Host attributes to initialize offers from.
     * @return Set of offers.
     */
    Set<HostOffer> build(Set<IHostAttributes> hostAttributes) {
      ImmutableSet.Builder<HostOffer> offers = ImmutableSet.builder();
      int id = 0;
      for (IHostAttributes attributes : hostAttributes) {
        Protos.Offer offer = Protos.Offer.newBuilder()
            .addAllResources(ImmutableSet.of(
                makeScalar(CPUS, cpu),
                makeScalar(RAM_MB, ram.as(Data.MB)),
                makeScalar(DISK_MB, disk.as(Data.MB)),
                makeRange(
                    PORTS,
                    IntStream.range(1, ports).boxed().collect(Collectors.toSet()))))
            .setId(Protos.OfferID.newBuilder().setValue(String.format(OFFER_ID_FORMAT, id++)))
            .setFrameworkId(Protos.FrameworkID.newBuilder().setValue(FRAMEWORK_ID))
            .setAgentId(Protos.AgentID.newBuilder().setValue(attributes.getSlaveId()))
            .setHostname(String.format(attributes.getHost()))
            .build();

        offers.add(new HostOffer(offer, attributes));
      }

      return offers.build();
    }

    private static Protos.Resource makeScalar(ResourceType type, double value) {
      return Protos.Resource.newBuilder()
          .setType(Protos.Value.Type.SCALAR)
          .setName(type.getMesosName())
          .setScalar(Protos.Value.Scalar.newBuilder().setValue(value).build())
          .build();
    }

    private static Protos.Resource makeRange(ResourceType type, Iterable<Integer> values) {
      return Protos.Resource.newBuilder()
          .setType(Protos.Value.Type.RANGES)
          .setName(type.getMesosName())
          .setRanges(Protos.Value.Ranges.newBuilder().addAllRange(
              Iterables.transform(Numbers.toRanges(values), Numbers.RANGE_TRANSFORM)))
          .build();
    }
  }
}