This is an automated email from the ASF dual-hosted git repository.
pvillard31 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new 2bdb2c901d3 NIFI-16380 Add Records Metrics to ConsumeKinesis (#11712)
2bdb2c901d3 is described below
commit 2bdb2c901d307407a618e037529d8d4641571f59
Author: David Handermann <[email protected]>
AuthorDate: Fri Sep 25 09:21:24 2026 -0500
NIFI-16380 Add Records Metrics to ConsumeKinesis (#11712)
---
.../processors/aws/kinesis/ConsumeKinesis.java | 51 ++++
.../processors/aws/kinesis/KinesisMetricName.java | 37 +++
.../processors/aws/kinesis/ConsumeKinesisTest.java | 257 ++++++++++++++++++++-
3 files changed, 333 insertions(+), 12 deletions(-)
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesis.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesis.java
index 83fdb502e46..56fa367de72 100644
---
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesis.java
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesis.java
@@ -46,6 +46,7 @@ import org.apache.nifi.processor.ProcessSession;
import org.apache.nifi.processor.Relationship;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.io.OutputStreamCallback;
+import org.apache.nifi.processor.metrics.CommitTiming;
import org.apache.nifi.processor.util.StandardValidators;
import
org.apache.nifi.processors.aws.credentials.provider.AwsCredentialsProviderService;
import org.apache.nifi.processors.aws.region.RegionUtil;
@@ -589,6 +590,11 @@ public class ConsumeKinesis extends AbstractProcessor
implements BacklogReportin
consumed = consumeRecords(claimedShards);
final List<ShardFetchResult> accepted =
discardRelinquishedResults(consumed, claimedShards);
+ for (final ShardFetchResult result : accepted) {
+ shardMillisBehindLatest.put(result.shardId(),
result.millisBehindLatest());
+ }
+ recordBehindMetrics(session, ownedShardIds);
+
if (accepted.isEmpty()) {
consumerClient.releaseShards(claimedShards);
context.yield();
@@ -895,6 +901,7 @@ public class ConsumeKinesis extends AbstractProcessor
implements BacklogReportin
if (processingStrategy == ProcessingStrategy.FLOW_FILE) {
final BatchAccumulator batch = new BatchAccumulator();
for (final List<ShardFetchResult> shardResults :
resultsByShard.values()) {
+ recordConsumptionMetrics(session, shardResults);
for (final ShardFetchResult result : shardResults) {
batch.updateMillisBehind(result.millisBehindLatest());
for (final UserRecord record : result.records()) {
@@ -910,6 +917,7 @@ public class ConsumeKinesis extends AbstractProcessor
implements BacklogReportin
for (final Map.Entry<String, List<ShardFetchResult>> entry :
resultsByShard.entrySet()) {
final BatchAccumulator batch = new BatchAccumulator();
batch.setLastShardId(entry.getKey());
+ recordConsumptionMetrics(session, entry.getValue());
for (final ShardFetchResult result : entry.getValue()) {
batch.updateMillisBehind(result.millisBehindLatest());
batch.updateSequenceRange(result);
@@ -946,6 +954,46 @@ public class ConsumeKinesis extends AbstractProcessor
implements BacklogReportin
return new WriteResult(produced, parseFailures, totalRecordCount,
totalBytesConsumed, maxMillisBehind);
}
+ private void recordConsumptionMetrics(final ProcessSession session, final
List<ShardFetchResult> shardResults) {
+ long recordCount = 0;
+ long bytesConsumed = 0;
+ String shardId = null;
+ for (final ShardFetchResult result : shardResults) {
+ shardId = result.shardId();
+ for (final UserRecord record : result.records()) {
+ recordCount++;
+ bytesConsumed += record.data().length;
+ }
+ }
+
+ if (recordCount == 0) {
+ return;
+ }
+
+ final Map<String, String> attributes = getMetricAttributes(streamName,
shardId);
+
session.adjustCounter(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
recordCount, attributes, CommitTiming.NOW);
+
session.adjustCounter(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
bytesConsumed, attributes, CommitTiming.NOW);
+ }
+
+ private void recordBehindMetrics(final ProcessSession session, final
Set<String> ownedShardIds) {
+ for (final String shardId : ownedShardIds) {
+ final Long millisBehind = shardMillisBehindLatest.get(shardId);
+ // Kinesis uses -1 when millisBehindLatest is not provided
+ if (millisBehind == null || millisBehind < 0) {
+ continue;
+ }
+ final Map<String, String> metricAttributes =
getMetricAttributes(streamName, shardId);
+
session.recordGauge(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
millisBehind, metricAttributes, CommitTiming.NOW);
+ }
+ }
+
+ private static Map<String, String> getMetricAttributes(final String
streamName, final String shardId) {
+ return Map.of(
+ ATTR_STREAM_NAME, streamName,
+ ATTR_SHARD_ID, shardId
+ );
+ }
+
private void writeFlowFilePerRecord(final ProcessSession session, final
List<ShardFetchResult> results,
final String streamName, final
BatchAccumulator batch, final List<FlowFile> output) {
for (final ShardFetchResult result : results) {
@@ -1360,6 +1408,9 @@ public class ConsumeKinesis extends AbstractProcessor
implements BacklogReportin
}
flowFile = session.putAllAttributes(flowFile, attributes);
parseFailureOutput.add(flowFile);
+
+ final Map<String, String> metricAttributes =
getMetricAttributes(streamName, record.shardId());
+
session.adjustCounter(KinesisMetricName.RECORDS_PARSED_ERRORS.getMetricName(),
1, metricAttributes, CommitTiming.NOW);
} catch (final Exception e) {
session.remove(flowFile);
throw e;
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/KinesisMetricName.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/KinesisMetricName.java
new file mode 100644
index 00000000000..cec30ea32d5
--- /dev/null
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/KinesisMetricName.java
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.processors.aws.kinesis;
+
+/**
+ * Metric names for ConsumeKinesis counters and gauges.
+ */
+public enum KinesisMetricName {
+ BYTES_CONSUMED("kinesis.bytes.consumed"),
+ RECORDS_CONSUMED("kinesis.records.consumed"),
+ CONSUMER_MILLISECONDS_BEHIND("kinesis.consumer.milliseconds.behind"),
+ RECORDS_PARSED_ERRORS("kinesis.records.parsed.errors");
+
+ private final String metricName;
+
+ KinesisMetricName(final String metricName) {
+ this.metricName = metricName;
+ }
+
+ public String getMetricName() {
+ return metricName;
+ }
+}
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisTest.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisTest.java
index ef538ed5fd1..67f263b2e84 100644
---
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisTest.java
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisTest.java
@@ -45,6 +45,7 @@ import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
@@ -52,6 +53,32 @@ import static org.mockito.Mockito.when;
class ConsumeKinesisTest {
+ private static final String TEST_STREAM_NAME = "test-stream";
+
+ private static final String DEFAULT_SHARD_ID = "shardId-000000000001";
+
+ private static final String FIRST_SHARD_ID = "shard-A";
+
+ private static final String SECOND_SHARD_ID = "shard-B";
+
+ private static final String FIRST_JSON_RECORD = "{\"id\":1}";
+
+ private static final String SECOND_JSON_RECORD = "{\"id\":2}";
+
+ private static final String THIRD_JSON_RECORD = "{\"id\":3}";
+
+ private static final String FOURTH_JSON_RECORD = "{\"id\":4}";
+
+ private static final String FIRST_FLOW_FILE_RECORD = "record-one";
+
+ private static final String SECOND_FLOW_FILE_RECORD = "record-two";
+
+ private static final long FIRST_SHARD_BEHIND_MS = 1000L;
+
+ private static final long SECOND_SHARD_BEHIND_MS = 2500L;
+
+ private static final long FLOW_FILE_BEHIND_MS = 5000L;
+
private TestRunner runner;
@BeforeEach
@@ -79,7 +106,7 @@ class ConsumeKinesisTest {
targetRunner.enableControllerService(credentialsService);
targetRunner.setProperty(ConsumeKinesis.APPLICATION_NAME, "test-app");
- targetRunner.setProperty(ConsumeKinesis.STREAM_NAME, "test-stream");
+ targetRunner.setProperty(ConsumeKinesis.STREAM_NAME, TEST_STREAM_NAME);
targetRunner.setProperty(ConsumeKinesis.AWS_CREDENTIALS_PROVIDER_SERVICE,
"creds");
}
@@ -168,6 +195,10 @@ class ConsumeKinesisTest {
runner.assertTransferCount(ConsumeKinesis.REL_PARSE_FAILURE, 0);
final MockFlowFile success =
runner.getFlowFilesForRelationship(ConsumeKinesis.REL_SUCCESS).getFirst();
success.assertAttributeEquals("record.count", "3");
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ assertEquals(3L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(List.of(0.0),
runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
attributes));
}
@Test
@@ -205,6 +236,10 @@ class ConsumeKinesisTest {
assertTrue(failureSequences.contains("1"));
assertTrue(failureSequences.contains("3"));
assertTrue(failureSequences.contains("5"));
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ assertEquals(5L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(3L,
runner.getCounterValue(KinesisMetricName.RECORDS_PARSED_ERRORS.getMetricName(),
attributes));
}
@Test
@@ -218,6 +253,10 @@ class ConsumeKinesisTest {
runner.assertTransferCount(ConsumeKinesis.REL_SUCCESS, 0);
runner.assertTransferCount(ConsumeKinesis.REL_PARSE_FAILURE, 3);
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ assertEquals(3L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(3L,
runner.getCounterValue(KinesisMetricName.RECORDS_PARSED_ERRORS.getMetricName(),
attributes));
}
@Test
@@ -234,7 +273,7 @@ class ConsumeKinesisTest {
final List<MockFlowFile> flowFiles =
runner.getFlowFilesForRelationship(ConsumeKinesis.REL_SUCCESS);
for (final MockFlowFile flowFile : flowFiles) {
flowFile.assertAttributeEquals("record.count", "1");
- flowFile.assertAttributeEquals(ConsumeKinesis.ATTR_STREAM_NAME,
"test-stream");
+ flowFile.assertAttributeEquals(ConsumeKinesis.ATTR_STREAM_NAME,
TEST_STREAM_NAME);
flowFile.assertAttributeEquals(ConsumeKinesis.ATTR_SHARD_ID,
"shardId-000000000001");
final String firstSequence =
flowFile.getAttribute(ConsumeKinesis.ATTR_FIRST_SEQUENCE);
final String lastSequence =
flowFile.getAttribute(ConsumeKinesis.ATTR_LAST_SEQUENCE);
@@ -284,10 +323,10 @@ class ConsumeKinesisTest {
@Test
void testMultipleShardsNoDataLoss() throws Exception {
- final ShardFetchResult shard1Result = new ShardFetchResult("shard-A",
- List.of(testRecord("10", "{\"id\":1}"), testRecord("20",
"{\"id\":2}")), 0L);
- final ShardFetchResult shard2Result = new ShardFetchResult("shard-B",
- List.of(testRecord("30", "{\"id\":3}"), testRecord("40",
"{\"id\":4}")), 0L);
+ final ShardFetchResult shard1Result = new
ShardFetchResult(FIRST_SHARD_ID,
+ List.of(testRecord("10", FIRST_JSON_RECORD), testRecord("20",
SECOND_JSON_RECORD)), FIRST_SHARD_BEHIND_MS);
+ final ShardFetchResult shard2Result = new
ShardFetchResult(SECOND_SHARD_ID,
+ List.of(testRecord("30", THIRD_JSON_RECORD), testRecord("40",
FOURTH_JSON_RECORD)), SECOND_SHARD_BEHIND_MS);
triggerWithResults(List.of(shard1Result, shard2Result), "RECORD");
@@ -301,8 +340,138 @@ class ConsumeKinesisTest {
shardsSeen.add(flowFile.getAttribute(ConsumeKinesis.ATTR_SHARD_ID));
totalRecords +=
Long.parseLong(flowFile.getAttribute("record.count"));
}
- assertEquals(Set.of("shard-A", "shard-B"), shardsSeen);
+ assertEquals(Set.of(FIRST_SHARD_ID, SECOND_SHARD_ID), shardsSeen);
assertEquals(4, totalRecords);
+
+ final Map<String, String> firstAttributes =
getMetricAttributes(FIRST_SHARD_ID);
+ final Map<String, String> secondAttributes =
getMetricAttributes(SECOND_SHARD_ID);
+ final long firstShardBytes = payloadBytes(FIRST_JSON_RECORD) +
payloadBytes(SECOND_JSON_RECORD);
+ final long secondShardBytes = payloadBytes(THIRD_JSON_RECORD) +
payloadBytes(FOURTH_JSON_RECORD);
+
+ assertEquals(2L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
firstAttributes));
+ assertEquals(2L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
secondAttributes));
+ assertEquals(firstShardBytes,
runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
firstAttributes));
+ assertEquals(secondShardBytes,
runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
secondAttributes));
+ assertEquals(List.of((double) FIRST_SHARD_BEHIND_MS),
+
runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
firstAttributes));
+ assertEquals(List.of((double) SECOND_SHARD_BEHIND_MS),
+
runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
secondAttributes));
+ }
+
+ @Test
+ void testEmptyConsumeDoesNotRecordMetrics() throws Exception {
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final TestableConsumeKinesis processor = new
TestableConsumeKinesis(mockShardManager, List.of());
+ runner = TestRunners.newTestRunner(processor);
+
+ setCommonProperties();
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, "FLOW_FILE");
+ runner.setProperty(ConsumeKinesis.CONSUMER_TYPE, "SHARED_THROUGHPUT");
+
+ runner.run();
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+
assertNull(runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+
assertNull(runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
attributes));
+
assertTrue(runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
attributes).isEmpty());
+ runner.assertTransferCount(ConsumeKinesis.REL_SUCCESS, 0);
+ }
+
+ @Test
+ void testEmptyPollRecordsBehindGaugeWhenStillLagging() throws Exception {
+ final Map<String, Long> shardLag = new LinkedHashMap<>();
+ shardLag.put(DEFAULT_SHARD_ID, FLOW_FILE_BEHIND_MS);
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final LagReportingConsumeKinesis processor = new
LagReportingConsumeKinesis(mockShardManager, shardLag);
+ runner = TestRunners.newTestRunner(processor);
+
+ setCommonProperties();
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, "FLOW_FILE");
+ runner.setProperty(ConsumeKinesis.CONSUMER_TYPE, "SHARED_THROUGHPUT");
+
+ runner.run(1, false, true);
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+
assertNull(runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+
assertNull(runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
attributes));
+ assertEquals(List.of((double) FLOW_FILE_BEHIND_MS),
+
runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
attributes));
+ runner.assertTransferCount(ConsumeKinesis.REL_SUCCESS, 0);
+ }
+
+ @Test
+ void testEmptyPollUpdatesBehindGaugeToCaughtUp() throws Exception {
+ final List<UserRecord> records = List.of(
+ testRecord("1", FIRST_FLOW_FILE_RECORD),
+ testRecord("2", SECOND_FLOW_FILE_RECORD));
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final ShardFetchResult firstFetch = new
ShardFetchResult(DEFAULT_SHARD_ID, records, FLOW_FILE_BEHIND_MS);
+ final SequentialFetchConsumeKinesis processor = new
SequentialFetchConsumeKinesis(mockShardManager, firstFetch, 0L);
+ runner = TestRunners.newTestRunner(processor);
+
+ setCommonProperties();
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, "FLOW_FILE");
+ runner.setProperty(ConsumeKinesis.CONSUMER_TYPE, "SHARED_THROUGHPUT");
+
+ runner.run(2, false, true);
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ final long expectedBytes = payloadBytes(FIRST_FLOW_FILE_RECORD) +
payloadBytes(SECOND_FLOW_FILE_RECORD);
+ assertEquals(2L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(expectedBytes,
runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
attributes));
+ assertEquals(List.of((double) FLOW_FILE_BEHIND_MS, 0.0),
+
runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
attributes));
+ runner.assertTransferCount(ConsumeKinesis.REL_SUCCESS, 2);
+ }
+
+ @Test
+ void testAbsentMillisBehindDoesNotRecordGauge() throws Exception {
+ final List<UserRecord> records = List.of(testRecord("1",
FIRST_FLOW_FILE_RECORD));
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final ShardFetchResult fetchResult = new
ShardFetchResult(DEFAULT_SHARD_ID, records, -1L);
+ final TestableConsumeKinesis processor = new
TestableConsumeKinesis(mockShardManager, fetchResult);
+ runner = TestRunners.newTestRunner(processor);
+
+ setCommonProperties();
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, "FLOW_FILE");
+ runner.setProperty(ConsumeKinesis.CONSUMER_TYPE, "SHARED_THROUGHPUT");
+
+ runner.run();
+
+ runner.assertTransferCount(ConsumeKinesis.REL_SUCCESS, 1);
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ assertEquals(1L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(payloadBytes(FIRST_FLOW_FILE_RECORD),
+
runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
attributes));
+
assertTrue(runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
attributes).isEmpty());
+ }
+
+ @Test
+ void testFlowFileStrategyRecordsBytesAndBehindMetrics() throws Exception {
+ final List<UserRecord> records = List.of(
+ testRecord("1", FIRST_FLOW_FILE_RECORD),
+ testRecord("2", SECOND_FLOW_FILE_RECORD));
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final ShardFetchResult fetchResult = new
ShardFetchResult(DEFAULT_SHARD_ID, records, FLOW_FILE_BEHIND_MS);
+ final TestableConsumeKinesis processor = new
TestableConsumeKinesis(mockShardManager, fetchResult);
+ runner = TestRunners.newTestRunner(processor);
+
+ setCommonProperties();
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, "FLOW_FILE");
+ runner.setProperty(ConsumeKinesis.CONSUMER_TYPE, "SHARED_THROUGHPUT");
+
+ runner.run();
+
+ runner.assertTransferCount(ConsumeKinesis.REL_SUCCESS, 2);
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ final long expectedBytes = payloadBytes(FIRST_FLOW_FILE_RECORD) +
payloadBytes(SECOND_FLOW_FILE_RECORD);
+ assertEquals(2L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(expectedBytes,
runner.getCounterValue(KinesisMetricName.BYTES_CONSUMED.getMetricName(),
attributes));
+ assertEquals(2L, runner.getCounterValue("Records Consumed"));
+ assertEquals(List.of((double) FLOW_FILE_BEHIND_MS),
+
runner.getGaugeValues(KinesisMetricName.CONSUMER_MILLISECONDS_BEHIND.getMetricName(),
attributes));
}
@Test
@@ -413,8 +582,8 @@ class ConsumeKinesisTest {
}
private void triggerWithRecords(final List<UserRecord> records) throws
Exception {
- final KinesisShardManager mockShardManager =
buildShardManager("shardId-000000000001");
- final ShardFetchResult fetchResult = new
ShardFetchResult("shardId-000000000001", records, 0L);
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final ShardFetchResult fetchResult = new
ShardFetchResult(DEFAULT_SHARD_ID, records, 0L);
final TestableConsumeKinesis processor = new
TestableConsumeKinesis(mockShardManager, fetchResult);
runner = TestRunners.newTestRunner(processor);
@@ -456,11 +625,15 @@ class ConsumeKinesisTest {
failure.assertContentEquals(expectedFailureContent);
failure.assertAttributeEquals(ConsumeKinesis.ATTR_FIRST_SEQUENCE,
expectedFailureSequence);
assertNotNull(failure.getAttribute(ConsumeKinesis.ATTR_RECORD_ERROR_MESSAGE));
+
+ final Map<String, String> attributes =
getMetricAttributes(DEFAULT_SHARD_ID);
+ assertEquals(3L,
runner.getCounterValue(KinesisMetricName.RECORDS_CONSUMED.getMetricName(),
attributes));
+ assertEquals(1L,
runner.getCounterValue(KinesisMetricName.RECORDS_PARSED_ERRORS.getMetricName(),
attributes));
}
private void triggerWithOutputStrategy(final List<UserRecord> records,
final String outputStrategy) throws Exception {
- final KinesisShardManager mockShardManager =
buildShardManager("shardId-000000000001");
- final ShardFetchResult fetchResult = new
ShardFetchResult("shardId-000000000001", records, 0L);
+ final KinesisShardManager mockShardManager =
buildShardManager(DEFAULT_SHARD_ID);
+ final ShardFetchResult fetchResult = new
ShardFetchResult(DEFAULT_SHARD_ID, records, 0L);
final TestableConsumeKinesis processor = new
TestableConsumeKinesis(mockShardManager, fetchResult);
runner = TestRunners.newTestRunner(processor);
@@ -538,7 +711,7 @@ class ConsumeKinesisTest {
private static UserRecord testRecord(final String sequenceNumber, final
String data, final Instant arrivalTimestamp) {
return new UserRecord(
- "shardId-000000000001",
+ DEFAULT_SHARD_ID,
sequenceNumber,
0,
"pk-" + sequenceNumber,
@@ -546,6 +719,17 @@ class ConsumeKinesisTest {
arrivalTimestamp);
}
+ private static Map<String, String> getMetricAttributes(final String
shardId) {
+ return Map.of(
+ ConsumeKinesis.ATTR_STREAM_NAME, TEST_STREAM_NAME,
+ ConsumeKinesis.ATTR_SHARD_ID, shardId
+ );
+ }
+
+ private static long payloadBytes(final String payload) {
+ return payload.getBytes(StandardCharsets.UTF_8).length;
+ }
+
static class TestableConsumeKinesis extends ConsumeKinesis {
private final KinesisShardManager mockShardManager;
private final List<ShardFetchResult> preloadedResults;
@@ -576,6 +760,55 @@ class ConsumeKinesisTest {
}
}
+ static class SequentialFetchConsumeKinesis extends ConsumeKinesis {
+ private final KinesisShardManager mockShardManager;
+ private final ShardFetchResult firstFetch;
+ private final long emptyPollMillisBehind;
+
+ SequentialFetchConsumeKinesis(final KinesisShardManager
mockShardManager, final ShardFetchResult firstFetch,
+ final long emptyPollMillisBehind) {
+ this.mockShardManager = mockShardManager;
+ this.firstFetch = firstFetch;
+ this.emptyPollMillisBehind = emptyPollMillisBehind;
+ }
+
+ @Override
+ protected KinesisShardManager createShardManager(final KinesisClient
kinesisClient, final DynamoDbClient dynamoDbClient,
+ final ComponentLog logger, final String checkpointTableName,
final String streamName) {
+ return mockShardManager;
+ }
+
+ @Override
+ protected KinesisConsumerClient createConsumerClient(final
KinesisClient kinesisClient, final ComponentLog logger,
+ final boolean efoMode) {
+ return new
SequentialFetchConsumerClient(mock(KinesisClient.class), logger, firstFetch,
emptyPollMillisBehind);
+ }
+ }
+
+ static class SequentialFetchConsumerClient extends StubConsumerClient {
+ private final ShardFetchResult firstFetch;
+ private final long emptyPollMillisBehind;
+ private int fetchCycles;
+
+ SequentialFetchConsumerClient(final KinesisClient kinesisClient, final
ComponentLog logger,
+ final ShardFetchResult firstFetch, final long
emptyPollMillisBehind) {
+ super(kinesisClient, logger);
+ this.firstFetch = firstFetch;
+ this.emptyPollMillisBehind = emptyPollMillisBehind;
+ }
+
+ @Override
+ void startFetches(final List<Shard> shards, final String streamName,
final int batchSize,
+ final String initialStreamPosition, final KinesisShardManager
shardManager) {
+ if (fetchCycles++ == 0) {
+ enqueueResult(firstFetch);
+ recordShardLag(firstFetch.shardId(),
firstFetch.millisBehindLatest());
+ } else {
+ recordShardLag(firstFetch.shardId(), emptyPollMillisBehind);
+ }
+ }
+ }
+
static class LagReportingConsumeKinesis extends ConsumeKinesis {
private final KinesisShardManager mockShardManager;
private final Map<String, Long> shardLag;
