@Override
public void setup(
StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) {
this.container = containingTask;
this.config = config;
this.output = output;
this.runtimeContext =
new StreamingRuntimeContext(
this, container.getEnvironment(), container.getAccumulatorMap());
stateKeySelector1 = config.getStatePartitioner(0, getUserCodeClassloader());
stateKeySelector2 = config.getStatePartitioner(1, getUserCodeClassloader());
try {
TypeSerializer<Object> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());
// if the keySerializer is null we still need to create the state backend
// for the non-partitioned state features it provides, such as the state output streams
String operatorIdentifier =
getClass().getSimpleName()
+ "_"
+ config.getVertexID()
+ "_"
+ runtimeContext.getIndexOfThisSubtask();
stateBackend = container.createStateBackend(operatorIdentifier, keySerializer);
} catch (Exception e) {
throw new RuntimeException("Could not initialize state backend. ", e);
}
}
/**
* Waits fro the task to be running. If this does not happen within the timeout, then a
* TimeoutException is thrown.
*
* @param timeout Timeout for the task to be running.
* @throws Exception
*/
public void waitForTaskRunning(long timeout) throws Exception {
if (taskThread == null) {
throw new IllegalStateException("Task thread was not started.");
} else {
if (taskThread.task instanceof StreamTask) {
StreamTask<?, ?> streamTask = (StreamTask<?, ?>) taskThread.task;
while (!streamTask.isRunning()) {
Thread.sleep(100);
if (!taskThread.isAlive()) {
if (taskThread.getError() != null) {
throw new Exception("Task Thread failed due to an error.", taskThread.getError());
} else {
throw new Exception("Task Thread unexpectedly shut down.");
}
}
}
} else {
throw new IllegalStateException("Not a StreamTask");
}
}
}
private static StreamTask<?, ?> createMockTask() {
Configuration configuration = new Configuration();
configuration.setString(ConfigConstants.STATE_BACKEND, "jobmanager");
StreamTask<?, ?> task = mock(StreamTask.class);
when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>());
when(task.getName()).thenReturn("Test task name");
when(task.getExecutionConfig()).thenReturn(new ExecutionConfig());
final TaskManagerRuntimeInfo mockTaskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
when(mockTaskManagerRuntimeInfo.getConfiguration()).thenReturn(configuration);
final Environment env = mock(Environment.class);
when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 1, 0, 1, 0));
when(env.getUserClassLoader())
.thenReturn(AggregatingAlignedProcessingTimeWindowOperatorTest.class.getClassLoader());
when(env.getMetricGroup()).thenReturn(new UnregisteredTaskMetricsGroup());
when(task.getEnvironment()).thenReturn(env);
return task;
}
@Override
public void run() {
synchronized (lock) {
try {
target.trigger(timestamp);
} catch (Throwable t) {
LOG.error("Caught exception while processing timer.", t);
if (task.timerException == null) {
task.timerException = new TimerException(t);
}
}
}
}
private static StreamTask<?, ?> createMockTaskWithTimer(
final ScheduledExecutorService timerService, final Object lock) {
StreamTask<?, ?> task = mock(StreamTask.class);
when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>());
when(task.getName()).thenReturn("Test task name");
when(task.getExecutionConfig()).thenReturn(new ExecutionConfig());
when(task.getEnvironment())
.thenReturn(
new MockEnvironment("MockTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024));
when(task.getCheckpointLock()).thenReturn(lock);
doAnswer(
new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
final Long timestamp = (Long) invocationOnMock.getArguments()[0];
final Triggerable target = (Triggerable) invocationOnMock.getArguments()[1];
timerService.schedule(
new Callable<Object>() {
@Override
public Object call() throws Exception {
synchronized (lock) {
target.trigger(timestamp);
}
return null;
}
},
timestamp - System.currentTimeMillis(),
TimeUnit.MILLISECONDS);
return null;
}
})
.when(task)
.registerTimer(anyLong(), any(Triggerable.class));
// ugly Java generic hacks to get the generic state backend into the mock
@SuppressWarnings("unchecked")
OngoingStubbing<StateBackend<?>> stubbing =
(OngoingStubbing<StateBackend<?>>) (OngoingStubbing<?>) when(task.getStateBackend());
stubbing.thenReturn(MemoryStateBackend.defaultInstance());
return task;
}
private static StreamTask<?, ?> createMockTask() {
StreamTask<?, ?> task = mock(StreamTask.class);
when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>());
when(task.getName()).thenReturn("Test task name");
when(task.getExecutionConfig()).thenReturn(new ExecutionConfig());
Environment env = mock(Environment.class);
when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 0, 1, 0));
when(env.getUserClassLoader())
.thenReturn(AggregatingAlignedProcessingTimeWindowOperatorTest.class.getClassLoader());
when(task.getEnvironment()).thenReturn(env);
// ugly java generic hacks to get the state backend into the mock
@SuppressWarnings("unchecked")
OngoingStubbing<StateBackend<?>> stubbing =
(OngoingStubbing<StateBackend<?>>) (OngoingStubbing<?>) when(task.getStateBackend());
stubbing.thenReturn(MemoryStateBackend.defaultInstance());
return task;
}
/**
* Register a timer callback. At the specified time the {@link Triggerable} will be invoked. This
* call is guaranteed to not happen concurrently with method calls on the operator.
*
* @param time The absolute time in milliseconds.
* @param target The target to be triggered.
*/
protected ScheduledFuture<?> registerTimer(long time, Triggerable target) {
return container.registerTimer(time, target);
}
public ClassLoader getUserCodeClassloader() {
return container.getUserCodeClassLoader();
}
/**
* Gets the execution config defined on the execution environment of the job to which this
* operator belongs.
*
* @return The job's execution config.
*/
public ExecutionConfig getExecutionConfig() {
return container.getExecutionConfig();
}