/*
* 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 com.facebook.presto.execution;
import com.facebook.presto.execution.TaskExecutor.TaskHandle;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import io.airlift.units.Duration;
import org.testng.annotations.Test;
import java.util.concurrent.Phaser;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import static com.google.common.collect.Iterables.getOnlyElement;
import static org.testng.Assert.assertEquals;
public class TaskExecutorTest
{
@Test(invocationCount = 100)
public void test()
throws Exception
{
TaskExecutor taskExecutor = new TaskExecutor(4);
taskExecutor.start();
try {
TaskHandle taskHandle = taskExecutor.addTask(new TaskId("test", "test", "test"));
final Phaser beginPhase = new Phaser();
beginPhase.register();
final Phaser verificationComplete = new Phaser();
verificationComplete.register();
// add two jobs
TestingJob driver1 = new TestingJob(beginPhase, verificationComplete, 10);
ListenableFuture<?> future1 = getOnlyElement(taskExecutor.enqueueSplits(taskHandle, true, ImmutableList.of(driver1)));
TestingJob driver2 = new TestingJob(beginPhase, verificationComplete, 10);
ListenableFuture<?> future2 = getOnlyElement(taskExecutor.enqueueSplits(taskHandle, true, ImmutableList.of(driver2)));
assertEquals(driver1.getCompletedPhases(), 0);
assertEquals(driver2.getCompletedPhases(), 0);
// verify worker have arrived but haven't processed yet
beginPhase.arriveAndAwaitAdvance();
assertEquals(driver1.getCompletedPhases(), 0);
assertEquals(driver2.getCompletedPhases(), 0);
verificationComplete.arriveAndAwaitAdvance();
// advance one phase and verify
beginPhase.arriveAndAwaitAdvance();
assertEquals(driver1.getCompletedPhases(), 1);
assertEquals(driver2.getCompletedPhases(), 1);
verificationComplete.arriveAndAwaitAdvance();
// add one more job
TestingJob driver3 = new TestingJob(beginPhase, verificationComplete, 10);
ListenableFuture<?> future3 = getOnlyElement(taskExecutor.enqueueSplits(taskHandle, false, ImmutableList.of(driver3)));
// advance one phase and verify
beginPhase.arriveAndAwaitAdvance();
assertEquals(driver1.getCompletedPhases(), 2);
assertEquals(driver2.getCompletedPhases(), 2);
assertEquals(driver3.getCompletedPhases(), 0);
verificationComplete.arriveAndAwaitAdvance();
// advance to the end of the first two task and verify
beginPhase.arriveAndAwaitAdvance();
for (int i = 0; i < 7; i++) {
verificationComplete.arriveAndAwaitAdvance();
beginPhase.arriveAndAwaitAdvance();
assertEquals(beginPhase.getPhase(), verificationComplete.getPhase() + 1);
}
assertEquals(driver1.getCompletedPhases(), 10);
assertEquals(driver2.getCompletedPhases(), 10);
assertEquals(driver3.getCompletedPhases(), 8);
future1.get(1, TimeUnit.SECONDS);
future2.get(1, TimeUnit.SECONDS);
verificationComplete.arriveAndAwaitAdvance();
// advance two more times and verify
beginPhase.arriveAndAwaitAdvance();
verificationComplete.arriveAndAwaitAdvance();
beginPhase.arriveAndAwaitAdvance();
assertEquals(driver1.getCompletedPhases(), 10);
assertEquals(driver2.getCompletedPhases(), 10);
assertEquals(driver3.getCompletedPhases(), 10);
future3.get(1, TimeUnit.SECONDS);
verificationComplete.arriveAndAwaitAdvance();
assertEquals(driver1.getFirstPhase(), 0);
assertEquals(driver2.getFirstPhase(), 0);
assertEquals(driver3.getFirstPhase(), 2);
assertEquals(driver1.getLastPhase(), 10);
assertEquals(driver2.getLastPhase(), 10);
assertEquals(driver3.getLastPhase(), 12);
}
finally {
taskExecutor.stop();
}
}
private static class TestingJob
implements SplitRunner
{
private final Phaser awaitWorkers;
private final Phaser awaitVerifiers;
private final int requiredPhases;
private final AtomicInteger completedPhases = new AtomicInteger();
private final AtomicInteger firstPhase = new AtomicInteger(-1);
private final AtomicInteger lastPhase = new AtomicInteger(-1);
public TestingJob(Phaser awaitWorkers, Phaser awaitVerifiers, int requiredPhases)
{
this.awaitWorkers = awaitWorkers;
this.awaitVerifiers = awaitVerifiers;
this.requiredPhases = requiredPhases;
awaitWorkers.register();
awaitVerifiers.register();
}
private int getFirstPhase()
{
return firstPhase.get();
}
private int getLastPhase()
{
return lastPhase.get();
}
private int getCompletedPhases()
{
return completedPhases.get();
}
@Override
public ListenableFuture<?> processFor(Duration duration)
throws Exception
{
int phase = awaitWorkers.arriveAndAwaitAdvance();
firstPhase.compareAndSet(-1, phase - 1);
lastPhase.set(phase);
awaitVerifiers.arriveAndAwaitAdvance();
completedPhases.getAndIncrement();
return Futures.immediateFuture(null);
}
@Override
public boolean isFinished()
{
boolean isFinished = completedPhases.get() >= requiredPhases;
if (isFinished) {
awaitVerifiers.arriveAndDeregister();
awaitWorkers.arriveAndDeregister();
}
return isFinished;
}
@Override
public void close()
{
}
}
}