org.apache.flink.api.common.io.ratelimiting.FlinkConnectorRateLimiter Java Examples

public TestKafkaConsumerThreadRateLimit(Logger log,
		Handover handover, Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue,
		String threadName, long pollTimeout,
		boolean useMetrics, MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		Consumer<byte[], byte[]> mockConsumer,
		FlinkConnectorRateLimiter rateLimiter) {
	super(log, handover, kafkaProperties, unassignedPartitionsQueue,
			threadName,
			pollTimeout, useMetrics, consumerMetricGroup, subtaskMetricGroup,
		rateLimiter);
	this.mockConsumer = mockConsumer;
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<T, TopicPartition>> unassignedPartitionsQueue,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		FlinkConnectorRateLimiter rateLimiter) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;

	if (rateLimiter != null) {
		this.rateLimiter = rateLimiter;
	}
}
 

public Kafka010Fetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<WatermarkStrategy<T>> watermarkStrategy,
		ProcessingTimeService processingTimeProvider,
		long autoWatermarkInterval,
		ClassLoader userCodeClassLoader,
		String taskNameWithSubtasks,
		KafkaDeserializationSchema<T> deserializer,
		Properties kafkaProperties,
		long pollTimeout,
		MetricGroup subtaskMetricGroup,
		MetricGroup consumerMetricGroup,
		boolean useMetrics,
		FlinkConnectorRateLimiter rateLimiter) throws Exception {
	super(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarkStrategy,
			processingTimeProvider,
			autoWatermarkInterval,
			userCodeClassLoader,
			consumerMetricGroup,
			useMetrics);

	this.deserializer = deserializer;
	this.handover = new Handover();

	this.consumerThread = new KafkaConsumerThread(
			LOG,
			handover,
			kafkaProperties,
			unassignedPartitionsQueue,
			"Kafka 0.10 Fetcher for " + taskNameWithSubtasks,
			pollTimeout,
			useMetrics,
			consumerMetricGroup,
			subtaskMetricGroup,
			rateLimiter);
}
 

public TestKafkaConsumerThreadRateLimit(Logger log,
		Handover handover, Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge, String threadName, long pollTimeout,
		boolean useMetrics, MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		KafkaConsumer mockConsumer,
		FlinkConnectorRateLimiter rateLimiter) {
	super(log, handover, kafkaProperties, unassignedPartitionsQueue, consumerCallBridge,
			threadName,
			pollTimeout, useMetrics, consumerMetricGroup, subtaskMetricGroup,
		rateLimiter);
	this.mockConsumer = mockConsumer;
}
 

public Kafka010Fetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<AssignerWithPeriodicWatermarks<T>> watermarksPeriodic,
		SerializedValue<AssignerWithPunctuatedWatermarks<T>> watermarksPunctuated,
		ProcessingTimeService processingTimeProvider,
		long autoWatermarkInterval,
		ClassLoader userCodeClassLoader,
		String taskNameWithSubtasks,
		KafkaDeserializationSchema<T> deserializer,
		Properties kafkaProperties,
		long pollTimeout,
		MetricGroup subtaskMetricGroup,
		MetricGroup consumerMetricGroup,
		boolean useMetrics,
		FlinkConnectorRateLimiter rateLimiter) throws Exception {
	super(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarksPeriodic,
			watermarksPunctuated,
			processingTimeProvider,
			autoWatermarkInterval,
			userCodeClassLoader,
			taskNameWithSubtasks,
			deserializer,
			kafkaProperties,
			pollTimeout,
			subtaskMetricGroup,
			consumerMetricGroup,
			useMetrics, rateLimiter);
}
 

@Override
protected AbstractFetcher<T, ?> createFetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<AssignerWithPeriodicWatermarks<T>> watermarksPeriodic,
		SerializedValue<AssignerWithPunctuatedWatermarks<T>> watermarksPunctuated,
		StreamingRuntimeContext runtimeContext,
		OffsetCommitMode offsetCommitMode,
		MetricGroup consumerMetricGroup,
		boolean useMetrics) throws Exception {

	// make sure that auto commit is disabled when our offset commit mode is ON_CHECKPOINTS;
	// this overwrites whatever setting the user configured in the properties
	adjustAutoCommitConfig(properties, offsetCommitMode);

	FlinkConnectorRateLimiter rateLimiter = super.getRateLimiter();
	// If a rateLimiter is set, then call rateLimiter.open() with the runtime context.
	if (rateLimiter != null) {
		rateLimiter.open(runtimeContext);
	}

	return new Kafka010Fetcher<>(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarksPeriodic,
			watermarksPunctuated,
			runtimeContext.getProcessingTimeService(),
			runtimeContext.getExecutionConfig().getAutoWatermarkInterval(),
			runtimeContext.getUserCodeClassLoader(),
			runtimeContext.getTaskNameWithSubtasks(),
			deserializer,
			properties,
			pollTimeout,
			runtimeContext.getMetricGroup(),
			consumerMetricGroup,
			useMetrics,
			rateLimiter);
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		FlinkConnectorRateLimiter rateLimiter) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);
	this.consumerCallBridge = checkNotNull(consumerCallBridge);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;

	if (rateLimiter != null) {
		this.rateLimiter = rateLimiter;
	}
}
 

PubSubSource(PubSubDeserializationSchema<OUT> deserializationSchema,
			PubSubSubscriberFactory pubSubSubscriberFactory,
			Credentials credentials,
			AcknowledgeOnCheckpointFactory acknowledgeOnCheckpointFactory,
			FlinkConnectorRateLimiter rateLimiter,
			int messagePerSecondRateLimit) {
	this.deserializationSchema = deserializationSchema;
	this.pubSubSubscriberFactory = pubSubSubscriberFactory;
	this.credentials = credentials;
	this.acknowledgeOnCheckpointFactory = acknowledgeOnCheckpointFactory;
	this.rateLimiter = rateLimiter;
	this.messagePerSecondRateLimit = messagePerSecondRateLimit;
}
 

@Override
protected AbstractFetcher<T, ?> createFetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<AssignerWithPeriodicWatermarks<T>> watermarksPeriodic,
		SerializedValue<AssignerWithPunctuatedWatermarks<T>> watermarksPunctuated,
		StreamingRuntimeContext runtimeContext,
		OffsetCommitMode offsetCommitMode,
		MetricGroup consumerMetricGroup,
		boolean useMetrics) throws Exception {

	// make sure that auto commit is disabled when our offset commit mode is ON_CHECKPOINTS;
	// this overwrites whatever setting the user configured in the properties
	adjustAutoCommitConfig(properties, offsetCommitMode);

	FlinkConnectorRateLimiter rateLimiter = super.getRateLimiter();
	// If a rateLimiter is set, then call rateLimiter.open() with the runtime context.
	if (rateLimiter != null) {
		rateLimiter.open(runtimeContext);
	}

	return new Kafka010Fetcher<>(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarksPeriodic,
			watermarksPunctuated,
			runtimeContext.getProcessingTimeService(),
			runtimeContext.getExecutionConfig().getAutoWatermarkInterval(),
			runtimeContext.getUserCodeClassLoader(),
			runtimeContext.getTaskNameWithSubtasks(),
			deserializer,
			properties,
			pollTimeout,
			runtimeContext.getMetricGroup(),
			consumerMetricGroup,
			useMetrics,
			rateLimiter);
}
 

public Kafka010Fetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<AssignerWithPeriodicWatermarks<T>> watermarksPeriodic,
		SerializedValue<AssignerWithPunctuatedWatermarks<T>> watermarksPunctuated,
		ProcessingTimeService processingTimeProvider,
		long autoWatermarkInterval,
		ClassLoader userCodeClassLoader,
		String taskNameWithSubtasks,
		KafkaDeserializationSchema<T> deserializer,
		Properties kafkaProperties,
		long pollTimeout,
		MetricGroup subtaskMetricGroup,
		MetricGroup consumerMetricGroup,
		boolean useMetrics,
		FlinkConnectorRateLimiter rateLimiter) throws Exception {
	super(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarksPeriodic,
			watermarksPunctuated,
			processingTimeProvider,
			autoWatermarkInterval,
			userCodeClassLoader,
			taskNameWithSubtasks,
			deserializer,
			kafkaProperties,
			pollTimeout,
			subtaskMetricGroup,
			consumerMetricGroup,
			useMetrics, rateLimiter);
}
 

public KafkaConsumerThread(
		Logger log,
		Handover handover,
		Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge,
		String threadName,
		long pollTimeout,
		boolean useMetrics,
		MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		FlinkConnectorRateLimiter rateLimiter) {

	super(threadName);
	setDaemon(true);

	this.log = checkNotNull(log);
	this.handover = checkNotNull(handover);
	this.kafkaProperties = checkNotNull(kafkaProperties);
	this.consumerMetricGroup = checkNotNull(consumerMetricGroup);
	this.subtaskMetricGroup = checkNotNull(subtaskMetricGroup);
	this.consumerCallBridge = checkNotNull(consumerCallBridge);

	this.unassignedPartitionsQueue = checkNotNull(unassignedPartitionsQueue);

	this.pollTimeout = pollTimeout;
	this.useMetrics = useMetrics;

	this.consumerReassignmentLock = new Object();
	this.nextOffsetsToCommit = new AtomicReference<>();
	this.running = true;

	if (rateLimiter != null) {
		this.rateLimiter = rateLimiter;
	}
}
 

public TestKafkaConsumerThreadRateLimit(Logger log,
		Handover handover, Properties kafkaProperties,
		ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue,
		KafkaConsumerCallBridge09 consumerCallBridge, String threadName, long pollTimeout,
		boolean useMetrics, MetricGroup consumerMetricGroup,
		MetricGroup subtaskMetricGroup,
		KafkaConsumer mockConsumer,
		FlinkConnectorRateLimiter rateLimiter) {
	super(log, handover, kafkaProperties, unassignedPartitionsQueue, consumerCallBridge,
			threadName,
			pollTimeout, useMetrics, consumerMetricGroup, subtaskMetricGroup,
		rateLimiter);
	this.mockConsumer = mockConsumer;
}
 

public FlinkConnectorRateLimiter getRateLimiter() {
	return rateLimiter;
}
 

@Test(timeout = 10000)
public void testRatelimiting() throws Exception {
	final String testTopic = "test-topic-ratelimit";

	// -------- setup mock KafkaConsumer with test data --------
	final int partition = 0;
	final byte[] payload = new byte[] {1};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(testTopic, partition, 15, payload, payload),
			new ConsumerRecord<>(testTopic, partition, 16, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(testTopic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// Sleep for one second in each consumer.poll() call to return 24 bytes / second
	final KafkaConsumer<byte[], byte[]> mockConsumer = mock(KafkaConsumer.class);
	PowerMockito.when(mockConsumer.poll(anyLong())).thenAnswer(
			invocationOnMock -> consumerRecords
	);

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<Object, TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
			new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<Object, TopicPartition>> newPartitions = new ArrayList<>(1);
	newPartitions.add(newPartition1);

	final ClosableBlockingQueue<KafkaTopicPartitionState<Object, TopicPartition>> unassignedPartitionsQueue =
			new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<Object, TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// --- ratelimiting properties ---
	StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
	Properties properties = new Properties();

	// -- mock Handover and logger ---
	Handover mockHandover = PowerMockito.mock(Handover.class);
	doNothing().when(mockHandover).produce(any());
	Logger mockLogger = mock(Logger.class);

	MetricGroup metricGroup = new UnregisteredMetricsGroup();
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(1L);
	rateLimiter.open(mockRuntimeContext);

	// -- Test Kafka Consumer thread ---

	KafkaConsumerThread testThread = new TestKafkaConsumerThreadRateLimit(
			mockLogger,
			mockHandover,
			properties,
			unassignedPartitionsQueue,
			"test",
			30L,
			false,
			metricGroup,
			metricGroup,
			mockConsumer,
			rateLimiter
	);

	testThread.start();
	// Wait for 4 seconds to ensure atleast 2 calls to consumer.poll()
	testThread.join(5000);
	assertNotNull(testThread.getRateLimiter());
	assertEquals(testThread.getRateLimiter().getRate(), 1, 0);

	// In a period of 5 seconds, no more than 3 calls to poll should be made.
	// The expected rate is 1 byte / second and we read 4 bytes in every consumer.poll()
	// call. The rate limiter should thus slow down the call by 4 seconds when the rate takes
	// effect.
	verify(mockConsumer, times(3)).poll(anyLong());
	testThread.shutdown();

}
 

public FlinkConnectorRateLimiter getRateLimiter() {
	return rateLimiter;
}
 

@VisibleForTesting
FlinkConnectorRateLimiter getRateLimiter() {
	return rateLimiter;
}
 

@Test(timeout = 10000)
public void testRatelimiting() throws Exception {
	final String testTopic = "test-topic-ratelimit";

	// -------- setup mock KafkaConsumer with test data --------
	final int partition = 0;
	final byte[] payload = new byte[] {1};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(testTopic, partition, 15, payload, payload),
			new ConsumerRecord<>(testTopic, partition, 16, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(testTopic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// Sleep for one second in each consumer.poll() call to return 24 bytes / second
	final KafkaConsumer<byte[], byte[]> mockConsumer = mock(KafkaConsumer.class);
	PowerMockito.when(mockConsumer.poll(anyLong())).thenAnswer(
			invocationOnMock -> consumerRecords
	);

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
			new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> newPartitions = new ArrayList<>(1);
	newPartitions.add(newPartition1);

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
			new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// --- ratelimiting properties ---
	StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
	Properties properties = new Properties();
	KafkaConsumerCallBridge09 mockBridge = mock(KafkaConsumerCallBridge09.class);

	// -- mock Handover and logger ---
	Handover mockHandover = PowerMockito.mock(Handover.class);
	doNothing().when(mockHandover).produce(any());
	Logger mockLogger = mock(Logger.class);

	MetricGroup metricGroup = new UnregisteredMetricsGroup();
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(1L);
	rateLimiter.open(mockRuntimeContext);

	// -- Test Kafka Consumer thread ---

	KafkaConsumerThread testThread = new TestKafkaConsumerThreadRateLimit(
			mockLogger,
			mockHandover,
			properties,
			unassignedPartitionsQueue,
			mockBridge,
			"test",
			30L,
			false,
			metricGroup,
			metricGroup,
			mockConsumer,
			rateLimiter
	);

	testThread.start();
	// Wait for 4 seconds to ensure atleast 2 calls to consumer.poll()
	testThread.join(5000);
	assertNotNull(testThread.getRateLimiter());
	assertEquals(testThread.getRateLimiter().getRate(), 1, 0);

	// In a period of 5 seconds, no more than 3 calls to poll should be made.
	// The expected rate is 1 byte / second and we read 4 bytes in every consumer.poll()
	// call. The rate limiter should thus slow down the call by 4 seconds when the rate takes
	// effect.
	verify(mockConsumer, times(3)).poll(anyLong());
	testThread.shutdown();

}
 

/**
 * Kafka09 specific RateLimiter test. This test produces 100 bytes of data to a test topic
 * and then runs a job with {@link FlinkKafkaConsumer09} as the source and a {@link GuavaFlinkConnectorRateLimiter} with
 * a desired rate of 3 bytes / second. Based on the execution time, the test asserts that this rate was not surpassed.
 * If no rate limiter is set on the consumer, the test should fail.
 */
@Test(timeout = 60000)
public void testRateLimitedConsumer() throws Exception {
	final String testTopic = "testRateLimitedConsumer";
	createTestTopic(testTopic, 3, 1);

	// ---------- Produce a stream into Kafka -------------------

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(1);
	env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
	env.getConfig().disableSysoutLogging();

	DataStream<String> stream = env.addSource(new SourceFunction<String>() {
		private static final long serialVersionUID = 1L;
		boolean running = true;

		@Override
		public void run(SourceContext<String> ctx) {
			long i = 0;
			while (running) {
				byte[] data = new byte[] {1};
				synchronized (ctx.getCheckpointLock()) {
					ctx.collect(new String(data)); // 1 byte
				}
				if (i++ == 100L) {
					running = false;
				}
			}
		}

		@Override
		public void cancel() {
			running = false;
		}
	});

	Properties producerProperties = new Properties();
	producerProperties.putAll(standardProps);
	producerProperties.putAll(secureProps);
	producerProperties.put("retries", 3);

	stream.addSink(new FlinkKafkaProducer09<>(testTopic, new SimpleStringSchema(), producerProperties));
	env.execute("Produce 100 bytes of data to test topic");

	// ---------- Consumer from Kafka in a ratelimited way -----------

	env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(1);
	env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
	env.getConfig().disableSysoutLogging();

	// ---------- RateLimiter config -------------
	final long globalRate = 10; // bytes/second
	FlinkKafkaConsumer09<String> consumer09 = new FlinkKafkaConsumer09<>(testTopic,
		new StringDeserializer(globalRate), standardProps);
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(globalRate);
	consumer09.setRateLimiter(rateLimiter);

	DataStream<String> stream1 = env.addSource(consumer09);
	stream1.addSink(new DiscardingSink<>());
	env.execute("Consume 100 bytes of data from test topic");

	// ------- Assertions --------------
	Assert.assertNotNull(consumer09.getRateLimiter());
	Assert.assertEquals(globalRate, consumer09.getRateLimiter().getRate());

	deleteTestTopic(testTopic);
}
 

public Kafka09Fetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<AssignerWithPeriodicWatermarks<T>> watermarksPeriodic,
		SerializedValue<AssignerWithPunctuatedWatermarks<T>> watermarksPunctuated,
		ProcessingTimeService processingTimeProvider,
		long autoWatermarkInterval,
		ClassLoader userCodeClassLoader,
		String taskNameWithSubtasks,
		KafkaDeserializationSchema<T> deserializer,
		Properties kafkaProperties,
		long pollTimeout,
		MetricGroup subtaskMetricGroup,
		MetricGroup consumerMetricGroup,
		boolean useMetrics,
		FlinkConnectorRateLimiter rateLimiter) throws Exception {
	super(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarksPeriodic,
			watermarksPunctuated,
			processingTimeProvider,
			autoWatermarkInterval,
			userCodeClassLoader,
			consumerMetricGroup,
			useMetrics);

	this.deserializer = deserializer;
	this.handover = new Handover();

	this.consumerThread = new KafkaConsumerThread(
			LOG,
			handover,
			kafkaProperties,
			unassignedPartitionsQueue,
			createCallBridge(),
			getFetcherName() + " for " + taskNameWithSubtasks,
			pollTimeout,
			useMetrics,
			consumerMetricGroup,
			subtaskMetricGroup,
			rateLimiter);
}
 

@VisibleForTesting
FlinkConnectorRateLimiter getRateLimiter() {
	return rateLimiter;
}
 

@Test(timeout = 10000)
public void testRatelimiting() throws Exception {
	final String testTopic = "test-topic-ratelimit";

	// -------- setup mock KafkaConsumer with test data --------
	final int partition = 0;
	final byte[] payload = new byte[] {1};

	final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
			new ConsumerRecord<>(testTopic, partition, 15, payload, payload),
			new ConsumerRecord<>(testTopic, partition, 16, payload, payload));

	final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
	data.put(new TopicPartition(testTopic, partition), records);

	final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

	// Sleep for one second in each consumer.poll() call to return 24 bytes / second
	final KafkaConsumer<byte[], byte[]> mockConsumer = mock(KafkaConsumer.class);
	PowerMockito.when(mockConsumer.poll(anyLong())).thenAnswer(
			invocationOnMock -> consumerRecords
	);

	whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

	// -------- new partitions with defined offsets --------

	KafkaTopicPartitionState<TopicPartition> newPartition1 = new KafkaTopicPartitionState<>(
			new KafkaTopicPartition(testTopic, 0), new TopicPartition(testTopic, 0));
	newPartition1.setOffset(KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	List<KafkaTopicPartitionState<TopicPartition>> newPartitions = new ArrayList<>(1);
	newPartitions.add(newPartition1);

	final ClosableBlockingQueue<KafkaTopicPartitionState<TopicPartition>> unassignedPartitionsQueue =
			new ClosableBlockingQueue<>();

	for (KafkaTopicPartitionState<TopicPartition> newPartition : newPartitions) {
		unassignedPartitionsQueue.add(newPartition);
	}

	// --- ratelimiting properties ---
	StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);
	Properties properties = new Properties();
	KafkaConsumerCallBridge09 mockBridge = mock(KafkaConsumerCallBridge09.class);

	// -- mock Handover and logger ---
	Handover mockHandover = PowerMockito.mock(Handover.class);
	doNothing().when(mockHandover).produce(any());
	Logger mockLogger = mock(Logger.class);

	MetricGroup metricGroup = new UnregisteredMetricsGroup();
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(1L);
	rateLimiter.open(mockRuntimeContext);

	// -- Test Kafka Consumer thread ---

	KafkaConsumerThread testThread = new TestKafkaConsumerThreadRateLimit(
			mockLogger,
			mockHandover,
			properties,
			unassignedPartitionsQueue,
			mockBridge,
			"test",
			30L,
			false,
			metricGroup,
			metricGroup,
			mockConsumer,
			rateLimiter
	);

	testThread.start();
	// Wait for 4 seconds to ensure atleast 2 calls to consumer.poll()
	testThread.join(5000);
	assertNotNull(testThread.getRateLimiter());
	assertEquals(testThread.getRateLimiter().getRate(), 1, 0);

	// In a period of 5 seconds, no more than 3 calls to poll should be made.
	// The expected rate is 1 byte / second and we read 4 bytes in every consumer.poll()
	// call. The rate limiter should thus slow down the call by 4 seconds when the rate takes
	// effect.
	verify(mockConsumer, times(3)).poll(anyLong());
	testThread.shutdown();

}
 

/**
 * Kafka09 specific RateLimiter test. This test produces 100 bytes of data to a test topic
 * and then runs a job with {@link FlinkKafkaConsumer09} as the source and a {@link GuavaFlinkConnectorRateLimiter} with
 * a desired rate of 3 bytes / second. Based on the execution time, the test asserts that this rate was not surpassed.
 * If no rate limiter is set on the consumer, the test should fail.
 */
@Test(timeout = 60000)
public void testRateLimitedConsumer() throws Exception {
	final String testTopic = "testRateLimitedConsumer";
	createTestTopic(testTopic, 3, 1);

	// ---------- Produce a stream into Kafka -------------------

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(1);
	env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
	env.getConfig().disableSysoutLogging();

	DataStream<String> stream = env.addSource(new SourceFunction<String>() {
		private static final long serialVersionUID = 1L;
		boolean running = true;

		@Override
		public void run(SourceContext<String> ctx) {
			long i = 0;
			while (running) {
				byte[] data = new byte[] {1};
				synchronized (ctx.getCheckpointLock()) {
					ctx.collect(new String(data)); // 1 byte
				}
				if (i++ == 100L) {
					running = false;
				}
			}
		}

		@Override
		public void cancel() {
			running = false;
		}
	});

	Properties producerProperties = new Properties();
	producerProperties.putAll(standardProps);
	producerProperties.putAll(secureProps);
	producerProperties.put("retries", 3);

	stream.addSink(new FlinkKafkaProducer09<>(testTopic, new SimpleStringSchema(), producerProperties));
	env.execute("Produce 100 bytes of data to test topic");

	// ---------- Consumer from Kafka in a ratelimited way -----------

	env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(1);
	env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
	env.getConfig().disableSysoutLogging();

	// ---------- RateLimiter config -------------
	final long globalRate = 10; // bytes/second
	FlinkKafkaConsumer09<String> consumer09 = new FlinkKafkaConsumer09<>(testTopic,
		new StringDeserializer(globalRate), standardProps);
	FlinkConnectorRateLimiter rateLimiter = new GuavaFlinkConnectorRateLimiter();
	rateLimiter.setRate(globalRate);
	consumer09.setRateLimiter(rateLimiter);

	DataStream<String> stream1 = env.addSource(consumer09);
	stream1.addSink(new DiscardingSink<>());
	env.execute("Consume 100 bytes of data from test topic");

	// ------- Assertions --------------
	Assert.assertNotNull(consumer09.getRateLimiter());
	Assert.assertEquals(globalRate, consumer09.getRateLimiter().getRate());

	deleteTestTopic(testTopic);
}
 

public FlinkConnectorRateLimiter getRateLimiter() {
	return rateLimiter;
}
 

public Kafka09Fetcher(
		SourceContext<T> sourceContext,
		Map<KafkaTopicPartition, Long> assignedPartitionsWithInitialOffsets,
		SerializedValue<AssignerWithPeriodicWatermarks<T>> watermarksPeriodic,
		SerializedValue<AssignerWithPunctuatedWatermarks<T>> watermarksPunctuated,
		ProcessingTimeService processingTimeProvider,
		long autoWatermarkInterval,
		ClassLoader userCodeClassLoader,
		String taskNameWithSubtasks,
		KafkaDeserializationSchema<T> deserializer,
		Properties kafkaProperties,
		long pollTimeout,
		MetricGroup subtaskMetricGroup,
		MetricGroup consumerMetricGroup,
		boolean useMetrics,
		FlinkConnectorRateLimiter rateLimiter) throws Exception {
	super(
			sourceContext,
			assignedPartitionsWithInitialOffsets,
			watermarksPeriodic,
			watermarksPunctuated,
			processingTimeProvider,
			autoWatermarkInterval,
			userCodeClassLoader,
			consumerMetricGroup,
			useMetrics);

	this.deserializer = deserializer;
	this.handover = new Handover();

	this.consumerThread = new KafkaConsumerThread(
			LOG,
			handover,
			kafkaProperties,
			unassignedPartitionsQueue,
			createCallBridge(),
			getFetcherName() + " for " + taskNameWithSubtasks,
			pollTimeout,
			useMetrics,
			consumerMetricGroup,
			subtaskMetricGroup,
			rateLimiter);
}
 

@VisibleForTesting
FlinkConnectorRateLimiter getRateLimiter() {
	return rateLimiter;
}
 

/**
 * Set a rate limiter to ratelimit bytes read from Kafka.
 * @param kafkaRateLimiter
 */
public void setRateLimiter(FlinkConnectorRateLimiter kafkaRateLimiter) {
	this.rateLimiter = kafkaRateLimiter;
}
 

/**
 * Set a rate limiter to ratelimit bytes read from Kafka.
 * @param kafkaRateLimiter
 */
public void setRateLimiter(FlinkConnectorRateLimiter kafkaRateLimiter) {
	this.rateLimiter = kafkaRateLimiter;
}
 

/**
 * Set a rate limiter to ratelimit bytes read from Kafka.
 * @param kafkaRateLimiter
 */
public void setRateLimiter(FlinkConnectorRateLimiter kafkaRateLimiter) {
	this.rateLimiter = kafkaRateLimiter;
}