/**
* 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.scheduling;
import java.util.Set;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.RateLimiter;
import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.common.util.BackoffStrategy;
import org.apache.aurora.gen.AssignedTask;
import org.apache.aurora.gen.JobKey;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.scheduler.async.DelayExecutor;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.events.PubsubEvent;
import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
import org.apache.aurora.scheduler.scheduling.TaskGroups.TaskGroupBatchWorker;
import org.apache.aurora.scheduler.scheduling.TaskGroups.TaskGroupsSettings;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
import org.apache.aurora.scheduler.testing.FakeScheduledExecutor;
import org.apache.aurora.scheduler.testing.FakeStatsProvider;
import org.junit.Before;
import org.junit.Test;
import static org.apache.aurora.gen.ScheduleStatus.ASSIGNED;
import static org.apache.aurora.gen.ScheduleStatus.INIT;
import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.junit.Assert.assertEquals;
public class TaskGroupsTest extends EasyMockTest {
private static final Amount<Long, Time> FIRST_SCHEDULE_DELAY = Amount.of(1L, Time.MILLISECONDS);
private static final Amount<Long, Time> RESCHEDULE_DELAY = FIRST_SCHEDULE_DELAY;
private static final IJobKey JOB_A = IJobKey.build(new JobKey("role", "test", "jobA"));
private static final String TASK_A_ID = "a";
private static final Set<String> SCHEDULED_RESULT = ImmutableSet.of(TASK_A_ID);
private BackoffStrategy backoffStrategy;
private TaskScheduler taskScheduler;
private RateLimiter rateLimiter;
private FakeScheduledExecutor clock;
private RescheduleCalculator rescheduleCalculator;
private TaskGroups taskGroups;
private TaskGroupBatchWorker batchWorker;
private StorageTestUtil storageUtil;
private FakeStatsProvider statsProvider;
@Before
public void setUp() throws Exception {
storageUtil = new StorageTestUtil(this);
storageUtil.expectOperations();
DelayExecutor executor = createMock(DelayExecutor.class);
clock = FakeScheduledExecutor.fromDelayExecutor(executor);
backoffStrategy = createMock(BackoffStrategy.class);
taskScheduler = createMock(TaskScheduler.class);
rateLimiter = createMock(RateLimiter.class);
rescheduleCalculator = createMock(RescheduleCalculator.class);
batchWorker = createMock(TaskGroupBatchWorker.class);
statsProvider = new FakeStatsProvider();
taskGroups = new TaskGroups(
executor,
new TaskGroupsSettings(FIRST_SCHEDULE_DELAY, backoffStrategy, rateLimiter, 2),
taskScheduler,
rescheduleCalculator,
batchWorker,
statsProvider);
}
@Test
public void testEvaluatedAfterFirstSchedulePenalty() throws Exception {
expect(rateLimiter.acquire()).andReturn(0D);
expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of(TASK_A_ID))))
.andReturn(SCHEDULED_RESULT);
expectBatchExecute(batchWorker, storageUtil.storage, control, SCHEDULED_RESULT)
.anyTimes();
control.replay();
taskGroups.taskChangedState(TaskStateChange.transition(makeTask(TASK_A_ID), INIT));
clock.advance(FIRST_SCHEDULE_DELAY);
assertEquals(0L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
}
@Test
public void testTaskDeletedBeforeEvaluating() throws Exception {
final IScheduledTask task = makeTask(TASK_A_ID);
expect(rateLimiter.acquire()).andReturn(0.5D);
expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of(TASK_A_ID))))
.andAnswer(() -> {
// Test a corner case where a task is deleted while it is being evaluated by the task
// scheduler. If not handled carefully, this could result in the scheduler trying again
// later to satisfy the deleted task.
taskGroups.tasksDeleted(new PubsubEvent.TasksDeleted(ImmutableSet.of(task)));
return ImmutableSet.of();
});
expectBatchExecute(batchWorker, storageUtil.storage, control, ImmutableSet.of())
.anyTimes();
expect(backoffStrategy.calculateBackoffMs(FIRST_SCHEDULE_DELAY.as(Time.MILLISECONDS)))
.andReturn(0L);
control.replay();
taskGroups.taskChangedState(TaskStateChange.transition(makeTask(Tasks.id(task)), INIT));
clock.advance(FIRST_SCHEDULE_DELAY);
assertEquals(1L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
}
@Test
public void testEvaluatedOnStartup() throws Exception {
expect(rateLimiter.acquire()).andReturn(0.000000001D);
expect(rescheduleCalculator.getStartupScheduleDelayMs(makeTask(TASK_A_ID))).andReturn(1L);
expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of(TASK_A_ID))))
.andReturn(ImmutableSet.of(TASK_A_ID));
expectBatchExecute(batchWorker, storageUtil.storage, control, SCHEDULED_RESULT)
.anyTimes();
control.replay();
taskGroups.taskChangedState(TaskStateChange.initialized(makeTask(TASK_A_ID)));
clock.advance(FIRST_SCHEDULE_DELAY);
clock.advance(RESCHEDULE_DELAY);
assertEquals(1L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
}
@Test
public void testMultipleTasksAndResistStarvation() throws Exception {
expect(rateLimiter.acquire()).andReturn(0.001D).times(2);
expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of("a0", "a1"))))
.andReturn(ImmutableSet.of("a0", "a1"));
expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of("b0"))))
.andReturn(ImmutableSet.of("b0"));
expectBatchExecute(
batchWorker,
storageUtil.storage,
control,
ImmutableSet.of("a0", "a1")).anyTimes();
expectBatchExecute(batchWorker, storageUtil.storage, control, ImmutableSet.of("b0"))
.anyTimes();
control.replay();
taskGroups.taskChangedState(TaskStateChange.transition(makeTask(JOB_A, "a0", 0), INIT));
taskGroups.taskChangedState(TaskStateChange.transition(makeTask(JOB_A, "a1", 1), INIT));
taskGroups.taskChangedState(TaskStateChange.transition(makeTask(JOB_A, "a2", 2), INIT));
taskGroups.taskChangedState(TaskStateChange.transition(
makeTask(IJobKey.build(JOB_A.newBuilder().setName("jobB")), "b0", 0), INIT));
clock.advance(FIRST_SCHEDULE_DELAY);
assertEquals(2L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
}
@Test
public void testNonPendingIgnored() {
control.replay();
IScheduledTask task =
IScheduledTask.build(makeTask(TASK_A_ID).newBuilder().setStatus(ASSIGNED));
taskGroups.taskChangedState(TaskStateChange.initialized(task));
}
private static IScheduledTask makeTask(String id) {
return makeTask(JOB_A, id, 0);
}
private static IScheduledTask makeTask(IJobKey jobKey, String id, int instanceId) {
return IScheduledTask.build(new ScheduledTask()
.setStatus(ScheduleStatus.PENDING)
.setAssignedTask(new AssignedTask()
.setInstanceId(instanceId)
.setTaskId(id)
.setTask(new TaskConfig()
.setJob(jobKey.newBuilder()))));
}
}