| /** |
| * 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.updater; |
| |
| import java.util.Optional; |
| import java.util.function.Consumer; |
| |
| import com.google.common.collect.ImmutableSet; |
| |
| import org.apache.aurora.common.testing.easymock.EasyMockTest; |
| import org.apache.aurora.gen.InstanceKey; |
| import org.apache.aurora.gen.InstanceTaskConfig; |
| 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.PercentageSlaPolicy; |
| import org.apache.aurora.gen.Range; |
| import org.apache.aurora.gen.ScheduleStatus; |
| import org.apache.aurora.gen.SlaPolicy; |
| import org.apache.aurora.scheduler.base.JobKeys; |
| import org.apache.aurora.scheduler.base.Query; |
| import org.apache.aurora.scheduler.base.TaskTestUtil; |
| import org.apache.aurora.scheduler.state.StateChangeResult; |
| import org.apache.aurora.scheduler.state.StateManager; |
| import org.apache.aurora.scheduler.storage.Storage; |
| import org.apache.aurora.scheduler.storage.entities.IInstanceKey; |
| import org.apache.aurora.scheduler.storage.entities.IJobKey; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdateInstructions; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey; |
| 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.updater.InstanceActionHandler.KillTask; |
| import org.easymock.Capture; |
| import org.junit.Before; |
| import org.junit.Test; |
| |
| import static org.easymock.EasyMock.anyObject; |
| import static org.easymock.EasyMock.capture; |
| import static org.easymock.EasyMock.eq; |
| import static org.easymock.EasyMock.expect; |
| import static org.easymock.EasyMock.expectLastCall; |
| import static org.junit.Assert.assertTrue; |
| import static org.junit.Assert.fail; |
| |
| public class KillTaskTest extends EasyMockTest { |
| private static final String TASK_ID = "task_id"; |
| private static final IJobKey JOB = JobKeys.from("role", "env", "job"); |
| private static final IInstanceKey INSTANCE = |
| IInstanceKey.build(new InstanceKey(JOB.newBuilder(), 0)); |
| private static final IJobUpdateKey UPDATE_ID = |
| IJobUpdateKey.build(new JobUpdateKey(JOB.newBuilder(), "update_id")); |
| private static final SlaPolicy TEST_SLA_POLICY_OLD = SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(95) |
| .setDurationSecs(3600)); |
| private static final SlaPolicy TEST_SLA_POLICY_NEW = SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(0) |
| .setDurationSecs(0)); |
| private static final ITaskConfig CONFIG_NO_SLA = ITaskConfig.build( |
| TaskTestUtil.makeConfig(JOB, true, Optional.empty()).newBuilder()); |
| private static final ITaskConfig OLD_CONFIG_WITH_SLA = |
| TaskTestUtil.makeConfig(JOB, true, Optional.of(TEST_SLA_POLICY_OLD)); |
| private static final ITaskConfig NEW_CONFIG_WITH_SLA = |
| TaskTestUtil.makeConfig(JOB, true, Optional.of(TEST_SLA_POLICY_NEW)); |
| private static final IJobUpdateInstructions INSTRUCTIONS = IJobUpdateInstructions.build( |
| new JobUpdateInstructions() |
| .setSettings( |
| new JobUpdateSettings() |
| .setMinWaitInInstanceRunningMs(1000))); |
| private static final IJobUpdateInstructions INSTRUCTIONS_SLA_AWARE = IJobUpdateInstructions.build( |
| new JobUpdateInstructions() |
| .setSettings( |
| new JobUpdateSettings() |
| .setMinWaitInInstanceRunningMs(1000) |
| .setSlaAware(true)) |
| .setInitialState(ImmutableSet.of(new InstanceTaskConfig( |
| OLD_CONFIG_WITH_SLA.newBuilder(), |
| ImmutableSet.of(new Range(0, 0))))) |
| .setDesiredState(new InstanceTaskConfig( |
| NEW_CONFIG_WITH_SLA.newBuilder(), |
| ImmutableSet.of(new Range(0, 0))))); |
| private static final IJobUpdateInstructions INSTRUCTIONS_SLA_AWARE_NO_POLICY |
| = IJobUpdateInstructions.build( |
| new JobUpdateInstructions() |
| .setSettings( |
| new JobUpdateSettings() |
| .setMinWaitInInstanceRunningMs(1000) |
| .setSlaAware(true)) |
| .setDesiredState(new InstanceTaskConfig( |
| CONFIG_NO_SLA.newBuilder(), |
| ImmutableSet.of(new Range(0, 0))))); |
| |
| private StorageTestUtil storageUtil; |
| private StateManager stateManager; |
| private InstanceActionHandler handler; |
| private UpdateAgentReserver updateAgentReserver; |
| private SlaKillController slaKillController; |
| |
| @Before |
| public void setUp() { |
| storageUtil = new StorageTestUtil(this); |
| storageUtil.expectOperations(); |
| stateManager = createMock(StateManager.class); |
| updateAgentReserver = createMock(UpdateAgentReserver.class); |
| handler = new KillTask(false); |
| slaKillController = createMock(SlaKillController.class); |
| } |
| |
| @Test |
| public void testInstanceKill() throws Exception { |
| storageUtil.expectTaskFetch( |
| Query.instanceScoped(INSTANCE).active(), |
| TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey())); |
| |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.KILLING, |
| Optional.of("Killed for job update " + UPDATE_ID.getId()))) |
| .andReturn(StateChangeResult.SUCCESS); |
| |
| control.replay(); |
| |
| handler.getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_BACK, |
| UPDATE_ID, |
| slaKillController); |
| } |
| |
| @Test |
| public void testKillForUpdateReservesAgentForInstance() throws Exception { |
| IScheduledTask task = TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey(), 1, "agent01"); |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active(), task); |
| |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.KILLING, |
| Optional.of("Killed for job update " + UPDATE_ID.getId()))) |
| .andReturn(StateChangeResult.SUCCESS); |
| updateAgentReserver.reserve(task.getAssignedTask().getSlaveId(), INSTANCE); |
| expectLastCall(); |
| |
| control.replay(); |
| |
| new KillTask(true).getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_BACK, |
| UPDATE_ID, |
| slaKillController); |
| } |
| |
| @Test |
| public void testInstanceNotFoundDoesNotThrow() throws Exception { |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active()); |
| |
| control.replay(); |
| |
| handler.getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_BACK, |
| UPDATE_ID, |
| slaKillController); |
| } |
| |
| /** |
| * Ensures that if an instance is killed with code {@code slaAware} option in the instructions, |
| * then the kill is sent to be handled by the {@link SlaKillController}. |
| */ |
| @Test |
| public void testInstanceKillSlaAware() throws Exception { |
| IScheduledTask task = TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey()); |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active(), task); |
| |
| slaKillController.slaKill( |
| eq(storageUtil.mutableStoreProvider), |
| eq(INSTANCE), |
| eq(task), |
| eq(UPDATE_ID), |
| eq(INSTRUCTIONS_SLA_AWARE.getDesiredState().getTask().getSlaPolicy()), |
| eq(JobUpdateStatus.ROLLING_FORWARD), |
| anyObject()); |
| expectLastCall(); |
| |
| control.replay(); |
| |
| handler.getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS_SLA_AWARE, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_FORWARD, |
| UPDATE_ID, |
| slaKillController); |
| } |
| |
| /** |
| * Ensures that if an instance killed while {@link JobUpdateStatus#ROLLING_BACK}, it uses the old |
| * configuration. |
| */ |
| @Test |
| public void testInstanceKillSlaAwareRollback() throws Exception { |
| IScheduledTask task = TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey(), 1, "agent01"); |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active(), task); |
| |
| Capture<Consumer<Storage.MutableStoreProvider>> killCommandCapture = createCapture(); |
| slaKillController.slaKill( |
| eq(storageUtil.mutableStoreProvider), |
| eq(INSTANCE), |
| eq(task), |
| eq(UPDATE_ID), |
| eq(INSTRUCTIONS_SLA_AWARE.getInitialState().asList().get(0).getTask().getSlaPolicy()), |
| eq(JobUpdateStatus.ROLLING_BACK), |
| capture(killCommandCapture)); |
| expectLastCall(); |
| |
| control.replay(); |
| |
| handler.getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS_SLA_AWARE, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_BACK, |
| UPDATE_ID, |
| slaKillController); |
| } |
| |
| /** |
| * Ensure the correct behavior of the consumer passed into |
| * {@link SlaKillController#slaKill}. It should behave as |
| * {@link KillTask#killAndMaybeReserve}. |
| */ |
| @Test |
| public void testInstanceKillSlaAwareKillCommand() throws Exception { |
| IScheduledTask task = TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey(), 1, "agent01"); |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active(), task); |
| |
| Capture<Consumer<Storage.MutableStoreProvider>> killCommandCapture = createCapture(); |
| slaKillController.slaKill( |
| eq(storageUtil.mutableStoreProvider), |
| eq(INSTANCE), |
| eq(task), |
| eq(UPDATE_ID), |
| eq(INSTRUCTIONS_SLA_AWARE.getDesiredState().getTask().getSlaPolicy()), |
| eq(JobUpdateStatus.ROLLING_FORWARD), |
| capture(killCommandCapture)); |
| expectLastCall(); |
| |
| // Ensure the correct kill command consumer has been passed in and expected behavior occurs |
| // when executed. |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.KILLING, |
| Optional.of("Killed for job update " + UPDATE_ID.getId()))) |
| .andReturn(StateChangeResult.SUCCESS); |
| updateAgentReserver.reserve(task.getAssignedTask().getSlaveId(), INSTANCE); |
| expectLastCall(); |
| |
| control.replay(); |
| |
| new KillTask(true).getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS_SLA_AWARE, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_FORWARD, |
| UPDATE_ID, |
| slaKillController); |
| |
| assertTrue(killCommandCapture.hasCaptured()); |
| killCommandCapture.getValue().accept(storageUtil.mutableStoreProvider); |
| } |
| |
| /** |
| * Ensures that if an instance is killed with code {@code slaAware} option in the instructions |
| * but no {@link org.apache.aurora.gen.SlaPolicy} with the task, then the kill does not fail but |
| * instead continues as a non-sla-aware kill. |
| */ |
| @Test |
| public void testInstanceKillSlaAwareMissingSlaPolicy() throws Exception { |
| IScheduledTask task = TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey(), 1, "agent01"); |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active(), task); |
| expect(stateManager.changeState( |
| storageUtil.mutableStoreProvider, |
| TASK_ID, |
| Optional.empty(), |
| ScheduleStatus.KILLING, |
| Optional.of("Killed for job update " + UPDATE_ID.getId()))) |
| .andReturn(StateChangeResult.SUCCESS); |
| |
| control.replay(); |
| |
| handler.getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS_SLA_AWARE_NO_POLICY, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLLING_FORWARD, |
| UPDATE_ID, |
| slaKillController); |
| } |
| |
| /** |
| * Ensures that if SLA-aware kill is called while not in an active state we throw an exception. |
| */ |
| @Test |
| public void testInstanceKillSlaAwareBadStatus() { |
| IScheduledTask task = TaskTestUtil.makeTask(TASK_ID, INSTANCE.getJobKey(), 1, "agent01"); |
| storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active(), task); |
| |
| control.replay(); |
| |
| try { |
| handler.getReevaluationDelay( |
| INSTANCE, |
| INSTRUCTIONS_SLA_AWARE_NO_POLICY, |
| storageUtil.mutableStoreProvider, |
| stateManager, |
| updateAgentReserver, |
| JobUpdateStatus.ROLL_FORWARD_PAUSED, |
| UPDATE_ID, |
| slaKillController); |
| fail(); |
| } catch (UpdateStateException e) { |
| // Expected |
| } |
| } |
| } |
