Examples with StateInitializationContext org.apache.flink.runtime.state.StateInitializationContext used on opensource projects

Example 1 with StateInitializationContext

use of org.apache.flink.runtime.state.StateInitializationContext in project flink by apache.

the class FlinkKafkaConsumerBaseTest method checkUseFetcherWhenNoCheckpoint.

/**
	 * Tests that on snapshots, states and offsets to commit to Kafka are correct
	 */
@SuppressWarnings("unchecked")
@Test
public void checkUseFetcherWhenNoCheckpoint() throws Exception {
    FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);
    List<KafkaTopicPartition> partitionList = new ArrayList<>(1);
    partitionList.add(new KafkaTopicPartition("test", 0));
    consumer.setSubscribedPartitions(partitionList);
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    TestingListState<Serializable> listState = new TestingListState<>();
    when(operatorStateStore.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
    // make the context signal that there is no restored state, then validate that
    when(initializationContext.isRestored()).thenReturn(false);
    consumer.initializeState(initializationContext);
    consumer.run(mock(SourceFunction.SourceContext.class));
}

Example 2 with StateInitializationContext

use of org.apache.flink.runtime.state.StateInitializationContext in project flink by apache.

the class FlinkKafkaConsumerBaseTest method checkRestoredCheckpointWhenFetcherNotReady.

/**
	 * Tests that no checkpoints happen when the fetcher is not running.
	 */
@Test
public void checkRestoredCheckpointWhenFetcherNotReady() throws Exception {
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    TestingListState<Serializable> listState = new TestingListState<>();
    listState.add(Tuple2.of(new KafkaTopicPartition("abc", 13), 16768L));
    listState.add(Tuple2.of(new KafkaTopicPartition("def", 7), 987654321L));
    FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);
    when(operatorStateStore.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
    when(initializationContext.isRestored()).thenReturn(true);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(17, 17));
    // ensure that the list was cleared and refilled. while this is an implementation detail, we use it here
    // to figure out that snapshotState() actually did something.
    Assert.assertTrue(listState.isClearCalled());
    Set<Serializable> expected = new HashSet<>();
    for (Serializable serializable : listState.get()) {
        expected.add(serializable);
    }
    int counter = 0;
    for (Serializable serializable : listState.get()) {
        assertTrue(expected.contains(serializable));
        counter++;
    }
    assertEquals(expected.size(), counter);
}

Example 3 with StateInitializationContext

use of org.apache.flink.runtime.state.StateInitializationContext in project flink by apache.

the class FlinkKafkaConsumerBaseTest method testSnapshotStateWithCommitOnCheckpointsEnabled.

@Test
@SuppressWarnings("unchecked")
public void testSnapshotStateWithCommitOnCheckpointsEnabled() throws Exception {
    // --------------------------------------------------------------------
    //   prepare fake states
    // --------------------------------------------------------------------
    final HashMap<KafkaTopicPartition, Long> state1 = new HashMap<>();
    state1.put(new KafkaTopicPartition("abc", 13), 16768L);
    state1.put(new KafkaTopicPartition("def", 7), 987654321L);
    final HashMap<KafkaTopicPartition, Long> state2 = new HashMap<>();
    state2.put(new KafkaTopicPartition("abc", 13), 16770L);
    state2.put(new KafkaTopicPartition("def", 7), 987654329L);
    final HashMap<KafkaTopicPartition, Long> state3 = new HashMap<>();
    state3.put(new KafkaTopicPartition("abc", 13), 16780L);
    state3.put(new KafkaTopicPartition("def", 7), 987654377L);
    // --------------------------------------------------------------------
    final AbstractFetcher<String, ?> fetcher = mock(AbstractFetcher.class);
    when(fetcher.snapshotCurrentState()).thenReturn(state1, state2, state3);
    final LinkedMap pendingOffsetsToCommit = new LinkedMap();
    FlinkKafkaConsumerBase<String> consumer = getConsumer(fetcher, pendingOffsetsToCommit, true);
    StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
    // enable checkpointing
    when(mockRuntimeContext.isCheckpointingEnabled()).thenReturn(true);
    consumer.setRuntimeContext(mockRuntimeContext);
    assertEquals(0, pendingOffsetsToCommit.size());
    OperatorStateStore backend = mock(OperatorStateStore.class);
    TestingListState<Serializable> listState = new TestingListState<>();
    when(backend.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(backend);
    when(initializationContext.isRestored()).thenReturn(false, true, true, true);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    // checkpoint 1
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(138, 138));
    HashMap<KafkaTopicPartition, Long> snapshot1 = new HashMap<>();
    for (Serializable serializable : listState.get()) {
        Tuple2<KafkaTopicPartition, Long> kafkaTopicPartitionLongTuple2 = (Tuple2<KafkaTopicPartition, Long>) serializable;
        snapshot1.put(kafkaTopicPartitionLongTuple2.f0, kafkaTopicPartitionLongTuple2.f1);
    }
    assertEquals(state1, snapshot1);
    assertEquals(1, pendingOffsetsToCommit.size());
    assertEquals(state1, pendingOffsetsToCommit.get(138L));
    // checkpoint 2
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(140, 140));
    HashMap<KafkaTopicPartition, Long> snapshot2 = new HashMap<>();
    for (Serializable serializable : listState.get()) {
        Tuple2<KafkaTopicPartition, Long> kafkaTopicPartitionLongTuple2 = (Tuple2<KafkaTopicPartition, Long>) serializable;
        snapshot2.put(kafkaTopicPartitionLongTuple2.f0, kafkaTopicPartitionLongTuple2.f1);
    }
    assertEquals(state2, snapshot2);
    assertEquals(2, pendingOffsetsToCommit.size());
    assertEquals(state2, pendingOffsetsToCommit.get(140L));
    // ack checkpoint 1
    consumer.notifyCheckpointComplete(138L);
    assertEquals(1, pendingOffsetsToCommit.size());
    assertTrue(pendingOffsetsToCommit.containsKey(140L));
    // checkpoint 3
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(141, 141));
    HashMap<KafkaTopicPartition, Long> snapshot3 = new HashMap<>();
    for (Serializable serializable : listState.get()) {
        Tuple2<KafkaTopicPartition, Long> kafkaTopicPartitionLongTuple2 = (Tuple2<KafkaTopicPartition, Long>) serializable;
        snapshot3.put(kafkaTopicPartitionLongTuple2.f0, kafkaTopicPartitionLongTuple2.f1);
    }
    assertEquals(state3, snapshot3);
    assertEquals(2, pendingOffsetsToCommit.size());
    assertEquals(state3, pendingOffsetsToCommit.get(141L));
    // ack checkpoint 3, subsumes number 2
    consumer.notifyCheckpointComplete(141L);
    assertEquals(0, pendingOffsetsToCommit.size());
    // invalid checkpoint
    consumer.notifyCheckpointComplete(666);
    assertEquals(0, pendingOffsetsToCommit.size());
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    listState = new TestingListState<>();
    when(operatorStateStore.getOperatorState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);
    // create 500 snapshots
    for (int i = 100; i < 600; i++) {
        consumer.snapshotState(new StateSnapshotContextSynchronousImpl(i, i));
        listState.clear();
    }
    assertEquals(FlinkKafkaConsumerBase.MAX_NUM_PENDING_CHECKPOINTS, pendingOffsetsToCommit.size());
    // commit only the second last
    consumer.notifyCheckpointComplete(598);
    assertEquals(1, pendingOffsetsToCommit.size());
    // access invalid checkpoint
    consumer.notifyCheckpointComplete(590);
    // and the last
    consumer.notifyCheckpointComplete(599);
    assertEquals(0, pendingOffsetsToCommit.size());
}

Example 4 with StateInitializationContext

use of org.apache.flink.runtime.state.StateInitializationContext in project flink by apache.

the class FlinkKafkaConsumerBaseTest method checkRestoredNullCheckpointWhenFetcherNotReady.

/**
	 * Tests that no checkpoints happen when the fetcher is not running.
	 */
@Test
public void checkRestoredNullCheckpointWhenFetcherNotReady() throws Exception {
    FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    TestingListState<Serializable> listState = new TestingListState<>();
    when(operatorStateStore.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
    when(initializationContext.isRestored()).thenReturn(false);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(17, 17));
    assertFalse(listState.get().iterator().hasNext());
}

Example 5 with StateInitializationContext

use of org.apache.flink.runtime.state.StateInitializationContext in project flink by apache.

the class FlinkKafkaConsumerBaseTest method testSnapshotStateWithCommitOnCheckpointsDisabled.

@Test
@SuppressWarnings("unchecked")
public void testSnapshotStateWithCommitOnCheckpointsDisabled() throws Exception {
    // --------------------------------------------------------------------
    //   prepare fake states
    // --------------------------------------------------------------------
    final HashMap<KafkaTopicPartition, Long> state1 = new HashMap<>();
    state1.put(new KafkaTopicPartition("abc", 13), 16768L);
    state1.put(new KafkaTopicPartition("def", 7), 987654321L);
    final HashMap<KafkaTopicPartition, Long> state2 = new HashMap<>();
    state2.put(new KafkaTopicPartition("abc", 13), 16770L);
    state2.put(new KafkaTopicPartition("def", 7), 987654329L);
    final HashMap<KafkaTopicPartition, Long> state3 = new HashMap<>();
    state3.put(new KafkaTopicPartition("abc", 13), 16780L);
    state3.put(new KafkaTopicPartition("def", 7), 987654377L);
    // --------------------------------------------------------------------
    final AbstractFetcher<String, ?> fetcher = mock(AbstractFetcher.class);
    when(fetcher.snapshotCurrentState()).thenReturn(state1, state2, state3);
    final LinkedMap pendingOffsetsToCommit = new LinkedMap();
    FlinkKafkaConsumerBase<String> consumer = getConsumer(fetcher, pendingOffsetsToCommit, true);
    StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
    // enable checkpointing
    when(mockRuntimeContext.isCheckpointingEnabled()).thenReturn(true);
    consumer.setRuntimeContext(mockRuntimeContext);
    // disable offset committing
    consumer.setCommitOffsetsOnCheckpoints(false);
    assertEquals(0, pendingOffsetsToCommit.size());
    OperatorStateStore backend = mock(OperatorStateStore.class);
    TestingListState<Serializable> listState = new TestingListState<>();
    when(backend.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(backend);
    when(initializationContext.isRestored()).thenReturn(false, true, true, true);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    // checkpoint 1
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(138, 138));
    HashMap<KafkaTopicPartition, Long> snapshot1 = new HashMap<>();
    for (Serializable serializable : listState.get()) {
        Tuple2<KafkaTopicPartition, Long> kafkaTopicPartitionLongTuple2 = (Tuple2<KafkaTopicPartition, Long>) serializable;
        snapshot1.put(kafkaTopicPartitionLongTuple2.f0, kafkaTopicPartitionLongTuple2.f1);
    }
    assertEquals(state1, snapshot1);
    // pending offsets to commit should not be updated
    assertEquals(0, pendingOffsetsToCommit.size());
    // checkpoint 2
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(140, 140));
    HashMap<KafkaTopicPartition, Long> snapshot2 = new HashMap<>();
    for (Serializable serializable : listState.get()) {
        Tuple2<KafkaTopicPartition, Long> kafkaTopicPartitionLongTuple2 = (Tuple2<KafkaTopicPartition, Long>) serializable;
        snapshot2.put(kafkaTopicPartitionLongTuple2.f0, kafkaTopicPartitionLongTuple2.f1);
    }
    assertEquals(state2, snapshot2);
    // pending offsets to commit should not be updated
    assertEquals(0, pendingOffsetsToCommit.size());
    // ack checkpoint 1
    consumer.notifyCheckpointComplete(138L);
    // not offsets should be committed
    verify(fetcher, never()).commitInternalOffsetsToKafka(anyMap());
    // checkpoint 3
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(141, 141));
    HashMap<KafkaTopicPartition, Long> snapshot3 = new HashMap<>();
    for (Serializable serializable : listState.get()) {
        Tuple2<KafkaTopicPartition, Long> kafkaTopicPartitionLongTuple2 = (Tuple2<KafkaTopicPartition, Long>) serializable;
        snapshot3.put(kafkaTopicPartitionLongTuple2.f0, kafkaTopicPartitionLongTuple2.f1);
    }
    assertEquals(state3, snapshot3);
    // pending offsets to commit should not be updated
    assertEquals(0, pendingOffsetsToCommit.size());
    // ack checkpoint 3, subsumes number 2
    consumer.notifyCheckpointComplete(141L);
    // not offsets should be committed
    verify(fetcher, never()).commitInternalOffsetsToKafka(anyMap());
    // invalid checkpoint
    consumer.notifyCheckpointComplete(666);
    // not offsets should be committed
    verify(fetcher, never()).commitInternalOffsetsToKafka(anyMap());
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    listState = new TestingListState<>();
    when(operatorStateStore.getOperatorState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);
    // create 500 snapshots
    for (int i = 100; i < 600; i++) {
        consumer.snapshotState(new StateSnapshotContextSynchronousImpl(i, i));
        listState.clear();
    }
    // pending offsets to commit should not be updated
    assertEquals(0, pendingOffsetsToCommit.size());
    // commit only the second last
    consumer.notifyCheckpointComplete(598);
    // not offsets should be committed
    verify(fetcher, never()).commitInternalOffsetsToKafka(anyMap());
    // access invalid checkpoint
    consumer.notifyCheckpointComplete(590);
    // not offsets should be committed
    verify(fetcher, never()).commitInternalOffsetsToKafka(anyMap());
    // and the last
    consumer.notifyCheckpointComplete(599);
    // not offsets should be committed
    verify(fetcher, never()).commitInternalOffsetsToKafka(anyMap());
}