| /** |
| * 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.thrift; |
| |
| import java.lang.reflect.InvocationTargetException; |
| import java.lang.reflect.Proxy; |
| import java.util.Map; |
| import java.util.Optional; |
| import java.util.Set; |
| |
| import com.google.common.base.Strings; |
| import com.google.common.base.Throwables; |
| import com.google.common.collect.ImmutableList; |
| import com.google.common.collect.ImmutableMap; |
| import com.google.common.collect.ImmutableSet; |
| import com.google.common.collect.Iterables; |
| import com.google.common.collect.Sets; |
| |
| import org.apache.aurora.common.testing.easymock.EasyMockTest; |
| import org.apache.aurora.gen.AssignedTask; |
| import org.apache.aurora.gen.AuroraAdmin; |
| import org.apache.aurora.gen.Constraint; |
| import org.apache.aurora.gen.Container; |
| import org.apache.aurora.gen.CoordinatorSlaPolicy; |
| import org.apache.aurora.gen.ExecutorConfig; |
| import org.apache.aurora.gen.ExplicitReconciliationSettings; |
| import org.apache.aurora.gen.HostStatus; |
| import org.apache.aurora.gen.Hosts; |
| import org.apache.aurora.gen.InstanceTaskConfig; |
| import org.apache.aurora.gen.JobConfiguration; |
| import org.apache.aurora.gen.JobUpdate; |
| import org.apache.aurora.gen.JobUpdateInstructions; |
| import org.apache.aurora.gen.JobUpdatePulseStatus; |
| import org.apache.aurora.gen.JobUpdateQuery; |
| import org.apache.aurora.gen.JobUpdateRequest; |
| import org.apache.aurora.gen.JobUpdateSettings; |
| import org.apache.aurora.gen.JobUpdateStrategy; |
| import org.apache.aurora.gen.JobUpdateSummary; |
| import org.apache.aurora.gen.LimitConstraint; |
| import org.apache.aurora.gen.ListBackupsResult; |
| import org.apache.aurora.gen.MaintenanceMode; |
| import org.apache.aurora.gen.MesosContainer; |
| import org.apache.aurora.gen.Metadata; |
| import org.apache.aurora.gen.PercentageSlaPolicy; |
| import org.apache.aurora.gen.PulseJobUpdateResult; |
| import org.apache.aurora.gen.QueryRecoveryResult; |
| import org.apache.aurora.gen.QueueJobUpdateStrategy; |
| import org.apache.aurora.gen.Range; |
| import org.apache.aurora.gen.ReadOnlyScheduler; |
| 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; |
| import org.apache.aurora.gen.Result; |
| import org.apache.aurora.gen.ScheduleStatus; |
| import org.apache.aurora.gen.ScheduledTask; |
| import org.apache.aurora.gen.SlaPolicy; |
| import org.apache.aurora.gen.StartJobUpdateResult; |
| import org.apache.aurora.gen.TaskConfig; |
| import org.apache.aurora.gen.TaskConstraint; |
| import org.apache.aurora.gen.TaskQuery; |
| import org.apache.aurora.gen.ValueConstraint; |
| import org.apache.aurora.gen.VariableBatchJobUpdateStrategy; |
| import org.apache.aurora.gen.apiConstants; |
| 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.configuration.ConfigurationManager; |
| import org.apache.aurora.scheduler.configuration.ConfigurationManager.TaskDescriptionException; |
| import org.apache.aurora.scheduler.configuration.SanitizedConfiguration; |
| import org.apache.aurora.scheduler.cron.CronException; |
| import org.apache.aurora.scheduler.cron.CronJobManager; |
| import org.apache.aurora.scheduler.cron.SanitizedCronJob; |
| import org.apache.aurora.scheduler.maintenance.MaintenanceController; |
| import org.apache.aurora.scheduler.quota.QuotaCheckResult; |
| import org.apache.aurora.scheduler.quota.QuotaManager; |
| import org.apache.aurora.scheduler.reconciliation.TaskReconciler; |
| import org.apache.aurora.scheduler.state.StateChangeResult; |
| import org.apache.aurora.scheduler.state.StateManager; |
| import org.apache.aurora.scheduler.state.UUIDGenerator; |
| import org.apache.aurora.scheduler.storage.SnapshotStore; |
| import org.apache.aurora.scheduler.storage.Storage.StorageException; |
| import org.apache.aurora.scheduler.storage.backup.Recovery; |
| import org.apache.aurora.scheduler.storage.backup.StorageBackup; |
| import org.apache.aurora.scheduler.storage.entities.IHostStatus; |
| import org.apache.aurora.scheduler.storage.entities.IJobConfiguration; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdate; |
| import org.apache.aurora.scheduler.storage.entities.IMetadata; |
| 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.testing.StorageTestUtil; |
| import org.apache.aurora.scheduler.testing.FakeStatsProvider; |
| import org.apache.aurora.scheduler.updater.JobUpdateController; |
| import org.apache.aurora.scheduler.updater.JobUpdateController.AuditData; |
| import org.apache.aurora.scheduler.updater.JobUpdateController.JobUpdatingException; |
| import org.apache.aurora.scheduler.updater.UpdateInProgressException; |
| import org.apache.aurora.scheduler.updater.UpdateStateException; |
| import org.apache.thrift.TException; |
| import org.easymock.EasyMock; |
| import org.easymock.IExpectationSetters; |
| import org.junit.Before; |
| import org.junit.Test; |
| |
| import static org.apache.aurora.gen.MaintenanceMode.DRAINING; |
| import static org.apache.aurora.gen.MaintenanceMode.NONE; |
| import static org.apache.aurora.gen.MaintenanceMode.SCHEDULED; |
| 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.gen.ResponseCode.INVALID_REQUEST; |
| import static org.apache.aurora.gen.ResponseCode.JOB_UPDATING_ERROR; |
| import static org.apache.aurora.gen.ResponseCode.OK; |
| import static org.apache.aurora.scheduler.configuration.ConfigurationManager.DEDICATED_ATTRIBUTE; |
| import static org.apache.aurora.scheduler.storage.backup.Recovery.RecoveryException; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.CRON_JOB; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.ENOUGH_QUOTA; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.IDENTITY; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.INSTANCE_KEY; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.INVALID_TASK_CONFIG; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.JOB_KEY; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.JOB_NAME; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.NOT_ENOUGH_QUOTA; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.ROLE; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.TASK_ID; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.UPDATE_KEY; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.USER; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.UU_ID; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.assertOkResponse; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.assertResponse; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.defaultTask; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.makeJob; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.makeProdJob; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.nonProductionTask; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.okResponse; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.productionTask; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.response; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.ADD_INSTANCES; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.CREATE_JOB; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.CREATE_OR_UPDATE_CRON; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.DRAIN_HOSTS; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.END_MAINTENANCE; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.INVALID_SLA_AWARE_UPDATE; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.KILL_TASKS; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.MAINTENANCE_STATUS; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.NOOP_JOB_UPDATE_MESSAGE; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.NO_CRON; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.PRUNE_TASKS; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.RESTART_SHARDS; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.SLA_DRAIN_HOSTS; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.START_JOB_UPDATE; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.START_MAINTENANCE; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.jobAlreadyExistsMessage; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.noCronScheduleMessage; |
| import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.notScheduledCronMessage; |
| import static org.easymock.EasyMock.anyInt; |
| import static org.easymock.EasyMock.anyObject; |
| import static org.easymock.EasyMock.eq; |
| import static org.easymock.EasyMock.expect; |
| import static org.easymock.EasyMock.expectLastCall; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertTrue; |
| import static org.junit.Assert.fail; |
| |
| public class SchedulerThriftInterfaceTest extends EasyMockTest { |
| |
| private static final String AUDIT_MESSAGE = "message"; |
| private static final AuditData AUDIT = new AuditData(USER, Optional.of(AUDIT_MESSAGE)); |
| private static final Thresholds THRESHOLDS = new Thresholds(1000, 2000); |
| private static final String EXECUTOR_NAME = apiConstants.AURORA_EXECUTOR_NAME; |
| private static final ImmutableSet<Metadata> METADATA = |
| ImmutableSet.of(new Metadata("k1", "v1"), new Metadata("k2", "v2")); |
| |
| private StorageTestUtil storageUtil; |
| private SnapshotStore snapshotStore; |
| private StorageBackup backup; |
| private Recovery recovery; |
| private MaintenanceController maintenance; |
| private AuroraAdmin.Iface thrift; |
| private CronJobManager cronJobManager; |
| private QuotaManager quotaManager; |
| private StateManager stateManager; |
| private UUIDGenerator uuidGenerator; |
| private JobUpdateController jobUpdateController; |
| private ReadOnlyScheduler.Iface readOnlyScheduler; |
| private AuditMessages auditMessages; |
| private TaskReconciler taskReconciler; |
| private FakeStatsProvider statsProvider; |
| |
| @Before |
| public void setUp() throws Exception { |
| storageUtil = new StorageTestUtil(this); |
| storageUtil.expectOperations(); |
| snapshotStore = createMock(SnapshotStore.class); |
| backup = createMock(StorageBackup.class); |
| recovery = createMock(Recovery.class); |
| maintenance = createMock(MaintenanceController.class); |
| cronJobManager = createMock(CronJobManager.class); |
| quotaManager = createMock(QuotaManager.class); |
| stateManager = createMock(StateManager.class); |
| uuidGenerator = createMock(UUIDGenerator.class); |
| jobUpdateController = createMock(JobUpdateController.class); |
| readOnlyScheduler = createMock(ReadOnlyScheduler.Iface.class); |
| auditMessages = createMock(AuditMessages.class); |
| taskReconciler = createMock(TaskReconciler.class); |
| statsProvider = new FakeStatsProvider(); |
| |
| thrift = getResponseProxy( |
| new SchedulerThriftInterface( |
| TaskTestUtil.CONFIGURATION_MANAGER, |
| THRESHOLDS, |
| storageUtil.storage, |
| snapshotStore, |
| backup, |
| recovery, |
| cronJobManager, |
| maintenance, |
| quotaManager, |
| stateManager, |
| uuidGenerator, |
| jobUpdateController, |
| readOnlyScheduler, |
| auditMessages, |
| taskReconciler, |
| statsProvider)); |
| } |
| |
| private static AuroraAdmin.Iface getResponseProxy(AuroraAdmin.Iface realThrift) { |
| // Capture all API method calls to validate response objects. |
| Class<AuroraAdmin.Iface> thriftClass = AuroraAdmin.Iface.class; |
| return (AuroraAdmin.Iface) Proxy.newProxyInstance( |
| thriftClass.getClassLoader(), |
| new Class<?>[] {thriftClass}, |
| (o, method, args) -> { |
| Response response; |
| try { |
| response = (Response) method.invoke(realThrift, args); |
| } catch (InvocationTargetException e) { |
| Throwables.propagateIfPossible(e.getTargetException(), TException.class); |
| throw e; |
| } |
| assertTrue(response.isSetResponseCode()); |
| assertNotNull(response.getDetails()); |
| return response; |
| }); |
| } |
| |
| private static SanitizedConfiguration fromUnsanitized(IJobConfiguration job) |
| throws TaskDescriptionException { |
| |
| return SanitizedConfiguration.fromUnsanitized(TaskTestUtil.CONFIGURATION_MANAGER, job); |
| } |
| |
| @Test |
| public void testCreateJob() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeProdJob()); |
| SanitizedConfiguration sanitized = fromUnsanitized(job); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| expectNoCronJob(); |
| expectInstanceQuotaCheck(sanitized, ENOUGH_QUOTA); |
| |
| stateManager.insertPendingTasks( |
| storageUtil.mutableStoreProvider, |
| sanitized.getJobConfig().getTaskConfig(), |
| sanitized.getInstanceIds()); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.createJob(makeProdJob())); |
| assertEquals(1L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsForCron() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeProdJob().setCronSchedule("")); |
| |
| control.replay(); |
| |
| assertEquals(invalidResponse(NO_CRON), thrift.createJob(job.newBuilder())); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsConfigCheck() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeJob(INVALID_TASK_CONFIG)); |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder())); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsJobExists() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeJob()); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), buildScheduledTask()); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder())); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsCronJobExists() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeJob()); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| expectCronJob(); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder())); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsInstanceCheck() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build( |
| makeJob(defaultTask(true), THRESHOLDS.getMaxTasksPerJob() + 1)); |
| |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| expectNoCronJob(); |
| expect(quotaManager.checkInstanceAddition( |
| anyObject(ITaskConfig.class), |
| anyInt(), |
| eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder())); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsQuotaCheck() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeProdJob()); |
| SanitizedConfiguration sanitized = fromUnsanitized(job); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| expectNoCronJob(); |
| expectInstanceQuotaCheck(sanitized, NOT_ENOUGH_QUOTA); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder())); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| private void assertMessageMatches(Response response, String string) { |
| // TODO(wfarner): This test coverage could be much better. Circle back to apply more thorough |
| // response contents testing throughout. |
| assertTrue(Iterables.any(response.getDetails(), detail -> detail.getMessage().equals(string))); |
| } |
| |
| @Test |
| public void testCreateEmptyJob() throws Exception { |
| control.replay(); |
| |
| JobConfiguration job = |
| new JobConfiguration().setKey(JOB_KEY.newBuilder()).setOwner(IDENTITY); |
| assertResponse(INVALID_REQUEST, thrift.createJob(job)); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobFailsNoExecutorOrContainerConfig() throws Exception { |
| JobConfiguration job = makeJob(); |
| job.getTaskConfig().unsetExecutorConfig(); |
| |
| control.replay(); |
| |
| Response response = thrift.createJob(job); |
| assertResponse(INVALID_REQUEST, response); |
| assertMessageMatches(response, ConfigurationManager.NO_EXECUTOR_OR_CONTAINER); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateHomogeneousJobNoInstances() throws Exception { |
| JobConfiguration job = makeJob(); |
| job.setInstanceCount(0); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job)); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobNegativeInstanceCount() throws Exception { |
| JobConfiguration job = makeJob(); |
| job.setInstanceCount(-1); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.createJob(job)); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobNoResources() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = productionTask(); |
| task.unsetResources(); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobBadCpu() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = productionTask() |
| .setResources(ImmutableSet.of( |
| numCpus(0.0), |
| ramMb(1024), |
| diskMb(1024))); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobBadRam() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = productionTask() |
| .setResources(ImmutableSet.of( |
| numCpus(1), |
| ramMb(-123), |
| diskMb(1024))); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateJobBadDisk() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = productionTask() |
| .setResources(ImmutableSet.of( |
| numCpus(1), |
| ramMb(1024), |
| diskMb(0))); |
| task.unsetResources(); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| } |
| |
| @Test |
| public void testCreateJobGoodResources() throws Exception { |
| |
| TaskConfig task = productionTask() |
| .setResources(ImmutableSet.of( |
| numCpus(1.0), |
| ramMb(1024), |
| diskMb(1024))); |
| |
| IJobConfiguration job = IJobConfiguration.build(makeJob(task)); |
| SanitizedConfiguration sanitized = fromUnsanitized(job); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| expectNoCronJob(); |
| expectInstanceQuotaCheck(sanitized, ENOUGH_QUOTA); |
| |
| stateManager.insertPendingTasks( |
| storageUtil.mutableStoreProvider, |
| sanitized.getJobConfig().getTaskConfig(), |
| sanitized.getInstanceIds()); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.createJob(job.newBuilder())); |
| } |
| |
| @Test |
| public void testCreateJobPopulateDefaults() throws Exception { |
| Set<Resource> resources = ImmutableSet.of(numCpus(1.0), ramMb(1024), diskMb(1024)); |
| TaskConfig task = new TaskConfig() |
| .setContactEmail("testing@twitter.com") |
| .setExecutorConfig( |
| new ExecutorConfig(EXECUTOR_NAME, "config")) // Arbitrary opaque data. |
| .setResources(resources) |
| .setIsService(true) |
| .setProduction(true) |
| .setTier(TaskTestUtil.PROD_TIER_NAME) |
| .setOwner(IDENTITY) |
| .setContainer(Container.mesos(new MesosContainer())) |
| .setJob(JOB_KEY.newBuilder()); |
| JobConfiguration job = makeJob(task); |
| |
| JobConfiguration sanitized = job.deepCopy(); |
| sanitized.getTaskConfig() |
| .setJob(JOB_KEY.newBuilder()) |
| .setPriority(0) |
| .setIsService(true) |
| .setProduction(true) |
| .setTaskLinks(ImmutableMap.of()) |
| .setConstraints(ImmutableSet.of()) |
| .setMaxTaskFailures(0) |
| .setResources(resources); |
| |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| expectNoCronJob(); |
| expectInstanceQuotaCheck(ITaskConfig.build(sanitized.getTaskConfig()), ENOUGH_QUOTA); |
| stateManager.insertPendingTasks( |
| storageUtil.mutableStoreProvider, |
| ITaskConfig.build(sanitized.getTaskConfig()), |
| ImmutableSet.of(0)); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.createJob(job)); |
| assertEquals(1L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testCreateUnauthorizedDedicatedJob() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = nonProductionTask(); |
| task.setConstraints(ImmutableSet.of(dedicatedConstraint(ImmutableSet.of("mesos")))); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testLimitConstraintForDedicatedJob() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = nonProductionTask(); |
| task.setConstraints(ImmutableSet.of(dedicatedConstraint(1))); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| @Test |
| public void testMultipleValueConstraintForDedicatedJob() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = nonProductionTask(); |
| task.setConstraints(ImmutableSet.of(dedicatedConstraint(ImmutableSet.of("mesos", "test")))); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| private IScheduledTask buildTaskForJobUpdate(int instanceId) { |
| return buildTaskForJobUpdate(instanceId, "data"); |
| } |
| |
| private IScheduledTask buildTaskForJobUpdate(int instanceId, String executorData) { |
| return buildTaskForJobUpdate( |
| instanceId, |
| executorData, |
| Optional.of(SlaPolicy.coordinatorSlaPolicy( |
| new CoordinatorSlaPolicy("test-url", "test-key")))); |
| } |
| |
| private IScheduledTask buildTaskForJobUpdate(int instanceId, |
| String executorData, |
| Optional<SlaPolicy> slaPolicy) { |
| |
| return IScheduledTask.build(new ScheduledTask() |
| .setAssignedTask(new AssignedTask() |
| .setInstanceId(instanceId) |
| .setTask(populatedTask() |
| .setIsService(true) |
| .setSlaPolicy(slaPolicy.orElse(null)) |
| .setExecutorConfig(new ExecutorConfig().setName(EXECUTOR_NAME) |
| .setData(executorData))))); |
| } |
| |
| private IScheduledTask buildScheduledTask() { |
| return buildScheduledTask(JOB_NAME, TASK_ID); |
| } |
| |
| private static IScheduledTask buildScheduledTask(String jobName, String taskId) { |
| return IScheduledTask.build(new ScheduledTask() |
| .setAssignedTask(new AssignedTask() |
| .setTaskId(taskId) |
| .setInstanceId(0) |
| .setTask(new TaskConfig() |
| .setJob(JOB_KEY.newBuilder().setName(jobName)) |
| .setOwner(IDENTITY)))); |
| } |
| |
| private void expectTransitionsToKilling() { |
| expectTransitionsToKilling(Optional.empty()); |
| } |
| |
| private void expectTransitionsToKilling(Optional<String> message) { |
| expect(auditMessages.killedByRemoteUser(message)).andReturn(Optional.of("test")); |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.KILLING, |
| Optional.of("test"))).andReturn(StateChangeResult.SUCCESS); |
| } |
| |
| @Test |
| public void testJobScopedKillsActive() throws Exception { |
| Query.Builder query = Query.unscoped().byJob(JOB_KEY).active(); |
| storageUtil.expectTaskFetch(query, buildScheduledTask()); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectTransitionsToKilling(); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.killTasks(JOB_KEY.newBuilder(), null, null)); |
| assertEquals(1L, statsProvider.getLongValue(KILL_TASKS)); |
| } |
| |
| @Test |
| public void testKillTasksInstanceScoped() throws Exception { |
| Query.Builder query = Query.instanceScoped(JOB_KEY, ImmutableSet.of(1)).active(); |
| storageUtil.expectTaskFetch(query, buildScheduledTask()); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectTransitionsToKilling(); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.killTasks(JOB_KEY.newBuilder(), ImmutableSet.of(1), null)); |
| assertEquals(1L, statsProvider.getLongValue(KILL_TASKS)); |
| } |
| |
| @Test |
| public void testKillTasksWhileJobUpdating() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectLastCall().andThrow(new JobUpdatingException("job is updating")); |
| |
| control.replay(); |
| |
| assertResponse(JOB_UPDATING_ERROR, thrift.killTasks(JOB_KEY.newBuilder(), null, null)); |
| assertEquals(0L, statsProvider.getLongValue(KILL_TASKS)); |
| } |
| |
| @Test |
| public void testKillNonExistentTasks() throws Exception { |
| Query.Builder query = Query.unscoped().byJob(JOB_KEY).active(); |
| storageUtil.expectTaskFetch(query); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| |
| control.replay(); |
| |
| Response response = thrift.killTasks(JOB_KEY.newBuilder(), null, null); |
| assertOkResponse(response); |
| assertMessageMatches(response, SchedulerThriftInterface.NO_TASKS_TO_KILL_MESSAGE); |
| assertEquals(0L, statsProvider.getLongValue(KILL_TASKS)); |
| } |
| |
| @Test |
| public void testKillTasksWithMessage() throws Exception { |
| String message = "Test message"; |
| Query.Builder query = Query.unscoped().byJob(JOB_KEY).active(); |
| storageUtil.expectTaskFetch(query, buildScheduledTask()); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectTransitionsToKilling(Optional.of(message)); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.killTasks(JOB_KEY.newBuilder(), null, message)); |
| } |
| |
| @Test |
| public void testPruneTasksRejectsActiveStates() throws Exception { |
| control.replay(); |
| Response rsp = thrift.pruneTasks(new TaskQuery().setStatuses(Tasks.ACTIVE_STATES)); |
| assertResponse(ResponseCode.ERROR, rsp); |
| assertEquals(0L, statsProvider.getLongValue(PRUNE_TASKS)); |
| } |
| |
| @Test |
| public void testPruneTasksRejectsMixedStates() throws Exception { |
| control.replay(); |
| Response rsp = thrift.pruneTasks(new TaskQuery().setStatuses( |
| ImmutableSet.of(ScheduleStatus.FINISHED, ScheduleStatus.KILLING))); |
| assertResponse(ResponseCode.ERROR, rsp); |
| assertEquals(0L, statsProvider.getLongValue(PRUNE_TASKS)); |
| } |
| |
| @Test |
| public void testPruneTasksAddsDefaultStatuses() throws Exception { |
| storageUtil.expectTaskFetch( |
| Query.arbitrary(new TaskQuery().setStatuses(Tasks.TERMINAL_STATES)), |
| buildScheduledTask()); |
| stateManager.deleteTasks( |
| storageUtil.mutableStoreProvider, |
| ImmutableSet.of(buildScheduledTask().getAssignedTask().getTaskId())); |
| control.replay(); |
| |
| assertOkResponse(thrift.pruneTasks(new TaskQuery())); |
| assertEquals(1L, statsProvider.getLongValue(PRUNE_TASKS)); |
| } |
| |
| @Test |
| public void testPruneTasksRespectsScopedTerminalState() throws Exception { |
| storageUtil.expectTaskFetch( |
| Query.arbitrary(new TaskQuery().setStatuses(ImmutableSet.of(ScheduleStatus.FAILED))), |
| buildScheduledTask()); |
| stateManager.deleteTasks( |
| storageUtil.mutableStoreProvider, |
| ImmutableSet.of(buildScheduledTask().getAssignedTask().getTaskId())); |
| control.replay(); |
| |
| assertOkResponse(thrift.pruneTasks( |
| new TaskQuery().setStatuses(ImmutableSet.of(ScheduleStatus.FAILED)))); |
| assertEquals(1L, statsProvider.getLongValue(PRUNE_TASKS)); |
| } |
| |
| @Test |
| public void testPruneTasksAppliesQueryLimit() throws Exception { |
| TaskQuery query = new TaskQuery().setLimit(3); |
| storageUtil.expectTaskFetch( |
| Query.arbitrary(query.setStatuses(Tasks.TERMINAL_STATES)), |
| buildScheduledTask("a/b/c", "task1"), |
| buildScheduledTask("a/b/c", "task2"), |
| buildScheduledTask("a/b/c", "task3"), |
| buildScheduledTask("a/b/c", "task4"), |
| buildScheduledTask("a/b/c", "task5")); |
| stateManager.deleteTasks( |
| storageUtil.mutableStoreProvider, |
| ImmutableSet.of("task1", "task2", "task3")); |
| control.replay(); |
| |
| assertOkResponse(thrift.pruneTasks(query)); |
| assertEquals(3L, statsProvider.getLongValue(PRUNE_TASKS)); |
| } |
| |
| @Test |
| public void testSetQuota() throws Exception { |
| ResourceAggregate resourceAggregate = new ResourceAggregate() |
| .setResources(ImmutableSet.of(numCpus(10.0), ramMb(200), diskMb(100))); |
| quotaManager.saveQuota( |
| ROLE, |
| IResourceAggregate.build(resourceAggregate), |
| storageUtil.mutableStoreProvider); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.setQuota(ROLE, resourceAggregate)); |
| } |
| |
| @Test |
| public void testSetQuotaFails() throws Exception { |
| ResourceAggregate resourceAggregate = new ResourceAggregate() |
| .setResources(ImmutableSet.of(numCpus(10.0), ramMb(200), diskMb(100))); |
| quotaManager.saveQuota( |
| ROLE, |
| IResourceAggregate.build(resourceAggregate.deepCopy() |
| .setResources(ImmutableSet.of(numCpus(10), ramMb(200), diskMb(100)))), |
| storageUtil.mutableStoreProvider); |
| |
| expectLastCall().andThrow(new QuotaManager.QuotaException("fail")); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.setQuota(ROLE, resourceAggregate)); |
| } |
| |
| @Test |
| public void testForceTaskState() throws Exception { |
| ScheduleStatus status = ScheduleStatus.FAILED; |
| |
| expect(auditMessages.transitionedBy()).andReturn(Optional.of("test")); |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.FAILED, |
| Optional.of("test"))).andReturn(StateChangeResult.SUCCESS); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.forceTaskState(TASK_ID, status)); |
| } |
| |
| @Test |
| public void testBackupControls() throws Exception { |
| backup.backupNow(); |
| |
| Set<String> backups = ImmutableSet.of("a", "b"); |
| expect(recovery.listBackups()).andReturn(backups); |
| |
| String backupId = "backup"; |
| recovery.stage(backupId); |
| |
| Query.Builder query = Query.taskScoped("taskId"); |
| Set<IScheduledTask> queryResult = ImmutableSet.of( |
| IScheduledTask.build(new ScheduledTask().setStatus(ScheduleStatus.RUNNING))); |
| expect(recovery.query(query)).andReturn(queryResult); |
| |
| recovery.deleteTasks(query); |
| |
| recovery.commit(); |
| |
| recovery.unload(); |
| |
| control.replay(); |
| |
| assertEquals(okEmptyResponse(), thrift.performBackup()); |
| |
| assertEquals( |
| okResponse(Result.listBackupsResult(new ListBackupsResult().setBackups(backups))), |
| thrift.listBackups()); |
| |
| assertEquals(okEmptyResponse(), thrift.stageRecovery(backupId)); |
| |
| assertEquals( |
| okResponse(Result.queryRecoveryResult( |
| new QueryRecoveryResult().setTasks(IScheduledTask.toBuildersSet(queryResult)))), |
| thrift.queryRecovery(query.get().newBuilder())); |
| |
| assertEquals(okEmptyResponse(), thrift.deleteRecoveryTasks(query.get().newBuilder())); |
| |
| assertEquals(okEmptyResponse(), thrift.commitRecovery()); |
| |
| assertEquals(okEmptyResponse(), thrift.unloadRecovery()); |
| } |
| |
| @Test |
| public void testRecoveryException() throws Exception { |
| Throwable recoveryException = new RecoveryException("Injected"); |
| |
| String backupId = "backup"; |
| recovery.stage(backupId); |
| expectLastCall().andThrow(recoveryException); |
| |
| control.replay(); |
| |
| try { |
| thrift.stageRecovery(backupId); |
| fail("No recovery exception thrown."); |
| } catch (RecoveryException e) { |
| assertEquals(recoveryException.getMessage(), e.getMessage()); |
| } |
| } |
| |
| @Test |
| public void testRestartShards() throws Exception { |
| Set<Integer> shards = ImmutableSet.of(0); |
| |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch( |
| Query.instanceScoped(JOB_KEY, shards).active(), |
| buildScheduledTask()); |
| |
| expect(auditMessages.restartedByRemoteUser()) |
| .andReturn(Optional.of("test")); |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.RESTARTING, |
| Optional.of("test"))).andReturn(StateChangeResult.SUCCESS); |
| |
| control.replay(); |
| |
| assertOkResponse( |
| thrift.restartShards(JOB_KEY.newBuilder(), shards)); |
| assertEquals(1L, statsProvider.getLongValue(RESTART_SHARDS)); |
| } |
| |
| @Test |
| public void testRestartShardsWhileJobUpdating() throws Exception { |
| Set<Integer> shards = ImmutableSet.of(1, 6); |
| |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectLastCall().andThrow(new JobUpdatingException("job is updating")); |
| |
| control.replay(); |
| |
| assertResponse( |
| JOB_UPDATING_ERROR, |
| thrift.restartShards(JOB_KEY.newBuilder(), shards)); |
| assertEquals(0L, statsProvider.getLongValue(KILL_TASKS)); |
| } |
| |
| @Test |
| public void testRestartShardsNotFoundTasksFailure() throws Exception { |
| Set<Integer> shards = ImmutableSet.of(1, 6); |
| |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.instanceScoped(JOB_KEY, shards).active()); |
| |
| control.replay(); |
| |
| assertResponse( |
| INVALID_REQUEST, |
| thrift.restartShards(JOB_KEY.newBuilder(), shards)); |
| assertEquals(0L, statsProvider.getLongValue(KILL_TASKS)); |
| } |
| |
| @Test |
| public void testReplaceCronTemplate() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB)); |
| |
| expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA); |
| cronJobManager.updateJob(anyObject(SanitizedCronJob.class)); |
| control.replay(); |
| |
| assertOkResponse(thrift.replaceCronTemplate(CRON_JOB)); |
| assertEquals(1L, statsProvider.getLongValue(CREATE_OR_UPDATE_CRON)); |
| } |
| |
| @Test |
| public void testReplaceCronTemplateDoesNotExist() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB)); |
| |
| expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA); |
| cronJobManager.updateJob(anyObject(SanitizedCronJob.class)); |
| expectLastCall().andThrow(new CronException("Nope")); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.replaceCronTemplate(CRON_JOB)); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_OR_UPDATE_CRON)); |
| } |
| |
| @Test |
| public void testStartCronJob() throws Exception { |
| cronJobManager.startJobNow(JOB_KEY); |
| control.replay(); |
| assertResponse(OK, thrift.startCronJob(JOB_KEY.newBuilder())); |
| } |
| |
| @Test |
| public void testStartCronJobFailsInCronManager() throws Exception { |
| cronJobManager.startJobNow(JOB_KEY); |
| expectLastCall().andThrow(new CronException("failed")); |
| control.replay(); |
| assertResponse(INVALID_REQUEST, thrift.startCronJob(JOB_KEY.newBuilder())); |
| } |
| |
| @Test |
| public void testScheduleCronCreatesJob() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB)); |
| |
| expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA); |
| |
| expectNoCronJob().times(2); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| cronJobManager.createJob(SanitizedCronJob.from(sanitized)); |
| control.replay(); |
| assertResponse(OK, thrift.scheduleCronJob(CRON_JOB)); |
| } |
| |
| @Test |
| public void testScheduleCronFailsCreationDueToExistingNonCron() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB)); |
| |
| expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA); |
| |
| expectNoCronJob(); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), buildScheduledTask()); |
| control.replay(); |
| assertEquals( |
| invalidResponse(jobAlreadyExistsMessage(JOB_KEY)), |
| thrift.scheduleCronJob(CRON_JOB)); |
| } |
| |
| @Test |
| public void testScheduleCronUpdatesJob() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB)); |
| |
| expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA); |
| |
| expectCronJob(); |
| cronJobManager.updateJob(SanitizedCronJob.from(sanitized)); |
| control.replay(); |
| |
| assertResponse(OK, thrift.scheduleCronJob(CRON_JOB)); |
| } |
| |
| @Test |
| public void testScheduleCronJobFailedTaskConfigValidation() throws Exception { |
| control.replay(); |
| IJobConfiguration job = IJobConfiguration.build(makeJob(INVALID_TASK_CONFIG)); |
| assertResponse(INVALID_REQUEST, thrift.scheduleCronJob(job.newBuilder())); |
| } |
| |
| @Test |
| public void testScheduleCronJobFailsWithNoCronSchedule() throws Exception { |
| control.replay(); |
| |
| assertEquals( |
| invalidResponse(noCronScheduleMessage(JOB_KEY)), |
| thrift.scheduleCronJob(makeJob())); |
| } |
| |
| @Test |
| public void testScheduleCronFailsQuotaCheck() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB)); |
| |
| expectCronQuotaCheck(sanitized.getJobConfig(), NOT_ENOUGH_QUOTA); |
| |
| control.replay(); |
| assertResponse(INVALID_REQUEST, thrift.scheduleCronJob(CRON_JOB)); |
| } |
| |
| @Test |
| public void testDescheduleCronJob() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expect(cronJobManager.deleteJob(JOB_KEY)).andReturn(true); |
| |
| control.replay(); |
| |
| assertResponse(OK, thrift.descheduleCronJob(CRON_JOB.getKey())); |
| } |
| |
| @Test |
| public void testDescheduleCronJobWhileJobUpdating() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectLastCall().andThrow(new JobUpdatingException("job is updating")); |
| control.replay(); |
| assertResponse(JOB_UPDATING_ERROR, thrift.descheduleCronJob(CRON_JOB.getKey())); |
| } |
| |
| @Test |
| public void testDescheduleNotACron() throws Exception { |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expect(cronJobManager.deleteJob(JOB_KEY)).andReturn(false); |
| control.replay(); |
| |
| assertEquals( |
| okEmptyResponse(notScheduledCronMessage(JOB_KEY)), |
| thrift.descheduleCronJob(JOB_KEY.newBuilder())); |
| } |
| |
| @Test |
| public void testUnauthorizedDedicatedJob() throws Exception { |
| control.replay(); |
| |
| TaskConfig task = nonProductionTask(); |
| task.setConstraints(ImmutableSet.of(dedicatedConstraint(ImmutableSet.of("mesos")))); |
| assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task))); |
| assertEquals(0L, statsProvider.getLongValue(CREATE_JOB)); |
| } |
| |
| private static Set<IHostStatus> status(String host, MaintenanceMode mode) { |
| return ImmutableSet.of(IHostStatus.build(new HostStatus(host, mode))); |
| } |
| |
| @Test |
| public void testHostMaintenance() throws Exception { |
| Set<String> hostnames = ImmutableSet.of("a"); |
| Set<IHostStatus> none = status("a", NONE); |
| Set<IHostStatus> scheduled = status("a", SCHEDULED); |
| Set<IHostStatus> draining = status("a", DRAINING); |
| Set<IHostStatus> drained = status("a", DRAINING); |
| expect(maintenance.getStatus(hostnames)).andReturn(none); |
| expect(maintenance.startMaintenance(hostnames)).andReturn(scheduled); |
| expect(maintenance.drain(hostnames)).andReturn(draining); |
| expect(maintenance.getStatus(hostnames)).andReturn(draining); |
| expect(maintenance.getStatus(hostnames)).andReturn(drained); |
| expect(maintenance.endMaintenance(hostnames)).andReturn(none); |
| |
| control.replay(); |
| |
| Hosts hosts = new Hosts(hostnames); |
| |
| assertEquals( |
| IHostStatus.toBuildersSet(none), |
| thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult() |
| .getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(MAINTENANCE_STATUS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(scheduled), |
| thrift.startMaintenance(hosts).getResult().getStartMaintenanceResult() |
| .getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(START_MAINTENANCE)); |
| assertEquals( |
| IHostStatus.toBuildersSet(draining), |
| thrift.drainHosts(hosts).getResult().getDrainHostsResult().getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(DRAIN_HOSTS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(draining), |
| thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult() |
| .getStatuses()); |
| assertEquals(2L, statsProvider.getLongValue(MAINTENANCE_STATUS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(drained), |
| thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult() |
| .getStatuses()); |
| assertEquals(3L, statsProvider.getLongValue(MAINTENANCE_STATUS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(none), |
| thrift.endMaintenance(hosts).getResult().getEndMaintenanceResult().getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(END_MAINTENANCE)); |
| } |
| |
| @Test |
| public void testSLAHostMaintenance() throws Exception { |
| Set<String> hostnames = ImmutableSet.of("a"); |
| SlaPolicy defaultSlaPolicy = SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy().setPercentage(95)); |
| Set<IHostStatus> none = status("a", NONE); |
| Set<IHostStatus> scheduled = status("a", SCHEDULED); |
| Set<IHostStatus> draining = status("a", DRAINING); |
| Set<IHostStatus> drained = status("a", DRAINING); |
| expect(maintenance.getStatus(hostnames)).andReturn(none); |
| expect(maintenance.startMaintenance(hostnames)).andReturn(scheduled); |
| expect(maintenance.slaDrain(hostnames, defaultSlaPolicy, 10)).andReturn(draining); |
| expect(maintenance.getStatus(hostnames)).andReturn(draining); |
| expect(maintenance.getStatus(hostnames)).andReturn(drained); |
| expect(maintenance.endMaintenance(hostnames)).andReturn(none); |
| |
| control.replay(); |
| |
| Hosts hosts = new Hosts(hostnames); |
| |
| assertEquals( |
| IHostStatus.toBuildersSet(none), |
| thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult() |
| .getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(MAINTENANCE_STATUS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(scheduled), |
| thrift.startMaintenance(hosts).getResult().getStartMaintenanceResult() |
| .getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(START_MAINTENANCE)); |
| assertEquals( |
| IHostStatus.toBuildersSet(draining), |
| thrift.slaDrainHosts(hosts, defaultSlaPolicy, 10) |
| .getResult() |
| .getDrainHostsResult() |
| .getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(SLA_DRAIN_HOSTS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(draining), |
| thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult() |
| .getStatuses()); |
| assertEquals(2L, statsProvider.getLongValue(MAINTENANCE_STATUS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(drained), |
| thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult() |
| .getStatuses()); |
| assertEquals(3L, statsProvider.getLongValue(MAINTENANCE_STATUS)); |
| assertEquals( |
| IHostStatus.toBuildersSet(none), |
| thrift.endMaintenance(hosts).getResult().getEndMaintenanceResult().getStatuses()); |
| assertEquals(1L, statsProvider.getLongValue(END_MAINTENANCE)); |
| } |
| |
| private static Response okEmptyResponse() { |
| return response(OK, Optional.empty()); |
| } |
| |
| private static Response okEmptyResponse(String message) { |
| return Responses.empty() |
| .setResponseCode(OK) |
| .setDetails(ImmutableList.of(new ResponseDetail(message))); |
| } |
| |
| private static Response invalidResponse(String message) { |
| return Responses.empty() |
| .setResponseCode(INVALID_REQUEST) |
| .setDetails(ImmutableList.of(new ResponseDetail(message))); |
| } |
| |
| @Test |
| public void testSnapshot() throws Exception { |
| snapshotStore.snapshot(); |
| expectLastCall(); |
| |
| snapshotStore.snapshot(); |
| expectLastCall().andThrow(new StorageException("mock error!")); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.snapshot()); |
| |
| try { |
| thrift.snapshot(); |
| fail("No StorageException thrown."); |
| } catch (StorageException e) { |
| // Expected. |
| } |
| } |
| |
| @Test |
| public void testExplicitTaskReconciliationWithNoBatchSize() throws Exception { |
| ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings(); |
| |
| taskReconciler.triggerExplicitReconciliation(Optional.empty()); |
| expectLastCall(); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.triggerExplicitTaskReconciliation(settings)); |
| } |
| |
| @Test |
| public void testExplicitTaskReconciliationWithValidBatchSize() throws Exception { |
| ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings(); |
| settings.setBatchSize(10); |
| |
| taskReconciler.triggerExplicitReconciliation(Optional.of(settings.getBatchSize())); |
| expectLastCall(); |
| |
| control.replay(); |
| assertOkResponse(thrift.triggerExplicitTaskReconciliation(settings)); |
| } |
| |
| @Test |
| public void testExplicitTaskReconciliationWithNegativeBatchSize() throws Exception { |
| ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings(); |
| settings.setBatchSize(-1000); |
| |
| control.replay(); |
| assertResponse(INVALID_REQUEST, thrift.triggerExplicitTaskReconciliation(settings)); |
| } |
| |
| @Test |
| public void testExplicitTaskReconciliationWithZeroBatchSize() throws Exception { |
| ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings(); |
| settings.setBatchSize(0); |
| |
| control.replay(); |
| assertResponse(INVALID_REQUEST, thrift.triggerExplicitTaskReconciliation(settings)); |
| } |
| |
| @Test |
| public void testImplicitTaskReconciliation() throws Exception { |
| taskReconciler.triggerImplicitReconciliation(); |
| expectLastCall(); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.triggerImplicitTaskReconciliation()); |
| } |
| |
| @Test |
| public void testAddInstancesWithInstanceKey() throws Exception { |
| expectNoCronJob(); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| IScheduledTask activeTask = buildScheduledTask(); |
| ITaskConfig task = activeTask.getAssignedTask().getTask(); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), activeTask); |
| expect(quotaManager.checkInstanceAddition( |
| task, |
| 2, |
| storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA); |
| stateManager.insertPendingTasks( |
| storageUtil.mutableStoreProvider, |
| task, |
| ImmutableSet.of(1, 2)); |
| |
| control.replay(); |
| |
| assertOkResponse(thrift.addInstances(INSTANCE_KEY, 2)); |
| assertEquals(2L, statsProvider.getLongValue(ADD_INSTANCES)); |
| } |
| |
| @Test |
| public void testAddInstancesWithInstanceKeyFailsWithNoInstance() throws Exception { |
| expectNoCronJob(); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active()); |
| |
| control.replay(); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_INSTANCE_ID), |
| thrift.addInstances(INSTANCE_KEY, 2)); |
| assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES)); |
| } |
| |
| @Test |
| public void testAddInstancesWithInstanceKeyFailsInvalidCount() throws Exception { |
| control.replay(); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_INSTANCE_COUNT), |
| thrift.addInstances(INSTANCE_KEY, 0)); |
| assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES)); |
| } |
| |
| @Test |
| public void testAddInstancesFailsCronJob() throws Exception { |
| expectCronJob(); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.addInstances(INSTANCE_KEY, 1)); |
| } |
| |
| @Test(expected = StorageException.class) |
| public void testAddInstancesFailsWithNonTransient() throws Exception { |
| expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andThrow(new StorageException("no retry")); |
| |
| control.replay(); |
| |
| thrift.addInstances(INSTANCE_KEY, 1); |
| } |
| |
| @Test |
| public void testAddInstancesWhileJobUpdating() throws Exception { |
| expectNoCronJob(); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| expectLastCall().andThrow(new JobUpdatingException("job is updating")); |
| |
| control.replay(); |
| |
| assertResponse(JOB_UPDATING_ERROR, thrift.addInstances(INSTANCE_KEY, 1)); |
| assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES)); |
| } |
| |
| @Test |
| public void testAddInstancesFailsQuotaCheck() throws Exception { |
| IScheduledTask activeTask = buildScheduledTask(); |
| ITaskConfig task = activeTask.getAssignedTask().getTask(); |
| expectNoCronJob(); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), activeTask); |
| expectInstanceQuotaCheck(task, NOT_ENOUGH_QUOTA); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.addInstances(INSTANCE_KEY, 1)); |
| assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES)); |
| } |
| |
| @Test |
| public void testAddInstancesInstanceCollisionFailure() throws Exception { |
| IScheduledTask activeTask = buildScheduledTask(); |
| ITaskConfig task = activeTask.getAssignedTask().getTask(); |
| expectNoCronJob(); |
| jobUpdateController.assertNotUpdating(JOB_KEY); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), activeTask); |
| expectInstanceQuotaCheck(task, ENOUGH_QUOTA); |
| stateManager.insertPendingTasks( |
| storageUtil.mutableStoreProvider, |
| task, |
| ImmutableSet.of(1)); |
| expectLastCall().andThrow(new IllegalArgumentException("instance collision")); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.addInstances(INSTANCE_KEY, 1)); |
| assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES)); |
| } |
| |
| @Test |
| public void testStartUpdate() throws Exception { |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| |
| ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); |
| |
| IScheduledTask oldTask1 = buildTaskForJobUpdate(0, "old"); |
| IScheduledTask oldTask2 = buildTaskForJobUpdate(1, "old"); |
| IScheduledTask oldTask3 = buildTaskForJobUpdate(2, "old2"); |
| IScheduledTask oldTask4 = buildTaskForJobUpdate(3, "old2"); |
| IScheduledTask oldTask5 = buildTaskForJobUpdate(4, "old"); |
| IScheduledTask oldTask6 = buildTaskForJobUpdate(5, "old"); |
| IScheduledTask oldTask7 = buildTaskForJobUpdate(6, "old"); |
| |
| IJobUpdate update = buildJobUpdate(6, newTask, ImmutableMap.of( |
| oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1), new Range(4, 6)), |
| oldTask3.getAssignedTask().getTask(), ImmutableSet.of(new Range(2, 3)) |
| )); |
| |
| expect(quotaManager.checkJobUpdate( |
| update, |
| storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA); |
| |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| storageUtil.expectTaskFetch( |
| Query.unscoped().byJob(JOB_KEY).active(), |
| oldTask1, |
| oldTask2, |
| oldTask3, |
| oldTask4, |
| oldTask5, |
| oldTask6, |
| oldTask7); |
| |
| jobUpdateController.start(update, AUDIT); |
| |
| control.replay(); |
| |
| Response response = |
| assertOkResponse(thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE)); |
| assertEquals( |
| new StartJobUpdateResult(UPDATE_KEY.newBuilder()).setUpdateSummary( |
| update.getSummary().newBuilder()), response.getResult().getStartJobUpdateResult()); |
| assertEquals(6L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| private void expectJobUpdateQuotaCheck(QuotaCheckResult result) { |
| expect(quotaManager.checkJobUpdate( |
| anyObject(IJobUpdate.class), |
| eq(storageUtil.mutableStoreProvider))).andReturn(result); |
| } |
| |
| @Test |
| public void testStartUpdateEmptyDesired() throws Exception { |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| |
| ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); |
| |
| IScheduledTask oldTask1 = buildTaskForJobUpdate(0); |
| IScheduledTask oldTask2 = buildTaskForJobUpdate(1); |
| |
| // Set instance count to 1 to generate empty desired state in diff. |
| IJobUpdate update = buildJobUpdate( |
| 1, |
| newTask, |
| ImmutableMap.of(oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1)))); |
| |
| expectJobUpdateQuotaCheck(ENOUGH_QUOTA); |
| |
| // Set diff-adjusted IJobUpdate expectations. |
| JobUpdate expected = update.newBuilder(); |
| expected.getInstructions().setInitialState(ImmutableSet.of( |
| new InstanceTaskConfig(newTask.newBuilder(), ImmutableSet.of(new Range(1, 1))))); |
| expected.getInstructions().unsetDesiredState(); |
| |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| storageUtil.expectTaskFetch( |
| Query.unscoped().byJob(JOB_KEY).active(), |
| oldTask1, |
| oldTask2); |
| |
| jobUpdateController.start(IJobUpdate.build(expected), AUDIT); |
| |
| control.replay(); |
| |
| Response response = assertOkResponse( |
| thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE)); |
| assertEquals( |
| new StartJobUpdateResult(UPDATE_KEY.newBuilder()).setUpdateSummary(expected.getSummary()), |
| response.getResult().getStartJobUpdateResult()); |
| assertEquals(1L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test(expected = NullPointerException.class) |
| public void testStartUpdateFailsNullRequest() throws Exception { |
| control.replay(); |
| thrift.startJobUpdate(null, AUDIT_MESSAGE); |
| } |
| |
| @Test(expected = NullPointerException.class) |
| public void testStartUpdateFailsNullTaskConfig() throws Exception { |
| control.replay(); |
| thrift.startJobUpdate( |
| new JobUpdateRequest(null, 5, buildJobUpdateSettings()), |
| AUDIT_MESSAGE); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidGroupSize() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setUpdateStrategy( |
| JobUpdateStrategy.queueStrategy(new QueueJobUpdateStrategy().setGroupSize(0))); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.NO_INSTANCES_MODIFIED), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidGroupSizeVariableBatch() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setUpdateStrategy( |
| JobUpdateStrategy.varBatchStrategy( |
| new VariableBatchJobUpdateStrategy().setGroupSizes(ImmutableList.of(1, 0, 4)))); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_GROUP_SIZE), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidGroupsSum() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setUpdateStrategy( |
| JobUpdateStrategy.varBatchStrategy( |
| new VariableBatchJobUpdateStrategy().setGroupSizes(ImmutableList.of()))); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.NO_INSTANCES_MODIFIED), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidMaxInstanceFailures() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setMaxPerInstanceFailures(-1); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_MAX_INSTANCE_FAILURES), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsTooManyPerInstanceFailures() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings() |
| .setMaxPerInstanceFailures(THRESHOLDS.getMaxUpdateInstanceFailures() + 10); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.TOO_MANY_POTENTIAL_FAILED_INSTANCES), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidMaxFailedInstances() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setMaxFailedInstances(-1); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_MAX_FAILED_INSTANCES), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidMinWaitInRunning() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setMinWaitInInstanceRunningMs(-1); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_MIN_WAIT_TO_RUNNING), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsNonServiceTask() throws Exception { |
| control.replay(); |
| JobUpdateRequest request = buildJobUpdateRequest(populatedTask().setIsService(false)); |
| assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInvalidPulseTimeout() throws Exception { |
| control.replay(); |
| |
| JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); |
| updateRequest.getSettings().setBlockIfNoPulsesAfterMs(-1); |
| |
| assertEquals( |
| invalidResponse(SchedulerThriftInterface.INVALID_PULSE_TIMEOUT), |
| thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsForCronJob() throws Exception { |
| JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); |
| expectCronJob(); |
| |
| control.replay(); |
| assertEquals(invalidResponse(NO_CRON), thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsConfigValidation() throws Exception { |
| JobUpdateRequest request = buildJobUpdateRequest(INVALID_TASK_CONFIG.setIsService(true)); |
| |
| control.replay(); |
| assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartNoopUpdate() throws Exception { |
| expectNoCronJob(); |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); |
| |
| IScheduledTask oldTask = buildTaskForJobUpdate(0); |
| storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); |
| |
| IJobUpdate update = buildJobUpdate( |
| 1, |
| newTask, |
| ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))); |
| |
| control.replay(); |
| JobUpdateRequest request = buildJobUpdateRequest(update); |
| Response response = thrift.startJobUpdate(request, AUDIT_MESSAGE); |
| assertResponse(OK, response); |
| assertEquals( |
| NOOP_JOB_UPDATE_MESSAGE, |
| Iterables.getOnlyElement(response.getDetails()).getMessage()); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateInvalidScope() throws Exception { |
| expectNoCronJob(); |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| |
| IScheduledTask oldTask = buildTaskForJobUpdate(0); |
| storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); |
| |
| ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); |
| JobUpdate builder = buildJobUpdate( |
| 1, |
| newTask, |
| ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))) |
| .newBuilder(); |
| builder.getInstructions().getSettings() |
| .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(100, 100))); |
| |
| control.replay(); |
| JobUpdateRequest request = buildJobUpdateRequest(IJobUpdate.build(builder)); |
| assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartScopeIncludesNoop() throws Exception { |
| // Test for regression of AURORA-1332: a scoped update should be allowed when unchanged |
| // instances are included in the scope. |
| |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| |
| ScheduledTask taskBuilder = buildTaskForJobUpdate(0).newBuilder(); |
| taskBuilder.getAssignedTask().getTask() |
| .setResources(ImmutableSet.of( |
| numCpus(100), |
| ramMb(1024), |
| diskMb(1024))); |
| IScheduledTask newTask = IScheduledTask.build(taskBuilder); |
| |
| IScheduledTask oldTask1 = buildTaskForJobUpdate(1); |
| IScheduledTask oldTask2 = buildTaskForJobUpdate(2); |
| storageUtil.expectTaskFetch( |
| Query.unscoped().byJob(JOB_KEY).active(), |
| newTask, oldTask1, oldTask2); |
| expectJobUpdateQuotaCheck(ENOUGH_QUOTA); |
| jobUpdateController.start(EasyMock.<IJobUpdate>anyObject(), eq(AUDIT)); |
| |
| ITaskConfig newConfig = newTask.getAssignedTask().getTask(); |
| JobUpdate builder = buildJobUpdate( |
| 3, |
| newConfig, |
| ImmutableMap.of( |
| newConfig, ImmutableSet.of(new Range(0, 0)), |
| oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(1, 2)))) |
| .newBuilder(); |
| builder.getInstructions().getSettings() |
| .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(0, 2))); |
| |
| control.replay(); |
| JobUpdateRequest request = buildJobUpdateRequest(IJobUpdate.build(builder)); |
| assertResponse(OK, thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(3L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInstanceCountCheck() throws Exception { |
| JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); |
| request.setInstanceCount(THRESHOLDS.getMaxTasksPerJob() + 1); |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active()); |
| expectJobUpdateQuotaCheck(ENOUGH_QUOTA); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsQuotaCheck() throws Exception { |
| JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| |
| IScheduledTask oldTask = buildTaskForJobUpdate(0); |
| storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); |
| |
| expectJobUpdateQuotaCheck(NOT_ENOUGH_QUOTA); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInController() throws Exception { |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| |
| IScheduledTask oldTask = buildTaskForJobUpdate(0, "old"); |
| ITaskConfig newTask = buildTaskForJobUpdate(0, "new").getAssignedTask().getTask(); |
| |
| IJobUpdate update = buildJobUpdate( |
| 1, |
| newTask, |
| ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))); |
| |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| expectJobUpdateQuotaCheck(ENOUGH_QUOTA); |
| |
| storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); |
| jobUpdateController.start(update, AUDIT); |
| expectLastCall().andThrow(new UpdateStateException("failed")); |
| |
| control.replay(); |
| |
| assertResponse( |
| INVALID_REQUEST, |
| thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE)); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsInControllerWhenUpdateInProgress() throws Exception { |
| expectGetRemoteUser(); |
| expectNoCronJob(); |
| |
| IScheduledTask oldTask = buildTaskForJobUpdate(0, "old"); |
| ITaskConfig newTask = buildTaskForJobUpdate(0, "new").getAssignedTask().getTask(); |
| |
| IJobUpdate update = buildJobUpdate( |
| 1, |
| newTask, |
| ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))); |
| |
| expect(uuidGenerator.createNew()).andReturn(UU_ID); |
| expectJobUpdateQuotaCheck(ENOUGH_QUOTA); |
| |
| storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); |
| jobUpdateController.start(update, AUDIT); |
| expectLastCall().andThrow(new UpdateInProgressException("failed", update.getSummary())); |
| |
| control.replay(); |
| |
| Response response = thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE); |
| |
| assertResponse(INVALID_REQUEST, response); |
| assertEquals( |
| new StartJobUpdateResult(UPDATE_KEY.newBuilder()).setUpdateSummary( |
| update.getSummary().newBuilder()), response.getResult().getStartJobUpdateResult()); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testStartUpdateFailsWhenSlaAwareWithNoPolicy() throws Exception { |
| IScheduledTask oldTask = buildTaskForJobUpdate(0, "old"); |
| ITaskConfig newTask = buildTaskForJobUpdate(0, "new", Optional.empty()) |
| .getAssignedTask() |
| .getTask(); |
| |
| IJobUpdate update = buildJobUpdate( |
| 1, |
| newTask, |
| ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))), |
| buildJobUpdateSettings().setSlaAware(true)); |
| |
| control.replay(); |
| |
| JobUpdateRequest request = buildJobUpdateRequest(update); |
| Response response = thrift.startJobUpdate(request, AUDIT_MESSAGE); |
| assertEquals(invalidResponse(INVALID_SLA_AWARE_UPDATE), response); |
| assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE)); |
| } |
| |
| @Test |
| public void testPauseJobUpdateByCoordinator() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.pause(UPDATE_KEY, AUDIT); |
| |
| control.replay(); |
| |
| assertResponse(OK, thrift.pauseJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testPauseJobUpdateByUser() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.pause(UPDATE_KEY, AUDIT); |
| |
| control.replay(); |
| |
| assertResponse(OK, thrift.pauseJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test(expected = IllegalArgumentException.class) |
| public void testPauseMessageTooLong() throws Exception { |
| expectGetRemoteUser(); |
| |
| control.replay(); |
| |
| assertResponse( |
| OK, |
| thrift.pauseJobUpdate( |
| UPDATE_KEY.newBuilder(), |
| Strings.repeat("*", AuditData.MAX_MESSAGE_LENGTH + 1))); |
| } |
| |
| @Test |
| public void testPauseJobUpdateFailsInController() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.pause(UPDATE_KEY, AUDIT); |
| expectLastCall().andThrow(new UpdateStateException("failed")); |
| |
| control.replay(); |
| |
| assertResponse( |
| INVALID_REQUEST, |
| thrift.pauseJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testResumeJobUpdate() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.resume(UPDATE_KEY, AUDIT); |
| |
| control.replay(); |
| |
| assertResponse(OK, thrift.resumeJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testResumeJobUpdateFailsInController() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.resume(UPDATE_KEY, AUDIT); |
| expectLastCall().andThrow(new UpdateStateException("failed")); |
| |
| control.replay(); |
| |
| assertResponse( |
| INVALID_REQUEST, |
| thrift.resumeJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testAbortJobUpdateByCoordinator() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.abort(UPDATE_KEY, AUDIT); |
| |
| control.replay(); |
| |
| assertResponse(OK, thrift.abortJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testAbortJobUpdateByUser() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.abort(UPDATE_KEY, AUDIT); |
| |
| control.replay(); |
| |
| assertResponse(OK, thrift.abortJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testAbortJobUpdateFailsInController() throws Exception { |
| expectGetRemoteUser(); |
| jobUpdateController.abort(UPDATE_KEY, AUDIT); |
| expectLastCall().andThrow(new UpdateStateException("failed")); |
| |
| control.replay(); |
| |
| assertResponse( |
| INVALID_REQUEST, |
| thrift.abortJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE)); |
| } |
| |
| @Test |
| public void testPulseJobUpdatePulsedAsCoordinator() throws Exception { |
| expect(jobUpdateController.pulse(UPDATE_KEY)).andReturn(JobUpdatePulseStatus.OK); |
| |
| control.replay(); |
| |
| assertEquals( |
| okResponse(Result.pulseJobUpdateResult( |
| new PulseJobUpdateResult(JobUpdatePulseStatus.OK))), |
| thrift.pulseJobUpdate(UPDATE_KEY.newBuilder())); |
| } |
| |
| @Test |
| public void testPulseJobUpdatePulsedAsUser() throws Exception { |
| expect(jobUpdateController.pulse(UPDATE_KEY)).andReturn(JobUpdatePulseStatus.OK); |
| |
| control.replay(); |
| |
| assertEquals( |
| okResponse(Result.pulseJobUpdateResult(new PulseJobUpdateResult(JobUpdatePulseStatus.OK))), |
| thrift.pulseJobUpdate(UPDATE_KEY.newBuilder())); |
| } |
| |
| @Test |
| public void testPulseJobUpdateFails() throws Exception { |
| expect(jobUpdateController.pulse(UPDATE_KEY)).andThrow(new UpdateStateException("failure")); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.pulseJobUpdate(UPDATE_KEY.newBuilder())); |
| } |
| |
| @Test |
| public void testGetJobUpdateSummaries() throws Exception { |
| Response updateSummary = Responses.empty() |
| .setResponseCode(OK) |
| .setDetails(ImmutableList.of(new ResponseDetail("summary"))); |
| |
| expect(readOnlyScheduler.getJobUpdateSummaries( |
| anyObject(JobUpdateQuery.class))).andReturn(updateSummary); |
| control.replay(); |
| |
| assertEquals(updateSummary, thrift.getJobUpdateSummaries(new JobUpdateQuery())); |
| } |
| |
| private IExpectationSetters<String> expectGetRemoteUser() { |
| return expect(auditMessages.getRemoteUserName()).andReturn(USER); |
| } |
| |
| private static TaskConfig populatedTask() { |
| return defaultTask(true).setConstraints(Sets.newHashSet( |
| new Constraint("host", TaskConstraint.limit(new LimitConstraint(1))))); |
| } |
| |
| private static Constraint dedicatedConstraint(int value) { |
| return new Constraint(DEDICATED_ATTRIBUTE, TaskConstraint.limit(new LimitConstraint(value))); |
| } |
| |
| private static Constraint dedicatedConstraint(Set<String> values) { |
| return new Constraint(DEDICATED_ATTRIBUTE, |
| TaskConstraint.value(new ValueConstraint(false, values))); |
| } |
| |
| private static JobUpdateRequest buildServiceJobUpdateRequest() { |
| return buildServiceJobUpdateRequest(defaultTask(true)); |
| } |
| |
| private static JobUpdateRequest buildServiceJobUpdateRequest(TaskConfig config) { |
| return buildJobUpdateRequest(config.setIsService(true)); |
| } |
| |
| private static JobUpdateRequest buildJobUpdateRequest(TaskConfig config) { |
| return new JobUpdateRequest() |
| .setInstanceCount(6) |
| .setSettings(buildJobUpdateSettings()) |
| .setTaskConfig(config); |
| } |
| |
| private static JobUpdateSettings buildJobUpdateSettings() { |
| return new JobUpdateSettings() |
| .setUpdateGroupSize(10) |
| .setUpdateStrategy( |
| JobUpdateStrategy.queueStrategy(new QueueJobUpdateStrategy().setGroupSize(10))) |
| .setMaxFailedInstances(2) |
| .setMaxPerInstanceFailures(1) |
| .setMinWaitInInstanceRunningMs(15000) |
| .setRollbackOnFailure(true); |
| } |
| |
| private static Integer rangesToInstanceCount(Set<IRange> ranges) { |
| int instanceCount = 0; |
| for (IRange range : ranges) { |
| instanceCount += range.getLast() - range.getFirst() + 1; |
| } |
| |
| return instanceCount; |
| } |
| |
| private static JobUpdateRequest buildJobUpdateRequest(IJobUpdate update) { |
| return new JobUpdateRequest() |
| .setInstanceCount(rangesToInstanceCount( |
| update.getInstructions().getDesiredState().getInstances())) |
| .setSettings(update.getInstructions().getSettings().newBuilder()) |
| .setTaskConfig(update.getInstructions().getDesiredState().getTask().newBuilder()) |
| .setMetadata(IMetadata.toBuildersSet(update.getSummary().getMetadata())); |
| } |
| |
| private static IJobUpdate buildJobUpdate( |
| int instanceCount, |
| ITaskConfig newConfig, |
| ImmutableMap<ITaskConfig, ImmutableSet<Range>> oldConfigMap) { |
| |
| return buildJobUpdate(instanceCount, newConfig, oldConfigMap, buildJobUpdateSettings()); |
| } |
| |
| private static IJobUpdate buildJobUpdate( |
| int instanceCount, |
| ITaskConfig newConfig, |
| ImmutableMap<ITaskConfig, ImmutableSet<Range>> oldConfigMap, |
| JobUpdateSettings jobUpdateSettings) { |
| |
| ImmutableSet.Builder<InstanceTaskConfig> builder = ImmutableSet.builder(); |
| for (Map.Entry<ITaskConfig, ImmutableSet<Range>> entry : oldConfigMap.entrySet()) { |
| builder.add(new InstanceTaskConfig(entry.getKey().newBuilder(), entry.getValue())); |
| } |
| |
| return IJobUpdate.build(new JobUpdate() |
| .setSummary(new JobUpdateSummary() |
| .setKey(UPDATE_KEY.newBuilder()) |
| .setUser(IDENTITY.getUser()) |
| .setMetadata(METADATA)) |
| .setInstructions(new JobUpdateInstructions() |
| .setSettings(jobUpdateSettings) |
| .setDesiredState(new InstanceTaskConfig() |
| .setTask(newConfig.newBuilder()) |
| .setInstances(ImmutableSet.of(new Range(0, instanceCount - 1)))) |
| .setInitialState(builder.build()))); |
| } |
| |
| private IExpectationSetters<?> expectCronJob() { |
| return expect(storageUtil.jobStore.fetchJob(JOB_KEY)) |
| .andReturn(Optional.of(IJobConfiguration.build(CRON_JOB))); |
| } |
| |
| private IExpectationSetters<?> expectNoCronJob() { |
| return expect(storageUtil.jobStore.fetchJob(JOB_KEY)) |
| .andReturn(Optional.empty()); |
| } |
| |
| private IExpectationSetters<?> expectInstanceQuotaCheck( |
| SanitizedConfiguration sanitized, |
| QuotaCheckResult result) { |
| |
| return expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), result); |
| } |
| |
| private IExpectationSetters<?> expectInstanceQuotaCheck( |
| ITaskConfig config, |
| QuotaCheckResult result) { |
| |
| return expect(quotaManager.checkInstanceAddition( |
| config, |
| 1, |
| storageUtil.mutableStoreProvider)).andReturn(result); |
| } |
| |
| private IExpectationSetters<?> expectCronQuotaCheck( |
| IJobConfiguration config, |
| QuotaCheckResult result) { |
| |
| return expect(quotaManager.checkCronUpdate(config, storageUtil.mutableStoreProvider)) |
| .andReturn(result); |
| } |
| } |
