Examples with TaskInfo org.apache.flink.api.common.TaskInfo used on opensource projects

Example 36 with TaskInfo

use of org.apache.flink.api.common.TaskInfo in project flink by apache.

the class StreamTaskStateInitializerImpl method streamOperatorStateContext.

// -----------------------------------------------------------------------------------------------------------------
@Override
public StreamOperatorStateContext streamOperatorStateContext(@Nonnull OperatorID operatorID, @Nonnull String operatorClassName, @Nonnull ProcessingTimeService processingTimeService, @Nonnull KeyContext keyContext, @Nullable TypeSerializer<?> keySerializer, @Nonnull CloseableRegistry streamTaskCloseableRegistry, @Nonnull MetricGroup metricGroup, double managedMemoryFraction, boolean isUsingCustomRawKeyedState) throws Exception {
    TaskInfo taskInfo = environment.getTaskInfo();
    OperatorSubtaskDescriptionText operatorSubtaskDescription = new OperatorSubtaskDescriptionText(operatorID, operatorClassName, taskInfo.getIndexOfThisSubtask(), taskInfo.getNumberOfParallelSubtasks());
    final String operatorIdentifierText = operatorSubtaskDescription.toString();
    final PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates = taskStateManager.prioritizedOperatorState(operatorID);
    CheckpointableKeyedStateBackend<?> keyedStatedBackend = null;
    OperatorStateBackend operatorStateBackend = null;
    CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
    CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
    InternalTimeServiceManager<?> timeServiceManager;
    try {
        // -------------- Keyed State Backend --------------
        keyedStatedBackend = keyedStatedBackend(keySerializer, operatorIdentifierText, prioritizedOperatorSubtaskStates, streamTaskCloseableRegistry, metricGroup, managedMemoryFraction);
        // -------------- Operator State Backend --------------
        operatorStateBackend = operatorStateBackend(operatorIdentifierText, prioritizedOperatorSubtaskStates, streamTaskCloseableRegistry);
        // -------------- Raw State Streams --------------
        rawKeyedStateInputs = rawKeyedStateInputs(prioritizedOperatorSubtaskStates.getPrioritizedRawKeyedState().iterator());
        streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);
        rawOperatorStateInputs = rawOperatorStateInputs(prioritizedOperatorSubtaskStates.getPrioritizedRawOperatorState().iterator());
        streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);
        // -------------- Internal Timer Service Manager --------------
        if (keyedStatedBackend != null) {
            // if the operator indicates that it is using custom raw keyed state,
            // then whatever was written in the raw keyed state snapshot was NOT written
            // by the internal timer services (because there is only ever one user of raw keyed
            // state);
            // in this case, timers should not attempt to restore timers from the raw keyed
            // state.
            final Iterable<KeyGroupStatePartitionStreamProvider> restoredRawKeyedStateTimers = (prioritizedOperatorSubtaskStates.isRestored() && !isUsingCustomRawKeyedState) ? rawKeyedStateInputs : Collections.emptyList();
            timeServiceManager = timeServiceManagerProvider.create(keyedStatedBackend, environment.getUserCodeClassLoader().asClassLoader(), keyContext, processingTimeService, restoredRawKeyedStateTimers);
        } else {
            timeServiceManager = null;
        }
        return new StreamOperatorStateContextImpl(prioritizedOperatorSubtaskStates.getRestoredCheckpointId(), operatorStateBackend, keyedStatedBackend, timeServiceManager, rawOperatorStateInputs, rawKeyedStateInputs);
    } catch (Exception ex) {
        // cleanup if something went wrong before results got published.
        if (keyedStatedBackend != null) {
            if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
                IOUtils.closeQuietly(keyedStatedBackend);
            }
            // release resource (e.g native resource)
            keyedStatedBackend.dispose();
        }
        if (operatorStateBackend != null) {
            if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
                IOUtils.closeQuietly(operatorStateBackend);
            }
            operatorStateBackend.dispose();
        }
        if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
            IOUtils.closeQuietly(rawKeyedStateInputs);
        }
        if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
            IOUtils.closeQuietly(rawOperatorStateInputs);
        }
        throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
    }
}

Example 37 with TaskInfo

use of org.apache.flink.api.common.TaskInfo in project flink by apache.

the class StreamingRuntimeContextTest method testReducingStateInstantiation.

@Test
public void testReducingStateInstantiation() throws Exception {
    final ExecutionConfig config = new ExecutionConfig();
    config.registerKryoType(Path.class);
    final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
    StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);
    @SuppressWarnings("unchecked") ReduceFunction<TaskInfo> reducer = (ReduceFunction<TaskInfo>) mock(ReduceFunction.class);
    ReducingStateDescriptor<TaskInfo> descr = new ReducingStateDescriptor<>("name", reducer, TaskInfo.class);
    context.getReducingState(descr);
    StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
    TypeSerializer<?> serializer = descrIntercepted.getSerializer();
    // check that the Path class is really registered, i.e., the execution config was applied
    assertTrue(serializer instanceof KryoSerializer);
    assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}

Example 38 with TaskInfo

use of org.apache.flink.api.common.TaskInfo in project flink by apache.

the class StreamingRuntimeContextTest method testMapStateInstantiation.

@Test
public void testMapStateInstantiation() throws Exception {
    final ExecutionConfig config = new ExecutionConfig();
    config.registerKryoType(Path.class);
    final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
    StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);
    MapStateDescriptor<String, TaskInfo> descr = new MapStateDescriptor<>("name", String.class, TaskInfo.class);
    context.getMapState(descr);
    MapStateDescriptor<?, ?> descrIntercepted = (MapStateDescriptor<?, ?>) descriptorCapture.get();
    TypeSerializer<?> valueSerializer = descrIntercepted.getValueSerializer();
    // check that the Path class is really registered, i.e., the execution config was applied
    assertTrue(valueSerializer instanceof KryoSerializer);
    assertTrue(((KryoSerializer<?>) valueSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}

Example 39 with TaskInfo

use of org.apache.flink.api.common.TaskInfo in project flink by apache.

the class StreamingRuntimeContextTest method testAggregatingStateInstantiation.

@Test
public void testAggregatingStateInstantiation() throws Exception {
    final ExecutionConfig config = new ExecutionConfig();
    config.registerKryoType(Path.class);
    final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
    StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);
    @SuppressWarnings("unchecked") AggregateFunction<String, TaskInfo, String> aggregate = (AggregateFunction<String, TaskInfo, String>) mock(AggregateFunction.class);
    AggregatingStateDescriptor<String, TaskInfo, String> descr = new AggregatingStateDescriptor<>("name", aggregate, TaskInfo.class);
    context.getAggregatingState(descr);
    AggregatingStateDescriptor<?, ?, ?> descrIntercepted = (AggregatingStateDescriptor<?, ?, ?>) descriptorCapture.get();
    TypeSerializer<?> serializer = descrIntercepted.getSerializer();
    // check that the Path class is really registered, i.e., the execution config was applied
    assertTrue(serializer instanceof KryoSerializer);
    assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}