content: /**
   * 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 com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Lists;
  
  import org.apache.aurora.common.quantity.Amount;
  import org.apache.aurora.common.quantity.Data;
  import org.apache.aurora.gen.Constraint;
  import org.apache.aurora.gen.JobKey;
  import org.apache.aurora.gen.LimitConstraint;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.gen.ScheduledTask;
  import org.apache.aurora.gen.TaskConstraint;
  import org.apache.aurora.gen.TaskEvent;
  import org.apache.aurora.gen.ValueConstraint;
  import org.apache.aurora.scheduler.base.TaskTestUtil;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  
  import static org.apache.aurora.gen.Resource.diskMb;
  import static org.apache.aurora.gen.Resource.numCpus;
  import static org.apache.aurora.gen.Resource.ramMb;
  
  /**
   * Task factory.
   */
  final class Tasks {
  
    private Tasks() {
      // Utility class.
    }
  
    /**
     * Builds tasks for the specified configuration.
     */
    static final class Builder {
      private JobKey jobKey = new JobKey("jmh", "dev", "benchmark");
      private String tier = TaskTestUtil.DEV_TIER_NAME;
      private double cpu = 6.0;
      private Amount<Long, Data> ram = Amount.of(8L, Data.GB);
      private Amount<Long, Data> disk = Amount.of(128L, Data.GB);
      private ScheduleStatus scheduleStatus = ScheduleStatus.PENDING;
      private ImmutableSet.Builder<Constraint> constraints = ImmutableSet.builder();
  
      private static int uuid = 0;
  
      Builder setRole(String newRole) {
        jobKey.setRole(newRole);
        return this;
      }
  
      Builder setEnv(String env) {
        jobKey.setEnvironment(env);
        return this;
      }
  
      Builder setJob(String job) {
        jobKey.setName(job);
        return this;
      }
  
      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 setScheduleStatus(ScheduleStatus newScheduleStatus) {
        scheduleStatus = newScheduleStatus;
        return this;
      }
  
      Builder setTier(String newTier) {
        tier = newTier;
        return this;
      }
  
      Builder addValueConstraint(String name, String value) {
        constraints.add(new Constraint()
            .setName(name)
            .setConstraint(TaskConstraint.value(new ValueConstraint()
                .setNegated(false)
                .setValues(ImmutableSet.of(value)))));
  
        return this;
      }
  
      Builder addLimitConstraint(String name, int limit) {
        constraints.add(new Constraint()
            .setName(name)
            .setConstraint(TaskConstraint.limit(new LimitConstraint()
                .setLimit(limit))));
  
        return this;
      }
  
      /**
       * Builds a set of {@link IScheduledTask} for the current configuration.
       *
       * @param count Number of tasks to build.
       * @return Set of tasks.
       */
      Set<IScheduledTask> build(int count) {
        ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder();
  
        for (int i = 0; i < count; i++) {
          String taskId =
              jobKey.getRole() + "-" + jobKey.getEnvironment() + "-" + i + "-" + uuid++;
  
          ScheduledTask builder = TaskTestUtil.makeTask(taskId, IJobKey.build(jobKey))
              .newBuilder()
              .setStatus(scheduleStatus)
              .setTaskEvents(Lists.newArrayList(
                  new TaskEvent(0, ScheduleStatus.PENDING),
                  new TaskEvent(1, scheduleStatus)));
          builder.getAssignedTask()
              .setInstanceId(i)
              .setTaskId(taskId);
          builder.getAssignedTask().setAssignedPorts(ImmutableMap.of());
          builder.getAssignedTask().getTask()
              .setConstraints(constraints.build())
              .setTier(tier)
              .setResources(ImmutableSet.of(
                  numCpus(cpu),
                  ramMb(ram.as(Data.MB)),
                  diskMb(disk.as(Data.MB))));
          tasks.add(IScheduledTask.build(builder));
        }
  
        return tasks.build();
      }
    }
  }
  