Java Code Examples for org.apache.flink.api.common.io.ratelimiting.FlinkConnectorRateLimiter#open()

The following examples show how to use org.apache.flink.api.common.io.ratelimiting.FlinkConnectorRateLimiter#open() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: FlinkKafkaConsumer010.java    From Flink-CEPplus with Apache License 2.0 5 votes vote down vote up 
@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);
}
Example 2
Source File: FlinkKafkaConsumer010.java    From flink with Apache License 2.0 5 votes vote down vote up 
@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);
}
Example 3
Source File: KafkaConsumerThreadTest.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up 
@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();

}
Example 4
Source File: KafkaConsumerThreadTest.java    From flink with Apache License 2.0 4 votes vote down vote up 
@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();

}
Example 5
Source File: KafkaConsumerThreadTest.java    From flink with Apache License 2.0 4 votes vote down vote up 
@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();

}