Search in sources :

Example 1 with TtlTimeProvider

use of org.apache.flink.runtime.state.ttl.TtlTimeProvider in project flink by apache.

the class StreamTaskStateInitializerImplTest method testWithRestore.

@SuppressWarnings("unchecked")
@Test
public void testWithRestore() throws Exception {
    StateBackend mockingBackend = spy(new StateBackend() {

        @Override
        public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(Environment env, JobID jobID, String operatorIdentifier, TypeSerializer<K> keySerializer, int numberOfKeyGroups, KeyGroupRange keyGroupRange, TaskKvStateRegistry kvStateRegistry, TtlTimeProvider ttlTimeProvider, MetricGroup metricGroup, @Nonnull Collection<KeyedStateHandle> stateHandles, CloseableRegistry cancelStreamRegistry) throws Exception {
            return mock(AbstractKeyedStateBackend.class);
        }

        @Override
        public OperatorStateBackend createOperatorStateBackend(Environment env, String operatorIdentifier, @Nonnull Collection<OperatorStateHandle> stateHandles, CloseableRegistry cancelStreamRegistry) throws Exception {
            return mock(OperatorStateBackend.class);
        }
    });
    OperatorID operatorID = new OperatorID(47L, 11L);
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    Random random = new Random(0x42);
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setManagedOperatorState(new OperatorStreamStateHandle(Collections.singletonMap("a", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 10 }, SPLIT_DISTRIBUTE)), CheckpointTestUtils.createDummyStreamStateHandle(random, null))).setRawOperatorState(new OperatorStreamStateHandle(Collections.singletonMap("_default_", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 20, 30 }, SPLIT_DISTRIBUTE)), CheckpointTestUtils.createDummyStreamStateHandle(random, null))).setManagedKeyedState(CheckpointTestUtils.createDummyKeyGroupStateHandle(random, null)).setRawKeyedState(CheckpointTestUtils.createDummyKeyGroupStateHandle(random, null)).setInputChannelState(singleton(createNewInputChannelStateHandle(10, random))).setResultSubpartitionState(singleton(createNewResultSubpartitionStateHandle(10, random))).build();
    taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(42L, taskStateSnapshot);
    StreamTaskStateInitializer streamTaskStateManager = streamTaskStateManager(mockingBackend, jobManagerTaskRestore, false);
    AbstractStreamOperator<?> streamOperator = mock(AbstractStreamOperator.class);
    when(streamOperator.getOperatorID()).thenReturn(operatorID);
    TypeSerializer<?> typeSerializer = new IntSerializer();
    CloseableRegistry closeableRegistry = new CloseableRegistry();
    StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(streamOperator.getOperatorID(), streamOperator.getClass().getSimpleName(), new TestProcessingTimeService(), streamOperator, typeSerializer, closeableRegistry, new UnregisteredMetricsGroup(), 1.0, false);
    OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
    CheckpointableKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
    InternalTimeServiceManager<?> timeServiceManager = stateContext.internalTimerServiceManager();
    CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
    CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();
    Assert.assertTrue("Expected the context to be restored", stateContext.isRestored());
    Assert.assertEquals(OptionalLong.of(42L), stateContext.getRestoredCheckpointId());
    Assert.assertNotNull(operatorStateBackend);
    Assert.assertNotNull(keyedStateBackend);
    // this is deactivated on purpose so that it does not attempt to consume the raw keyed
    // state.
    Assert.assertNull(timeServiceManager);
    Assert.assertNotNull(keyedStateInputs);
    Assert.assertNotNull(operatorStateInputs);
    int count = 0;
    for (KeyGroupStatePartitionStreamProvider keyedStateInput : keyedStateInputs) {
        ++count;
    }
    Assert.assertEquals(1, count);
    count = 0;
    for (StatePartitionStreamProvider operatorStateInput : operatorStateInputs) {
        ++count;
    }
    Assert.assertEquals(3, count);
    checkCloseablesRegistered(closeableRegistry, operatorStateBackend, keyedStateBackend, keyedStateInputs, operatorStateInputs);
}
Also used : IntSerializer(org.apache.flink.api.common.typeutils.base.IntSerializer) UnregisteredMetricsGroup(org.apache.flink.metrics.groups.UnregisteredMetricsGroup) KeyGroupRange(org.apache.flink.runtime.state.KeyGroupRange) MetricGroup(org.apache.flink.metrics.MetricGroup) JobManagerTaskRestore(org.apache.flink.runtime.checkpoint.JobManagerTaskRestore) TaskKvStateRegistry(org.apache.flink.runtime.query.TaskKvStateRegistry) OperatorID(org.apache.flink.runtime.jobgraph.OperatorID) KeyedStateHandle(org.apache.flink.runtime.state.KeyedStateHandle) CloseableRegistry(org.apache.flink.core.fs.CloseableRegistry) StateBackend(org.apache.flink.runtime.state.StateBackend) OperatorStateBackend(org.apache.flink.runtime.state.OperatorStateBackend) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) AbstractKeyedStateBackend(org.apache.flink.runtime.state.AbstractKeyedStateBackend) CheckpointableKeyedStateBackend(org.apache.flink.runtime.state.CheckpointableKeyedStateBackend) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) KeyGroupStatePartitionStreamProvider(org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider) KeyGroupStatePartitionStreamProvider(org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider) StatePartitionStreamProvider(org.apache.flink.runtime.state.StatePartitionStreamProvider) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) Random(java.util.Random) TtlTimeProvider(org.apache.flink.runtime.state.ttl.TtlTimeProvider) AbstractKeyedStateBackend(org.apache.flink.runtime.state.AbstractKeyedStateBackend) OperatorStreamStateHandle(org.apache.flink.runtime.state.OperatorStreamStateHandle) OperatorStateBackend(org.apache.flink.runtime.state.OperatorStateBackend) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) Environment(org.apache.flink.runtime.execution.Environment) TestProcessingTimeService(org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService) OperatorStateHandle(org.apache.flink.runtime.state.OperatorStateHandle) JobID(org.apache.flink.api.common.JobID) TaskStateManagerImplTest(org.apache.flink.runtime.state.TaskStateManagerImplTest) Test(org.junit.Test)

Aggregations

Random (java.util.Random)1 JobID (org.apache.flink.api.common.JobID)1 IntSerializer (org.apache.flink.api.common.typeutils.base.IntSerializer)1 CloseableRegistry (org.apache.flink.core.fs.CloseableRegistry)1 MetricGroup (org.apache.flink.metrics.MetricGroup)1 UnregisteredMetricsGroup (org.apache.flink.metrics.groups.UnregisteredMetricsGroup)1 JobManagerTaskRestore (org.apache.flink.runtime.checkpoint.JobManagerTaskRestore)1 OperatorSubtaskState (org.apache.flink.runtime.checkpoint.OperatorSubtaskState)1 TaskStateSnapshot (org.apache.flink.runtime.checkpoint.TaskStateSnapshot)1 Environment (org.apache.flink.runtime.execution.Environment)1 OperatorID (org.apache.flink.runtime.jobgraph.OperatorID)1 DummyEnvironment (org.apache.flink.runtime.operators.testutils.DummyEnvironment)1 TaskKvStateRegistry (org.apache.flink.runtime.query.TaskKvStateRegistry)1 AbstractKeyedStateBackend (org.apache.flink.runtime.state.AbstractKeyedStateBackend)1 CheckpointableKeyedStateBackend (org.apache.flink.runtime.state.CheckpointableKeyedStateBackend)1 KeyGroupRange (org.apache.flink.runtime.state.KeyGroupRange)1 KeyGroupStatePartitionStreamProvider (org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider)1 KeyedStateHandle (org.apache.flink.runtime.state.KeyedStateHandle)1 OperatorStateBackend (org.apache.flink.runtime.state.OperatorStateBackend)1 OperatorStateHandle (org.apache.flink.runtime.state.OperatorStateHandle)1