Flink FLIP-27

// org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator
private final Set<TopicPartition> assignedPartitions;
public KafkaSourceEnumState snapshotState(long checkpointId) throws Exception {
    return new KafkaSourceEnumState(assignedPartitions);
}

// org.apache.flink.streaming.api.operators.SourceOperator#snapshotState
public void snapshotState(StateSnapshotContext context) throws Exception {
    long checkpointId = context.getCheckpointId();
    LOG.debug("Taking a snapshot for checkpoint {}", checkpointId);
    readerState.update(sourceReader.snapshotState(checkpointId));
}

// org.apache.flink.connector.base.source.reader.SourceReaderBase
public List<SplitT> snapshotState(long checkpointId) {
    List<SplitT> splits = new ArrayList<>();
    splitStates.forEach((id, context) -> splits.add(toSplitType(id, context.state)));
    return splits;
}

public void addSplits(List<SplitT> splits) {
    LOG.info("Adding split(s) to reader: {}", splits);
    // Initialize the state for each split.
    splits.forEach(
            s ->
                    splitStates.put(
                            s.splitId(), new SplitContext<>(s.splitId(), initializedState(s))));
    // Hand over the splits to the split fetcher to start fetch.
    splitFetcherManager.addSplits(splits);
}

currentSplitContext = splitStates.get(nextSplitId);
currentSplitOutput = currentSplitContext.getOrCreateSplitOutput(output);
recordEmitter.emitRecord(record, currentSplitOutput, currentSplitContext.state);

public void emitRecord(
        ConsumerRecord<byte[], byte[]> consumerRecord,
        SourceOutput<T> output,
        KafkaPartitionSplitState splitState)
        throws Exception {
    try {
        sourceOutputWrapper.setSourceOutput(output);
        sourceOutputWrapper.setTimestamp(consumerRecord.timestamp());
        deserializationSchema.deserialize(consumerRecord, sourceOutputWrapper);
        splitState.setCurrentOffset(consumerRecord.offset() + 1);
    } catch (Exception e) {
        throw new IOException("Failed to deserialize consumer record due to", e);
    }
}

// org.apache.flink.connector.kafka.source.KafkaSource
public SplitEnumerator<KafkaPartitionSplit, KafkaSourceEnumState> restoreEnumerator(
        SplitEnumeratorContext<KafkaPartitionSplit> enumContext,
        KafkaSourceEnumState checkpoint)
        throws IOException {
    return new KafkaSourceEnumerator(
            subscriber,
            startingOffsetsInitializer,
            stoppingOffsetsInitializer,
            props,
            enumContext,
            boundedness,
            checkpoint.assignedPartitions());
}

// org.apache.flink.streaming.api.operators.SourceOperator
public void initializeState(StateInitializationContext context) throws Exception {
    super.initializeState(context);
    final ListState<byte[]> rawState =
            context.getOperatorStateStore().getListState(SPLITS_STATE_DESC);
    readerState = new SimpleVersionedListState<>(rawState, splitSerializer);
}

public void open() throws Exception {
    ...
    // restore the state if necessary.
    final List<SplitT> splits = CollectionUtil.iterableToList(readerState.get());
    if (!splits.isEmpty()) {
        sourceReader.addSplits(splits);
    }
    ...
}