| /** |
| * 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.durability; |
| |
| import java.util.Optional; |
| import java.util.Set; |
| import java.util.concurrent.atomic.AtomicBoolean; |
| import java.util.concurrent.locks.ReentrantLock; |
| import java.util.stream.Collectors; |
| import java.util.stream.Stream; |
| |
| import com.google.common.base.Function; |
| import com.google.common.base.Functions; |
| import com.google.common.collect.ImmutableList; |
| import com.google.common.collect.ImmutableSet; |
| |
| import org.apache.aurora.common.testing.easymock.EasyMockTest; |
| import org.apache.aurora.gen.AssignedTask; |
| import org.apache.aurora.gen.Attribute; |
| import org.apache.aurora.gen.HostAttributes; |
| import org.apache.aurora.gen.HostMaintenanceRequest; |
| import org.apache.aurora.gen.InstanceTaskConfig; |
| import org.apache.aurora.gen.JobConfiguration; |
| import org.apache.aurora.gen.JobInstanceUpdateEvent; |
| import org.apache.aurora.gen.JobUpdate; |
| import org.apache.aurora.gen.JobUpdateAction; |
| import org.apache.aurora.gen.JobUpdateEvent; |
| import org.apache.aurora.gen.JobUpdateInstructions; |
| import org.apache.aurora.gen.JobUpdateKey; |
| import org.apache.aurora.gen.JobUpdateSettings; |
| import org.apache.aurora.gen.JobUpdateStatus; |
| import org.apache.aurora.gen.JobUpdateStrategy; |
| import org.apache.aurora.gen.JobUpdateSummary; |
| import org.apache.aurora.gen.MaintenanceMode; |
| import org.apache.aurora.gen.PercentageSlaPolicy; |
| import org.apache.aurora.gen.QueueJobUpdateStrategy; |
| import org.apache.aurora.gen.Range; |
| import org.apache.aurora.gen.ResourceAggregate; |
| import org.apache.aurora.gen.ScheduleStatus; |
| import org.apache.aurora.gen.ScheduledTask; |
| import org.apache.aurora.gen.SlaPolicy; |
| import org.apache.aurora.gen.TaskConfig; |
| import org.apache.aurora.gen.storage.Op; |
| import org.apache.aurora.gen.storage.PruneJobUpdateHistory; |
| import org.apache.aurora.gen.storage.RemoveHostMaintenanceRequest; |
| import org.apache.aurora.gen.storage.RemoveJob; |
| import org.apache.aurora.gen.storage.RemoveJobUpdates; |
| import org.apache.aurora.gen.storage.RemoveLock; |
| import org.apache.aurora.gen.storage.RemoveQuota; |
| import org.apache.aurora.gen.storage.RemoveTasks; |
| import org.apache.aurora.gen.storage.SaveCronJob; |
| import org.apache.aurora.gen.storage.SaveFrameworkId; |
| import org.apache.aurora.gen.storage.SaveHostAttributes; |
| import org.apache.aurora.gen.storage.SaveHostMaintenanceRequest; |
| import org.apache.aurora.gen.storage.SaveJobInstanceUpdateEvent; |
| import org.apache.aurora.gen.storage.SaveJobUpdate; |
| import org.apache.aurora.gen.storage.SaveJobUpdateEvent; |
| import org.apache.aurora.gen.storage.SaveLock; |
| import org.apache.aurora.gen.storage.SaveQuota; |
| import org.apache.aurora.gen.storage.SaveTasks; |
| import org.apache.aurora.scheduler.base.JobKeys; |
| import org.apache.aurora.scheduler.base.TaskTestUtil; |
| import org.apache.aurora.scheduler.base.Tasks; |
| import org.apache.aurora.scheduler.events.EventSink; |
| import org.apache.aurora.scheduler.events.PubsubEvent; |
| import org.apache.aurora.scheduler.resources.ResourceTestUtil; |
| import org.apache.aurora.scheduler.storage.AttributeStore; |
| import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider; |
| import org.apache.aurora.scheduler.storage.Storage.MutateWork; |
| import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult; |
| import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult.Quiet; |
| import org.apache.aurora.scheduler.storage.durability.Persistence.Edit; |
| import org.apache.aurora.scheduler.storage.entities.IHostAttributes; |
| import org.apache.aurora.scheduler.storage.entities.IHostMaintenanceRequest; |
| import org.apache.aurora.scheduler.storage.entities.IJobConfiguration; |
| import org.apache.aurora.scheduler.storage.entities.IJobInstanceUpdateEvent; |
| import org.apache.aurora.scheduler.storage.entities.IJobKey; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdate; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdateEvent; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey; |
| import org.apache.aurora.scheduler.storage.entities.IResourceAggregate; |
| import org.apache.aurora.scheduler.storage.entities.IScheduledTask; |
| import org.apache.aurora.scheduler.storage.testing.StorageTestUtil; |
| import org.easymock.Capture; |
| import org.easymock.EasyMock; |
| import org.easymock.IAnswer; |
| 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.anyObject; |
| import static org.easymock.EasyMock.capture; |
| import static org.easymock.EasyMock.expect; |
| import static org.easymock.EasyMock.expectLastCall; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertTrue; |
| |
| public class DurableStorageTest extends EasyMockTest { |
| |
| private static final IJobKey JOB_KEY = JobKeys.from("role", "env", "name"); |
| private static final IJobUpdateKey UPDATE_ID = |
| IJobUpdateKey.build(new JobUpdateKey(JOB_KEY.newBuilder(), "testUpdateId")); |
| |
| private DurableStorage durableStorage; |
| private Persistence persistence; |
| private StorageTestUtil storageUtil; |
| private EventSink eventSink; |
| |
| @Before |
| public void setUp() { |
| persistence = createMock(Persistence.class); |
| storageUtil = new StorageTestUtil(this); |
| eventSink = createMock(EventSink.class); |
| |
| durableStorage = new DurableStorage( |
| persistence, |
| storageUtil.storage, |
| storageUtil.schedulerStore, |
| storageUtil.jobStore, |
| storageUtil.taskStore, |
| storageUtil.quotaStore, |
| storageUtil.attributeStore, |
| storageUtil.jobUpdateStore, |
| storageUtil.hostMaintenanceStore, |
| eventSink, |
| new ReentrantLock(), |
| TaskTestUtil.THRIFT_BACKFILL); |
| |
| storageUtil.storage.prepare(); |
| } |
| |
| @Test |
| public void testStart() throws Exception { |
| // We should initialize persistence. |
| persistence.prepare(); |
| |
| // Our start should recover persistence and then forward to the underlying storage start of the |
| // supplied initialization logic. |
| AtomicBoolean initialized = new AtomicBoolean(false); |
| MutateWork.NoResult.Quiet initializationLogic = provider -> { |
| // Creating a mock and expecting apply(storeProvider) does not work here for whatever |
| // reason. |
| initialized.set(true); |
| }; |
| |
| Capture<MutateWork.NoResult.Quiet> recoverAndInitializeWork = createCapture(); |
| storageUtil.storage.write(capture(recoverAndInitializeWork)); |
| expectLastCall().andAnswer(() -> { |
| recoverAndInitializeWork.getValue().apply(storageUtil.mutableStoreProvider); |
| return null; |
| }); |
| |
| Capture<MutateWork<Void, RuntimeException>> initializationWork = createCapture(); |
| expect(storageUtil.storage.write(capture(initializationWork))).andAnswer( |
| () -> { |
| initializationWork.getValue().apply(storageUtil.mutableStoreProvider); |
| return null; |
| }); |
| |
| // Populate all Op types. |
| buildReplayOps(); |
| |
| storageUtil.expectStoreAccesses(); |
| |
| control.replay(); |
| |
| durableStorage.prepare(); |
| durableStorage.start(initializationLogic); |
| assertTrue(initialized.get()); |
| } |
| |
| private void buildReplayOps() throws Exception { |
| ImmutableSet.Builder<Edit> builder = ImmutableSet.builder(); |
| |
| builder.add(Edit.op(Op.saveFrameworkId(new SaveFrameworkId("bob")))); |
| storageUtil.schedulerStore.saveFrameworkId("bob"); |
| |
| JobConfiguration actualJob = new JobConfiguration().setTaskConfig(nonBackfilledConfig()); |
| JobConfiguration expectedJob = |
| new JobConfiguration().setTaskConfig(makeConfig(JOB_KEY).newBuilder()); |
| SaveCronJob cronJob = new SaveCronJob().setJobConfig(actualJob); |
| builder.add(Edit.op(Op.saveCronJob(cronJob))); |
| storageUtil.jobStore.saveAcceptedJob(IJobConfiguration.build(expectedJob)); |
| |
| RemoveJob removeJob = new RemoveJob(JOB_KEY.newBuilder()); |
| builder.add(Edit.op(Op.removeJob(removeJob))); |
| storageUtil.jobStore.removeJob(JOB_KEY); |
| |
| ScheduledTask actualTask = makeTask("id", JOB_KEY).newBuilder(); |
| actualTask.getAssignedTask().setTask(nonBackfilledConfig()); |
| IScheduledTask expectedTask = makeTask("id", JOB_KEY); |
| SaveTasks saveTasks = new SaveTasks(ImmutableSet.of(actualTask)); |
| builder.add(Edit.op(Op.saveTasks(saveTasks))); |
| storageUtil.taskStore.saveTasks(ImmutableSet.of(expectedTask)); |
| |
| // Side-effects from a storage reset, caused by a snapshot. |
| builder.add(Edit.deleteAll()); |
| storageUtil.jobStore.deleteJobs(); |
| storageUtil.taskStore.deleteAllTasks(); |
| storageUtil.quotaStore.deleteQuotas(); |
| storageUtil.attributeStore.deleteHostAttributes(); |
| storageUtil.jobUpdateStore.deleteAllUpdates(); |
| storageUtil.hostMaintenanceStore.deleteHostMaintenanceRequests(); |
| |
| RemoveTasks removeTasks = new RemoveTasks(ImmutableSet.of("taskId1")); |
| builder.add(Edit.op(Op.removeTasks(removeTasks))); |
| storageUtil.taskStore.deleteTasks(removeTasks.getTaskIds()); |
| |
| ResourceAggregate resourceAggregate = new ResourceAggregate() |
| .setResources(ImmutableSet.of(numCpus(1.0), ramMb(32), diskMb(64))); |
| SaveQuota saveQuota = new SaveQuota(JOB_KEY.getRole(), resourceAggregate); |
| builder.add(Edit.op(Op.saveQuota(saveQuota))); |
| storageUtil.quotaStore.saveQuota( |
| saveQuota.getRole(), |
| IResourceAggregate.build(resourceAggregate)); |
| |
| builder.add(Edit.op(Op.removeQuota(new RemoveQuota(JOB_KEY.getRole())))); |
| storageUtil.quotaStore.removeQuota(JOB_KEY.getRole()); |
| |
| // This entry lacks a slave ID, and should therefore be discarded. |
| SaveHostAttributes hostAttributes1 = new SaveHostAttributes(new HostAttributes() |
| .setHost("host1") |
| .setMode(MaintenanceMode.DRAINED)); |
| builder.add(Edit.op(Op.saveHostAttributes(hostAttributes1))); |
| |
| SaveHostAttributes hostAttributes2 = new SaveHostAttributes(new HostAttributes() |
| .setHost("host2") |
| .setSlaveId("slave2") |
| .setMode(MaintenanceMode.DRAINED)); |
| builder.add(Edit.op(Op.saveHostAttributes(hostAttributes2))); |
| expect(storageUtil.attributeStore.saveHostAttributes( |
| IHostAttributes.build(hostAttributes2.getHostAttributes()))).andReturn(true); |
| |
| builder.add(Edit.op(Op.saveLock(new SaveLock()))); |
| // TODO(jly): Deprecated, this is a no-op to be removed in 0.21. See AURORA-1959. |
| |
| builder.add(Edit.op(Op.removeLock(new RemoveLock()))); |
| // TODO(jly): Deprecated, this is a no-op to be removed in 0.21. See AURORA-1959. |
| |
| JobUpdate actualUpdate = new JobUpdate() |
| .setSummary(new JobUpdateSummary().setKey(UPDATE_ID.newBuilder())) |
| .setInstructions(new JobUpdateInstructions() |
| .setInitialState( |
| ImmutableSet.of(new InstanceTaskConfig().setTask(nonBackfilledConfig()))) |
| .setDesiredState(new InstanceTaskConfig().setTask(nonBackfilledConfig())) |
| .setSettings(new JobUpdateSettings())); |
| JobUpdate expectedUpdate = actualUpdate.deepCopy(); |
| expectedUpdate.getInstructions().getDesiredState().setTask(makeConfig(JOB_KEY).newBuilder()); |
| expectedUpdate.getInstructions().getInitialState() |
| .forEach(e -> e.setTask(makeConfig(JOB_KEY).newBuilder())); |
| expectedUpdate.getInstructions() |
| .getSettings() |
| .setUpdateStrategy( |
| JobUpdateStrategy.queueStrategy(new QueueJobUpdateStrategy().setGroupSize(0))); |
| SaveJobUpdate saveUpdate = new SaveJobUpdate().setJobUpdate(actualUpdate); |
| builder.add(Edit.op(Op.saveJobUpdate(saveUpdate))); |
| storageUtil.jobUpdateStore.saveJobUpdate(IJobUpdate.build(expectedUpdate)); |
| |
| SaveJobUpdateEvent saveUpdateEvent = |
| new SaveJobUpdateEvent(new JobUpdateEvent(), UPDATE_ID.newBuilder()); |
| builder.add(Edit.op(Op.saveJobUpdateEvent(saveUpdateEvent))); |
| storageUtil.jobUpdateStore.saveJobUpdateEvent( |
| UPDATE_ID, |
| IJobUpdateEvent.build(saveUpdateEvent.getEvent())); |
| |
| SaveJobInstanceUpdateEvent saveInstanceEvent = new SaveJobInstanceUpdateEvent( |
| new JobInstanceUpdateEvent(), |
| UPDATE_ID.newBuilder()); |
| builder.add(Edit.op(Op.saveJobInstanceUpdateEvent(saveInstanceEvent))); |
| storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent( |
| UPDATE_ID, |
| IJobInstanceUpdateEvent.build(saveInstanceEvent.getEvent())); |
| |
| builder.add(Edit.op(Op.pruneJobUpdateHistory(new PruneJobUpdateHistory(5, 10L)))); |
| // No expectation - this op is ignored. |
| |
| builder.add(Edit.op(Op.removeJobUpdate( |
| new RemoveJobUpdates().setKeys(ImmutableSet.of(UPDATE_ID.newBuilder()))))); |
| storageUtil.jobUpdateStore.removeJobUpdates(ImmutableSet.of(UPDATE_ID)); |
| |
| expect(persistence.recover()).andReturn(builder.build().stream()); |
| } |
| |
| private TaskConfig nonBackfilledConfig() { |
| // When more fields have to be backfilled |
| // modify this method. |
| return makeConfig(JOB_KEY).newBuilder(); |
| } |
| |
| abstract class AbstractStorageFixture { |
| private final AtomicBoolean runCalled = new AtomicBoolean(false); |
| |
| AbstractStorageFixture() { |
| // Prevent otherwise silent noop tests that forget to call run(). |
| addTearDown(new TearDown() { |
| @Override |
| public void tearDown() { |
| assertTrue(runCalled.get()); |
| } |
| }); |
| } |
| |
| void run() throws Exception { |
| runCalled.set(true); |
| |
| // Expect basic start operations. |
| |
| // Initialize persistence. |
| persistence.prepare(); |
| |
| // Replay the ops and perform and supplied initializationWork. |
| // Simulate NOOP initialization work |
| // Creating a mock and expecting apply(storeProvider) does not work here for whatever |
| // reason. |
| MutateWork.NoResult.Quiet initializationLogic = storeProvider -> { |
| // No-op. |
| }; |
| |
| Capture<MutateWork.NoResult.Quiet> recoverAndInitializeWork = createCapture(); |
| storageUtil.storage.write(capture(recoverAndInitializeWork)); |
| expectLastCall().andAnswer(() -> { |
| recoverAndInitializeWork.getValue().apply(storageUtil.mutableStoreProvider); |
| return null; |
| }); |
| |
| expect(persistence.recover()).andReturn(Stream.empty()); |
| Capture<MutateWork<Void, RuntimeException>> recoveryWork = createCapture(); |
| expect(storageUtil.storage.write(capture(recoveryWork))).andAnswer( |
| () -> { |
| recoveryWork.getValue().apply(storageUtil.mutableStoreProvider); |
| return null; |
| }); |
| |
| // Setup custom test expectations. |
| setupExpectations(); |
| |
| control.replay(); |
| |
| // Start the system. |
| durableStorage.prepare(); |
| durableStorage.start(initializationLogic); |
| |
| // Exercise the system. |
| runTest(); |
| } |
| |
| protected void setupExpectations() throws Exception { |
| // Default to no expectations. |
| } |
| |
| protected abstract void runTest(); |
| } |
| |
| abstract class AbstractMutationFixture extends AbstractStorageFixture { |
| @Override |
| protected void runTest() { |
| durableStorage.write((Quiet) AbstractMutationFixture.this::performMutations); |
| } |
| |
| protected abstract void performMutations(MutableStoreProvider storeProvider); |
| } |
| |
| private void expectPersist(Op op, Op... ops) { |
| try { |
| // Workaround for comparing streams. |
| persistence.persist(anyObject()); |
| expectLastCall().andAnswer((IAnswer<Void>) () -> { |
| assertEquals( |
| ImmutableList.<Op>builder().add(op).add(ops).build(), |
| ((Stream<Op>) EasyMock.getCurrentArguments()[0]).collect(Collectors.toList())); |
| |
| return null; |
| }); |
| } catch (Persistence.PersistenceException e) { |
| throw new RuntimeException(e); |
| } |
| } |
| |
| @Test |
| public void testSaveFrameworkId() throws Exception { |
| String frameworkId = "bob"; |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.schedulerStore.saveFrameworkId(frameworkId); |
| expectPersist(Op.saveFrameworkId(new SaveFrameworkId(frameworkId))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getSchedulerStore().saveFrameworkId(frameworkId); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveAcceptedJob() throws Exception { |
| IJobConfiguration jobConfig = |
| IJobConfiguration.build(new JobConfiguration().setKey(JOB_KEY.newBuilder())); |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.jobStore.saveAcceptedJob(jobConfig); |
| expectPersist(Op.saveCronJob(new SaveCronJob(jobConfig.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getCronJobStore().saveAcceptedJob(jobConfig); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testRemoveJob() throws Exception { |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.jobStore.removeJob(JOB_KEY); |
| expectPersist(Op.removeJob(new RemoveJob().setJobKey(JOB_KEY.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getCronJobStore().removeJob(JOB_KEY); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveTasks() throws Exception { |
| Set<IScheduledTask> tasks = ImmutableSet.of(task("a", ScheduleStatus.INIT)); |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.taskStore.saveTasks(tasks); |
| expectPersist(Op.saveTasks(new SaveTasks(IScheduledTask.toBuildersSet(tasks)))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getUnsafeTaskStore().saveTasks(tasks); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testMutateTasks() throws Exception { |
| String taskId = "fred"; |
| Function<IScheduledTask, IScheduledTask> mutation = Functions.identity(); |
| Optional<IScheduledTask> mutated = Optional.of(task("a", ScheduleStatus.STARTING)); |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| expect(storageUtil.taskStore.mutateTask(taskId, mutation)).andReturn(mutated); |
| expectPersist(Op.saveTasks(new SaveTasks(ImmutableSet.of(mutated.get().newBuilder())))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| assertEquals(mutated, storeProvider.getUnsafeTaskStore().mutateTask(taskId, mutation)); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testNestedTransactions() throws Exception { |
| String taskId = "fred"; |
| Function<IScheduledTask, IScheduledTask> mutation = Functions.identity(); |
| Optional<IScheduledTask> mutated = Optional.of(task("a", ScheduleStatus.STARTING)); |
| ImmutableSet<String> tasksToRemove = ImmutableSet.of("b"); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| expect(storageUtil.taskStore.mutateTask(taskId, mutation)).andReturn(mutated); |
| |
| storageUtil.taskStore.deleteTasks(tasksToRemove); |
| |
| expectPersist( |
| Op.saveTasks(new SaveTasks(ImmutableSet.of(mutated.get().newBuilder()))), |
| Op.removeTasks(new RemoveTasks(tasksToRemove))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| assertEquals(mutated, storeProvider.getUnsafeTaskStore().mutateTask(taskId, mutation)); |
| |
| durableStorage.write((NoResult.Quiet) |
| innerProvider -> innerProvider.getUnsafeTaskStore().deleteTasks(tasksToRemove)); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveAndMutateTasks() throws Exception { |
| String taskId = "fred"; |
| Function<IScheduledTask, IScheduledTask> mutation = Functions.identity(); |
| Set<IScheduledTask> saved = ImmutableSet.of(task("a", ScheduleStatus.INIT)); |
| Optional<IScheduledTask> mutated = Optional.of(task("a", ScheduleStatus.PENDING)); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.taskStore.saveTasks(saved); |
| |
| // Nested transaction with result. |
| expect(storageUtil.taskStore.mutateTask(taskId, mutation)).andReturn(mutated); |
| |
| // Resulting stream operation. |
| expectPersist(Op.saveTasks(new SaveTasks(ImmutableSet.of(mutated.get().newBuilder())))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getUnsafeTaskStore().saveTasks(saved); |
| assertEquals(mutated, storeProvider.getUnsafeTaskStore().mutateTask(taskId, mutation)); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveAndMutateTasksNoCoalesceUniqueIds() throws Exception { |
| String taskId = "fred"; |
| Function<IScheduledTask, IScheduledTask> mutation = Functions.identity(); |
| Set<IScheduledTask> saved = ImmutableSet.of(task("b", ScheduleStatus.INIT)); |
| Optional<IScheduledTask> mutated = Optional.of(task("a", ScheduleStatus.PENDING)); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.taskStore.saveTasks(saved); |
| |
| // Nested transaction with result. |
| expect(storageUtil.taskStore.mutateTask(taskId, mutation)).andReturn(mutated); |
| |
| // Resulting stream operation. |
| expectPersist(Op.saveTasks(new SaveTasks( |
| ImmutableSet.<ScheduledTask>builder() |
| .addAll(IScheduledTask.toBuildersList(saved)) |
| .add(mutated.get().newBuilder()) |
| .build()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getUnsafeTaskStore().saveTasks(saved); |
| assertEquals(mutated, storeProvider.getUnsafeTaskStore().mutateTask(taskId, mutation)); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testRemoveTasksQuery() throws Exception { |
| IScheduledTask task = task("a", ScheduleStatus.FINISHED); |
| Set<String> taskIds = Tasks.ids(task); |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.taskStore.deleteTasks(taskIds); |
| expectPersist(Op.removeTasks(new RemoveTasks(taskIds))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getUnsafeTaskStore().deleteTasks(taskIds); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testRemoveTasksIds() throws Exception { |
| Set<String> taskIds = ImmutableSet.of("42"); |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.taskStore.deleteTasks(taskIds); |
| expectPersist(Op.removeTasks(new RemoveTasks(taskIds))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getUnsafeTaskStore().deleteTasks(taskIds); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveQuota() throws Exception { |
| String role = "role"; |
| IResourceAggregate quota = ResourceTestUtil.aggregate(1.0, 128L, 1024L); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.quotaStore.saveQuota(role, quota); |
| expectPersist(Op.saveQuota(new SaveQuota(role, quota.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getQuotaStore().saveQuota(role, quota); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testRemoveQuota() throws Exception { |
| String role = "role"; |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.quotaStore.removeQuota(role); |
| expectPersist(Op.removeQuota(new RemoveQuota(role))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getQuotaStore().removeQuota(role); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveHostAttributes() throws Exception { |
| String host = "hostname"; |
| Set<Attribute> attributes = |
| ImmutableSet.of(new Attribute().setName("attr").setValues(ImmutableSet.of("value"))); |
| Optional<IHostAttributes> hostAttributes = Optional.of( |
| IHostAttributes.build(new HostAttributes() |
| .setHost(host) |
| .setAttributes(attributes))); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| expect(storageUtil.attributeStore.getHostAttributes(host)) |
| .andReturn(Optional.empty()); |
| |
| expect(storageUtil.attributeStore.getHostAttributes(host)).andReturn(hostAttributes); |
| |
| expect(storageUtil.attributeStore.saveHostAttributes(hostAttributes.get())).andReturn(true); |
| eventSink.post(new PubsubEvent.HostAttributesChanged(hostAttributes.get())); |
| expectPersist( |
| Op.saveHostAttributes(new SaveHostAttributes(hostAttributes.get().newBuilder()))); |
| |
| expect(storageUtil.attributeStore.saveHostAttributes(hostAttributes.get())) |
| .andReturn(false); |
| |
| expect(storageUtil.attributeStore.getHostAttributes(host)).andReturn(hostAttributes); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| AttributeStore.Mutable store = storeProvider.getAttributeStore(); |
| assertEquals(Optional.empty(), store.getHostAttributes(host)); |
| |
| assertTrue(store.saveHostAttributes(hostAttributes.get())); |
| |
| assertEquals(hostAttributes, store.getHostAttributes(host)); |
| |
| assertFalse(store.saveHostAttributes(hostAttributes.get())); |
| |
| assertEquals(hostAttributes, store.getHostAttributes(host)); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveUpdate() throws Exception { |
| IJobUpdate update = IJobUpdate.build(new JobUpdate() |
| .setSummary(new JobUpdateSummary() |
| .setKey(UPDATE_ID.newBuilder()) |
| .setUser("user")) |
| .setInstructions(new JobUpdateInstructions() |
| .setDesiredState(new InstanceTaskConfig() |
| .setTask(new TaskConfig()) |
| .setInstances(ImmutableSet.of(new Range(0, 3)))) |
| .setInitialState(ImmutableSet.of(new InstanceTaskConfig() |
| .setTask(new TaskConfig()) |
| .setInstances(ImmutableSet.of(new Range(0, 3))))) |
| .setSettings(new JobUpdateSettings()))); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.jobUpdateStore.saveJobUpdate(update); |
| expectPersist(Op.saveJobUpdate(new SaveJobUpdate().setJobUpdate(update.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getJobUpdateStore().saveJobUpdate(update); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveJobUpdateEvent() throws Exception { |
| IJobUpdateEvent event = IJobUpdateEvent.build(new JobUpdateEvent() |
| .setStatus(JobUpdateStatus.ROLLING_BACK) |
| .setTimestampMs(12345L)); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.jobUpdateStore.saveJobUpdateEvent(UPDATE_ID, event); |
| expectPersist(Op.saveJobUpdateEvent(new SaveJobUpdateEvent( |
| event.newBuilder(), |
| UPDATE_ID.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getJobUpdateStore().saveJobUpdateEvent(UPDATE_ID, event); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveJobInstanceUpdateEvent() throws Exception { |
| IJobInstanceUpdateEvent event = IJobInstanceUpdateEvent.build(new JobInstanceUpdateEvent() |
| .setAction(JobUpdateAction.INSTANCE_ROLLING_BACK) |
| .setTimestampMs(12345L) |
| .setInstanceId(0)); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(UPDATE_ID, event); |
| expectPersist(Op.saveJobInstanceUpdateEvent( |
| new SaveJobInstanceUpdateEvent( |
| event.newBuilder(), |
| UPDATE_ID.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getJobUpdateStore().saveJobInstanceUpdateEvent(UPDATE_ID, event); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testRemoveJobUpdates() throws Exception { |
| IJobUpdateKey key = IJobUpdateKey.build(new JobUpdateKey() |
| .setJob(JOB_KEY.newBuilder()) |
| .setId("update-id")); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() throws Exception { |
| storageUtil.expectWrite(); |
| storageUtil.jobUpdateStore.removeJobUpdates(ImmutableSet.of(key)); |
| |
| // No transaction is generated since this version is currently in 'read-only' |
| // compatibility mode for this operation type. |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getJobUpdateStore().removeJobUpdates(ImmutableSet.of(key)); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testSaveHostMaintenanceRequest() throws Exception { |
| String host = "hostname"; |
| IHostMaintenanceRequest hostMaintenanceRequest = IHostMaintenanceRequest.build( |
| new HostMaintenanceRequest() |
| .setHost(host) |
| .setDefaultSlaPolicy(SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(95) |
| .setDurationSecs(1800))) |
| .setTimeoutSecs(1800) |
| ); |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() { |
| storageUtil.expectWrite(); |
| storageUtil.hostMaintenanceStore.saveHostMaintenanceRequest(hostMaintenanceRequest); |
| expectPersist(Op.saveHostMaintenanceRequest( |
| new SaveHostMaintenanceRequest(hostMaintenanceRequest.newBuilder()))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getHostMaintenanceStore().saveHostMaintenanceRequest( |
| hostMaintenanceRequest); |
| } |
| }.run(); |
| } |
| |
| @Test |
| public void testRemoveHostMaintenanceRequest() throws Exception { |
| String host = "hostname"; |
| |
| new AbstractMutationFixture() { |
| @Override |
| protected void setupExpectations() { |
| storageUtil.expectWrite(); |
| storageUtil.hostMaintenanceStore.removeHostMaintenanceRequest(host); |
| expectPersist(Op.removeHostMaintenanceRequest(new RemoveHostMaintenanceRequest(host))); |
| } |
| |
| @Override |
| protected void performMutations(MutableStoreProvider storeProvider) { |
| storeProvider.getHostMaintenanceStore().removeHostMaintenanceRequest(host); |
| } |
| }.run(); |
| } |
| |
| private static IScheduledTask task(String id, ScheduleStatus status) { |
| return IScheduledTask.build(new ScheduledTask() |
| .setStatus(status) |
| .setAssignedTask(new AssignedTask() |
| .setTaskId(id) |
| .setTask(new TaskConfig()))); |
| } |
| } |
