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: Pushing transactions up in QuotaManager.
Date Tue, 24 Feb 2015 21:37:28 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master 2aa90e7ff -> 012059148


Pushing transactions up in QuotaManager.

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


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

Branch: refs/heads/master
Commit: 0120591488642a36b705145a891062d8a837be08
Parents: 2aa90e7
Author: Maxim Khutornenko <maxim@apache.org>
Authored: Tue Feb 24 13:33:33 2015 -0800
Committer: Maxim Khutornenko <maxim@apache.org>
Committed: Tue Feb 24 13:33:33 2015 -0800

----------------------------------------------------------------------
 .../aurora/scheduler/quota/QuotaManager.java    | 114 +++++++++--------
 .../scheduler/thrift/ReadOnlySchedulerImpl.java |  16 ++-
 .../thrift/SchedulerThriftInterface.java        |  25 +++-
 .../scheduler/quota/QuotaManagerImplTest.java   | 127 ++++++++++++-------
 .../thrift/ReadOnlySchedulerImplTest.java       |   2 +-
 .../thrift/SchedulerThriftInterfaceTest.java    | 115 +++++++++++------
 .../aurora/scheduler/thrift/ThriftIT.java       |   3 +-
 7 files changed, 245 insertions(+), 157 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/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 bb72721..199dc11 100644
--- a/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java
+++ b/src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java
@@ -31,7 +31,6 @@ import com.google.common.collect.Multimap;
 import com.google.common.collect.Range;
 import com.google.common.collect.RangeSet;
 import com.google.common.collect.Sets;
-import com.google.inject.Inject;
 
 import org.apache.aurora.gen.JobUpdateQuery;
 import org.apache.aurora.gen.ResourceAggregate;
@@ -39,9 +38,8 @@ 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.storage.JobUpdateStore;
-import org.apache.aurora.scheduler.storage.Storage;
+import org.apache.aurora.scheduler.storage.QuotaStore;
 import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
-import org.apache.aurora.scheduler.storage.Storage.Work;
 import org.apache.aurora.scheduler.storage.entities.IInstanceTaskConfig;
 import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
@@ -71,17 +69,20 @@ public interface QuotaManager {
    *
    * @param role Quota owner.
    * @param quota Quota to save.
+   * @param quoteStore A quota store.
    * @throws QuotaException If provided quota specification is invalid.
    */
-  void saveQuota(String role, IResourceAggregate quota) throws QuotaException;
+  void saveQuota(String role, IResourceAggregate quota, QuotaStore.Mutable quoteStore)
+      throws QuotaException;
 
   /**
    * Gets {@code QuotaInfo} for the specified role.
    *
    * @param role Quota owner.
+   * @param storeProvider A store provider to access quota data.
    * @return {@code QuotaInfo} instance.
    */
-  QuotaInfo getQuotaInfo(String role);
+  QuotaInfo getQuotaInfo(String role, StoreProvider storeProvider);
 
   /**
    * Checks if there is enough resource quota available for adding {@code instances} of
@@ -90,18 +91,23 @@ public interface QuotaManager {
    *
    * @param template Task resource requirement.
    * @param instances Number of additional instances requested.
+   * @param storeProvider A store provider to access quota data.
    * @return {@code QuotaComparisonResult} instance with quota check result details.
    */
-  QuotaCheckResult checkInstanceAddition(ITaskConfig template, int instances);
+  QuotaCheckResult checkInstanceAddition(
+      ITaskConfig template,
+      int instances,
+      StoreProvider storeProvider);
 
   /**
    * Checks if there is enough resource quota available for performing a job update represented
    * by the {@code jobUpdate}. The quota is defined at the task owner (role) level.
    *
    * @param jobUpdate Job update to check quota for.
+   * @param storeProvider A store provider to access quota data.
    * @return {@code QuotaComparisonResult} instance with quota check result details.
    */
-  QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate);
+  QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate, StoreProvider storeProvider);
 
   /**
    * Thrown when quota related operation failed.
@@ -113,19 +119,15 @@ public interface QuotaManager {
   }
 
   /**
-   * Quota provider that stores quotas in the canonical {@link Storage} system.
+   * Quota provider that stores quotas in the canonical store.
    */
   class QuotaManagerImpl implements QuotaManager {
-    private final Storage storage;
-
-    @Inject
-    QuotaManagerImpl(Storage storage) {
-      this.storage = requireNonNull(storage);
-    }
 
     @Override
-    public void saveQuota(final String ownerRole, final IResourceAggregate quota)
-        throws QuotaException {
+    public void saveQuota(
+        final String ownerRole,
+        final IResourceAggregate quota,
+        QuotaStore.Mutable quoteStore) throws QuotaException {
 
       if (!quota.isSetNumCpus() || !quota.isSetRamMb() || !quota.isSetDiskMb()) {
         throw new QuotaException("Missing quota specification(s) in: " + quota.toString());
@@ -135,27 +137,26 @@ public interface QuotaManager {
         throw new QuotaException("Negative values in: " + quota.toString());
       }
 
-      storage.write(new Storage.MutateWork.NoResult.Quiet() {
-        @Override
-        protected void execute(Storage.MutableStoreProvider storeProvider) {
-          storeProvider.getQuotaStore().saveQuota(ownerRole, quota);
-        }
-      });
+      quoteStore.saveQuota(ownerRole, quota);
     }
 
     @Override
-    public QuotaInfo getQuotaInfo(final String role) {
-      return getQuotaInfo(role, Optional.<IJobUpdate>absent());
+    public QuotaInfo getQuotaInfo(String role, StoreProvider storeProvider) {
+      return getQuotaInfo(role, Optional.<IJobUpdate>absent(), storeProvider);
     }
 
     @Override
-    public QuotaCheckResult checkInstanceAddition(ITaskConfig template, int instances) {
+    public QuotaCheckResult checkInstanceAddition(
+        ITaskConfig template,
+        int instances,
+        StoreProvider storeProvider) {
+
       Preconditions.checkArgument(instances >= 0);
       if (!template.isProduction()) {
         return new QuotaCheckResult(SUFFICIENT_QUOTA);
       }
 
-      QuotaInfo quotaInfo = getQuotaInfo(template.getJob().getRole());
+      QuotaInfo quotaInfo = getQuotaInfo(template.getJob().getRole(), storeProvider);
       IResourceAggregate requestedTotal = add(
           quotaInfo.getProdConsumption(),
           ResourceAggregates.scale(fromTasks(ImmutableSet.of(template)), instances));
@@ -164,7 +165,7 @@ public interface QuotaManager {
     }
 
     @Override
-    public QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate) {
+    public QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate, StoreProvider storeProvider) {
       requireNonNull(jobUpdate);
       if (!jobUpdate.getInstructions().isSetDesiredState()
           || !jobUpdate.getInstructions().getDesiredState().getTask().isProduction()) {
@@ -174,7 +175,8 @@ public interface QuotaManager {
 
       QuotaInfo quotaInfo = getQuotaInfo(
           jobUpdate.getSummary().getJobKey().getRole(),
-          Optional.of(jobUpdate));
+          Optional.of(jobUpdate),
+          storeProvider);
 
       return QuotaCheckResult.greaterOrEqual(quotaInfo.getQuota(), quotaInfo.getProdConsumption());
     }
@@ -187,41 +189,41 @@ public interface QuotaManager {
      *
      * @param role Role to get quota info for.
      * @param requestedUpdate An optional {@code IJobUpdate} to forecast the consumption.
+     * @param storeProvider A store provider to access quota data.
      * @return {@code QuotaInfo} with quota and consumption details.
      */
-    private QuotaInfo getQuotaInfo(final String role, final Optional<IJobUpdate> requestedUpdate) {
-      return storage.read(new Work.Quiet<QuotaInfo>() {
-        @Override
-        public QuotaInfo apply(StoreProvider storeProvider) {
-          FluentIterable<IScheduledTask> tasks = FluentIterable.from(
-              storeProvider.getTaskStore().fetchTasks(Query.roleScoped(role).active()));
-
-          Map<IJobKey, IJobUpdate> updates = Maps.newHashMap(
-              fetchActiveJobUpdates(storeProvider.getJobUpdateStore(), role)
-                  .uniqueIndex(UPDATE_TO_JOB_KEY));
-
-          // Mix in a requested job update (if present) to correctly calculate consumption.
-          // This would be an update that is not saved in the store yet (i.e. the one quota is
-          // checked for).
-          if (requestedUpdate.isPresent()) {
-            updates.put(requestedUpdate.get().getSummary().getJobKey(), requestedUpdate.get());
-          }
+    private QuotaInfo getQuotaInfo(
+        String role,
+        Optional<IJobUpdate> requestedUpdate,
+        StoreProvider storeProvider) {
+
+      FluentIterable<IScheduledTask> tasks = FluentIterable.from(
+          storeProvider.getTaskStore().fetchTasks(Query.roleScoped(role).active()));
+
+      Map<IJobKey, IJobUpdate> updates = Maps.newHashMap(
+          fetchActiveJobUpdates(storeProvider.getJobUpdateStore(), role)
+              .uniqueIndex(UPDATE_TO_JOB_KEY));
+
+      // Mix in a requested job update (if present) to correctly calculate consumption.
+      // This would be an update that is not saved in the store yet (i.e. the one quota is
+      // checked for).
+      if (requestedUpdate.isPresent()) {
+        updates.put(requestedUpdate.get().getSummary().getJobKey(), requestedUpdate.get());
+      }
 
-          Map<IJobKey, IJobConfiguration> cronTemplates =
-              FluentIterable.from(storeProvider.getJobStore().fetchJobs())
-                  .filter(Predicates.compose(Predicates.equalTo(role), JobKeys.CONFIG_TO_ROLE))
-                  .uniqueIndex(JobKeys.FROM_CONFIG);
+      Map<IJobKey, IJobConfiguration> cronTemplates =
+          FluentIterable.from(storeProvider.getJobStore().fetchJobs())
+              .filter(Predicates.compose(Predicates.equalTo(role), JobKeys.CONFIG_TO_ROLE))
+              .uniqueIndex(JobKeys.FROM_CONFIG);
 
-          IResourceAggregate prodConsumed = getConsumption(tasks, updates, cronTemplates, true);
+      IResourceAggregate prodConsumed = getConsumption(tasks, updates, cronTemplates, true);
 
-          IResourceAggregate nonProdConsumed = getConsumption(tasks, updates, cronTemplates, false);
+      IResourceAggregate nonProdConsumed = getConsumption(tasks, updates, cronTemplates, false);
 
-          IResourceAggregate quota =
-              storeProvider.getQuotaStore().fetchQuota(role).or(ResourceAggregates.none());
+      IResourceAggregate quota =
+          storeProvider.getQuotaStore().fetchQuota(role).or(ResourceAggregates.none());
 
-          return new QuotaInfo(quota, prodConsumed, nonProdConsumed);
-        }
-      });
+      return new QuotaInfo(quota, prodConsumed, nonProdConsumed);
     }
 
     private IResourceAggregate getConsumption(

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java b/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
index 3476385..b9dce16 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
@@ -289,13 +289,17 @@ class ReadOnlySchedulerImpl implements ReadOnlyScheduler.Iface {
   @Override
   public Response getQuota(final String ownerRole) {
     MorePreconditions.checkNotBlank(ownerRole);
+    return storage.read(new Quiet<Response>() {
+      @Override
+      public Response apply(StoreProvider storeProvider) {
+        QuotaInfo quotaInfo = quotaManager.getQuotaInfo(ownerRole, storeProvider);
+        GetQuotaResult result = new GetQuotaResult(quotaInfo.getQuota().newBuilder())
+            .setProdConsumption(quotaInfo.getProdConsumption().newBuilder())
+            .setNonProdConsumption(quotaInfo.getNonProdConsumption().newBuilder());
 
-    QuotaInfo quotaInfo = quotaManager.getQuotaInfo(ownerRole);
-    GetQuotaResult result = new GetQuotaResult(quotaInfo.getQuota().newBuilder())
-        .setProdConsumption(quotaInfo.getProdConsumption().newBuilder())
-        .setNonProdConsumption(quotaInfo.getNonProdConsumption().newBuilder());
-
-    return ok(Result.getQuotaResult(result));
+        return ok(Result.getQuotaResult(result));
+      }
+    });
   }
 
   @Override

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/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 d24b31b..059626f 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
@@ -39,6 +39,7 @@ import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Multimaps;
 import com.google.common.collect.Range;
+
 import com.twitter.common.args.Arg;
 import com.twitter.common.args.CmdLine;
 import com.twitter.common.args.constraints.Positive;
@@ -279,7 +280,10 @@ class SchedulerThriftInterface implements AuroraAdmin.Iface {
           ITaskConfig template = sanitized.getJobConfig().getTaskConfig();
           int count = sanitized.getJobConfig().getInstanceCount();
 
-          validateTaskLimits(template, count, quotaManager.checkInstanceAddition(template, count));
+          validateTaskLimits(
+              template,
+              count,
+              quotaManager.checkInstanceAddition(template, count, storeProvider));
 
           LOG.info("Launching " + count + " tasks.");
           stateManager.insertPendingTasks(
@@ -346,7 +350,10 @@ class SchedulerThriftInterface implements AuroraAdmin.Iface {
           ITaskConfig template = sanitized.getJobConfig().getTaskConfig();
           int count = sanitized.getJobConfig().getInstanceCount();
 
-          validateTaskLimits(template, count, quotaManager.checkInstanceAddition(template, count));
+          validateTaskLimits(
+              template,
+              count,
+              quotaManager.checkInstanceAddition(template, count, storeProvider));
 
           // TODO(mchucarroll): Merge CronJobManager.createJob/updateJob
           if (updateOnly || getCronJob(storeProvider, jobKey).isPresent()) {
@@ -664,7 +671,15 @@ class SchedulerThriftInterface implements AuroraAdmin.Iface {
     requireNonNull(session);
 
     try {
-      quotaManager.saveQuota(ownerRole, IResourceAggregate.build(resourceAggregate));
+      storage.write(new MutateWork.NoResult<QuotaException>() {
+        @Override
+        protected void execute(MutableStoreProvider store) throws QuotaException {
+          quotaManager.saveQuota(
+              ownerRole,
+              IResourceAggregate.build(resourceAggregate),
+              store.getQuotaStore());
+        }
+      });
       return ok();
     } catch (QuotaException e) {
       return error(INVALID_REQUEST, e);
@@ -941,7 +956,7 @@ class SchedulerThriftInterface implements AuroraAdmin.Iface {
           validateTaskLimits(
               task,
               currentTasks.size() + config.getInstanceIdsSize(),
-              quotaManager.checkInstanceAddition(task, config.getInstanceIdsSize()));
+              quotaManager.checkInstanceAddition(task, config.getInstanceIdsSize(), storeProvider));
 
           storage.write(new NoResult.Quiet() {
             @Override
@@ -1207,7 +1222,7 @@ class SchedulerThriftInterface implements AuroraAdmin.Iface {
           validateTaskLimits(
               request.getTaskConfig(),
               request.getInstanceCount(),
-              quotaManager.checkJobUpdate(update));
+              quotaManager.checkJobUpdate(update, storeProvider));
 
           jobUpdateController.start(update, context.getIdentity());
           return ok(Result.startJobUpdateResult(new StartJobUpdateResult(updateId)));

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/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 eac2033..b79b07c 100644
--- a/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
@@ -37,6 +37,7 @@ import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.quota.QuotaManager.QuotaException;
 import org.apache.aurora.scheduler.quota.QuotaManager.QuotaManagerImpl;
 import org.apache.aurora.scheduler.storage.JobUpdateStore;
+import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
 import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
@@ -69,12 +70,14 @@ public class QuotaManagerImplTest extends EasyMockTest {
   private StorageTestUtil storageUtil;
   private JobUpdateStore jobUpdateStore;
   private QuotaManagerImpl quotaManager;
+  private StoreProvider storeProvider;
 
   @Before
   public void setUp() throws Exception {
     storageUtil = new StorageTestUtil(this);
+    storeProvider = storageUtil.storeProvider;
     jobUpdateStore = storageUtil.jobUpdateStore;
-    quotaManager = new QuotaManagerImpl(storageUtil.storage);
+    quotaManager = new QuotaManagerImpl();
     storageUtil.expectOperations();
   }
 
@@ -95,7 +98,7 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     assertEquals(
         new QuotaInfo(from(4, 4, 4), from(6, 6, 6), from(9, 9, 9)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -130,7 +133,7 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     assertEquals(
         new QuotaInfo(from(4, 4, 4), from(7, 7, 7), from(10, 10, 10)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -151,7 +154,7 @@ public class QuotaManagerImplTest extends EasyMockTest {
     // Expected consumption from: prodTask + updatingProdTask + job update.
     assertEquals(
         new QuotaInfo(from(4, 4, 4), from(7, 7, 7), from(2, 2, 2)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -167,7 +170,7 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     assertEquals(
         new QuotaInfo(from(4, 4, 4), from(0, 0, 0), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -179,7 +182,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -192,7 +196,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -205,7 +210,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -218,7 +224,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -232,7 +239,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -240,7 +248,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
   public void testCheckQuotaSkippedForNonProdRequest() {
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, false), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, false), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -255,7 +264,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -268,7 +278,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(2, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(2, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertTrue(checkQuota.getDetails().get().contains("CPU"));
   }
@@ -282,7 +293,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 2, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 2, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertTrue(checkQuota.getDetails().get().contains("RAM"));
   }
@@ -296,7 +308,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 2, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 2, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertTrue(checkQuota.getDetails().get().contains("DISK"));
   }
@@ -312,11 +325,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(2, 2, 2, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(2, 2, 2, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(5, 5, 5), from(4, 4, 4), from(7, 7, 7)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -330,11 +344,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(5, 5, 5), from(4, 4, 4), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -347,11 +362,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(5, 5, 5), from(4, 4, 4), from(8, 8, 8)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -364,11 +380,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 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));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -381,11 +398,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(5, 5, 5), from(5, 5, 5), from(1, 1, 1)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -397,11 +415,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(6, 6, 6), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -413,11 +432,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(6, 6, 6), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -429,11 +449,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(6, 6, 6), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -458,11 +479,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(4, 4, 4), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -487,11 +509,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(4, 4, 4), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -516,11 +539,12 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1);
+    QuotaCheckResult checkQuota =
+        quotaManager.checkInstanceAddition(taskConfig(1, 1, 1, true), 1, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(6, 6, 6), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -544,11 +568,11 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update);
+    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(6, 6, 6), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -569,11 +593,11 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update);
+    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider);
     assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(4, 4, 4), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -597,11 +621,11 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update);
+    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
     assertEquals(
         new QuotaInfo(from(6, 6, 6), from(6, 6, 6), from(0, 0, 0)),
-        quotaManager.getQuotaInfo(ROLE));
+        quotaManager.getQuotaInfo(ROLE, storeProvider));
   }
 
   @Test
@@ -616,7 +640,7 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update);
+    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -634,7 +658,8 @@ public class QuotaManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(IJobUpdate.build(updateBuilder));
+    QuotaCheckResult checkQuota =
+        quotaManager.checkJobUpdate(IJobUpdate.build(updateBuilder), storeProvider);
     assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult());
   }
 
@@ -643,19 +668,25 @@ public class QuotaManagerImplTest extends EasyMockTest {
     storageUtil.quotaStore.saveQuota(ROLE, QUOTA);
 
     control.replay();
-    quotaManager.saveQuota(ROLE, QUOTA);
+    quotaManager.saveQuota(ROLE, QUOTA, storageUtil.mutableStoreProvider.getQuotaStore());
   }
 
   @Test(expected = QuotaException.class)
   public void testSaveQuotaFailsMissingSpecs() throws Exception {
     control.replay();
-    quotaManager.saveQuota(ROLE, IResourceAggregate.build(new ResourceAggregate()));
+    quotaManager.saveQuota(
+        ROLE,
+        IResourceAggregate.build(new ResourceAggregate()),
+        storageUtil.mutableStoreProvider.getQuotaStore());
   }
 
   @Test(expected = QuotaException.class)
   public void testSaveQuotaFailsNegativeValues() throws Exception {
     control.replay();
-    quotaManager.saveQuota(ROLE, IResourceAggregate.build(new ResourceAggregate(-2.0, 4, 5)));
+    quotaManager.saveQuota(
+        ROLE,
+        IResourceAggregate.build(new ResourceAggregate(-2.0, 4, 5)),
+        storageUtil.mutableStoreProvider.getQuotaStore());
   }
 
   private IExpectationSetters<?> expectTasks(IScheduledTask... tasks) {

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java b/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
index daaa759..e969b97 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
@@ -281,7 +281,7 @@ public class ReadOnlySchedulerImplTest extends EasyMockTest {
   @Test
   public void testGetQuota() throws Exception {
     QuotaInfo infoMock = createMock(QuotaInfo.class);
-    expect(quotaManager.getQuotaInfo(ROLE)).andReturn(infoMock);
+    expect(quotaManager.getQuotaInfo(ROLE, storageUtil.storeProvider)).andReturn(infoMock);
     expect(infoMock.getQuota()).andReturn(QUOTA);
     expect(infoMock.getProdConsumption()).andReturn(CONSUMED);
     IResourceAggregate nonProdConsumed = IResourceAggregate.build(new ResourceAggregate(1, 0, 0));

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/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 5b0565e..d364188 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
@@ -92,6 +92,7 @@ import org.apache.aurora.scheduler.cron.CronJobManager;
 import org.apache.aurora.scheduler.cron.CronPredictor;
 import org.apache.aurora.scheduler.cron.SanitizedCronJob;
 import org.apache.aurora.scheduler.metadata.NearestFit;
+import org.apache.aurora.scheduler.quota.QuotaCheckResult;
 import org.apache.aurora.scheduler.quota.QuotaManager;
 import org.apache.aurora.scheduler.state.LockManager;
 import org.apache.aurora.scheduler.state.LockManager.LockException;
@@ -281,8 +282,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized, ENOUGH_QUOTA);
 
     stateManager.insertPendingTasks(
         storageUtil.mutableStoreProvider,
@@ -304,8 +304,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized, ENOUGH_QUOTA);
 
     stateManager.insertPendingTasks(
         storageUtil.mutableStoreProvider,
@@ -397,8 +396,10 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     lockManager.validateIfLocked(LOCK_KEY, Optional.<ILock>absent());
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
     expectNoCronJob();
-    expect(quotaManager.checkInstanceAddition(anyObject(ITaskConfig.class), anyInt()))
-        .andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkInstanceAddition(
+        anyObject(ITaskConfig.class),
+        anyInt(),
+        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);
 
     control.replay();
 
@@ -413,8 +414,11 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     lockManager.validateIfLocked(LOCK_KEY, Optional.<ILock>absent());
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
     expectNoCronJob();
-    expect(quotaManager.checkInstanceAddition(anyObject(ITaskConfig.class), anyInt()))
-        .andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkInstanceAddition(
+        anyObject(ITaskConfig.class),
+        anyInt(),
+        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);
+
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(Strings.repeat("a", MAX_TASK_ID_LENGTH + 1));
 
@@ -433,8 +437,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(NOT_ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized, NOT_ENOUGH_QUOTA);
 
     control.replay();
 
@@ -582,8 +585,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     expect(taskIdGenerator.generate(ITaskConfig.build(sanitized.getTaskConfig()), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(ITaskConfig.build(sanitized.getTaskConfig()), 1))
-        .andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(ITaskConfig.build(sanitized.getTaskConfig()), ENOUGH_QUOTA);
     stateManager.insertPendingTasks(
         storageUtil.mutableStoreProvider,
         ITaskConfig.build(sanitized.getTaskConfig()),
@@ -855,7 +857,10 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
         .setNumCpus(10)
         .setDiskMb(100)
         .setRamMb(200);
-    quotaManager.saveQuota(ROLE, IResourceAggregate.build(resourceAggregate));
+    quotaManager.saveQuota(
+        ROLE,
+        IResourceAggregate.build(resourceAggregate),
+        storageUtil.quotaStore);
 
     control.replay();
 
@@ -868,7 +873,11 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
         .setNumCpus(10)
         .setDiskMb(100)
         .setRamMb(200);
-    quotaManager.saveQuota(ROLE, IResourceAggregate.build(resourceAggregate));
+    quotaManager.saveQuota(
+        ROLE,
+        IResourceAggregate.build(resourceAggregate),
+        storageUtil.quotaStore);
+
     expectLastCall().andThrow(new QuotaManager.QuotaException("fail"));
 
     control.replay();
@@ -1041,9 +1050,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
-
+    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
     cronJobManager.updateJob(anyObject(SanitizedCronJob.class));
     control.replay();
 
@@ -1081,9 +1088,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
-
+    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
     cronJobManager.updateJob(anyObject(SanitizedCronJob.class));
     expectLastCall().andThrow(new CronException("Nope"));
 
@@ -1125,8 +1130,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
 
     expectNoCronJob().times(2);
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
@@ -1144,8 +1148,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
 
     expectNoCronJob();
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), buildScheduledTask());
@@ -1164,8 +1167,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), ENOUGH_QUOTA);
 
     expectCronJob();
     cronJobManager.updateJob(SanitizedCronJob.from(sanitized));
@@ -1222,8 +1224,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(taskIdGenerator.generate(sanitized.getJobConfig().getTaskConfig(), 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(sanitized.getJobConfig().getTaskConfig(), 1))
-        .andReturn(NOT_ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), NOT_ENOUGH_QUOTA);
 
     control.replay();
     assertResponse(INVALID_REQUEST, thrift.scheduleCronJob(CRON_JOB, null, SESSION));
@@ -1566,7 +1567,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
     expect(taskIdGenerator.generate(populatedTask, 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(populatedTask, 1)).andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(populatedTask, ENOUGH_QUOTA);
     stateManager.insertPendingTasks(
         storageUtil.mutableStoreProvider,
         populatedTask,
@@ -1590,7 +1591,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
     expect(taskIdGenerator.generate(populatedTask, 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(populatedTask, 1)).andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(populatedTask, ENOUGH_QUOTA);
     stateManager.insertPendingTasks(
         storageUtil.mutableStoreProvider,
         populatedTask,
@@ -1664,8 +1665,10 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     lockManager.validateIfLocked(LOCK_KEY, Optional.of(LOCK));
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
-    expect(quotaManager.checkInstanceAddition(anyObject(ITaskConfig.class), anyInt()))
-        .andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkInstanceAddition(
+        anyObject(ITaskConfig.class),
+        anyInt(),
+        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);
     expect(taskIdGenerator.generate(populatedTask, 1))
         .andReturn(Strings.repeat("a", MAX_TASK_ID_LENGTH + 1));
 
@@ -1684,7 +1687,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
     expect(taskIdGenerator.generate(populatedTask, 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(populatedTask, 1)).andReturn(NOT_ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(populatedTask, NOT_ENOUGH_QUOTA);
 
     control.replay();
 
@@ -1701,7 +1704,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
     expect(taskIdGenerator.generate(populatedTask, 1))
         .andReturn(TASK_ID);
-    expect(quotaManager.checkInstanceAddition(populatedTask, 1)).andReturn(ENOUGH_QUOTA);
+    expectInstanceQuotaCheck(populatedTask, ENOUGH_QUOTA);
     stateManager.insertPendingTasks(
         storageUtil.mutableStoreProvider,
         populatedTask,
@@ -1818,7 +1821,10 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
         oldTask3.getAssignedTask().getTask(), ImmutableSet.of(new Range(2, 3))
     ));
 
-    expect(quotaManager.checkJobUpdate(update)).andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkJobUpdate(
+        update,
+        storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA);
+
     expect(uuidGenerator.createNew()).andReturn(UU_ID);
     storageUtil.expectTaskFetch(
         Query.unscoped().byJob(JOB_KEY).active(),
@@ -1859,7 +1865,9 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
         newTask,
         ImmutableMap.of(oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1))));
 
-    expect(quotaManager.checkJobUpdate(anyObject(IJobUpdate.class))).andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkJobUpdate(
+        anyObject(IJobUpdate.class),
+        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);
 
     // Set diff-adjusted IJobUpdate expectations.
     JobUpdate expected = update.newBuilder();
@@ -2071,7 +2079,9 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     expect(uuidGenerator.createNew()).andReturn(UU_ID);
     storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active());
-    expect(quotaManager.checkJobUpdate(anyObject(IJobUpdate.class))).andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkJobUpdate(
+        anyObject(IJobUpdate.class),
+        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);
 
     control.replay();
 
@@ -2085,7 +2095,10 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     expectNoCronJob();
     expect(uuidGenerator.createNew()).andReturn(UU_ID);
     storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active());
-    expect(quotaManager.checkJobUpdate(anyObject(IJobUpdate.class))).andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkJobUpdate(
+        anyObject(IJobUpdate.class),
+        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);
+
     expect(taskIdGenerator.generate(ITaskConfig.build(request.getTaskConfig()), 6))
         .andReturn(Strings.repeat("a", MAX_TASK_ID_LENGTH + 1));
 
@@ -2110,7 +2123,9 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     IJobUpdate update = buildJobUpdate(6, config, ImmutableMap.of(
         oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))));
 
-    expect(quotaManager.checkJobUpdate(update)).andReturn(NOT_ENOUGH_QUOTA);
+    expect(quotaManager.checkJobUpdate(
+        update,
+        storageUtil.mutableStoreProvider)).andReturn(NOT_ENOUGH_QUOTA);
 
     control.replay();
 
@@ -2132,7 +2147,10 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     expect(uuidGenerator.createNew()).andReturn(UU_ID);
     expect(taskIdGenerator.generate(newTask, 1)).andReturn(TASK_ID);
-    expect(quotaManager.checkJobUpdate(update)).andReturn(ENOUGH_QUOTA);
+    expect(quotaManager.checkJobUpdate(
+        update,
+        storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA);
+
     storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask);
     jobUpdateController.start(update, USER);
     expectLastCall().andThrow(new UpdateStateException("failed"));
@@ -2474,4 +2492,21 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     return expect(storageUtil.jobStore.fetchJob(JOB_KEY))
         .andReturn(Optional.<IJobConfiguration>absent());
   }
+
+  private IExpectationSetters<?> expectInstanceQuotaCheck(
+      SanitizedConfiguration sanitized,
+      QuotaCheckResult result) {
+
+    return expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), result);
+  }
+
+  private IExpectationSetters<?> expectInstanceQuotaCheck(
+      ITaskConfig config,
+      QuotaCheckResult result) {
+
+    return expect(quotaManager.checkInstanceAddition(
+        config,
+        1,
+        storageUtil.mutableStoreProvider)).andReturn(result);
+  }
 }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/01205914/src/test/java/org/apache/aurora/scheduler/thrift/ThriftIT.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/ThriftIT.java b/src/test/java/org/apache/aurora/scheduler/thrift/ThriftIT.java
index 2b8a9e4..bb19fc7 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/ThriftIT.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/ThriftIT.java
@@ -182,7 +182,8 @@ public class ThriftIT extends EasyMockTest {
 
   @Test
   public void testProvisionAccess() throws Exception {
-    quotaManager.saveQuota(USER, QUOTA);
+    storageTestUtil.expectOperations();
+    quotaManager.saveQuota(USER, QUOTA, storageTestUtil.quotaStore);
     expectLastCall().times(2);
 
     control.replay();


Mime
View raw message