ResourceAggregate schema changes for resource refactoring. Reviewed at https://reviews.apache.org/r/46716/
Project: http://git-wip-us.apache.org/repos/asf/aurora/repo Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/7e30ebe3 Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/7e30ebe3 Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/7e30ebe3 Branch: refs/heads/master Commit: 7e30ebe34444b2947f37061144ee0e25b12e5956 Parents: a5b9603 Author: Maxim Khutornenko <[email protected]> Authored: Wed Apr 27 12:05:11 2016 -0700 Committer: Maxim Khutornenko <[email protected]> Committed: Wed Apr 27 12:05:11 2016 -0700 ---------------------------------------------------------------------- .../thrift/org/apache/aurora/gen/api.thrift | 20 +-- .../scheduler/resources/ResourceAggregates.java | 30 +++- .../scheduler/storage/db/DbQuotaStore.java | 24 ++-- .../scheduler/storage/db/DbTaskStore.java | 5 +- .../scheduler/storage/db/QuotaMapper.java | 28 +++- .../scheduler/storage/db/TaskConfigManager.java | 15 +- .../scheduler/storage/db/TaskConfigMapper.java | 8 +- .../aurora/scheduler/storage/db/TaskMapper.java | 6 +- .../V005_CreateQuotaResourceTable.java | 68 +++++++++ .../storage/db/views/DBResourceAggregate.java | 57 ++++++++ .../scheduler/storage/db/views/DBSaveQuota.java | 29 ++++ .../scheduler/storage/db/views/Pairs.java | 8 -- .../scheduler/storage/log/LogStorage.java | 3 +- .../scheduler/storage/log/ThriftBackfill.java | 31 +++- .../thrift/SchedulerThriftInterface.java | 4 +- .../python/apache/aurora/client/api/__init__.py | 8 +- .../aurora/scheduler/storage/db/QuotaMapper.xml | 58 ++++++-- .../scheduler/storage/db/TaskConfigMapper.xml | 12 +- .../aurora/scheduler/storage/db/TaskMapper.xml | 6 +- .../aurora/scheduler/storage/db/schema.sql | 29 ++-- .../scheduler/quota/QuotaManagerImplTest.java | 141 +++++++++---------- .../scheduler/resources/ResourceTestUtil.java | 48 +++++++ .../scheduler/stats/AsyncStatsModuleTest.java | 4 +- .../scheduler/stats/ResourceCounterTest.java | 9 +- .../scheduler/stats/SlotSizeCounterTest.java | 13 +- .../storage/backup/StorageBackupTest.java | 4 +- .../scheduler/storage/db/DbQuotaStoreTest.java | 8 +- .../scheduler/storage/log/LogStorageTest.java | 15 +- .../storage/log/SnapshotStoreImplIT.java | 3 +- .../storage/log/ThriftBackfillTest.java | 28 ++++ .../storage/mem/StorageTransactionTest.java | 8 +- .../aurora/scheduler/thrift/Fixtures.java | 5 +- .../thrift/SchedulerThriftInterfaceTest.java | 16 ++- .../aurora/scheduler/thrift/ThriftIT.java | 5 +- .../python/apache/aurora/client/api/test_api.py | 10 ++ 35 files changed, 545 insertions(+), 221 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/api/src/main/thrift/org/apache/aurora/gen/api.thrift ---------------------------------------------------------------------- diff --git a/api/src/main/thrift/org/apache/aurora/gen/api.thrift b/api/src/main/thrift/org/apache/aurora/gen/api.thrift index b1b2f1d..3847095 100644 --- a/api/src/main/thrift/org/apache/aurora/gen/api.thrift +++ b/api/src/main/thrift/org/apache/aurora/gen/api.thrift @@ -41,15 +41,6 @@ struct Identity { 2: string user } -struct ResourceAggregate { - /** Number of CPU cores allotted. */ - 1: double numCpus - /** Megabytes of RAM allotted. */ - 2: i64 ramMb - /** Megabytes of disk space allotted. */ - 3: i64 diskMb -} - /** A single host attribute. */ struct Attribute { 1: string name @@ -278,6 +269,17 @@ struct TaskConfig { 29: Container container = { "mesos": {} } } +struct ResourceAggregate { + /** Number of CPU cores allotted. */ + 1: double numCpus + /** Megabytes of RAM allotted. */ + 2: i64 ramMb + /** Megabytes of disk space allotted. */ + 3: i64 diskMb + /** Aggregated resource values. */ + 4: set<Resource> resources +} + /** Defines the policy for launching a new cron job when one is already running. */ enum CronCollisionPolicy { /** Kills the existing job with the colliding name, and runs the new cron job. */ http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/resources/ResourceAggregates.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/resources/ResourceAggregates.java b/src/main/java/org/apache/aurora/scheduler/resources/ResourceAggregates.java index 302eb87..1d19b32 100644 --- a/src/main/java/org/apache/aurora/scheduler/resources/ResourceAggregates.java +++ b/src/main/java/org/apache/aurora/scheduler/resources/ResourceAggregates.java @@ -24,19 +24,39 @@ import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; public final class ResourceAggregates { public static final IResourceAggregate EMPTY = - IResourceAggregate.build(new ResourceAggregate(0, 0, 0)); + IResourceAggregate.build(new ResourceAggregate() + .setNumCpus(0) + .setRamMb(0) + .setDiskMb(0) + ); public static final IResourceAggregate SMALL = - IResourceAggregate.build(new ResourceAggregate(1.0, 1024, 4096)); + IResourceAggregate.build(new ResourceAggregate() + .setNumCpus(1.0) + .setRamMb(1024) + .setDiskMb(4096) + ); public static final IResourceAggregate MEDIUM = - IResourceAggregate.build(new ResourceAggregate(4.0, 8192, 16384)); + IResourceAggregate.build(new ResourceAggregate() + .setNumCpus(4.0) + .setRamMb(8192) + .setDiskMb(16384) + ); public static final IResourceAggregate LARGE = - IResourceAggregate.build(new ResourceAggregate(8.0, 16384, 32768)); + IResourceAggregate.build(new ResourceAggregate() + .setNumCpus(8.0) + .setRamMb(16384) + .setDiskMb(32768) + ); public static final IResourceAggregate XLARGE = - IResourceAggregate.build(new ResourceAggregate(16.0, 32768, 65536)); + IResourceAggregate.build(new ResourceAggregate() + .setNumCpus(16.0) + .setRamMb(32768) + .setDiskMb(65536) + ); private ResourceAggregates() { // Utility class. http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/DbQuotaStore.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/DbQuotaStore.java b/src/main/java/org/apache/aurora/scheduler/storage/db/DbQuotaStore.java index ac63e67..e7afbad 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/DbQuotaStore.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/DbQuotaStore.java @@ -14,13 +14,15 @@ package org.apache.aurora.scheduler.storage.db; import java.util.Map; +import java.util.stream.Collectors; import com.google.common.base.Optional; -import com.google.common.collect.ImmutableMap; import com.google.inject.Inject; -import org.apache.aurora.gen.storage.SaveQuota; import org.apache.aurora.scheduler.storage.QuotaStore; +import org.apache.aurora.scheduler.storage.db.views.DBResourceAggregate; +import org.apache.aurora.scheduler.storage.db.views.DBSaveQuota; +import org.apache.aurora.scheduler.storage.db.views.Pairs; import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import static java.util.Objects.requireNonNull; @@ -43,17 +45,15 @@ class DbQuotaStore implements QuotaStore.Mutable { @Override public Optional<IResourceAggregate> fetchQuota(String role) { return Optional.fromNullable(mapper.select(role)) - .transform(IResourceAggregate::build); + .transform(DBResourceAggregate::toImmutable); } @Timed("quota_store_fetch_quotas") @Override public Map<String, IResourceAggregate> fetchQuotas() { - ImmutableMap.Builder<String, IResourceAggregate> results = ImmutableMap.builder(); - for (SaveQuota result : mapper.selectAll()) { - results.put(result.getRole(), IResourceAggregate.build(result.getQuota())); - } - return results.build(); + return Pairs.toMap(mapper.selectAll().stream() + .map(DBSaveQuota::toImmutable) + .collect(Collectors.toList())); } @Timed("quota_store_delete_quotas") @@ -71,6 +71,12 @@ class DbQuotaStore implements QuotaStore.Mutable { @Timed("quota_store_save_quota") @Override public void saveQuota(String role, IResourceAggregate quota) { - mapper.merge(role, quota.newBuilder()); + mapper.delete(role); + InsertResult quotaInsert = new InsertResult(); + mapper.insert(role, quota.newBuilder(), quotaInsert); + mapper.insertResources( + quotaInsert.getId(), + DBResourceAggregate.mapFromResources(quota.getResources())); + } } http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java b/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java index b9ed417..a649a6e 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java @@ -38,7 +38,6 @@ import org.apache.aurora.scheduler.base.Query.Builder; import org.apache.aurora.scheduler.base.Tasks; import org.apache.aurora.scheduler.storage.TaskStore; import org.apache.aurora.scheduler.storage.db.views.DbScheduledTask; -import org.apache.aurora.scheduler.storage.db.views.Pairs; import org.apache.aurora.scheduler.storage.entities.IJobKey; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import org.apache.aurora.scheduler.storage.entities.ITaskConfig; @@ -147,9 +146,7 @@ class DbTaskStore implements TaskStore.Mutable { taskMapper.insertTaskEvents(result.getId(), task.getTaskEvents()); } if (!task.getAssignedTask().getAssignedPorts().isEmpty()) { - taskMapper.insertPorts( - result.getId(), - Pairs.fromMap(task.getAssignedTask().getAssignedPorts())); + taskMapper.insertPorts(result.getId(), task.getAssignedTask().getAssignedPorts()); } } } http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/QuotaMapper.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/QuotaMapper.java b/src/main/java/org/apache/aurora/scheduler/storage/db/QuotaMapper.java index 3e5c434..7ebbb43 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/QuotaMapper.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/QuotaMapper.java @@ -14,11 +14,13 @@ package org.apache.aurora.scheduler.storage.db; import java.util.List; +import java.util.Map; import javax.annotation.Nullable; import org.apache.aurora.gen.ResourceAggregate; -import org.apache.aurora.gen.storage.SaveQuota; +import org.apache.aurora.scheduler.storage.db.views.DBResourceAggregate; +import org.apache.aurora.scheduler.storage.db.views.DBSaveQuota; import org.apache.ibatis.annotations.Param; /** @@ -26,12 +28,26 @@ import org.apache.ibatis.annotations.Param; */ interface QuotaMapper { /** - * Saves the quota for the given {@code role}, updating the existing value if it exists. + * Inserts the quota for the given {@code role}. * - * @param role Role to save quota for. + * @param role Role to insert quota for. * @param quota Quota value to store. + * @param result Container for auto-generated ID of the inserted row. */ - void merge(@Param("role") String role, @Param("quota") ResourceAggregate quota); + void insert( + @Param("role") String role, + @Param("quota") ResourceAggregate quota, + @Param("result") InsertResult result); + + /** + * Insert quota resources. + * + * @param quotaId Quota ID to merge resources for. + * @param values Resources to merge. + */ + void insertResources( + @Param("quotaId") long quotaId, + @Param("values") Map<Integer, String> values); /** * Gets the quota assigned to a role. @@ -40,14 +56,14 @@ interface QuotaMapper { * @return The previously-saved quota for the role, if it exists. */ @Nullable - ResourceAggregate select(String role); + DBResourceAggregate select(String role); /** * Gets all saved quotas. * * @return All quotas stored in the database. */ - List<SaveQuota> selectAll(); + List<DBSaveQuota> selectAll(); /** * Removes the quota stored for a role. http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigManager.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigManager.java b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigManager.java index 17b2490..07bdf22 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigManager.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigManager.java @@ -20,11 +20,8 @@ import javax.inject.Inject; import com.google.common.collect.Maps; -import org.apache.aurora.GuavaUtils; -import org.apache.aurora.common.collections.Pair; -import org.apache.aurora.scheduler.resources.ResourceType; +import org.apache.aurora.scheduler.storage.db.views.DBResourceAggregate; import org.apache.aurora.scheduler.storage.db.views.DbTaskConfig; -import org.apache.aurora.scheduler.storage.db.views.Pairs; import org.apache.aurora.scheduler.storage.entities.IAppcImage; import org.apache.aurora.scheduler.storage.entities.IConstraint; import org.apache.aurora.scheduler.storage.entities.IDockerContainer; @@ -102,11 +99,7 @@ class TaskConfigManager { if (!config.getResources().isEmpty()) { configMapper.insertResources( configInsert.getId(), - config.getResources().stream() - .map(e -> Pair.of( - ResourceType.fromResource(e).getValue(), - ResourceType.fromResource(e).getTypeConverter().stringify(e.getRawValue()))) - .collect(GuavaUtils.toImmutableList())); + DBResourceAggregate.mapFromResources(config.getResources())); } if (!config.getRequestedPorts().isEmpty()) { @@ -114,9 +107,7 @@ class TaskConfigManager { } if (!config.getTaskLinks().isEmpty()) { - configMapper.insertTaskLinks( - configInsert.getId(), - Pairs.fromMap(config.getTaskLinks())); + configMapper.insertTaskLinks(configInsert.getId(), config.getTaskLinks()); } if (!config.getMetadata().isEmpty()) { http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.java b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.java index 04666ad..6e06af4 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.java @@ -14,9 +14,9 @@ package org.apache.aurora.scheduler.storage.db; import java.util.List; +import java.util.Map; import java.util.Set; -import org.apache.aurora.common.collections.Pair; import org.apache.aurora.scheduler.storage.db.views.DbTaskConfig; import org.apache.aurora.scheduler.storage.entities.IConstraint; import org.apache.aurora.scheduler.storage.entities.IDockerContainer; @@ -111,9 +111,7 @@ interface TaskConfigMapper extends GarbageCollectedTableMapper { * @param configId Task config ID. * @param links Task links to insert. */ - void insertTaskLinks( - @Param("configId") long configId, - @Param("links") List<Pair<String, String>> links); + void insertTaskLinks(@Param("configId") long configId, @Param("links") Map<String, String> links); /** * Inserts the container association within an {@link ITaskConfig}. @@ -185,5 +183,5 @@ interface TaskConfigMapper extends GarbageCollectedTableMapper { */ void insertResources( @Param("configId") long configId, - @Param("values") List<Pair<Integer, String>> values); + @Param("values") Map<Integer, String> values); } http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java index ea12165..cbcef84 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java @@ -14,11 +14,11 @@ package org.apache.aurora.scheduler.storage.db; import java.util.List; +import java.util.Map; import java.util.Set; import javax.annotation.Nullable; -import org.apache.aurora.common.collections.Pair; import org.apache.aurora.gen.JobKey; import org.apache.aurora.scheduler.storage.db.views.DbScheduledTask; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; @@ -83,9 +83,7 @@ interface TaskMapper { * @param taskRowId Task row ID. * @param ports Assigned ports to insert. */ - void insertPorts( - @Param("taskRowId") long taskRowId, - @Param("ports") List<Pair<String, Integer>> ports); + void insertPorts(@Param("taskRowId") long taskRowId, @Param("ports") Map<String, Integer> ports); /** * Deletes all task rows. http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/migration/V005_CreateQuotaResourceTable.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/migration/V005_CreateQuotaResourceTable.java b/src/main/java/org/apache/aurora/scheduler/storage/db/migration/V005_CreateQuotaResourceTable.java new file mode 100644 index 0000000..e9e5de1 --- /dev/null +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/migration/V005_CreateQuotaResourceTable.java @@ -0,0 +1,68 @@ +/** + * 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.scheduler.storage.db.migration; + +import java.math.BigDecimal; + +import org.apache.aurora.scheduler.resources.ResourceType; +import org.apache.ibatis.migration.MigrationScript; + +public class V005_CreateQuotaResourceTable implements MigrationScript { + @Override + public BigDecimal getId() { + return BigDecimal.valueOf(5L); + } + + @Override + public String getDescription() { + return "Create the quota_resource table."; + } + + @Override + public String getUpScript() { + return "CREATE TABLE IF NOT EXISTS quota_resource(" + + "id IDENTITY," + + "quota_id BIGINT NOT NULL REFERENCES quotas(id) ON DELETE CASCADE," + + "type_id INT NOT NULL REFERENCES resource_types(id)," + + "value VARCHAR NOT NULL," + + "UNIQUE(quota_id, type_id, value)" + + ");\n" + + migrateScript(); + } + + @Override + public String getDownScript() { + return "DROP TABLE IF EXISTS quota_resource;"; + } + + private static String migrateScript() { + return "MERGE INTO quota_resource(quota_id, type_id, value) " + + "KEY(quota_id, type_id, value) " + + "SELECT q.id, rt.id, q.num_cpus FROM quotas q " + + "JOIN resource_types rt ON rt.name = " + + String.format("'%s';%n", ResourceType.CPUS.name()) + + + "MERGE INTO quota_resource(quota_id, type_id, value) " + + "KEY(quota_id, type_id, value) " + + "SELECT q.id, rt.id, q.ram_mb FROM quotas q " + + "JOIN resource_types rt ON rt.name = " + + String.format("'%s';%n", ResourceType.RAM_MB.name()) + + + "MERGE INTO quota_resource(quota_id, type_id, value) " + + "KEY(quota_id, type_id, value) " + + "SELECT q.id, rt.id, q.disk_mb FROM quotas q " + + "JOIN resource_types rt ON rt.name = " + + String.format("'%s';%n", ResourceType.DISK_MB.name()); + } +} http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBResourceAggregate.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBResourceAggregate.java b/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBResourceAggregate.java new file mode 100644 index 0000000..4f5c07e --- /dev/null +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBResourceAggregate.java @@ -0,0 +1,57 @@ +/** + * 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.scheduler.storage.db.views; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import org.apache.aurora.common.collections.Pair; +import org.apache.aurora.gen.ResourceAggregate; +import org.apache.aurora.scheduler.resources.ResourceType; +import org.apache.aurora.scheduler.storage.entities.IResource; +import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; + +import static org.apache.aurora.GuavaUtils.toImmutableSet; + +public final class DBResourceAggregate { + private double numCpus; + private long ramMb; + private long diskMb; + private List<DBResource> resources; + + private DBResourceAggregate() { + } + + public ResourceAggregate toThrift() { + return new ResourceAggregate() + .setNumCpus(numCpus) + .setRamMb(ramMb) + .setDiskMb(diskMb) + .setResources(resources.stream().map(DBResource::toThrift).collect(toImmutableSet())); + } + + public static Map<Integer, String> mapFromResources(Set<IResource> resources) { + return Pairs.toMap(resources.stream() + .map(e -> Pair.of( + ResourceType.fromResource(e).getValue(), + ResourceType.fromResource(e).getTypeConverter().stringify(e.getRawValue()))) + .collect(Collectors.toList())); + } + + public IResourceAggregate toImmutable() { + return IResourceAggregate.build(toThrift()); + } +} http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBSaveQuota.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBSaveQuota.java b/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBSaveQuota.java new file mode 100644 index 0000000..3d48592 --- /dev/null +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/views/DBSaveQuota.java @@ -0,0 +1,29 @@ +/** + * 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.scheduler.storage.db.views; + +import org.apache.aurora.common.collections.Pair; +import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; + +public final class DBSaveQuota { + private String role; + private DBResourceAggregate quota; + + private DBSaveQuota() { + } + + public Pair<String, IResourceAggregate> toImmutable() { + return Pair.of(role, quota.toImmutable()); + } +} http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/db/views/Pairs.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/views/Pairs.java b/src/main/java/org/apache/aurora/scheduler/storage/db/views/Pairs.java index 922578f..106b7af 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/db/views/Pairs.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/db/views/Pairs.java @@ -13,12 +13,10 @@ */ package org.apache.aurora.scheduler.storage.db.views; -import java.util.List; import java.util.Map; import com.google.common.collect.ImmutableMap; -import org.apache.aurora.GuavaUtils; import org.apache.aurora.common.collections.Pair; /** @@ -37,10 +35,4 @@ public final class Pairs { } return map.build(); } - - public static <K, V> List<Pair<K, V>> fromMap(Map<K, V> map) { - return map.entrySet().stream() - .map((e) -> Pair.of(e.getKey(), e.getValue())) - .collect(GuavaUtils.toImmutableList()); - } } http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java b/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java index 39b6567..1ddc117 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java @@ -67,7 +67,6 @@ import org.apache.aurora.scheduler.storage.entities.IJobUpdateEvent; import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey; import org.apache.aurora.scheduler.storage.entities.ILock; import org.apache.aurora.scheduler.storage.entities.ILockKey; -import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -348,7 +347,7 @@ public class LogStorage implements NonVolatileStorage, DistributedSnapshotStore SaveQuota saveQuota = op.getSaveQuota(); writeBehindQuotaStore.saveQuota( saveQuota.getRole(), - IResourceAggregate.build(saveQuota.getQuota())); + ThriftBackfill.backfillResourceAggregate(saveQuota.getQuota())); }) .put( Op._Fields.REMOVE_QUOTA, http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/java/org/apache/aurora/scheduler/storage/log/ThriftBackfill.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/storage/log/ThriftBackfill.java b/src/main/java/org/apache/aurora/scheduler/storage/log/ThriftBackfill.java index 9c86aa0..ad45085 100644 --- a/src/main/java/org/apache/aurora/scheduler/storage/log/ThriftBackfill.java +++ b/src/main/java/org/apache/aurora/scheduler/storage/log/ThriftBackfill.java @@ -20,12 +20,14 @@ import org.apache.aurora.gen.JobConfiguration; import org.apache.aurora.gen.JobUpdate; import org.apache.aurora.gen.JobUpdateInstructions; import org.apache.aurora.gen.Resource; +import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.gen.ScheduledTask; import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.scheduler.resources.ResourceType; import org.apache.aurora.scheduler.storage.entities.IJobConfiguration; import org.apache.aurora.scheduler.storage.entities.IJobUpdate; import org.apache.aurora.scheduler.storage.entities.IResource; +import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import static org.apache.aurora.scheduler.resources.ResourceType.CPUS; @@ -41,8 +43,8 @@ public final class ThriftBackfill { // Utility class. } - private static Resource getResource(TaskConfig config, ResourceType type) { - return config.getResources().stream() + private static Resource getResource(Set<Resource> resources, ResourceType type) { + return resources.stream() .filter(e -> ResourceType.fromResource(IResource.build(e)).equals(type)) .findFirst() .orElseThrow(() -> new IllegalArgumentException("Missing resource definition for " + type)); @@ -65,9 +67,9 @@ public final class ThriftBackfill { } } } else { - config.setNumCpus(getResource(config, CPUS).getNumCpus()); - config.setRamMb(getResource(config, RAM_MB).getRamMb()); - config.setDiskMb(getResource(config, DISK_MB).getDiskMb()); + config.setNumCpus(getResource(config.getResources(), CPUS).getNumCpus()); + config.setRamMb(getResource(config.getResources(), RAM_MB).getRamMb()); + config.setDiskMb(getResource(config.getResources(), DISK_MB).getDiskMb()); Set<String> ports = config.getResources().stream() .filter(e -> ResourceType.fromResource(IResource.build(e)).equals(PORTS)) .map(Resource::getNamedPort) @@ -103,6 +105,25 @@ public final class ThriftBackfill { .collect(GuavaUtils.toImmutableSet()); } + /** + * Ensures ResourceAggregate.resources and correspondent deprecated fields are all populated. + * + * @param aggregate ResourceAggregate to backfill. + * @return Backfilled IResourceAggregate. + */ + public static IResourceAggregate backfillResourceAggregate(ResourceAggregate aggregate) { + if (!aggregate.isSetResources() || aggregate.getResources().isEmpty()) { + aggregate.addToResources(Resource.numCpus(aggregate.getNumCpus())); + aggregate.addToResources(Resource.ramMb(aggregate.getRamMb())); + aggregate.addToResources(Resource.diskMb(aggregate.getDiskMb())); + } else { + aggregate.setNumCpus(getResource(aggregate.getResources(), CPUS).getNumCpus()); + aggregate.setRamMb(getResource(aggregate.getResources(), RAM_MB).getRamMb()); + aggregate.setDiskMb(getResource(aggregate.getResources(), DISK_MB).getDiskMb()); + } + return IResourceAggregate.build(aggregate); + } + private static ScheduledTask backfillScheduledTask(ScheduledTask task) { backfillTask(task.getAssignedTask().getTask()); return task; http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/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 14abebb..4d032b9 100644 --- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java +++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java @@ -106,9 +106,9 @@ import org.apache.aurora.scheduler.storage.entities.IJobUpdateRequest; import org.apache.aurora.scheduler.storage.entities.IJobUpdateSettings; import org.apache.aurora.scheduler.storage.entities.ILockKey; import org.apache.aurora.scheduler.storage.entities.IRange; -import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import org.apache.aurora.scheduler.storage.entities.ITaskConfig; +import org.apache.aurora.scheduler.storage.log.ThriftBackfill; import org.apache.aurora.scheduler.thrift.aop.AnnotatedAuroraAdmin; import org.apache.aurora.scheduler.thrift.auth.DecoratedThrift; import org.apache.aurora.scheduler.updater.JobDiff; @@ -493,7 +493,7 @@ class SchedulerThriftInterface implements AnnotatedAuroraAdmin { try { storage.write((NoResult<QuotaException>) store -> quotaManager.saveQuota( ownerRole, - IResourceAggregate.build(resourceAggregate), + ThriftBackfill.backfillResourceAggregate(resourceAggregate), store)); return ok(); } catch (QuotaException e) { http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/python/apache/aurora/client/api/__init__.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/api/__init__.py b/src/main/python/apache/aurora/client/api/__init__.py index bbbe67b..68baf8f 100644 --- a/src/main/python/apache/aurora/client/api/__init__.py +++ b/src/main/python/apache/aurora/client/api/__init__.py @@ -31,6 +31,7 @@ from gen.apache.aurora.api.ttypes import ( JobUpdateKey, JobUpdateQuery, JobUpdateRequest, + Resource, ResourceAggregate, TaskQuery ) @@ -286,7 +287,12 @@ class AuroraClientAPI(object): def set_quota(self, role, cpu, ram, disk): log.info("Setting quota for user:%s cpu:%f ram:%d disk: %d" % (role, cpu, ram, disk)) - return self._scheduler_proxy.setQuota(role, ResourceAggregate(cpu, ram, disk)) + return self._scheduler_proxy.setQuota( + role, + ResourceAggregate(cpu, ram, disk, frozenset([ + Resource(numCpus=cpu), + Resource(ramMb=ram), + Resource(diskMb=disk)]))) def force_task_state(self, task_id, status): log.info("Requesting that task %s transition to state %s" % (task_id, status)) http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/resources/org/apache/aurora/scheduler/storage/db/QuotaMapper.xml ---------------------------------------------------------------------- diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/QuotaMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/QuotaMapper.xml index 0283ec1..0f46968 100644 --- a/src/main/resources/org/apache/aurora/scheduler/storage/db/QuotaMapper.xml +++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/QuotaMapper.xml @@ -16,15 +16,13 @@ PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd";> <mapper namespace="org.apache.aurora.scheduler.storage.db.QuotaMapper"> - <insert id="merge"> - MERGE INTO quotas ( - id, + <insert id="insert" useGeneratedKeys="true" keyColumn="id" keyProperty="result.id"> + INSERT INTO quotas ( role, num_cpus, ram_mb, disk_mb - ) KEY(role) VALUES ( - null, + ) VALUES ( #{role}, #{quota.numCpus}, #{quota.ramMb}, @@ -32,18 +30,54 @@ ) </insert> - <select id="select" resultType="org.apache.aurora.gen.ResourceAggregate"> - SELECT * FROM quotas - WHERE role = #{id} - </select> + <insert id="insertResources"> + INSERT INTO quota_resource ( + quota_id, + type_id, + value + ) VALUES ( + <foreach index="type" item="value" collection="values" separator="),("> + #{quotaId}, + #{type}, + #{value} + </foreach> + ) + </insert> - <resultMap id="quotaResultMap" type="org.apache.aurora.gen.storage.SaveQuota"> + <resultMap id="quotaMap" type="org.apache.aurora.scheduler.storage.db.views.DBResourceAggregate"> <id column="id" /> - <association property="quota" javaType="org.apache.aurora.gen.ResourceAggregate" /> + <collection + property="resources" + columnPrefix="qr_" + resultMap="org.apache.aurora.scheduler.storage.db.TaskConfigMapper.resourceMap"/> </resultMap> + <resultMap id="quotaResultMap" type="org.apache.aurora.scheduler.storage.db.views.DBSaveQuota"> + <id column="id" /> + <association property="quota" resultMap="quotaMap" /> + </resultMap> + + <sql id="unscopedSelect"> + SELECT + q.id, + q.role, + q.num_cpus, + q.ram_mb, + q.disk_mb, + qr.id as qr_id, + qr.type_id as qr_type_id, + qr.value as qr_value + FROM quotas AS q + INNER JOIN quota_resource AS qr ON qr.quota_id = q.id + </sql> + + <select id="select" resultMap="quotaMap"> + <include refid="unscopedSelect"/> + WHERE role = #{id} + </select> + <select id="selectAll" resultMap="quotaResultMap"> - SELECT * FROM quotas + <include refid="unscopedSelect"/> </select> <delete id="delete"> http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml ---------------------------------------------------------------------- diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml index 4162797..6b5e130 100644 --- a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml +++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskConfigMapper.xml @@ -284,10 +284,10 @@ type_id, value ) VALUES ( - <foreach item="value" collection="values" separator="),("> + <foreach index="type" item="value" collection="values" separator="),("> #{configId}, - #{value.first}, - #{value.second} + #{type}, + #{value} </foreach> ) </insert> @@ -298,10 +298,10 @@ label, url ) VALUES ( - <foreach item="link" collection="links" separator="),("> + <foreach index="label" item="url" collection="links" separator="),("> #{configId}, - #{link.first}, - #{link.second} + #{label}, + #{url} </foreach> ) </insert> http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml ---------------------------------------------------------------------- diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml index b75e2b0..0ecffab 100644 --- a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml +++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml @@ -207,10 +207,10 @@ name, port ) VALUES ( - <foreach item="port" collection="ports" separator="),("> + <foreach index="name" item="port" collection="ports" separator="),("> #{taskRowId}, - #{port.first}, - #{port.second} + #{name}, + #{port} </foreach> ) </insert> http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql ---------------------------------------------------------------------- diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql b/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql index 45ec1bf..5d905db 100644 --- a/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql +++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql @@ -43,16 +43,6 @@ CREATE TABLE locks( UNIQUE(token) ); -CREATE TABLE quotas( - id IDENTITY, - role VARCHAR NOT NULL, - num_cpus DOUBLE NOT NULL, - ram_mb BIGINT NOT NULL, - disk_mb BIGINT NOT NULL, - - UNIQUE(role) -); - CREATE TABLE maintenance_modes( id INT PRIMARY KEY, name VARCHAR NOT NULL, @@ -116,6 +106,25 @@ CREATE TABLE task_resource( UNIQUE(task_config_id, type_id, value) ); +CREATE TABLE quotas( + id IDENTITY, + role VARCHAR NOT NULL, + num_cpus DOUBLE NOT NULL, + ram_mb BIGINT NOT NULL, + disk_mb BIGINT NOT NULL, + + UNIQUE(role) +); + +CREATE TABLE quota_resource( + id IDENTITY, + quota_id BIGINT NOT NULL REFERENCES quotas(id) ON DELETE CASCADE, + type_id INT NOT NULL REFERENCES resource_types(id), + value VARCHAR NOT NULL, + + UNIQUE(quota_id, type_id, value) +); + CREATE TABLE task_constraints( id IDENTITY, task_config_id BIGINT NOT NULL REFERENCES task_configs(id) ON DELETE CASCADE, http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/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 fc01821..d591d0f 100644 --- a/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java +++ b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java @@ -29,7 +29,6 @@ import org.apache.aurora.gen.JobUpdateInstructions; import org.apache.aurora.gen.JobUpdateKey; import org.apache.aurora.gen.JobUpdateSummary; import org.apache.aurora.gen.Range; -import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.gen.ScheduledTask; import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.gen.TaskConstraint; @@ -39,6 +38,7 @@ import org.apache.aurora.scheduler.base.Query; import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.quota.QuotaManager.QuotaException; import org.apache.aurora.scheduler.quota.QuotaManager.QuotaManagerImpl; +import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.storage.JobUpdateStore; import org.apache.aurora.scheduler.storage.Storage.StoreProvider; import org.apache.aurora.scheduler.storage.entities.IJobConfiguration; @@ -69,10 +69,7 @@ public class QuotaManagerImplTest extends EasyMockTest { private static final String JOB_NAME = "job"; private static final IJobUpdateKey UPDATE_KEY = IJobUpdateKey.build(new JobUpdateKey(JobKeys.from(ROLE, ENV, JOB_NAME).newBuilder(), "u1")); - private static final IResourceAggregate QUOTA = IResourceAggregate.build(new ResourceAggregate() - .setNumCpus(1.0) - .setRamMb(100L) - .setDiskMb(200L)); + private static final IResourceAggregate QUOTA = aggregate(1.0, 100L, 200L); private static final Query.Builder ACTIVE_QUERY = Query.roleScoped(ROLE).active(); private StorageTestUtil storageUtil; @@ -95,7 +92,7 @@ public class QuotaManagerImplTest extends EasyMockTest { IScheduledTask prodDedicatedTask = prodDedicatedTask("foo2", 5, 5, 5); IScheduledTask nonProdSharedTask = nonProdTask("bar1", 2, 2, 2); IScheduledTask nonProdDedicatedTask = nonProdDedicatedTask("bar2", 7, 7, 7); - IResourceAggregate quota = IResourceAggregate.build(new ResourceAggregate(4, 4, 4)); + IResourceAggregate quota = aggregate(4, 4, 4); expectQuota(quota); expectTasks(prodSharedTask, nonProdSharedTask, prodDedicatedTask, nonProdDedicatedTask); @@ -107,7 +104,12 @@ public class QuotaManagerImplTest extends EasyMockTest { control.replay(); assertEquals( - new QuotaInfo(from(4, 4, 4), from(6, 6, 6), from(5, 5, 5), from(9, 9, 9), from(7, 7, 7)), + new QuotaInfo( + aggregate(4, 4, 4), + aggregate(6, 6, 6), + aggregate(5, 5, 5), + aggregate(9, 9, 9), + aggregate(7, 7, 7)), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @@ -115,7 +117,7 @@ public class QuotaManagerImplTest extends EasyMockTest { public void testGetQuotaInfoWithCronTasks() { IScheduledTask prodTask = prodTask("pc", 6, 6, 6); IScheduledTask nonProdTask = prodTask("npc", 7, 7, 7); - IResourceAggregate quota = IResourceAggregate.build(new ResourceAggregate(4, 4, 4)); + IResourceAggregate quota = aggregate(4, 4, 4); expectQuota(quota); expectTasks(prodTask, nonProdTask); @@ -140,7 +142,7 @@ public class QuotaManagerImplTest extends EasyMockTest { control.replay(); assertEquals( - new QuotaInfo(from(4, 4, 4), from(7, 7, 7), EMPTY, from(10, 10, 10), EMPTY), + new QuotaInfo(aggregate(4, 4, 4), aggregate(7, 7, 7), EMPTY, aggregate(10, 10, 10), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @@ -150,7 +152,7 @@ public class QuotaManagerImplTest extends EasyMockTest { IScheduledTask updatingProdTask = createTask(JOB_NAME, "id1", 3, 3, 3, true, 1); IScheduledTask updatingFilteredProdTask = createTask(JOB_NAME, "id0", 3, 3, 3, true, 0); IScheduledTask nonProdTask = createTask("bar", "id1", 2, 2, 2, false, 0); - IResourceAggregate quota = IResourceAggregate.build(new ResourceAggregate(4, 4, 4)); + IResourceAggregate quota = aggregate(4, 4, 4); expectQuota(quota); expectTasks(prodTask, updatingProdTask, updatingFilteredProdTask, nonProdTask); @@ -161,13 +163,13 @@ public class QuotaManagerImplTest extends EasyMockTest { // Expected consumption from: prodTask + updatingProdTask + job update. assertEquals( - new QuotaInfo(from(4, 4, 4), from(7, 7, 7), EMPTY, from(2, 2, 2), EMPTY), + new QuotaInfo(aggregate(4, 4, 4), aggregate(7, 7, 7), EMPTY, aggregate(2, 2, 2), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testGetQuotaInfoNoTasksNoUpdatesNoCronJobs() { - IResourceAggregate quota = IResourceAggregate.build(new ResourceAggregate(4, 4, 4)); + IResourceAggregate quota = aggregate(4, 4, 4); expectQuota(quota); expectNoTasks(); @@ -177,13 +179,13 @@ public class QuotaManagerImplTest extends EasyMockTest { control.replay(); assertEquals( - new QuotaInfo(from(4, 4, 4), from(0, 0, 0), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(4, 4, 4), aggregate(0, 0, 0), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaPasses() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectTasks(prodTask("foo", 2, 2, 2)); expectJobUpdates(taskConfig(1, 1, 1, true), taskConfig(1, 1, 1, true)); expectNoCronJobs(); @@ -197,7 +199,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaPassesNoTasks() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectNoTasks(); expectJobUpdates(taskConfig(1, 1, 1, true), taskConfig(1, 1, 1, true)); expectNoCronJobs(); @@ -211,7 +213,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaPassesNoUpdates() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectTasks(prodTask("foo", 2, 2, 2)); expectNoJobUpdates(); expectNoCronJobs(); @@ -225,7 +227,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaPassesNoTasksNoUpdates() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectNoTasks(); expectNoJobUpdates(); expectNoCronJobs(); @@ -239,7 +241,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaPassesNonProdUnaccounted() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectTasks(prodTask("foo", 2, 2, 2), createTask("bar", "id2", 5, 5, 5, false, 0)); expectNoJobUpdates(); @@ -290,7 +292,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaExceedsCpu() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectTasks(prodTask("foo", 3, 3, 3)); expectNoJobUpdates(); expectNoCronJobs(); @@ -305,7 +307,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaExceedsRam() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectTasks(prodTask("foo", 3, 3, 3)); expectNoJobUpdates(); expectNoCronJobs(); @@ -320,7 +322,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaExceedsDisk() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(4, 4, 4))); + expectQuota(aggregate(4, 4, 4)); expectTasks(prodTask("foo", 3, 3, 3)); expectNoJobUpdates(); expectNoCronJobs(); @@ -335,7 +337,7 @@ public class QuotaManagerImplTest extends EasyMockTest { @Test public void testCheckQuotaExceedsCron() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectNoTasks().times(2); expectNoJobUpdates().times(2); expectCronJobs( @@ -348,13 +350,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(7, 7, 7), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, aggregate(7, 7, 7), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaUpdatingTasksFilteredOut() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectTasks(prodTask("foo", 2, 2, 2), createTask(JOB_NAME, "id2", 3, 3, 3, true, 0)) .times(2); @@ -367,13 +369,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaNonProdUpdatesUnaccounted() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 2, 2, 2)).times(2); expectJobUpdates(taskConfig(8, 8, 8, false), taskConfig(4, 4, 4, false), 2); @@ -385,13 +387,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(8, 8, 8), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, aggregate(8, 8, 8), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaProdToNonUpdateUnaccounted() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 1, 1, 1)).times(2); expectJobUpdates(taskConfig(1, 1, 1, true), taskConfig(7, 7, 7, false), 2); @@ -403,13 +405,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(7, 7, 7), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, aggregate(7, 7, 7), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaNonToProdUpdateExceedsQuota() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 2, 2, 2)).times(2); expectJobUpdates(taskConfig(1, 1, 1, false), taskConfig(1, 1, 1, true), 2); @@ -421,13 +423,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(1, 1, 1), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(5, 5, 5), EMPTY, aggregate(1, 1, 1), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaOldJobUpdateConfigMatters() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 2, 2, 2)).times(2); expectJobUpdates(taskConfig(2, 2, 2, true), taskConfig(1, 1, 1, true), 2); expectNoCronJobs().times(2); @@ -438,13 +440,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(6, 6, 6), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaUpdateAddsInstances() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 2, 2, 2)).times(2); expectJobUpdates(taskConfig(1, 1, 1, true), 1, taskConfig(1, 1, 1, true), 2, 2); expectNoCronJobs().times(2); @@ -455,13 +457,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(6, 6, 6), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaUpdateRemovesInstances() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 2, 2, 2)).times(2); expectJobUpdates(taskConfig(1, 1, 1, true), 2, taskConfig(1, 1, 1, true), 1, 2); expectNoCronJobs().times(2); @@ -472,13 +474,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(6, 6, 6), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaUpdateInitialConfigsUsedForFiltering() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask(JOB_NAME, 2, 2, 2)).times(2); ITaskConfig config = taskConfig(2, 2, 2, true); @@ -501,13 +503,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(4, 4, 4), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaUpdateDesiredConfigsUsedForFiltering() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask(JOB_NAME, 2, 2, 2)).times(2); ITaskConfig config = taskConfig(2, 2, 2, true); @@ -530,13 +532,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(4, 4, 4), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaNoDesiredState() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask("bar", 2, 2, 2)).times(2); ITaskConfig config = taskConfig(2, 2, 2, true); @@ -559,13 +561,13 @@ public class QuotaManagerImplTest extends EasyMockTest { 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), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(6, 6, 6), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaNewInPlaceUpdate() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks( prodTask("foo", 2, 2, 2), createTask(JOB_NAME, "id1", 2, 2, 2, true, 0), @@ -587,13 +589,13 @@ public class QuotaManagerImplTest extends EasyMockTest { QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider); assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(6, 6, 6), from(6, 6, 6), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(6, 6, 6), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaNewUpdateAddsInstances() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks(prodTask("foo", 2, 2, 2), prodTask(JOB_NAME, 2, 2, 2)).times(2); expectNoJobUpdates().times(2); @@ -612,13 +614,13 @@ public class QuotaManagerImplTest extends EasyMockTest { QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider); assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(6, 6, 6), from(4, 4, 4), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(4, 4, 4), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaNewUpdateRemovesInstances() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(6, 6, 6))).times(2); + expectQuota(aggregate(6, 6, 6)).times(2); expectTasks( prodTask("foo", 2, 2, 2), createTask(JOB_NAME, "id1", 2, 2, 2, true, 0), @@ -640,7 +642,7 @@ public class QuotaManagerImplTest extends EasyMockTest { QuotaCheckResult checkQuota = quotaManager.checkJobUpdate(update, storeProvider); assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(6, 6, 6), from(6, 6, 6), EMPTY, from(0, 0, 0), EMPTY), + new QuotaInfo(aggregate(6, 6, 6), aggregate(6, 6, 6), EMPTY, aggregate(0, 0, 0), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @@ -701,7 +703,7 @@ public class QuotaManagerImplTest extends EasyMockTest { expectNoCronJobs(); IScheduledTask prodTask = prodTask("foo", 1, 1, 1); expectTasks(prodTask); - expectQuota(IResourceAggregate.build(new ResourceAggregate(1, 1, 1))); + expectQuota(aggregate(1, 1, 1)); storageUtil.quotaStore.saveQuota(ROLE, QUOTA); @@ -717,24 +719,18 @@ public class QuotaManagerImplTest extends EasyMockTest { expectNoJobUpdates(); expectNoCronJobs(); expectNoTasks(); - expectQuota(IResourceAggregate.build(new ResourceAggregate(1, 1, 1))); + expectQuota(aggregate(1, 1, 1)); storageUtil.quotaStore.saveQuota(ROLE, EMPTY); control.replay(); - quotaManager.saveQuota( - ROLE, - IResourceAggregate.build(new ResourceAggregate()), - storageUtil.mutableStoreProvider); + quotaManager.saveQuota(ROLE, aggregate(0, 0, 0), storageUtil.mutableStoreProvider); } @Test(expected = QuotaException.class) public void testSaveQuotaFailsNegativeValues() throws Exception { control.replay(); - quotaManager.saveQuota( - ROLE, - IResourceAggregate.build(new ResourceAggregate(-2.0, 4, 5)), - storageUtil.mutableStoreProvider); + quotaManager.saveQuota(ROLE, aggregate(-2.0, 4, 5), storageUtil.mutableStoreProvider); } @Test(expected = QuotaException.class) @@ -743,19 +739,16 @@ public class QuotaManagerImplTest extends EasyMockTest { expectNoCronJobs(); IScheduledTask prodTask = prodTask("foo", 10, 100, 100); expectTasks(prodTask); - expectQuota(IResourceAggregate.build(new ResourceAggregate(20, 200, 200))); + expectQuota(aggregate(20, 200, 200)); control.replay(); - quotaManager.saveQuota( - ROLE, - IResourceAggregate.build(new ResourceAggregate(1, 1, 1)), - storageUtil.mutableStoreProvider); + quotaManager.saveQuota(ROLE, aggregate(1, 1, 1), storageUtil.mutableStoreProvider); } @Test public void testCheckQuotaCronUpdateDownsize() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectNoTasks().times(2); expectNoJobUpdates().times(2); @@ -769,13 +762,13 @@ public class QuotaManagerImplTest extends EasyMockTest { quotaManager.checkCronUpdate(createJob(prodTask("pc", 1, 1, 1), 2), storeProvider); assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(5, 5, 5), from(4, 4, 4), EMPTY, from(7, 7, 7), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, aggregate(7, 7, 7), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaCronUpdateUpsize() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectNoTasks().times(2); expectNoJobUpdates().times(2); @@ -789,13 +782,13 @@ public class QuotaManagerImplTest extends EasyMockTest { quotaManager.checkCronUpdate(createJob(prodTask("pc", 5, 5, 5), 1), storeProvider); assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(5, 5, 5), from(4, 4, 4), EMPTY, from(7, 7, 7), EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, aggregate(7, 7, 7), EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaCronUpdateFails() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectNoTasks().times(2); expectNoJobUpdates().times(2); @@ -809,13 +802,13 @@ public class QuotaManagerImplTest extends EasyMockTest { quotaManager.checkCronUpdate(createJob(prodTask("pc", 5, 5, 5), 2), storeProvider); assertEquals(INSUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(5, 5, 5), from(4, 4, 4), EMPTY, EMPTY, EMPTY), + new QuotaInfo(aggregate(5, 5, 5), aggregate(4, 4, 4), EMPTY, EMPTY, EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @Test public void testCheckQuotaCronCreate() { - expectQuota(IResourceAggregate.build(new ResourceAggregate(5, 5, 5))).times(2); + expectQuota(aggregate(5, 5, 5)).times(2); expectNoTasks().times(2); expectNoJobUpdates().times(2); expectNoCronJobs().times(2); @@ -827,7 +820,7 @@ public class QuotaManagerImplTest extends EasyMockTest { quotaManager.checkCronUpdate(createJob(prodTask("pc", 5, 5, 5), 1), storeProvider); assertEquals(SUFFICIENT_QUOTA, checkQuota.getResult()); assertEquals( - new QuotaInfo(from(5, 5, 5), EMPTY, EMPTY, EMPTY, EMPTY), + new QuotaInfo(aggregate(5, 5, 5), EMPTY, EMPTY, EMPTY, EMPTY), quotaManager.getQuotaInfo(ROLE, storeProvider)); } @@ -990,7 +983,7 @@ public class QuotaManagerImplTest extends EasyMockTest { .setInstanceCount(instanceCount)); } - private static IResourceAggregate from(double cpu, int ramMb, int diskMb) { - return IResourceAggregate.build(new ResourceAggregate(cpu, ramMb, diskMb)); + private static IResourceAggregate aggregate(double numCpus, long ramMb, long diskMb) { + return ResourceTestUtil.nonBackfilledAggregate(numCpus, ramMb, diskMb); } } http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/resources/ResourceTestUtil.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/resources/ResourceTestUtil.java b/src/test/java/org/apache/aurora/scheduler/resources/ResourceTestUtil.java new file mode 100644 index 0000000..04a8238 --- /dev/null +++ b/src/test/java/org/apache/aurora/scheduler/resources/ResourceTestUtil.java @@ -0,0 +1,48 @@ +/** + * 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.scheduler.resources; + +import com.google.common.collect.ImmutableSet; + +import org.apache.aurora.gen.ResourceAggregate; +import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; + +import static org.apache.aurora.gen.Resource.diskMb; +import static org.apache.aurora.gen.Resource.numCpus; +import static org.apache.aurora.gen.Resource.ramMb; + +/** + * Convenience methods for working with resources. + */ +public final class ResourceTestUtil { + + private ResourceTestUtil() { + // Utility class. + } + + public static IResourceAggregate aggregate(double numCpus, long ramMb, long diskMb) { + return IResourceAggregate.build(new ResourceAggregate(numCpus, ramMb, diskMb, ImmutableSet.of( + numCpus(numCpus), + ramMb(ramMb), + diskMb(diskMb) + ))); + } + + public static IResourceAggregate nonBackfilledAggregate(double numCpus, long ramMb, long diskMb) { + return IResourceAggregate.build(new ResourceAggregate() + .setNumCpus(numCpus) + .setRamMb(ramMb) + .setDiskMb(diskMb)); + } +} http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/stats/AsyncStatsModuleTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/stats/AsyncStatsModuleTest.java b/src/test/java/org/apache/aurora/scheduler/stats/AsyncStatsModuleTest.java index fdd42b7..7fcf47a 100644 --- a/src/test/java/org/apache/aurora/scheduler/stats/AsyncStatsModuleTest.java +++ b/src/test/java/org/apache/aurora/scheduler/stats/AsyncStatsModuleTest.java @@ -55,7 +55,9 @@ public class AsyncStatsModuleTest extends EasyMockTest { private static MachineResource resource(boolean revocable, double cpu) { return new MachineResource( - IResourceAggregate.build(new ResourceAggregate(cpu, 0, 0)), false, revocable); + IResourceAggregate.build(new ResourceAggregate().setNumCpus(cpu)), + false, + revocable); } private static Protos.Resource getCpuResource(boolean revocable, double value) { http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java b/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java index 845de41..fd4b434 100644 --- a/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java +++ b/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java @@ -22,7 +22,6 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import org.apache.aurora.gen.Constraint; -import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.gen.ScheduleStatus; import org.apache.aurora.gen.ScheduledTask; import org.apache.aurora.gen.TaskConfig; @@ -32,11 +31,11 @@ import org.apache.aurora.scheduler.base.JobKeys; import org.apache.aurora.scheduler.base.Query; import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.configuration.ConfigurationManager; +import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.storage.Storage; import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult; import org.apache.aurora.scheduler.storage.db.DbUtil; import org.apache.aurora.scheduler.storage.entities.IJobKey; -import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import org.junit.Before; @@ -119,10 +118,8 @@ public class ResourceCounterTest { @Test public void testComputeQuotaAllocationTotals() { storage.write((NoResult.Quiet) storeProvider -> { - storeProvider.getQuotaStore() - .saveQuota("a", IResourceAggregate.build(new ResourceAggregate(1, 1, 1))); - storeProvider.getQuotaStore() - .saveQuota("b", IResourceAggregate.build(new ResourceAggregate(2, 3, 4))); + storeProvider.getQuotaStore().saveQuota("a", ResourceTestUtil.aggregate(1, 1, 1)); + storeProvider.getQuotaStore().saveQuota("b", ResourceTestUtil.aggregate(2, 3, 4)); }); assertEquals(new Metric(3, 4, 5), resourceCounter.computeQuotaAllocationTotals()); http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/stats/SlotSizeCounterTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/stats/SlotSizeCounterTest.java b/src/test/java/org/apache/aurora/scheduler/stats/SlotSizeCounterTest.java index bbf1097..b1c8f75 100644 --- a/src/test/java/org/apache/aurora/scheduler/stats/SlotSizeCounterTest.java +++ b/src/test/java/org/apache/aurora/scheduler/stats/SlotSizeCounterTest.java @@ -21,7 +21,6 @@ import com.google.common.collect.ImmutableMap; import org.apache.aurora.common.stats.StatsProvider; import org.apache.aurora.common.testing.easymock.EasyMockTest; -import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.scheduler.resources.ResourceAggregates; import org.apache.aurora.scheduler.stats.SlotSizeCounter.MachineResource; import org.apache.aurora.scheduler.stats.SlotSizeCounter.MachineResourceProvider; @@ -29,13 +28,13 @@ import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.junit.Before; import org.junit.Test; +import static org.apache.aurora.scheduler.resources.ResourceTestUtil.aggregate; import static org.easymock.EasyMock.expect; import static org.junit.Assert.assertEquals; public class SlotSizeCounterTest extends EasyMockTest { - private static final IResourceAggregate SMALL = - IResourceAggregate.build(new ResourceAggregate(1.0, 1024, 4096)); + private static final IResourceAggregate SMALL = aggregate(1.0, 1024, 4096); private static final IResourceAggregate LARGE = ResourceAggregates.scale(SMALL, 4); private static final Map<String, IResourceAggregate> SLOT_SIZES = ImmutableMap.of( @@ -106,8 +105,7 @@ public class SlotSizeCounterTest extends EasyMockTest { @Test public void testTinyOffers() { expectStatExport(); - expectGetSlots(new MachineResource( - IResourceAggregate.build(new ResourceAggregate(0.1, 1, 1)), false, false)); + expectGetSlots(new MachineResource(aggregate(0.1, 1, 1), false, false)); control.replay(); @@ -126,8 +124,7 @@ public class SlotSizeCounterTest extends EasyMockTest { public void testStarvedResourceVector() { expectStatExport(); expectGetSlots( - new MachineResource( - IResourceAggregate.build(new ResourceAggregate(1000, 16384, 1)), false, false)); + new MachineResource(aggregate(1000, 16384, 1), false, false)); control.replay(); @@ -152,7 +149,7 @@ public class SlotSizeCounterTest extends EasyMockTest { new MachineResource(LARGE, false, true), new MachineResource(LARGE, true, true), new MachineResource(ResourceAggregates.scale(LARGE, 4), false, false), - new MachineResource(IResourceAggregate.build(new ResourceAggregate(1, 1, 1)), false, false), + new MachineResource(aggregate(1, 1, 1), false, false), new MachineResource(SMALL, true, false), new MachineResource(SMALL, true, false), new MachineResource(ResourceAggregates.scale(SMALL, 2), true, false)); http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/storage/backup/StorageBackupTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/backup/StorageBackupTest.java b/src/test/java/org/apache/aurora/scheduler/storage/backup/StorageBackupTest.java index 7b3f185..fff376f 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/backup/StorageBackupTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/backup/StorageBackupTest.java @@ -34,7 +34,6 @@ import org.apache.aurora.gen.Attribute; import org.apache.aurora.gen.HostAttributes; import org.apache.aurora.gen.JobConfiguration; import org.apache.aurora.gen.JobKey; -import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.gen.ScheduledTask; import org.apache.aurora.gen.storage.QuotaConfiguration; import org.apache.aurora.gen.storage.SchedulerMetadata; @@ -49,6 +48,7 @@ import org.junit.Rule; import org.junit.Test; import org.junit.rules.TemporaryFolder; +import static org.apache.aurora.scheduler.resources.ResourceTestUtil.aggregate; import static org.easymock.EasyMock.expect; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -179,7 +179,7 @@ public class StorageBackupTest extends EasyMockTest { snapshot.setCronJobs(ImmutableSet.of( new StoredCronJob(new JobConfiguration().setKey(new JobKey("owner", "env", "jobA"))))); snapshot.setQuotaConfigurations( - ImmutableSet.of(new QuotaConfiguration("roleA", new ResourceAggregate(10, 1024, 1024)))); + ImmutableSet.of(new QuotaConfiguration("roleA", aggregate(10, 1024, 1024).newBuilder()))); snapshot.setSchedulerMetadata(new SchedulerMetadata().setFrameworkId("frameworkId")); snapshot.setTasks(ImmutableSet.of(new ScheduledTask())); return snapshot; http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/storage/db/DbQuotaStoreTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/db/DbQuotaStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/db/DbQuotaStoreTest.java index e0ec995..275d0fd 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/db/DbQuotaStoreTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/db/DbQuotaStoreTest.java @@ -19,7 +19,7 @@ import java.util.Map; import com.google.common.base.Optional; import com.google.common.collect.ImmutableMap; -import org.apache.aurora.gen.ResourceAggregate; +import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.storage.Storage; import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult; import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; @@ -33,10 +33,8 @@ public class DbQuotaStoreTest { private static final String ROLE_A = "roleA"; private static final String ROLE_B = "roleB"; - private static final IResourceAggregate QUOTA_A = - IResourceAggregate.build(new ResourceAggregate(1.0D, 2, 3)); - private static final IResourceAggregate QUOTA_B = - IResourceAggregate.build(new ResourceAggregate(2.0D, 4, 6)); + private static final IResourceAggregate QUOTA_A = ResourceTestUtil.aggregate(1.0, 2, 3); + private static final IResourceAggregate QUOTA_B = ResourceTestUtil.aggregate(2.0, 4, 6); private Storage storage; http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java b/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java index 7ac0b58..be1132b 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java @@ -86,6 +86,7 @@ import org.apache.aurora.scheduler.log.Log; import org.apache.aurora.scheduler.log.Log.Entry; import org.apache.aurora.scheduler.log.Log.Position; import org.apache.aurora.scheduler.log.Log.Stream; +import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.storage.AttributeStore; import org.apache.aurora.scheduler.storage.SnapshotStore; import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider; @@ -113,6 +114,9 @@ import org.easymock.Capture; import org.junit.Before; import org.junit.Test; +import static org.apache.aurora.gen.Resource.diskMb; +import static org.apache.aurora.gen.Resource.numCpus; +import static org.apache.aurora.gen.Resource.ramMb; import static org.apache.aurora.scheduler.base.TaskTestUtil.makeConfig; import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask; import static org.easymock.EasyMock.capture; @@ -302,11 +306,16 @@ public class LogStorageTest extends EasyMockTest { builder.add(createTransaction(Op.removeTasks(removeTasks))); storageUtil.taskStore.deleteTasks(removeTasks.getTaskIds()); - SaveQuota saveQuota = new SaveQuota(JOB_KEY.getRole(), new ResourceAggregate()); + ResourceAggregate nonBackfilled = new ResourceAggregate() + .setNumCpus(1.0) + .setRamMb(32) + .setDiskMb(64); + SaveQuota saveQuota = new SaveQuota(JOB_KEY.getRole(), nonBackfilled); builder.add(createTransaction(Op.saveQuota(saveQuota))); storageUtil.quotaStore.saveQuota( saveQuota.getRole(), - IResourceAggregate.build(saveQuota.getQuota())); + IResourceAggregate.build(nonBackfilled.deepCopy() + .setResources(ImmutableSet.of(numCpus(1.0), ramMb(32), diskMb(64))))); builder.add(createTransaction(Op.removeQuota(new RemoveQuota(JOB_KEY.getRole())))); storageUtil.quotaStore.removeQuota(JOB_KEY.getRole()); @@ -730,7 +739,7 @@ public class LogStorageTest extends EasyMockTest { @Test public void testSaveQuota() throws Exception { String role = "role"; - IResourceAggregate quota = IResourceAggregate.build(new ResourceAggregate(1.0, 128L, 1024L)); + IResourceAggregate quota = ResourceTestUtil.aggregate(1.0, 128L, 1024L); new AbstractMutationFixture() { @Override http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplIT.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplIT.java b/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplIT.java index 05168ba..aeab07d 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplIT.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplIT.java @@ -193,7 +193,8 @@ public class SnapshotStoreImplIT { .setInstanceCount(1) .setTaskConfig(TASK_CONFIG.newBuilder())); private static final String ROLE = "role"; - private static final IResourceAggregate QUOTA = ResourceAggregates.LARGE; + private static final IResourceAggregate QUOTA = + ThriftBackfill.backfillResourceAggregate(ResourceAggregates.LARGE.newBuilder()); private static final IHostAttributes ATTRIBUTES = IHostAttributes.build( new HostAttributes("host", ImmutableSet.of(new Attribute("attr", ImmutableSet.of("value")))) .setMode(MaintenanceMode.NONE) http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/storage/log/ThriftBackfillTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/ThriftBackfillTest.java b/src/test/java/org/apache/aurora/scheduler/storage/log/ThriftBackfillTest.java index cc039fb..f213b7f 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/log/ThriftBackfillTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/log/ThriftBackfillTest.java @@ -15,7 +15,9 @@ package org.apache.aurora.scheduler.storage.log; import com.google.common.collect.ImmutableSet; +import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.gen.TaskConfig; +import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.junit.Test; import static org.apache.aurora.gen.Resource.diskMb; @@ -93,4 +95,30 @@ public class ThriftBackfillTest { TaskConfig config = new TaskConfig().setResources(ImmutableSet.of(numCpus(1.0), ramMb(32))); ThriftBackfill.backfillTask(config); } + + @Test + public void testResourceAggregateFieldsToSet() { + ResourceAggregate aggregate = new ResourceAggregate() + .setNumCpus(1.0) + .setRamMb(32) + .setDiskMb(64); + + IResourceAggregate expected = IResourceAggregate.build(aggregate.deepCopy() + .setResources(ImmutableSet.of(numCpus(1.0), ramMb(32), diskMb(64)))); + + assertEquals(expected, ThriftBackfill.backfillResourceAggregate(aggregate)); + } + + @Test + public void testResourceAggregateSetToFields() { + ResourceAggregate aggregate = new ResourceAggregate() + .setResources(ImmutableSet.of(numCpus(1.0), ramMb(32), diskMb(64))); + + IResourceAggregate expected = IResourceAggregate.build(aggregate.deepCopy() + .setNumCpus(1.0) + .setRamMb(32) + .setDiskMb(64)); + + assertEquals(expected, ThriftBackfill.backfillResourceAggregate(aggregate)); + } } http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java b/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java index ba5969a..25f34e2 100644 --- a/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java +++ b/src/test/java/org/apache/aurora/scheduler/storage/mem/StorageTransactionTest.java @@ -27,15 +27,14 @@ import com.google.common.util.concurrent.MoreExecutors; import com.google.common.util.concurrent.ThreadFactoryBuilder; import org.apache.aurora.common.testing.TearDownTestCase; -import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.scheduler.base.Query; import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.base.Tasks; +import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.storage.Storage; import org.apache.aurora.scheduler.storage.Storage.MutateWork; import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult; import org.apache.aurora.scheduler.storage.db.DbUtil; -import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import org.junit.Before; import org.junit.Test; @@ -114,12 +113,9 @@ public class StorageTransactionTest extends TearDownTestCase { @Test public void testWritesUnderTransaction() { - IResourceAggregate quota = IResourceAggregate - .build(new ResourceAggregate().setDiskMb(100).setNumCpus(2.0).setRamMb(512)); - try { storage.write(storeProvider -> { - storeProvider.getQuotaStore().saveQuota("a", quota); + storeProvider.getQuotaStore().saveQuota("a", ResourceTestUtil.aggregate(2.0, 512, 100)); throw new CustomException(); }); fail("Expected CustomException to be thrown."); http://git-wip-us.apache.org/repos/asf/aurora/blob/7e30ebe3/src/test/java/org/apache/aurora/scheduler/thrift/Fixtures.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/Fixtures.java b/src/test/java/org/apache/aurora/scheduler/thrift/Fixtures.java index 54585a9..a2e2395 100644 --- a/src/test/java/org/apache/aurora/scheduler/thrift/Fixtures.java +++ b/src/test/java/org/apache/aurora/scheduler/thrift/Fixtures.java @@ -37,7 +37,6 @@ import org.apache.aurora.gen.JobUpdateKey; import org.apache.aurora.gen.LockKey; import org.apache.aurora.gen.MesosContainer; import org.apache.aurora.gen.Resource; -import org.apache.aurora.gen.ResourceAggregate; import org.apache.aurora.gen.Response; import org.apache.aurora.gen.ResponseCode; import org.apache.aurora.gen.ResponseDetail; @@ -46,6 +45,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.quota.QuotaCheckResult; +import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.storage.entities.IJobKey; import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey; import org.apache.aurora.scheduler.storage.entities.ILockKey; @@ -74,8 +74,7 @@ final class Fixtures { private static final Function<String, ResponseDetail> MESSAGE_TO_DETAIL = message -> new ResponseDetail().setMessage(message); static final String CRON_SCHEDULE = "0 * * * *"; - static final IResourceAggregate QUOTA = - IResourceAggregate.build(new ResourceAggregate(10.0, 1024, 2048)); + static final IResourceAggregate QUOTA = ResourceTestUtil.aggregate(10.0, 1024, 2048); static final QuotaCheckResult ENOUGH_QUOTA = new QuotaCheckResult(SUFFICIENT_QUOTA); static final QuotaCheckResult NOT_ENOUGH_QUOTA = new QuotaCheckResult(INSUFFICIENT_QUOTA); static final InstanceKey INSTANCE_KEY = new InstanceKey(JOB_KEY.newBuilder(), 0);
