@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);
}
}
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() {
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;
}
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;
}
