aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject incubator-aurora git commit: Fixing cron update quota checking.
Date Fri, 27 Feb 2015 22:02:06 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master b436ec521 -> 388a5b3b1


Fixing cron update quota checking.

Bugs closed: AURORA-1134

Reviewed at https://reviews.apache.org/r/31334/


Project: http://git-wip-us.apache.org/repos/asf/incubator-aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-aurora/commit/388a5b3b
Tree: http://git-wip-us.apache.org/repos/asf/incubator-aurora/tree/388a5b3b
Diff: http://git-wip-us.apache.org/repos/asf/incubator-aurora/diff/388a5b3b

Branch: refs/heads/master
Commit: 388a5b3b10c217798108612b2cd8a58c0d6cfa14
Parents: b436ec5
Author: Maxim Khutornenko <maxim@apache.org>
Authored: Fri Feb 27 14:01:22 2015 -0800
Committer: Maxim Khutornenko <maxim@apache.org>
Committed: Fri Feb 27 14:01:22 2015 -0800

----------------------------------------------------------------------
 .../aurora/scheduler/quota/QuotaManager.java    | 79 ++++++++++++----
 .../thrift/SchedulerThriftInterface.java        |  2 +-
 .../scheduler/quota/QuotaManagerImplTest.java   | 99 ++++++++++++++++++++
 .../thrift/SchedulerThriftInterfaceTest.java    | 20 ++--
 4 files changed, 177 insertions(+), 23 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/388a5b3b/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java b/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java
index 8955df1..39e930c 100644
--- a/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java
+++ b/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java
@@ -55,6 +55,7 @@ import org.apache.aurora.scheduler.updater.Updates;
 
 import static java.util.Objects.requireNonNull;
 
+import static org.apache.aurora.scheduler.base.ResourceAggregates.EMPTY;
 import static org.apache.aurora.scheduler.base.Tasks.INFO_TO_JOB_KEY;
 import static org.apache.aurora.scheduler.base.Tasks.IS_PRODUCTION;
 import static org.apache.aurora.scheduler.base.Tasks.SCHEDULED_TO_INFO;
@@ -110,6 +111,16 @@ public interface QuotaManager {
   QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate, StoreProvider storeProvider);
 
   /**
+   * Check if there is enough resource quota available for creating or updating a cron job
+   * represented by the {@code cronConfig}. The quota is defined at the task owner (role)
level.
+   *
+   * @param cronConfig Cron job configuration.
+   * @param storeProvider A store provider to access quota data.
+   * @return{@code QuotaComparisonResult} instance with quota check result details.
+   */
+  QuotaCheckResult checkCronUpdate(IJobConfiguration cronConfig, StoreProvider storeProvider);
+
+  /**
    * Thrown when quota related operation failed.
    */
   class QuotaException extends Exception {
@@ -157,9 +168,8 @@ public interface QuotaManager {
       }
 
       QuotaInfo quotaInfo = getQuotaInfo(template.getJob().getRole(), storeProvider);
-      IResourceAggregate requestedTotal = add(
-          quotaInfo.getProdConsumption(),
-          ResourceAggregates.scale(fromTasks(ImmutableSet.of(template)), instances));
+      IResourceAggregate requestedTotal =
+          add(quotaInfo.getProdConsumption(), scale(template, instances));
 
       return QuotaCheckResult.greaterOrEqual(quotaInfo.getQuota(), requestedTotal);
     }
@@ -181,6 +191,32 @@ public interface QuotaManager {
       return QuotaCheckResult.greaterOrEqual(quotaInfo.getQuota(), quotaInfo.getProdConsumption());
     }
 
+    @Override
+    public QuotaCheckResult checkCronUpdate(
+        IJobConfiguration cronConfig,
+        StoreProvider storeProvider) {
+
+      if (!cronConfig.getTaskConfig().isProduction()) {
+        return new QuotaCheckResult(SUFFICIENT_QUOTA);
+      }
+
+      QuotaInfo quotaInfo =
+          getQuotaInfo(cronConfig.getKey().getRole(), Optional.<IJobUpdate>absent(),
storeProvider);
+
+      Optional<IJobConfiguration> oldCron =
+          storeProvider.getCronJobStore().fetchJob(cronConfig.getKey());
+
+      IResourceAggregate oldResource = oldCron.isPresent() ? scale(oldCron.get()) : EMPTY;
+
+      // Calculate requested total as a sum of current prod consumption and a delta between
+      // new and old cron templates.
+      IResourceAggregate requestedTotal = add(
+          quotaInfo.getProdConsumption(),
+          subtract(scale(cronConfig), oldResource));
+
+      return QuotaCheckResult.greaterOrEqual(quotaInfo.getQuota(), requestedTotal);
+    }
+
     /**
      * Gets QuotaInfo with currently allocated quota and actual consumption data.
      * <p>
@@ -268,7 +304,7 @@ public interface QuotaManager {
           .filter(buildNonUpdatingTasksFilter(updatesByKey))
           .transform(SCHEDULED_TO_INFO));
 
-      IResourceAggregate updateConsumption = ResourceAggregates.EMPTY;
+      IResourceAggregate updateConsumption = EMPTY;
       for (IJobUpdate update : updatesByKey.values()) {
         updateConsumption =
             add(updateConsumption, instructionsToResources(update.getInstructions(), isProd));
@@ -293,11 +329,11 @@ public interface QuotaManager {
       Multimap<IJobKey, ITaskConfig> taskConfigsByKey =
           tasks.transform(SCHEDULED_TO_INFO).index(INFO_TO_JOB_KEY);
 
-      IResourceAggregate totalConsumption = ResourceAggregates.EMPTY;
+      IResourceAggregate totalConsumption = EMPTY;
       for (IJobConfiguration config : cronTemplates.values()) {
         if (isProd == config.getTaskConfig().isProduction()) {
-          IResourceAggregate templateConsumption = ResourceAggregates.scale(
-              fromTasks(ImmutableSet.of(config.getTaskConfig())), config.getInstanceCount());
+          IResourceAggregate templateConsumption =
+              scale(config.getTaskConfig(), config.getInstanceCount());
 
           IResourceAggregate taskConsumption = fromTasks(taskConfigsByKey.get(config.getKey()));
 
@@ -396,14 +432,12 @@ public interface QuotaManager {
         final boolean isProd) {
 
       // Calculate initial state consumption.
-      IResourceAggregate initial = ResourceAggregates.EMPTY;
+      IResourceAggregate initial = EMPTY;
       for (IInstanceTaskConfig group : instructions.getInitialState()) {
         ITaskConfig task = group.getTask();
         if (isProd == task.isProduction()) {
           for (IRange range : group.getInstances()) {
-            initial = add(initial, ResourceAggregates.scale(
-                fromTasks(ImmutableSet.of(task)),
-                instanceCountFromRange(range)));
+            initial = add(initial, scale(task, instanceCountFromRange(range)));
           }
         }
       }
@@ -414,12 +448,10 @@ public interface QuotaManager {
             @Override
             public IResourceAggregate apply(IInstanceTaskConfig input) {
               return isProd == input.getTask().isProduction()
-                  ? ResourceAggregates.scale(
-                  fromTasks(ImmutableSet.of(input.getTask())),
-                  getUpdateInstanceCount(input.getInstances()))
-                  : ResourceAggregates.EMPTY;
+                  ? scale(input.getTask(), getUpdateInstanceCount(input.getInstances()))
+                  : EMPTY;
             }
-          }).or(ResourceAggregates.EMPTY);
+          }).or(EMPTY);
 
       // Calculate result as max(existing, desired) per resource type.
       return max(initial, desired);
@@ -432,6 +464,13 @@ public interface QuotaManager {
           .setDiskMb(a.getDiskMb() + b.getDiskMb()));
     }
 
+    private static IResourceAggregate subtract(IResourceAggregate a, IResourceAggregate b)
{
+      return IResourceAggregate.build(new ResourceAggregate()
+          .setNumCpus(a.getNumCpus() - b.getNumCpus())
+          .setRamMb(a.getRamMb() - b.getRamMb())
+          .setDiskMb(a.getDiskMb() - b.getDiskMb()));
+    }
+
     private static IResourceAggregate max(IResourceAggregate a, IResourceAggregate b) {
       return IResourceAggregate.build(new ResourceAggregate()
           .setNumCpus(Math.max(a.getNumCpus(), b.getNumCpus()))
@@ -439,6 +478,14 @@ public interface QuotaManager {
           .setDiskMb(Math.max(a.getDiskMb(), b.getDiskMb())));
     }
 
+    private static IResourceAggregate scale(ITaskConfig taskConfig, int instanceCount) {
+      return ResourceAggregates.scale(fromTasks(ImmutableSet.of(taskConfig)), instanceCount);
+    }
+
+    private static IResourceAggregate scale(IJobConfiguration jobConfiguration) {
+      return scale(jobConfiguration.getTaskConfig(), jobConfiguration.getInstanceCount());
+    }
+
     private static IResourceAggregate fromTasks(Iterable<ITaskConfig> tasks) {
       double cpu = 0;
       int ramMb = 0;

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/388a5b3b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
index b62adba..8ec5f9a 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
@@ -353,7 +353,7 @@ class SchedulerThriftInterface implements AuroraAdmin.Iface {
           validateTaskLimits(
               template,
               count,
-              quotaManager.checkInstanceAddition(template, count, storeProvider));
+              quotaManager.checkCronUpdate(sanitized.getJobConfig(), storeProvider));
 
           // TODO(mchucarroll): Merge CronJobManager.createJob/updateJob
           if (updateOnly || getCronJob(storeProvider, jobKey).isPresent()) {

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/388a5b3b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
index 97e60d8..0480aca 100644
--- a/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
@@ -36,6 +36,7 @@ import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
+import org.apache.aurora.scheduler.base.ResourceAggregates;
 import org.apache.aurora.scheduler.quota.QuotaManager.QuotaException;
 import org.apache.aurora.scheduler.quota.QuotaManager.QuotaManagerImpl;
 import org.apache.aurora.scheduler.storage.JobUpdateStore;
@@ -55,6 +56,7 @@ import org.junit.Test;
 import static org.apache.aurora.scheduler.quota.QuotaCheckResult.Result.INSUFFICIENT_QUOTA;
 import static org.apache.aurora.scheduler.quota.QuotaCheckResult.Result.SUFFICIENT_QUOTA;
 import static org.apache.aurora.scheduler.quota.QuotaManager.QuotaManagerImpl.updateQuery;
+import static org.easymock.EasyMock.anyObject;
 import static org.easymock.EasyMock.expect;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -690,6 +692,94 @@ public class QuotaManagerImplTest extends EasyMockTest {
         storageUtil.mutableStoreProvider.getQuotaStore());
   }
 
+  @Test
+  public void testCheckQuotaCronUpdateDownsize() {
+    expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2);
+    expectNoTasks().times(2);
+    expectNoJobUpdates().times(2);
+
+    IJobConfiguration job = createJob(createProdTask("pc", 4, 4, 4), 1);
+    expectCronJobs(job, createJob(createNonProdTask("npc", 7, 7, 7), 1)).times(2);
+    expectCronJob(job);
+
+    control.replay();
+
+    QuotaCheckResult checkQuota =
+        quotaManager.checkCronUpdate(createJob(createProdTask("pc", 1, 1, 1), 2), storeProvider);
+    assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
+    assertEquals(
+        new QuotaInfo(from(5, 5, 5), from(4, 4, 4), from(7, 7, 7)),
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
+  }
+
+  @Test
+  public void testCheckQuotaCronUpdateUpsize() {
+    expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2);
+    expectNoTasks().times(2);
+    expectNoJobUpdates().times(2);
+
+    IJobConfiguration job = createJob(createProdTask("pc", 4, 4, 4), 1);
+    expectCronJobs(job, createJob(createNonProdTask("npc", 7, 7, 7), 1)).times(2);
+    expectCronJob(job);
+
+    control.replay();
+
+    QuotaCheckResult checkQuota =
+        quotaManager.checkCronUpdate(createJob(createProdTask("pc", 5, 5, 5), 1), storeProvider);
+    assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
+    assertEquals(
+        new QuotaInfo(from(5, 5, 5), from(4, 4, 4), from(7, 7, 7)),
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
+  }
+
+  @Test
+  public void testCheckQuotaCronUpdateFails() {
+    expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2);
+    expectNoTasks().times(2);
+    expectNoJobUpdates().times(2);
+
+    IJobConfiguration job = createJob(createProdTask("pc", 4, 4, 4), 1);
+    expectCronJobs(job).times(2);
+    expectCronJob(job);
+
+    control.replay();
+
+    QuotaCheckResult checkQuota =
+        quotaManager.checkCronUpdate(createJob(createProdTask("pc", 5, 5, 5), 2), storeProvider);
+    assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
+    assertEquals(
+        new QuotaInfo(from(5, 5, 5), from(4, 4, 4), ResourceAggregates.EMPTY),
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
+  }
+
+  @Test
+  public void testCheckQuotaCronCreate() {
+    expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2);
+    expectNoTasks().times(2);
+    expectNoJobUpdates().times(2);
+    expectNoCronJobs().times(2);
+    expectNoCronJob();
+
+    control.replay();
+
+    QuotaCheckResult checkQuota =
+        quotaManager.checkCronUpdate(createJob(createProdTask("pc", 5, 5, 5), 1), storeProvider);
+    assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
+    assertEquals(
+        new QuotaInfo(from(5, 5, 5), ResourceAggregates.EMPTY, ResourceAggregates.EMPTY),
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
+  }
+
+  @Test
+  public void testCheckQuotaNonProdCron() {
+    control.replay();
+
+    QuotaCheckResult checkQuota =
+        quotaManager.checkCronUpdate(createJob(createNonProdTask("np", 5, 5, 5), 1), storeProvider);
+
+    assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
+  }
+
   private IExpectationSetters<?> expectTasks(IScheduledTask... tasks) {
     return storageUtil.expectTaskFetch(ACTIVE_QUERY, tasks);
   }
@@ -767,6 +857,15 @@ public class QuotaManagerImplTest extends EasyMockTest {
     return expect(storageUtil.jobStore.fetchJobs()).andReturn(builder.build());
   }
 
+  private IExpectationSetters<?> expectCronJob(IJobConfiguration job) {
+    return expect(storageUtil.jobStore.fetchJob(job.getKey())).andReturn(Optional.of(job));
+  }
+
+  private IExpectationSetters<?> expectNoCronJob() {
+    return expect(storageUtil.jobStore.fetchJob(anyObject(IJobKey.class)))
+        .andReturn(Optional.<IJobConfiguration>absent());
+  }
+
   private IExpectationSetters<Optional<IResourceAggregate>> expectQuota(IResourceAggregate
quota) {
     return expect(storageUtil.quotaStore.fetchQuota(ROLE))
         .andReturn(Optional.of(quota));

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/388a5b3b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
index 98cb25a..8247408 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
@@ -1051,7 +1051,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
+    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
     cronJobManager.updateJob(anyObject(SanitizedCronJob.class));
     control.replay();
 
@@ -1089,7 +1089,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
+    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
     cronJobManager.updateJob(anyObject(SanitizedCronJob.class));
     expectLastCall().andThrow(new CronException("Nope"));
 
@@ -1131,7 +1131,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
+    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
 
     expectNoCronJob().times(2);
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
@@ -1149,7 +1149,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
+    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
 
     expectNoCronJob();
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), buildScheduledTask());
@@ -1168,7 +1168,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
+    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
 
     expectCronJob();
     cronJobManager.updateJob(SanitizedCronJob.from(sanitized));
@@ -1225,7 +1225,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), NOT_ENOUGH_QUOTA);
+    expectCronQuotaCheck(sanitized.getJobConfig(), NOT_ENOUGH_QUOTA);
 
     control.replay();
     assertResponse(INVALID_REQUEST, thrift.scheduleCronJob(CRON_JOB, null, SESSION));
@@ -2514,4 +2514,12 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
         1,
         storageUtil.mutableStoreProvider)).andReturn(result);
   }
+
+  private IExpectationSetters<?> expectCronQuotaCheck(
+      IJobConfiguration config,
+      QuotaCheckResult result) {
+
+    return expect(quotaManager.checkCronUpdate(config, storageUtil.mutableStoreProvider))
+        .andReturn(result);
+  }
 }


Mime
View raw message