org.apache.flink.streaming.connectors.kinesis.model.StreamShardMetadata Java Examples

/**
 * Utility function to convert {@link StreamShardHandle} into {@link StreamShardMetadata}.
 *
 * @param streamShardHandle the {@link StreamShardHandle} to be converted
 * @return a {@link StreamShardMetadata} object
 */
public static StreamShardMetadata convertToStreamShardMetadata(StreamShardHandle streamShardHandle) {
	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();

	streamShardMetadata.setStreamName(streamShardHandle.getStreamName());
	streamShardMetadata.setShardId(streamShardHandle.getShard().getShardId());
	streamShardMetadata.setParentShardId(streamShardHandle.getShard().getParentShardId());
	streamShardMetadata.setAdjacentParentShardId(streamShardHandle.getShard().getAdjacentParentShardId());

	if (streamShardHandle.getShard().getHashKeyRange() != null) {
		streamShardMetadata.setStartingHashKey(streamShardHandle.getShard().getHashKeyRange().getStartingHashKey());
		streamShardMetadata.setEndingHashKey(streamShardHandle.getShard().getHashKeyRange().getEndingHashKey());
	}

	if (streamShardHandle.getShard().getSequenceNumberRange() != null) {
		streamShardMetadata.setStartingSequenceNumber(streamShardHandle.getShard().getSequenceNumberRange().getStartingSequenceNumber());
		streamShardMetadata.setEndingSequenceNumber(streamShardHandle.getShard().getSequenceNumberRange().getEndingSequenceNumber());
	}

	return streamShardMetadata;
}
 

/**
 * Utility function to convert {@link StreamShardMetadata} into {@link StreamShardHandle}.
 *
 * @param streamShardMetadata the {@link StreamShardMetadata} to be converted
 * @return a {@link StreamShardHandle} object
 */
public static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata) {
	Shard shard = new Shard();
	shard.withShardId(streamShardMetadata.getShardId());
	shard.withParentShardId(streamShardMetadata.getParentShardId());
	shard.withAdjacentParentShardId(streamShardMetadata.getAdjacentParentShardId());

	HashKeyRange hashKeyRange = new HashKeyRange();
	hashKeyRange.withStartingHashKey(streamShardMetadata.getStartingHashKey());
	hashKeyRange.withEndingHashKey(streamShardMetadata.getEndingHashKey());
	shard.withHashKeyRange(hashKeyRange);

	SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.withStartingSequenceNumber(streamShardMetadata.getStartingSequenceNumber());
	sequenceNumberRange.withEndingSequenceNumber(streamShardMetadata.getEndingSequenceNumber());
	shard.withSequenceNumberRange(sequenceNumberRange);

	return new StreamShardHandle(streamShardMetadata.getStreamName(), shard);
}
 

/**
 * Utility function to convert {@link StreamShardHandle} into {@link StreamShardMetadata}.
 *
 * @param streamShardHandle the {@link StreamShardHandle} to be converted
 * @return a {@link StreamShardMetadata} object
 */
public static StreamShardMetadata convertToStreamShardMetadata(StreamShardHandle streamShardHandle) {
	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();

	streamShardMetadata.setStreamName(streamShardHandle.getStreamName());
	streamShardMetadata.setShardId(streamShardHandle.getShard().getShardId());
	streamShardMetadata.setParentShardId(streamShardHandle.getShard().getParentShardId());
	streamShardMetadata.setAdjacentParentShardId(streamShardHandle.getShard().getAdjacentParentShardId());

	if (streamShardHandle.getShard().getHashKeyRange() != null) {
		streamShardMetadata.setStartingHashKey(streamShardHandle.getShard().getHashKeyRange().getStartingHashKey());
		streamShardMetadata.setEndingHashKey(streamShardHandle.getShard().getHashKeyRange().getEndingHashKey());
	}

	if (streamShardHandle.getShard().getSequenceNumberRange() != null) {
		streamShardMetadata.setStartingSequenceNumber(streamShardHandle.getShard().getSequenceNumberRange().getStartingSequenceNumber());
		streamShardMetadata.setEndingSequenceNumber(streamShardHandle.getShard().getSequenceNumberRange().getEndingSequenceNumber());
	}

	return streamShardMetadata;
}
 

/**
 * Utility function to convert {@link StreamShardMetadata} into {@link StreamShardHandle}.
 *
 * @param streamShardMetadata the {@link StreamShardMetadata} to be converted
 * @return a {@link StreamShardHandle} object
 */
public static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata) {
	Shard shard = new Shard();
	shard.withShardId(streamShardMetadata.getShardId());
	shard.withParentShardId(streamShardMetadata.getParentShardId());
	shard.withAdjacentParentShardId(streamShardMetadata.getAdjacentParentShardId());

	HashKeyRange hashKeyRange = new HashKeyRange();
	hashKeyRange.withStartingHashKey(streamShardMetadata.getStartingHashKey());
	hashKeyRange.withEndingHashKey(streamShardMetadata.getEndingHashKey());
	shard.withHashKeyRange(hashKeyRange);

	SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.withStartingSequenceNumber(streamShardMetadata.getStartingSequenceNumber());
	sequenceNumberRange.withEndingSequenceNumber(streamShardMetadata.getEndingSequenceNumber());
	shard.withSequenceNumberRange(sequenceNumberRange);

	return new StreamShardHandle(streamShardMetadata.getStreamName(), shard);
}
 

/**
 * Utility function to convert {@link StreamShardMetadata} into {@link StreamShardHandle}.
 *
 * @param streamShardMetadata the {@link StreamShardMetadata} to be converted
 * @return a {@link StreamShardHandle} object
 */
public static StreamShardHandle convertToStreamShardHandle(StreamShardMetadata streamShardMetadata) {
	Shard shard = new Shard();
	shard.withShardId(streamShardMetadata.getShardId());
	shard.withParentShardId(streamShardMetadata.getParentShardId());
	shard.withAdjacentParentShardId(streamShardMetadata.getAdjacentParentShardId());

	HashKeyRange hashKeyRange = new HashKeyRange();
	hashKeyRange.withStartingHashKey(streamShardMetadata.getStartingHashKey());
	hashKeyRange.withEndingHashKey(streamShardMetadata.getEndingHashKey());
	shard.withHashKeyRange(hashKeyRange);

	SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.withStartingSequenceNumber(streamShardMetadata.getStartingSequenceNumber());
	sequenceNumberRange.withEndingSequenceNumber(streamShardMetadata.getEndingSequenceNumber());
	shard.withSequenceNumberRange(sequenceNumberRange);

	return new StreamShardHandle(streamShardMetadata.getStreamName(), shard);
}
 

/**
 * Utility function to convert {@link StreamShardHandle} into {@link StreamShardMetadata}.
 *
 * @param streamShardHandle the {@link StreamShardHandle} to be converted
 * @return a {@link StreamShardMetadata} object
 */
public static StreamShardMetadata convertToStreamShardMetadata(StreamShardHandle streamShardHandle) {
	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();

	streamShardMetadata.setStreamName(streamShardHandle.getStreamName());
	streamShardMetadata.setShardId(streamShardHandle.getShard().getShardId());
	streamShardMetadata.setParentShardId(streamShardHandle.getShard().getParentShardId());
	streamShardMetadata.setAdjacentParentShardId(streamShardHandle.getShard().getAdjacentParentShardId());

	if (streamShardHandle.getShard().getHashKeyRange() != null) {
		streamShardMetadata.setStartingHashKey(streamShardHandle.getShard().getHashKeyRange().getStartingHashKey());
		streamShardMetadata.setEndingHashKey(streamShardHandle.getShard().getHashKeyRange().getEndingHashKey());
	}

	if (streamShardHandle.getShard().getSequenceNumberRange() != null) {
		streamShardMetadata.setStartingSequenceNumber(streamShardHandle.getShard().getSequenceNumberRange().getStartingSequenceNumber());
		streamShardMetadata.setEndingSequenceNumber(streamShardHandle.getShard().getSequenceNumberRange().getEndingSequenceNumber());
	}

	return streamShardMetadata;
}
 

@Test
public void testLegacyKinesisStreamShardToStreamShardMetadataConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();
	streamShardMetadata.setStreamName(streamName);
	streamShardMetadata.setShardId(shardId);
	streamShardMetadata.setParentShardId(parentShardId);
	streamShardMetadata.setAdjacentParentShardId(adjacentParentShardId);
	streamShardMetadata.setStartingHashKey(startingHashKey);
	streamShardMetadata.setEndingHashKey(endingHashKey);
	streamShardMetadata.setStartingSequenceNumber(startingSequenceNumber);
	streamShardMetadata.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	KinesisStreamShard kinesisStreamShard = new KinesisStreamShard(streamName, shard);

	assertEquals(streamShardMetadata, KinesisStreamShard.convertToStreamShardMetadata(kinesisStreamShard));
}
 

/**
 * Creates a snapshot of the current last processed sequence numbers of each subscribed shard.
 *
 * @return state snapshot
 */
public HashMap<StreamShardMetadata, SequenceNumber> snapshotState() {
	// this method assumes that the checkpoint lock is held
	assert Thread.holdsLock(checkpointLock);

	HashMap<StreamShardMetadata, SequenceNumber> stateSnapshot = new HashMap<>();
	for (KinesisStreamShardState shardWithState : subscribedShardsState) {
		stateSnapshot.put(shardWithState.getStreamShardMetadata(), shardWithState.getLastProcessedSequenceNum());
	}
	return stateSnapshot;
}
 

private static Map<StreamShardMetadata, SequenceNumber> removeEquivalenceWrappers(
		Map<StreamShardMetadata.EquivalenceWrapper, SequenceNumber> equivalenceWrappedMap) {

	Map<StreamShardMetadata, SequenceNumber> unwrapped = new HashMap<>();
	for (Map.Entry<StreamShardMetadata.EquivalenceWrapper, SequenceNumber> wrapped : equivalenceWrappedMap.entrySet()) {
		unwrapped.put(wrapped.getKey().getShardMetadata(), wrapped.getValue());
	}

	return unwrapped;
}
 

public TestFetcher(
		List<String> streams,
		SourceFunction.SourceContext<T> sourceContext,
		RuntimeContext runtimeContext,
		Properties configProps,
		KinesisDeserializationSchema<T> deserializationSchema,
		HashMap<StreamShardMetadata, SequenceNumber> testStateSnapshot,
		List<StreamShardHandle> testInitialDiscoveryShards) {

	super(streams, sourceContext, runtimeContext, configProps, deserializationSchema, DEFAULT_SHARD_ASSIGNER, null, null);

	this.testStateSnapshot = testStateSnapshot;
	this.testInitialDiscoveryShards = testInitialDiscoveryShards;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	TypeInformation<Tuple2<StreamShardMetadata, SequenceNumber>> shardsStateTypeInfo = new TupleTypeInfo<>(
		TypeInformation.of(StreamShardMetadata.class),
		TypeInformation.of(SequenceNumber.class));

	sequenceNumsStateForCheckpoint = context.getOperatorStateStore().getUnionListState(
		new ListStateDescriptor<>(sequenceNumsStateStoreName, shardsStateTypeInfo));

	if (context.isRestored()) {
		if (sequenceNumsToRestore == null) {
			sequenceNumsToRestore = new HashMap<>();
			for (Tuple2<StreamShardMetadata, SequenceNumber> kinesisSequenceNumber : sequenceNumsStateForCheckpoint.get()) {
				sequenceNumsToRestore.put(
					// we wrap the restored metadata inside an equivalence wrapper that checks only stream name and shard id,
					// so that if a shard had been closed (due to a Kinesis reshard operation, for example) since
					// the savepoint and has a different metadata than what we last stored,
					// we will still be able to match it in sequenceNumsToRestore. Please see FLINK-8484 for details.
					new StreamShardMetadata.EquivalenceWrapper(kinesisSequenceNumber.f0),
					kinesisSequenceNumber.f1);
			}

			LOG.info("Setting restore state in the FlinkKinesisConsumer. Using the following offsets: {}",
				sequenceNumsToRestore);
		}
	} else {
		LOG.info("No restore state for FlinkKinesisConsumer.");
	}
}
 

@Test
public void testStreamShardMetadataAndHandleConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata kinesisStreamShard = new StreamShardMetadata();
	kinesisStreamShard.setStreamName(streamName);
	kinesisStreamShard.setShardId(shardId);
	kinesisStreamShard.setParentShardId(parentShardId);
	kinesisStreamShard.setAdjacentParentShardId(adjacentParentShardId);
	kinesisStreamShard.setStartingHashKey(startingHashKey);
	kinesisStreamShard.setEndingHashKey(endingHashKey);
	kinesisStreamShard.setStartingSequenceNumber(startingSequenceNumber);
	kinesisStreamShard.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	StreamShardHandle streamShardHandle = new StreamShardHandle(streamName, shard);

	assertEquals(kinesisStreamShard, KinesisDataFetcher.convertToStreamShardMetadata(streamShardHandle));
	assertEquals(streamShardHandle, KinesisDataFetcher.convertToStreamShardHandle(kinesisStreamShard));
}
 

@Test
public void testStreamShardMetadataAndHandleConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata kinesisStreamShard = new StreamShardMetadata();
	kinesisStreamShard.setStreamName(streamName);
	kinesisStreamShard.setShardId(shardId);
	kinesisStreamShard.setParentShardId(parentShardId);
	kinesisStreamShard.setAdjacentParentShardId(adjacentParentShardId);
	kinesisStreamShard.setStartingHashKey(startingHashKey);
	kinesisStreamShard.setEndingHashKey(endingHashKey);
	kinesisStreamShard.setStartingSequenceNumber(startingSequenceNumber);
	kinesisStreamShard.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	StreamShardHandle streamShardHandle = new StreamShardHandle(streamName, shard);

	assertEquals(kinesisStreamShard, KinesisDataFetcher.convertToStreamShardMetadata(streamShardHandle));
	assertEquals(streamShardHandle, KinesisDataFetcher.convertToStreamShardHandle(kinesisStreamShard));
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	TypeInformation<Tuple2<StreamShardMetadata, SequenceNumber>> shardsStateTypeInfo = new TupleTypeInfo<>(
		TypeInformation.of(StreamShardMetadata.class),
		TypeInformation.of(SequenceNumber.class));

	sequenceNumsStateForCheckpoint = context.getOperatorStateStore().getUnionListState(
		new ListStateDescriptor<>(sequenceNumsStateStoreName, shardsStateTypeInfo));

	if (context.isRestored()) {
		if (sequenceNumsToRestore == null) {
			sequenceNumsToRestore = new HashMap<>();
			for (Tuple2<StreamShardMetadata, SequenceNumber> kinesisSequenceNumber : sequenceNumsStateForCheckpoint.get()) {
				sequenceNumsToRestore.put(
					// we wrap the restored metadata inside an equivalence wrapper that checks only stream name and shard id,
					// so that if a shard had been closed (due to a Kinesis reshard operation, for example) since
					// the savepoint and has a different metadata than what we last stored,
					// we will still be able to match it in sequenceNumsToRestore. Please see FLINK-8484 for details.
					new StreamShardMetadata.EquivalenceWrapper(kinesisSequenceNumber.f0),
					kinesisSequenceNumber.f1);
			}

			LOG.info("Setting restore state in the FlinkKinesisConsumer. Using the following offsets: {}",
				sequenceNumsToRestore);
		}
	} else {
		LOG.info("No restore state for FlinkKinesisConsumer.");
	}
}
 

/**
 * Creates a snapshot of the current last processed sequence numbers of each subscribed shard.
 *
 * @return state snapshot
 */
public HashMap<StreamShardMetadata, SequenceNumber> snapshotState() {
	// this method assumes that the checkpoint lock is held
	assert Thread.holdsLock(checkpointLock);

	HashMap<StreamShardMetadata, SequenceNumber> stateSnapshot = new HashMap<>();
	for (KinesisStreamShardState shardWithState : subscribedShardsState) {
		stateSnapshot.put(shardWithState.getStreamShardMetadata(), shardWithState.getLastProcessedSequenceNum());
	}
	return stateSnapshot;
}
 

@Test
public void testLegacyKinesisStreamShardToStreamShardMetadataConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();
	streamShardMetadata.setStreamName(streamName);
	streamShardMetadata.setShardId(shardId);
	streamShardMetadata.setParentShardId(parentShardId);
	streamShardMetadata.setAdjacentParentShardId(adjacentParentShardId);
	streamShardMetadata.setStartingHashKey(startingHashKey);
	streamShardMetadata.setEndingHashKey(endingHashKey);
	streamShardMetadata.setStartingSequenceNumber(startingSequenceNumber);
	streamShardMetadata.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	KinesisStreamShard kinesisStreamShard = new KinesisStreamShard(streamName, shard);

	assertEquals(streamShardMetadata, KinesisStreamShard.convertToStreamShardMetadata(kinesisStreamShard));
}
 

public TestFetcher(
		List<String> streams,
		SourceFunction.SourceContext<T> sourceContext,
		RuntimeContext runtimeContext,
		Properties configProps,
		KinesisDeserializationSchema<T> deserializationSchema,
		HashMap<StreamShardMetadata, SequenceNumber> testStateSnapshot,
		List<StreamShardHandle> testInitialDiscoveryShards) {

	super(streams, sourceContext, runtimeContext, configProps, deserializationSchema, DEFAULT_SHARD_ASSIGNER, null, null);

	this.testStateSnapshot = testStateSnapshot;
	this.testInitialDiscoveryShards = testInitialDiscoveryShards;
}
 

private static Map<StreamShardMetadata, SequenceNumber> removeEquivalenceWrappers(
		Map<StreamShardMetadata.EquivalenceWrapper, SequenceNumber> equivalenceWrappedMap) {

	Map<StreamShardMetadata, SequenceNumber> unwrapped = new HashMap<>();
	for (Map.Entry<StreamShardMetadata.EquivalenceWrapper, SequenceNumber> wrapped : equivalenceWrappedMap.entrySet()) {
		unwrapped.put(wrapped.getKey().getShardMetadata(), wrapped.getValue());
	}

	return unwrapped;
}
 

private static Map<StreamShardMetadata, SequenceNumber> removeEquivalenceWrappers(
		Map<StreamShardMetadata.EquivalenceWrapper, SequenceNumber> equivalenceWrappedMap) {

	Map<StreamShardMetadata, SequenceNumber> unwrapped = new HashMap<>();
	for (Map.Entry<StreamShardMetadata.EquivalenceWrapper, SequenceNumber> wrapped : equivalenceWrappedMap.entrySet()) {
		unwrapped.put(wrapped.getKey().getShardMetadata(), wrapped.getValue());
	}

	return unwrapped;
}
 

/**
 * Creates a snapshot of the current last processed sequence numbers of each subscribed shard.
 *
 * @return state snapshot
 */
public HashMap<StreamShardMetadata, SequenceNumber> snapshotState() {
	// this method assumes that the checkpoint lock is held
	assert Thread.holdsLock(checkpointLock);

	HashMap<StreamShardMetadata, SequenceNumber> stateSnapshot = new HashMap<>();
	for (KinesisStreamShardState shardWithState : subscribedShardsState) {
		stateSnapshot.put(shardWithState.getStreamShardMetadata(), shardWithState.getLastProcessedSequenceNum());
	}
	return stateSnapshot;
}
 

public TestFetcher(
		List<String> streams,
		SourceFunction.SourceContext<T> sourceContext,
		RuntimeContext runtimeContext,
		Properties configProps,
		KinesisDeserializationSchema<T> deserializationSchema,
		HashMap<StreamShardMetadata, SequenceNumber> testStateSnapshot,
		List<StreamShardHandle> testInitialDiscoveryShards) {

	super(streams, sourceContext, runtimeContext, configProps, deserializationSchema, DEFAULT_SHARD_ASSIGNER, null, null);

	this.testStateSnapshot = testStateSnapshot;
	this.testInitialDiscoveryShards = testInitialDiscoveryShards;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	TypeInformation<Tuple2<StreamShardMetadata, SequenceNumber>> shardsStateTypeInfo = new TupleTypeInfo<>(
		TypeInformation.of(StreamShardMetadata.class),
		TypeInformation.of(SequenceNumber.class));

	sequenceNumsStateForCheckpoint = context.getOperatorStateStore().getUnionListState(
		new ListStateDescriptor<>(sequenceNumsStateStoreName, shardsStateTypeInfo));

	if (context.isRestored()) {
		if (sequenceNumsToRestore == null) {
			sequenceNumsToRestore = new HashMap<>();
			for (Tuple2<StreamShardMetadata, SequenceNumber> kinesisSequenceNumber : sequenceNumsStateForCheckpoint.get()) {
				sequenceNumsToRestore.put(
					// we wrap the restored metadata inside an equivalence wrapper that checks only stream name and shard id,
					// so that if a shard had been closed (due to a Kinesis reshard operation, for example) since
					// the savepoint and has a different metadata than what we last stored,
					// we will still be able to match it in sequenceNumsToRestore. Please see FLINK-8484 for details.
					new StreamShardMetadata.EquivalenceWrapper(kinesisSequenceNumber.f0),
					kinesisSequenceNumber.f1);
			}

			LOG.info("Setting restore state in the FlinkKinesisConsumer. Using the following offsets: {}",
				sequenceNumsToRestore);
		}
	} else {
		LOG.info("No restore state for FlinkKinesisConsumer.");
	}
}
 

@Test
public void testStreamShardMetadataAndHandleConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata kinesisStreamShard = new StreamShardMetadata();
	kinesisStreamShard.setStreamName(streamName);
	kinesisStreamShard.setShardId(shardId);
	kinesisStreamShard.setParentShardId(parentShardId);
	kinesisStreamShard.setAdjacentParentShardId(adjacentParentShardId);
	kinesisStreamShard.setStartingHashKey(startingHashKey);
	kinesisStreamShard.setEndingHashKey(endingHashKey);
	kinesisStreamShard.setStartingSequenceNumber(startingSequenceNumber);
	kinesisStreamShard.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	StreamShardHandle streamShardHandle = new StreamShardHandle(streamName, shard);

	assertEquals(kinesisStreamShard, KinesisDataFetcher.convertToStreamShardMetadata(streamShardHandle));
	assertEquals(streamShardHandle, KinesisDataFetcher.convertToStreamShardHandle(kinesisStreamShard));
}
 

@Test
public void testLegacyKinesisStreamShardToStreamShardMetadataConversion() {
	String streamName = "fakeStream1";
	String shardId = "shard-000001";
	String parentShardId = "shard-000002";
	String adjacentParentShardId = "shard-000003";
	String startingHashKey = "key-000001";
	String endingHashKey = "key-000010";
	String startingSequenceNumber = "seq-0000021";
	String endingSequenceNumber = "seq-00000031";

	StreamShardMetadata streamShardMetadata = new StreamShardMetadata();
	streamShardMetadata.setStreamName(streamName);
	streamShardMetadata.setShardId(shardId);
	streamShardMetadata.setParentShardId(parentShardId);
	streamShardMetadata.setAdjacentParentShardId(adjacentParentShardId);
	streamShardMetadata.setStartingHashKey(startingHashKey);
	streamShardMetadata.setEndingHashKey(endingHashKey);
	streamShardMetadata.setStartingSequenceNumber(startingSequenceNumber);
	streamShardMetadata.setEndingSequenceNumber(endingSequenceNumber);

	Shard shard = new Shard()
		.withShardId(shardId)
		.withParentShardId(parentShardId)
		.withAdjacentParentShardId(adjacentParentShardId)
		.withHashKeyRange(new HashKeyRange()
			.withStartingHashKey(startingHashKey)
			.withEndingHashKey(endingHashKey))
		.withSequenceNumberRange(new SequenceNumberRange()
			.withStartingSequenceNumber(startingSequenceNumber)
			.withEndingSequenceNumber(endingSequenceNumber));
	KinesisStreamShard kinesisStreamShard = new KinesisStreamShard(streamName, shard);

	assertEquals(streamShardMetadata, KinesisStreamShard.convertToStreamShardMetadata(kinesisStreamShard));
}
 

/**
 * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link StreamShardMetadata#getShardId()} or
 * {@link StreamShardMetadata#getStreamName()} does not result in the shard not being able to be restored.
 * This handles the corner case where the stored shard metadata is open (no ending sequence number), but after the
 * job restore, the shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
 * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection of restored shard metadata.
 * <p></p>
 * Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard before attempting to find back
 * the sequence number to restore.
 */
@Test
public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored() throws Exception {
	// ----------------------------------------------------------------------
	// setup initial state
	// ----------------------------------------------------------------------

	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

	// ----------------------------------------------------------------------
	// mock operator state backend and initial state for initializeState()
	// ----------------------------------------------------------------------

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
		listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
	}

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	// ----------------------------------------------------------------------
	// mock fetcher
	// ----------------------------------------------------------------------

	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
	List<StreamShardHandle> shards = new ArrayList<>();

	// create a fake stream shard handle based on the first entry in the restored state
	final StreamShardHandle originalStreamShardHandle = fakeRestoredState.keySet().iterator().next();
	final StreamShardHandle closedStreamShardHandle = new StreamShardHandle(originalStreamShardHandle.getStreamName(), originalStreamShardHandle.getShard());
	// close the shard handle by setting an ending sequence number
	final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.setEndingSequenceNumber("1293844");
	closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);

	shards.add(closedStreamShardHandle);

	when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

	// assume the given config is correct
	PowerMockito.mockStatic(KinesisConfigUtil.class);
	PowerMockito.doNothing().when(KinesisConfigUtil.class);

	// ----------------------------------------------------------------------
	// start to test fetcher's initial state seeding
	// ----------------------------------------------------------------------

	TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer(
		"fakeStream", new Properties(), 10, 2);
	consumer.initializeState(initializationContext);
	consumer.open(new Configuration());
	consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

	Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
		new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(closedStreamShardHandle),
			closedStreamShardHandle, fakeRestoredState.get(closedStreamShardHandle)));
}
 

@SuppressWarnings("unchecked")
private void writeSnapshot(String path, HashMap<StreamShardMetadata, SequenceNumber> state) throws Exception {
	final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(state.size());
	for (StreamShardMetadata shardMetadata : state.keySet()) {
		Shard shard = new Shard();
		shard.setShardId(shardMetadata.getShardId());

		SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
		sequenceNumberRange.withStartingSequenceNumber("1");
		shard.setSequenceNumberRange(sequenceNumberRange);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), shard));
	}

	final TestFetcher<String> fetcher = new TestFetcher<>(
		Collections.singletonList(TEST_STREAM_NAME),
		new TestSourceContext<>(),
		new TestRuntimeContext(true, 1, 0),
		TestUtils.getStandardProperties(),
		new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
		state,
		initialDiscoveryShards);

	final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(
		fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));

	StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumer);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();
	testHarness.open();

	final AtomicReference<Throwable> error = new AtomicReference<>();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				consumer.run(new TestSourceContext<>());
			} catch (Throwable t) {
				t.printStackTrace();
				error.set(t);
			}
		}
	};
	runner.start();

	fetcher.waitUntilRun();

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(snapshot, path);

	consumerOperator.close();
	runner.join();
}
 

@Test
public void testStreamShardMetadataSerializedUsingPojoSerializer() {
	TypeInformation<StreamShardMetadata> typeInformation = TypeInformation.of(StreamShardMetadata.class);
	assertTrue(typeInformation.createSerializer(new ExecutionConfig()) instanceof PojoSerializer);
}
 

/**
 * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link StreamShardMetadata#getShardId()} or
 * {@link StreamShardMetadata#getStreamName()} does not result in the shard not being able to be restored.
 * This handles the corner case where the stored shard metadata is open (no ending sequence number), but after the
 * job restore, the shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
 * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection of restored shard metadata.
 * <p></p>
 * Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard before attempting to find back
 * the sequence number to restore.
 */
@Test
public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored() throws Exception {
	// ----------------------------------------------------------------------
	// setup initial state
	// ----------------------------------------------------------------------

	HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

	// ----------------------------------------------------------------------
	// mock operator state backend and initial state for initializeState()
	// ----------------------------------------------------------------------

	TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
	for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
		listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
	}

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);
	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(true);

	// ----------------------------------------------------------------------
	// mock fetcher
	// ----------------------------------------------------------------------

	KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
	List<StreamShardHandle> shards = new ArrayList<>();

	// create a fake stream shard handle based on the first entry in the restored state
	final StreamShardHandle originalStreamShardHandle = fakeRestoredState.keySet().iterator().next();
	final StreamShardHandle closedStreamShardHandle = new StreamShardHandle(originalStreamShardHandle.getStreamName(), originalStreamShardHandle.getShard());
	// close the shard handle by setting an ending sequence number
	final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
	sequenceNumberRange.setEndingSequenceNumber("1293844");
	closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);

	shards.add(closedStreamShardHandle);

	when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

	// assume the given config is correct
	PowerMockito.mockStatic(KinesisConfigUtil.class);
	PowerMockito.doNothing().when(KinesisConfigUtil.class);

	// ----------------------------------------------------------------------
	// start to test fetcher's initial state seeding
	// ----------------------------------------------------------------------

	TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer(
		"fakeStream", new Properties(), 10, 2);
	consumer.initializeState(initializationContext);
	consumer.open(new Configuration());
	consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

	Mockito.verify(mockedFetcher).registerNewSubscribedShardState(
		new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(closedStreamShardHandle),
			closedStreamShardHandle, fakeRestoredState.get(closedStreamShardHandle)));
}
 

@Test
public void testStreamShardMetadataSerializedUsingPojoSerializer() {
	TypeInformation<StreamShardMetadata> typeInformation = TypeInformation.of(StreamShardMetadata.class);
	assertTrue(typeInformation.createSerializer(new ExecutionConfig()) instanceof PojoSerializer);
}
 

@Override
public void run(SourceContext<T> sourceContext) throws Exception {

	// all subtasks will run a fetcher, regardless of whether or not the subtask will initially have
	// shards to subscribe to; fetchers will continuously poll for changes in the shard list, so all subtasks
	// can potentially have new shards to subscribe to later on
	KinesisDataFetcher<T> fetcher = createFetcher(streams, sourceContext, getRuntimeContext(), configProps, deserializer);

	// initial discovery
	List<StreamShardHandle> allShards = fetcher.discoverNewShardsToSubscribe();

	for (StreamShardHandle shard : allShards) {
		StreamShardMetadata.EquivalenceWrapper kinesisStreamShard =
			new StreamShardMetadata.EquivalenceWrapper(KinesisDataFetcher.convertToStreamShardMetadata(shard));

		if (sequenceNumsToRestore != null) {

			if (sequenceNumsToRestore.containsKey(kinesisStreamShard)) {
				// if the shard was already seen and is contained in the state,
				// just use the sequence number stored in the state
				fetcher.registerNewSubscribedShardState(
					new KinesisStreamShardState(kinesisStreamShard.getShardMetadata(), shard, sequenceNumsToRestore.get(kinesisStreamShard)));

				if (LOG.isInfoEnabled()) {
					LOG.info("Subtask {} is seeding the fetcher with restored shard {}," +
							" starting state set to the restored sequence number {}",
						getRuntimeContext().getIndexOfThisSubtask(), shard.toString(), sequenceNumsToRestore.get(kinesisStreamShard));
				}
			} else {
				// the shard wasn't discovered in the previous run, therefore should be consumed from the beginning
				fetcher.registerNewSubscribedShardState(
					new KinesisStreamShardState(kinesisStreamShard.getShardMetadata(), shard, SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get()));

				if (LOG.isInfoEnabled()) {
					LOG.info("Subtask {} is seeding the fetcher with new discovered shard {}," +
							" starting state set to the SENTINEL_EARLIEST_SEQUENCE_NUM",
						getRuntimeContext().getIndexOfThisSubtask(), shard.toString());
				}
			}
		} else {
			// we're starting fresh; use the configured start position as initial state
			SentinelSequenceNumber startingSeqNum =
				InitialPosition.valueOf(configProps.getProperty(
					ConsumerConfigConstants.STREAM_INITIAL_POSITION,
					ConsumerConfigConstants.DEFAULT_STREAM_INITIAL_POSITION)).toSentinelSequenceNumber();

			fetcher.registerNewSubscribedShardState(
				new KinesisStreamShardState(kinesisStreamShard.getShardMetadata(), shard, startingSeqNum.get()));

			if (LOG.isInfoEnabled()) {
				LOG.info("Subtask {} will be seeded with initial shard {}, starting state set as sequence number {}",
					getRuntimeContext().getIndexOfThisSubtask(), shard.toString(), startingSeqNum.get());
			}
		}
	}

	// check that we are running before starting the fetcher
	if (!running) {
		return;
	}

	// expose the fetcher from this point, so that state
	// snapshots can be taken from the fetcher's state holders
	this.fetcher = fetcher;

	// start the fetcher loop. The fetcher will stop running only when cancel() or
	// close() is called, or an error is thrown by threads created by the fetcher
	fetcher.runFetcher();

	// check that the fetcher has terminated before fully closing
	fetcher.awaitTermination();
	sourceContext.close();
}