@Override public void open() throws Exception { super.open(); outTuple = outSerializer.createInstance(); } }
@Override public void close() throws Exception { super.close(); functionsClosed = true; FunctionUtils.closeFunction(userFunction); }
@Override public void dispose() throws Exception { super.dispose(); if (!functionsClosed) { functionsClosed = true; FunctionUtils.closeFunction(userFunction); } }
public StreamingRuntimeContext(AbstractStreamOperator<?> operator, Environment env, Map<String, Accumulator<?, ?>> accumulators) { super(env.getTaskInfo(), env.getUserClassLoader(), operator.getExecutionConfig(), accumulators, env.getDistributedCacheEntries(), operator.getMetricGroup()); this.operator = operator; this.taskEnvironment = env; this.streamConfig = new StreamConfig(env.getTaskConfiguration()); this.operatorUniqueID = operator.getOperatorID().toString(); }
@Override public final void initializeState() throws Exception { final TypeSerializer<?> keySerializer = config.getStateKeySerializer(getUserCodeClassloader()); Preconditions.checkNotNull(getContainingTask()); final CloseableRegistry streamTaskCloseableRegistry = Preconditions.checkNotNull(containingTask.getCancelables()); getOperatorID(), getClass().getSimpleName(), this, this.keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getExecutionConfig()); initializeState(initializationContext); } finally { closeFromRegistry(operatorStateInputs, streamTaskCloseableRegistry); closeFromRegistry(keyedStateInputs, streamTaskCloseableRegistry);
@SuppressWarnings("unchecked") private static AbstractStreamOperator<?> createDescriptorCapturingMockOp( final AtomicReference<Object> ref, final ExecutionConfig config) throws Exception { AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class); KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class); DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config); when(operatorMock.getExecutionConfig()).thenReturn(config); doAnswer(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocationOnMock) throws Throwable { ref.set(invocationOnMock.getArguments()[2]); return null; } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(StateDescriptor.class)); when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore); when(operatorMock.getOperatorID()).thenReturn(new OperatorID()); return operatorMock; }
when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenCallRealMethod(); doCallRealMethod().when(operator).close(); doCallRealMethod().when(operator).dispose(); doReturn(containingTask).when(operator).getContainingTask(); operator.snapshotState( checkpointId, timestamp, verify(futureKeyedStateHandle).cancel(anyBoolean()); operator.close(); operator.dispose();
@Override public void initializeState(StateInitializationContext context) throws Exception { super.initializeState(context); StreamingFunctionUtils.restoreFunctionState(context, userFunction); }
@Override public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) { super.setup(containingTask, config, output); FunctionUtils.setFunctionRuntimeContext(userFunction, getRuntimeContext()); }
@Override public void snapshotState(StateSnapshotContext context) throws Exception { super.snapshotState(context); StreamingFunctionUtils.snapshotFunctionState(context, getOperatorStateBackend(), userFunction); }
factory, keyGroupRange, getContainingTask().getCancelables())) { snapshotState(snapshotContext); getOperatorName() + "."; if (!getContainingTask().isCanceled()) { LOG.info(snapshotFailMessage, snapshotException);
when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenCallRealMethod(); doReturn(containingTask).when(operator).getContainingTask(); doThrow(failingException).when(operator).snapshotState(eq(context)); operator.snapshotState( checkpointId, timestamp,
public void processWatermark2(Watermark mark) throws Exception { input2Watermark = mark.getTimestamp(); long newMin = Math.min(input1Watermark, input2Watermark); if (newMin > combinedWatermark) { combinedWatermark = newMin; processWatermark(new Watermark(combinedWatermark)); } }
@Override public <K, V> KeyedValueState<K, V> getKeyedValueState( ValueStateDescriptor<V> descriptor) throws Exception { descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig()); return operator.getKeyedState( new KeyedValueStateDescriptor<>( descriptor.getName(), (TypeSerializer<K>) operator.getKeySerializer(), descriptor.getSerializer() ) ); }
public StreamingRuntimeContext(AbstractStreamOperator<?> operator, Environment env, Map<String, Accumulator<?, ?>> accumulators) { super(env.getTaskInfo(), env.getUserClassLoader(), operator.getExecutionConfig(), accumulators, env.getDistributedCacheEntries(), operator.getMetricGroup()); this.operator = operator; this.taskEnvironment = env; this.streamConfig = new StreamConfig(env.getTaskConfiguration()); }
private static <T extends Serializable> List<T> runRichSourceFunction(SourceFunction<T> sourceFunction) throws Exception { try (MockEnvironment environment = new MockEnvironmentBuilder() .setTaskName("MockTask") .setMemorySize(3 * 1024 * 1024) .setInputSplitProvider(new MockInputSplitProvider()) .setBufferSize(1024) .build()) { AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class); when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig()); RuntimeContext runtimeContext = new StreamingRuntimeContext( operator, environment, new HashMap<>()); ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext); ((RichFunction) sourceFunction).open(new Configuration()); return runNonRichSourceFunction(sourceFunction); } }
@Override public void initializeState(StateInitializationContext context) throws Exception { super.initializeState(context); checkState(checkpointedState == null, "The reader state has already been initialized."); checkpointedState = context.getOperatorStateStore().getSerializableListState("splits"); int subtaskIdx = getRuntimeContext().getIndexOfThisSubtask(); if (context.isRestored()) { LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIdx); // this may not be null in case we migrate from a previous Flink version. if (restoredReaderState == null) { restoredReaderState = new ArrayList<>(); for (TimestampedFileInputSplit split : checkpointedState.get()) { restoredReaderState.add(split); } if (LOG.isDebugEnabled()) { LOG.debug("{} (taskIdx={}) restored {}.", getClass().getSimpleName(), subtaskIdx, restoredReaderState); } } } else { LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIdx); } }
@Override public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) { super.setup(containingTask, config, output); if (priorityQueue == null) { priorityQueue = new PriorityQueue<>(INITIAL_PRIORITY_QUEUE_CAPACITY, new StreamRecordComparator<IN>()); } if (siddhiRuntimeHandlers == null) { siddhiRuntimeHandlers = new ConcurrentHashMap<>(); } }
@Override public void snapshotState(StateSnapshotContext context) throws Exception { super.snapshotState(context);
@SuppressWarnings("unchecked") private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception { AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class); ExecutionConfig config = new ExecutionConfig(); KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class); DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config); when(operatorMock.getExecutionConfig()).thenReturn(config); doAnswer(new Answer<MapState<Integer, String>>() { @Override public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable { MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class)); when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore); when(operatorMock.getOperatorID()).thenReturn(new OperatorID()); return operatorMock; }