| /** |
| * 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.sla; |
| |
| import java.io.IOException; |
| import java.net.URLEncoder; |
| import java.util.Map; |
| import java.util.Optional; |
| import java.util.concurrent.CountDownLatch; |
| import java.util.concurrent.ScheduledExecutorService; |
| import java.util.concurrent.Semaphore; |
| import java.util.concurrent.TimeUnit; |
| import java.util.stream.Collectors; |
| import javax.servlet.http.HttpServletRequest; |
| import javax.servlet.http.HttpServletResponse; |
| |
| import com.google.common.base.Joiner; |
| import com.google.common.collect.ImmutableMap; |
| import com.google.common.collect.ImmutableSet; |
| import com.google.gson.Gson; |
| import com.google.gson.JsonObject; |
| import com.google.gson.JsonParser; |
| import com.google.inject.AbstractModule; |
| import com.google.inject.Guice; |
| import com.google.inject.Injector; |
| import com.google.inject.TypeLiteral; |
| |
| import org.apache.aurora.common.stats.StatsProvider; |
| import org.apache.aurora.common.testing.easymock.EasyMockTest; |
| import org.apache.aurora.gen.CoordinatorSlaPolicy; |
| import org.apache.aurora.gen.CountSlaPolicy; |
| import org.apache.aurora.gen.PercentageSlaPolicy; |
| import org.apache.aurora.gen.ScheduleStatus; |
| import org.apache.aurora.gen.ScheduledTask; |
| import org.apache.aurora.gen.ServerInfo; |
| import org.apache.aurora.gen.SlaPolicy; |
| import org.apache.aurora.scheduler.TierManager; |
| import org.apache.aurora.scheduler.base.AsyncUtil; |
| 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.state.StateManager; |
| import org.apache.aurora.scheduler.storage.Storage; |
| import org.apache.aurora.scheduler.storage.entities.IScheduledTask; |
| import org.apache.aurora.scheduler.storage.entities.IServerInfo; |
| import org.apache.aurora.scheduler.storage.entities.ISlaPolicy; |
| import org.apache.aurora.scheduler.storage.testing.StorageTestUtil; |
| import org.apache.aurora.scheduler.testing.FakeStatsProvider; |
| import org.apache.thrift.TException; |
| import org.apache.thrift.TSerializer; |
| import org.apache.thrift.protocol.TSimpleJSONProtocol; |
| import org.asynchttpclient.AsyncHttpClient; |
| import org.asynchttpclient.DefaultAsyncHttpClient; |
| import org.eclipse.jetty.server.Request; |
| import org.eclipse.jetty.server.Server; |
| import org.eclipse.jetty.server.handler.AbstractHandler; |
| import org.junit.Before; |
| import org.junit.Test; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import static org.apache.aurora.gen.ScheduleStatus.PENDING; |
| import static org.apache.aurora.gen.ScheduleStatus.RUNNING; |
| import static org.apache.aurora.scheduler.base.TaskTestUtil.TIER_MANAGER; |
| import static org.easymock.EasyMock.expect; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNotEquals; |
| import static org.junit.Assert.fail; |
| |
| public class SlaManagerTest extends EasyMockTest { |
| |
| private static final Logger LOG = LoggerFactory.getLogger(SlaManagerTest.class); |
| private static final String CLUSTER_NAME = "test_cluster"; |
| private static final String STATS_URL_PREFIX = "fake_url"; |
| private static final String HOST_A = "a"; |
| |
| private static final ISlaPolicy PERCENTAGE_SLA_POLICY = ISlaPolicy.build( |
| SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(66) |
| .setDurationSecs(1800))); |
| |
| private static final ISlaPolicy COUNT_SLA_POLICY = ISlaPolicy.build( |
| SlaPolicy.countSlaPolicy( |
| new CountSlaPolicy() |
| .setCount(2) |
| .setDurationSecs(1800))); |
| |
| private AsyncHttpClient httpClient; |
| private SlaManager slaManager; |
| private StorageTestUtil storageUtil; |
| private StateManager stateManager; |
| private IServerInfo serverInfo; |
| private Server jettyServer; |
| private CountDownLatch coordinatorResponded; |
| |
| @Before |
| public void setUp() { |
| jettyServer = new Server(0); // Start Jetty server with ephemeral port |
| storageUtil = new StorageTestUtil(this); |
| storageUtil.expectOperations(); |
| stateManager = createMock(StateManager.class); |
| httpClient = new DefaultAsyncHttpClient(); |
| coordinatorResponded = new CountDownLatch(1); |
| |
| serverInfo = IServerInfo.build( |
| new ServerInfo() |
| .setClusterName(CLUSTER_NAME) |
| .setStatsUrlPrefix(STATS_URL_PREFIX)); |
| |
| Injector injector = Guice.createInjector( |
| new AbstractModule() { |
| @Override |
| protected void configure() { |
| bind(Storage.class).toInstance(storageUtil.storage); |
| bind(StateManager.class).toInstance(stateManager); |
| bind(StatsProvider.class).toInstance(new FakeStatsProvider()); |
| bind(TierManager.class).toInstance(TIER_MANAGER); |
| bind(AsyncHttpClient.class) |
| .annotatedWith(SlaManager.HttpClient.class) |
| .toInstance(httpClient); |
| |
| bind(new TypeLiteral<Integer>() { }) |
| .annotatedWith(SlaManager.MinRequiredInstances.class) |
| .toInstance(2); |
| |
| bind(new TypeLiteral<Integer>() { }) |
| .annotatedWith(SlaManager.MaxParallelCoordinators.class) |
| .toInstance(10); |
| |
| bind(ScheduledExecutorService.class) |
| .annotatedWith(SlaManager.SlaManagerExecutor.class) |
| .toInstance(AsyncUtil.loggingScheduledExecutor( |
| 10, "SlaManagerTest-%d", LOG)); |
| |
| bind(IServerInfo.class).toInstance(serverInfo); |
| } |
| } |
| ); |
| slaManager = injector.getInstance(SlaManager.class); |
| |
| addTearDown(() -> jettyServer.stop()); |
| } |
| |
| private static IScheduledTask makeTask( |
| String taskId, |
| int instanceId, |
| ScheduleStatus status) { |
| return makeTask(taskId, instanceId, status, 1000, true); |
| } |
| |
| private static IScheduledTask makeTask( |
| String taskId, |
| int instanceId, |
| long runningSince) { |
| return makeTask(taskId, instanceId, RUNNING, runningSince, true); |
| } |
| |
| private static IScheduledTask makeTask( |
| String taskId, |
| int instanceId, |
| ScheduleStatus status, |
| long runningSince, |
| boolean prod) { |
| ScheduledTask builder = TaskTestUtil.addStateTransition( |
| TaskTestUtil.makeTask(taskId, TaskTestUtil.JOB, instanceId, prod), |
| status, |
| runningSince).newBuilder(); |
| builder.getAssignedTask().setSlaveHost(SlaManagerTest.HOST_A); |
| return IScheduledTask.build(builder); |
| } |
| |
| /** |
| * Verifies that SLA check passes and the supplied {@link Storage.MutateWork} gets executed |
| * for a job that has {@link CountSlaPolicy#count} + 1 tasks that have been RUNNING for |
| * the required {@link CountSlaPolicy#durationSecs} with {@link CountSlaPolicy#count} set to 2. |
| */ |
| @Test |
| public void testCheckCountSlaPassesThenActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, RUNNING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| COUNT_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check gets overridden, passes and the supplied {@link Storage.MutateWork} is |
| * executed for a job that has and aggressive {@link CountSlaPolicy#count} > total number of |
| * instances. |
| */ |
| @Test |
| public void testCheckCountSlaOverridesAggressiveCountPassesThenActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, RUNNING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.countSlaPolicy( |
| new CountSlaPolicy() |
| .setCount(3) |
| .setDurationSecs(1800))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check gets overridden, fails and the supplied {@link Storage.MutateWork} is |
| * not executed for a job that has and aggressive {@link CountSlaPolicy#count} > total number of |
| * instances. |
| */ |
| @Test |
| public void testCheckCountSlaOverridesAggressiveCountFailsAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, PENDING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.countSlaPolicy( |
| new CountSlaPolicy() |
| .setCount(3) |
| .setDurationSecs(1800))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check fails and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job that has {@link CountSlaPolicy#count} tasks that have been RUNNING for |
| * the required {@link CountSlaPolicy#durationSecs} and 1 task in PENDING with |
| * {@link CountSlaPolicy#count} set to 2. |
| */ |
| @Test |
| public void testCheckCountSlaFailsDoesNotAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, PENDING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| COUNT_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check is skipped and the supplied {@link Storage.MutateWork} gets |
| * executed for a job that does not meet {@link SlaManager#minRequiredInstances}. |
| */ |
| @Test |
| public void testCheckCountSlaFailsDueToMinRequiredInstancesActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.countSlaPolicy( |
| new CountSlaPolicy() |
| .setCount(0) |
| .setDurationSecs(0))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check is skipped and the supplied {@link Storage.MutateWork} gets |
| * executed for a job is not the correct tier. |
| */ |
| @Test |
| public void testCheckCountSlaFailsDueToTierActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING, 1000, false); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.countSlaPolicy( |
| new CountSlaPolicy() |
| .setCount(0) |
| .setDurationSecs(0))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check fails and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job that has {@link CountSlaPolicy#count} tasks that have been RUNNING for |
| * the required {@link CountSlaPolicy#durationSecs} and 1 task in RUNNING for less than |
| * {@link CountSlaPolicy#durationSecs} with {@link CountSlaPolicy#count} set to 2. |
| */ |
| @Test |
| public void testCheckCountSlaFailsDueToRunningTimeDoesNotAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, System.currentTimeMillis()); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| COUNT_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that when {@code force} is {@code True} supplied {@link Storage.MutateWork} gets |
| * executed for a job that has 1 task in RUNNING for less than {@link CountSlaPolicy#durationSecs} |
| * with {@link CountSlaPolicy#count} set to 2, without checking SLA. |
| */ |
| @Test |
| public void testCheckCountSlaCheckForceAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, System.currentTimeMillis()); |
| |
| // expect that the fetchTask is the work is called after force |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| COUNT_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| true); |
| } |
| |
| /** |
| * Verifies that SLA check passes and the supplied {@link Storage.MutateWork} gets executed |
| * for a job that has more {@link PercentageSlaPolicy#percentage} tasks that have been RUNNING for |
| * the required {@link PercentageSlaPolicy#durationSecs} with |
| * {@link PercentageSlaPolicy#percentage} set to 66. |
| */ |
| @Test |
| public void testCheckPercentageSlaPassesThenActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, RUNNING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| PERCENTAGE_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check overridden and the supplied {@link Storage.MutateWork} gets executed |
| * for a job that requires aggressive {@link PercentageSlaPolicy#percentage} tasks to be RUNNING |
| * for the required {@link PercentageSlaPolicy#durationSecs} with |
| * {@link PercentageSlaPolicy#percentage} set to 67. |
| */ |
| @Test |
| public void testCheckPercentageSlaWithAggressivePercentagePassesThenActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, RUNNING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| // We set the required running percentage to 67 which will allow 0 instances to be in |
| // maintenance is overridden to allow exactly 1 instance to be in maintenance. |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(67) |
| .setDurationSecs(1800))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check is overridden and the supplied {@link Storage.MutateWork} gets |
| * executed for a job that requires aggressive {@link PercentageSlaPolicy#percentage} tasks to |
| * be RUNNING for the required {@link PercentageSlaPolicy#durationSecs} with |
| * {@link PercentageSlaPolicy#percentage} set to 67. |
| */ |
| @Test |
| public void testCheckPercentageSlaWithAggressivePercentageFailsAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, PENDING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| control.replay(); |
| |
| // We set the required running percentage to 67 which will allow 0 instances to be in |
| // maintenance is not overridden to allow 1 instance to be in maintenance. |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(67) |
| .setDurationSecs(1800))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check fails and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job that has {@link PercentageSlaPolicy#percentage} tasks that have been |
| * RUNNING for the required {@link PercentageSlaPolicy#durationSecs} and 1 task in PENDING with |
| * {@link PercentageSlaPolicy#percentage} set to 66. |
| */ |
| @Test |
| public void testCheckPercentageSlaFailsDoesNotAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, PENDING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| PERCENTAGE_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check is skipped and the supplied {@link Storage.MutateWork} gets |
| * executed for a job that does not meet {@link SlaManager#minRequiredInstances}. |
| */ |
| @Test |
| public void testCheckPercentageSlaFailsDueToMinRequiredInstancesActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(0) |
| .setDurationSecs(0))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check is skipped and the supplied {@link Storage.MutateWork} gets |
| * executed for a job is not correct tier. |
| */ |
| @Test |
| public void testCheckPercentageSlaFailsDueToTierActs() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING, 1000, false); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1)); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build( |
| SlaPolicy.percentageSlaPolicy( |
| new PercentageSlaPolicy() |
| .setPercentage(0) |
| .setDurationSecs(0))), |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that SLA check fails and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job that has {@link PercentageSlaPolicy#percentage} tasks that have been |
| * RUNNING for the required {@link PercentageSlaPolicy#durationSecs} and 1 task in RUNNING for |
| * less than {@link PercentageSlaPolicy#durationSecs} with {@link PercentageSlaPolicy#percentage} |
| * set to 66. |
| */ |
| @Test |
| public void testCheckPercentageSlaFailsDueToRunningTimeDoesNotAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| IScheduledTask task3 = makeTask("taskC", 3, System.currentTimeMillis()); |
| |
| // mock calls to fetch all active tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).active())) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| // mock calls to fetch all RUNNING tasks for the job for sla calculation |
| expect(storageUtil.taskStore.fetchTasks(Query.jobScoped(Tasks.getJob(task1)).byStatus(RUNNING))) |
| .andReturn(ImmutableSet.of(task1, task2, task3)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| PERCENTAGE_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| false); |
| } |
| |
| /** |
| * Verifies that when {@code force} is {@code True} the supplied {@link Storage.MutateWork} gets |
| * executed for a job that has 1 task in RUNNING for less than |
| * {@link PercentageSlaPolicy#durationSecs} |
| * with {@link PercentageSlaPolicy#percentage} set to 66, without checking SLA. |
| */ |
| @Test |
| public void testCheckPercentageSlaCheckForceAct() { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| |
| // expect that the fetchTask is the work is called after force |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| PERCENTAGE_SLA_POLICY, |
| storeProvider -> storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()), |
| ImmutableMap.of(), |
| true); |
| } |
| |
| /** |
| * Verifies that SLA check passes and the supplied {@link Storage.MutateWork} gets executed |
| * for a job when {@link CoordinatorSlaPolicy#coordinatorUrl} responds |
| * with {@code {"drain": true}}. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaPassesThenActs() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorResponse(task1, "{\"drain\": true}")); |
| jettyServer.start(); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| while (!coordinatorResponded.await(100, TimeUnit.MILLISECONDS)) { |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| |
| // wait until we are sure that the server has responded |
| coordinatorResponded.await(); |
| workCalled.await(); |
| |
| assertEquals(0, coordinatorResponded.getCount()); |
| // check the work was called |
| assertEquals(0, workCalled.getCount()); |
| } |
| |
| /** |
| * Verifies that SLA check passes and the supplied {@link Storage.MutateWork} gets executed |
| * for a job when {@link CoordinatorSlaPolicy#coordinatorUrl} responds |
| * with {@code {"drain": true}} when supplied with extra params. |
| */ |
| @Test |
| public void testCheckCoordinatorWithExtraParamsSlaPassesThenActs() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| Map<String, String> extraParams = ImmutableMap.of("someKey", "someValue"); |
| jettyServer.setHandler(mockCoordinatorResponse(task1, "{\"drain\": true}", extraParams)); |
| jettyServer.start(); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| while (!coordinatorResponded.await(100, TimeUnit.MILLISECONDS)) { |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| extraParams, |
| false); |
| } |
| |
| // wait until we are sure that the server has responded |
| coordinatorResponded.await(); |
| workCalled.await(); |
| |
| assertEquals(0, coordinatorResponded.getCount()); |
| // check the work was called |
| assertEquals(0, workCalled.getCount()); |
| } |
| |
| /** |
| * Verifies that SLA check passes and the supplied {@link Storage.MutateWork} gets executed |
| * for a job when {@link CoordinatorSlaPolicy#coordinatorUrl} responds |
| * with {@code {"drain": true}}. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaPassesThenActsCustomStatusKey() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorResponse(task1, "{\"custom-key\": true}")); |
| jettyServer.start(); |
| |
| // expect that the fetchTask in the work is called, after sla check passes |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| while (!coordinatorResponded.await(100, TimeUnit.MILLISECONDS)) { |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy("custom-key"), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| |
| workCalled.await(); |
| |
| assertEquals(0, coordinatorResponded.getCount()); |
| // check the work was called |
| assertEquals(0, workCalled.getCount()); |
| } |
| |
| /** |
| * Verifies that SLA check passes and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job when {@link CoordinatorSlaPolicy#coordinatorUrl} responds |
| * with {@code {"drain": false}}. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaFailsDoesNotAct() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorResponse(task1, "{\"drain\": false}")); |
| jettyServer.start(); |
| |
| control.replay(); |
| |
| while (!coordinatorResponded.await(100, TimeUnit.MILLISECONDS)) { |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| |
| try { |
| workCalled.await(5, TimeUnit.SECONDS); |
| } catch (InterruptedException e) { |
| // expected |
| } |
| |
| assertEquals(0, coordinatorResponded.getCount()); |
| // check the work was not called |
| assertEquals(1, workCalled.getCount()); |
| } |
| |
| /** |
| * Verifies that when {@code force} is {@code True} the supplied {@link Storage.MutateWork} gets |
| * executed without checking the SLA. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaCheckForceAct() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorResponse(task1, "{\"drain\": false}")); |
| jettyServer.start(); |
| |
| // expect that the fetchTask in the work is called after force |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| ImmutableMap.of(), |
| true); |
| |
| workCalled.await(); |
| |
| // coordinator is not contacted |
| assertEquals(1, coordinatorResponded.getCount()); |
| // check the work was called |
| assertEquals(0, workCalled.getCount()); |
| } |
| |
| /** |
| * Verifies that SLA check fails and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job when {@link CoordinatorSlaPolicy#coordinatorUrl} responds |
| * with error. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaErrorDoesNotAct() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorError()); |
| jettyServer.start(); |
| |
| control.replay(); |
| |
| while (!coordinatorResponded.await(100, TimeUnit.MILLISECONDS)) { |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| |
| try { |
| workCalled.await(1, TimeUnit.SECONDS); |
| } catch (InterruptedException e) { |
| // expected |
| } |
| |
| assertEquals(0, coordinatorResponded.getCount()); |
| // check the work was not called |
| assertEquals(1, workCalled.getCount()); |
| } |
| |
| /** |
| * Verifies that SLA check fails and the supplied {@link Storage.MutateWork} does not get |
| * executed for a job when {@link CoordinatorSlaPolicy#coordinatorUrl} throws. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaThrowsDoesNotAct() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| CountDownLatch workCalled = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorError()); |
| jettyServer.start(); |
| |
| control.replay(); |
| |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build(SlaPolicy.coordinatorSlaPolicy( |
| new CoordinatorSlaPolicy() |
| .setCoordinatorUrl("http://does.not.exist.com"))), |
| storeProvider -> { |
| // set the marker to indicate that we performed the work |
| workCalled.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| |
| try { |
| workCalled.await(1, TimeUnit.SECONDS); |
| } catch (InterruptedException e) { |
| // expected |
| } |
| |
| // check the work was not called |
| assertEquals(1, workCalled.getCount()); |
| } |
| |
| /** |
| * Start 2 actions simultaneously and make the first one to enter the coordinator lock to block |
| * on a semaphore. Only when the semaphore is released will the next action be able acquire the |
| * lock. At the end both actions should have completed. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaNoLockDoesNotAct() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| |
| CountDownLatch blocked1 = new CountDownLatch(1); |
| CountDownLatch blocked2 = new CountDownLatch(1); |
| CountDownLatch action1 = new CountDownLatch(1); |
| CountDownLatch action2 = new CountDownLatch(1); |
| CountDownLatch finishAction1 = new CountDownLatch(1); |
| CountDownLatch finishAction2 = new CountDownLatch(1); |
| |
| jettyServer.setHandler(mockCoordinatorResponses("{\"drain\": true}")); |
| jettyServer.start(); |
| |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| expect(storageUtil.taskStore.fetchTask(task2.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task2)); |
| |
| control.replay(); |
| |
| // kick off a new thread to attempt starting action1 |
| new Thread(() -> { |
| try { |
| // wait until action1 enters the lock |
| while (!action1.await(100, TimeUnit.MILLISECONDS)) { |
| // start action1 |
| slaManager.checkSlaThenAct( |
| task1, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| LOG.info("Starting action1 for task:{}", slaManager.getTaskKey(task1)); |
| // set the marker to indicate that we started the work |
| action1.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| // we will block here to make sure that the lock is held |
| blocked1.await(); |
| finishAction1.countDown(); |
| LOG.info("Finished action1 for task:{}", slaManager.getTaskKey(task1)); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| } catch (InterruptedException e) { |
| fail(); |
| } |
| }).start(); |
| |
| // wait for action to ge the lock |
| action1.await(); |
| |
| // kick off a new thread to attempt starting action2 |
| new Thread(() -> { |
| try { |
| // wait until action2 enters the lock |
| while (!action2.await(100, TimeUnit.MILLISECONDS)) { |
| // start action2 |
| slaManager.checkSlaThenAct( |
| task2, |
| createCoordinatorSlaPolicy(), |
| storeProvider -> { |
| LOG.info("Starting action2 for task:{}", slaManager.getTaskKey(task2)); |
| // set the marker to indicate that we started the work |
| action2.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task2.getAssignedTask().getTaskId()); |
| // we will block here to make sure that the lock is held |
| blocked2.await(); |
| finishAction2.countDown(); |
| LOG.info("Finished action2 for task:{}", slaManager.getTaskKey(task2)); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| } catch (InterruptedException e) { |
| fail(); |
| } |
| }).start(); |
| |
| // both actions should not have entered critical section |
| assertNotEquals(action1.getCount(), action2.getCount()); |
| // action1 should have entered the lock |
| assertEquals(0, action1.getCount()); |
| // action2 should not have entered the lock |
| assertEquals(1, action2.getCount()); |
| |
| // unblock blocked action1 |
| blocked1.countDown(); |
| // wait for both of the actions to finish |
| finishAction1.await(); |
| |
| // wait for the second action to enter the lock |
| action2.await(); |
| |
| // action2 should have entered the lock |
| assertEquals(0, action2.getCount()); |
| |
| // unblock blocked action2 |
| blocked2.countDown(); |
| |
| finishAction2.await(); |
| } |
| |
| /** |
| * Start 2 actions simultaneously with different coordinator urls, so they can both enter the |
| * appropriate coordinator locks and try to acquire the semaphore. Both the actions must be able |
| * to enter the critical area and start the work simultaneously. |
| * At the end only one action would have fully completed the work, since the semaphore is not |
| * released. |
| */ |
| @Test |
| public void testCheckCoordinatorSlaParallelAct() throws Exception { |
| IScheduledTask task1 = makeTask("taskA", 1, RUNNING); |
| IScheduledTask task2 = makeTask("taskB", 2, RUNNING); |
| |
| Semaphore blocked = new Semaphore(1); |
| CountDownLatch action1 = new CountDownLatch(1); |
| CountDownLatch action2 = new CountDownLatch(1); |
| CountDownLatch finished = new CountDownLatch(2); |
| |
| jettyServer.setHandler(mockCoordinatorResponses("{\"drain\": true}")); |
| jettyServer.start(); |
| |
| expect(storageUtil.taskStore.fetchTask(task1.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task1)); |
| |
| expect(storageUtil.taskStore.fetchTask(task2.getAssignedTask().getTaskId())) |
| .andReturn(Optional.of(task2)); |
| |
| control.replay(); |
| |
| // kick off a new thread to attempt starting action1 |
| new Thread(() -> { |
| try { |
| // wait until action1 enters the lock |
| while (!action1.await(100, TimeUnit.MILLISECONDS)) { |
| // start action1 |
| slaManager.checkSlaThenAct( |
| task1, |
| ISlaPolicy.build(SlaPolicy.coordinatorSlaPolicy( |
| new CoordinatorSlaPolicy() |
| .setCoordinatorUrl( |
| // Note that the url is different although referring to the same server |
| String.format("http://localhost:%d", jettyServer.getURI().getPort())))), |
| storeProvider -> { |
| LOG.info("Starting action for task:{}", slaManager.getTaskKey(task1)); |
| // set the marker to indicate that we started the work |
| action1.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task1.getAssignedTask().getTaskId()); |
| // we will block here to make sure that the lock is held |
| blocked.acquire(); |
| finished.countDown(); |
| LOG.info("Finished action for task:{}", slaManager.getTaskKey(task1)); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| } catch (InterruptedException e) { |
| fail(); |
| } |
| }).start(); |
| |
| // kick off a new thread to attempt starting action2 |
| new Thread(() -> { |
| try { |
| // wait until action2 enters the lock |
| while (!action2.await(100, TimeUnit.MILLISECONDS)) { |
| // start action2 |
| slaManager.checkSlaThenAct( |
| task2, |
| ISlaPolicy.build(SlaPolicy.coordinatorSlaPolicy( |
| new CoordinatorSlaPolicy() |
| .setCoordinatorUrl( |
| // Note that the url is different although referring to the same server |
| String.format("http://0.0.0.0:%d", jettyServer.getURI().getPort())))), |
| storeProvider -> { |
| LOG.info("Starting action for task:{}", slaManager.getTaskKey(task2)); |
| // set the marker to indicate that we started the work |
| action2.countDown(); |
| storeProvider |
| .getUnsafeTaskStore() |
| .fetchTask(task2.getAssignedTask().getTaskId()); |
| // we will block here to make sure that the lock is held |
| blocked.acquire(); |
| finished.countDown(); |
| LOG.info("Finished action for task:{}", slaManager.getTaskKey(task2)); |
| return null; |
| }, |
| ImmutableMap.of(), |
| false); |
| } |
| } catch (InterruptedException e) { |
| fail(); |
| } |
| }).start(); |
| |
| // wait for both of the actions to get the lock |
| action1.await(); |
| action2.await(); |
| |
| assertEquals(0, action1.getCount()); |
| assertEquals(0, action2.getCount()); |
| |
| // only 1 action has fully completed |
| assertEquals(1, finished.getCount()); |
| |
| // unblock the blocked action |
| blocked.release(); |
| |
| finished.await(); |
| |
| // both actions have fully completed |
| assertEquals(0, finished.getCount()); |
| } |
| |
| private ISlaPolicy createCoordinatorSlaPolicy() { |
| return createCoordinatorSlaPolicy("drain"); |
| } |
| |
| private ISlaPolicy createCoordinatorSlaPolicy(String statusKey) { |
| return ISlaPolicy.build(SlaPolicy.coordinatorSlaPolicy( |
| new CoordinatorSlaPolicy() |
| .setCoordinatorUrl(String.format("http://localhost:%d", jettyServer.getURI().getPort())) |
| .setStatusKey(statusKey) |
| )); |
| } |
| |
| private AbstractHandler mockCoordinatorResponse( |
| IScheduledTask task, |
| String pollResponse) { |
| return mockCoordinatorResponse(task, pollResponse, ImmutableMap.of()); |
| } |
| |
| private AbstractHandler mockCoordinatorResponse( |
| IScheduledTask task, |
| String pollResponse, |
| Map<String, String> params) { |
| |
| return new AbstractHandler() { |
| @Override |
| public void handle( |
| String target, |
| Request baseRequest, |
| HttpServletRequest request, |
| HttpServletResponse response) throws IOException { |
| try { |
| String taskKey = slaManager.getTaskKey(task); |
| String query = Joiner |
| .on("=") |
| .join(SlaManager.TASK_PARAM, URLEncoder.encode(taskKey, "UTF-8")); |
| |
| String taskConfig = new TSerializer(new TSimpleJSONProtocol.Factory()) |
| .toString(task.newBuilder()); |
| JsonObject jsonBody = new JsonObject(); |
| jsonBody.add("taskConfig", new JsonParser().parse(taskConfig)); |
| jsonBody.addProperty(SlaManager.TASK_PARAM, taskKey); |
| params.forEach(jsonBody::addProperty); |
| String body = new Gson().toJson(jsonBody); |
| |
| if (request.getQueryString().equals(query) |
| && request.getReader().lines().collect(Collectors.joining()).equals(body)) { |
| createResponse(baseRequest, response, pollResponse); |
| } |
| coordinatorResponded.countDown(); |
| } catch (TException e) { |
| fail(); |
| } |
| } |
| }; |
| } |
| |
| private AbstractHandler mockCoordinatorResponses(String mockResponse) { |
| return new AbstractHandler() { |
| @Override |
| public void handle(String target, Request baseRequest, HttpServletRequest request, |
| HttpServletResponse response) throws IOException { |
| createResponse(baseRequest, response, mockResponse); |
| } |
| }; |
| } |
| |
| private void createResponse( |
| Request baseRequest, |
| HttpServletResponse response, |
| String mockResponse) throws IOException { |
| response.setStatus(HttpServletResponse.SC_OK); |
| response.getWriter().write(mockResponse); |
| response.getWriter().flush(); |
| baseRequest.setHandled(true); |
| } |
| |
| private AbstractHandler mockCoordinatorError() { |
| return new AbstractHandler() { |
| @Override |
| public void handle(String target, Request baseRequest, HttpServletRequest request, |
| HttpServletResponse response) { |
| response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR); |
| baseRequest.setHandled(true); |
| coordinatorResponded.countDown(); |
| } |
| }; |
| } |
| } |
