You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@nifi.apache.org by ex...@apache.org on 2025/10/17 18:55:19 UTC
(nifi) branch main updated: NIFI-14727 Added new ConsumeKinesis Processor (#10077)
This is an automated email from the ASF dual-hosted git repository.
exceptionfactory 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 939931fc2f NIFI-14727 Added new ConsumeKinesis Processor (#10077)
939931fc2f is described below
commit 939931fc2fc9245f31dbf69b8b8903f5eb31e17c
Author: Alaksiej Ščarbaty <65...@users.noreply.github.com>
AuthorDate: Fri Oct 17 20:55:13 2025 +0200
NIFI-14727 Added new ConsumeKinesis Processor (#10077)
Co-authored-by: ForThoseWhoComeAfter <mb...@gmail.com>
Co-authored-by: Mark Payne <ma...@hotmail.com>
Signed-off-by: David Handermann <ex...@apache.org>
---
nifi-assembly/pom.xml | 6 +
.../nifi-aws-abstract-processors/pom.xml | 5 +
.../processors/aws/v2/AbstractAwsProcessor.java | 1 +
.../nifi-aws-bundle/nifi-aws-kinesis-nar/pom.xml | 252 +++++++
.../src/main/resources/META-INF/LICENSE | 209 ++++++
.../src/main/resources/META-INF/NOTICE | 386 +++++++++++
.../nifi-aws-bundle/nifi-aws-kinesis/pom.xml | 115 ++++
.../processors/aws/kinesis/ConsumeKinesis.java | 749 +++++++++++++++++++++
.../aws/kinesis/ConsumeKinesisAttributes.java | 80 +++
.../aws/kinesis/MemoryBoundRecordBuffer.java | 684 +++++++++++++++++++
.../aws/kinesis/ReaderRecordProcessor.java | 265 ++++++++
.../nifi/processors/aws/kinesis/RecordBuffer.java | 96 +++
.../services/org.apache.nifi.processor.Processor | 16 +
.../additionalDetails.md | 136 ++++
.../processors/aws/kinesis/ConsumeKinesisIT.java | 642 ++++++++++++++++++
.../processors/aws/kinesis/ConsumeKinesisTest.java | 60 ++
.../processors/aws/kinesis/JsonRecordAssert.java | 73 ++
.../aws/kinesis/KinesisRecordPayload.java | 36 +
.../aws/kinesis/MemoryBoundRecordBufferTest.java | 734 ++++++++++++++++++++
.../aws/kinesis/ReaderRecordProcessorTest.java | 410 +++++++++++
.../nifi-aws-bundle/nifi-aws-nar/pom.xml | 6 -
.../nifi-aws-bundle/nifi-aws-processors/pom.xml | 7 +-
.../pom.xml | 20 +-
.../nifi/processors/aws/region}/RegionUtilV2.java | 2 +-
.../nifi-aws-service-api-nar/pom.xml | 32 +
.../nifi-aws-bundle/nifi-aws-service-api/pom.xml | 6 -
nifi-extension-bundles/nifi-aws-bundle/pom.xml | 53 +-
27 files changed, 5037 insertions(+), 44 deletions(-)
diff --git a/nifi-assembly/pom.xml b/nifi-assembly/pom.xml
index 536ef5789f..7528009443 100644
--- a/nifi-assembly/pom.xml
+++ b/nifi-assembly/pom.xml
@@ -423,6 +423,12 @@ language governing permissions and limitations under the License. -->
<version>2.7.0-SNAPSHOT</version>
<type>nar</type>
</dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-kinesis-nar</artifactId>
+ <version>2.7.0-SNAPSHOT</version>
+ <type>nar</type>
+ </dependency>
<dependency>
<groupId>org.apache.nifi</groupId>
<artifactId>nifi-network-processors-nar</artifactId>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/pom.xml
index d771f985ff..946cfb6541 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/pom.xml
@@ -94,6 +94,11 @@
<artifactId>commons-lang3</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-regions</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.nifi</groupId>
<artifactId>nifi-aws-service-api</artifactId>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/AbstractAwsProcessor.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/AbstractAwsProcessor.java
index 11214aee72..707f150c18 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/AbstractAwsProcessor.java
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/AbstractAwsProcessor.java
@@ -37,6 +37,7 @@ import org.apache.nifi.processor.VerifiableProcessor;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.util.StandardValidators;
import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService;
+import org.apache.nifi.processors.aws.region.RegionUtilV2;
import org.apache.nifi.proxy.ProxyConfiguration;
import org.apache.nifi.proxy.ProxySpec;
import org.apache.nifi.ssl.SSLContextProvider;
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/pom.xml
new file mode 100644
index 0000000000..38e23f0014
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/pom.xml
@@ -0,0 +1,252 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ 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.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+
+ <parent>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-bundle</artifactId>
+ <version>2.7.0-SNAPSHOT</version>
+ </parent>
+
+ <artifactId>nifi-aws-kinesis-nar</artifactId>
+ <packaging>nar</packaging>
+
+ <dependencyManagement>
+ <dependencies>
+ <!-- Mark AWS SDK v2 dependencies as provided from nifi-aws-service-api-nar -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>sdk-core</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>aws-core</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>auth</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>regions</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>profiles</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>utils</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>http-client-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>annotations</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>json-utils</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>metrics-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>endpoints-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>apache-client</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>aws-query-protocol</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>protocol-core</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>aws-json-protocol</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>sts</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>third-party-jackson-core</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Additional AWS SDK v2 core components -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>utils-lite</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>http-auth-aws</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>http-auth-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>http-auth</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>http-auth-aws-eventstream</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>identity-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>retries-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>retries</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>checksums-spi</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>checksums</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>netty-nio-client</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- AWS EventStream -->
+ <dependency>
+ <groupId>software.amazon.eventstream</groupId>
+ <artifactId>eventstream</artifactId>
+ <version>1.0.1</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Apache HTTP Components -->
+ <dependency>
+ <groupId>org.apache.httpcomponents</groupId>
+ <artifactId>httpclient</artifactId>
+ <version>${org.apache.httpcomponents.httpclient.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.httpcomponents</groupId>
+ <artifactId>httpcore</artifactId>
+ <version>${org.apache.httpcomponents.httpcore.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Google Guava and related -->
+ <dependency>
+ <groupId>com.google.guava</groupId>
+ <artifactId>guava</artifactId>
+ <version>${guava.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Reactive Streams -->
+ <dependency>
+ <groupId>org.reactivestreams</groupId>
+ <artifactId>reactive-streams</artifactId>
+ <version>1.0.4</version>
+ <scope>provided</scope>
+ </dependency>
+ </dependencies>
+ </dependencyManagement>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-kinesis</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-service-api-nar</artifactId>
+ <version>${project.version}</version>
+ <type>nar</type>
+ </dependency>
+ </dependencies>
+</project>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/src/main/resources/META-INF/LICENSE b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/src/main/resources/META-INF/LICENSE
new file mode 100644
index 0000000000..44893cdb29
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/src/main/resources/META-INF/LICENSE
@@ -0,0 +1,209 @@
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed 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.
+
+APACHE NIFI SUBCOMPONENTS:
+
+The Apache NiFi project contains subcomponents with separate copyright
+notices and license terms. Your use of the source code for the these
+subcomponents is subject to the terms and conditions of the following
+licenses.
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/src/main/resources/META-INF/NOTICE b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/src/main/resources/META-INF/NOTICE
new file mode 100644
index 0000000000..504e484ed0
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis-nar/src/main/resources/META-INF/NOTICE
@@ -0,0 +1,386 @@
+nifi-aws-kinesis-3-nar
+Copyright 2015-2025 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).
+
+******************
+Apache Software License v2
+******************
+
+The following binary components are provided under the Apache Software License v2
+
+ (ASLv2) Amazon Event Stream
+ The following NOTICE information applies:
+ AWS EventStream for Java
+ Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+ (ASLv2) Amazon Kinesis Client Library
+ The following NOTICE information applies:
+ AmazonKinesisClientLibrary
+ Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+ (ASLv2) Apache Commons Collections
+ The following NOTICE information applies:
+ Apache Commons Collections
+ Copyright 2001-2015 The Apache Software Foundation
+
+ This product includes software developed by
+ The Apache Software Foundation (http://www.apache.org/).
+
+ (ASLv2) Apache Commons Collections 4
+ The following NOTICE information applies:
+ Apache Commons Collections
+ Copyright 2001-2025 The Apache Software Foundation
+
+ This product includes software developed at
+ The Apache Software Foundation (https://www.apache.org/).
+
+ (ASLv2) AWS SDK for Java V2
+ The following NOTICE information applies:
+ AWS SDK for Java 2.0
+ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+ This product includes software developed by
+ Amazon Technologies, Inc (http://www.amazon.com/).
+
+ **********************
+ THIRD PARTY COMPONENTS
+ **********************
+ This software includes third party software subject to the following copyrights:
+ - XML parsing and utility functions from JetS3t - Copyright 2006-2009 James Murty.
+ - PKCS#1 PEM encoded private key parsing and utility functions from oauth.googlecode.com - Copyright 1998-2010 AOL Inc.
+ - Apache Commons Lang - https://github.com/apache/commons-lang
+ - Netty Reactive Streams - https://github.com/playframework/netty-reactive-streams
+ - Jackson-core - https://github.com/FasterXML/jackson-core
+ - Jackson-dataformat-cbor - https://github.com/FasterXML/jackson-dataformats-binary
+
+ The licenses for these third party components are included in LICENSE.txt
+
+ - For Apache Commons Lang see also this required NOTICE:
+ Apache Commons Lang
+ Copyright 2001-2020 The Apache Software Foundation
+
+ This product includes software developed at
+ The Apache Software Foundation (https://www.apache.org/).
+
+ (ASLv2) DynamoDB Lock Client
+ The following NOTICE information applies:
+ DynamoDB Lock Client
+ Copyright 2013 - 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+ (ASLv2) Google Guava
+ The following NOTICE information applies:
+ Google Guava
+ Copyright (C) 2024 The Guava Authors
+
+ (ASLv2) JSpecify
+ The following NOTICE information applies:
+ JSpecify
+ Copyright 2022 The JSpecify Authors
+
+ (ASLv2) RxJava
+ The following NOTICE information applies:
+ RxJava
+ Copyright (c) 2016-present, RxJava Contributors.
+
+ (ASLv2) The Netty Project
+ The following NOTICE information applies:
+
+ The Netty Project
+ =================
+
+ Please visit the Netty web site for more information:
+
+ * https://netty.io/
+
+ Copyright 2014 The Netty Project
+
+ The Netty Project 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:
+
+ https://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.
+
+ Also, please refer to each LICENSE.<component>.txt file, which is located in
+ the 'license' directory of the distribution file, for the license terms of the
+ components that this product depends on.
+
+ -------------------------------------------------------------------------------
+ This product contains the extensions to Java Collections Framework which has
+ been derived from the works by JSR-166 EG, Doug Lea, and Jason T. Greene:
+
+ * LICENSE:
+ * license/LICENSE.jsr166y.txt (Public Domain)
+ * HOMEPAGE:
+ * http://gee.cs.oswego.edu/cgi-bin/viewcvs.cgi/jsr166/
+ * http://viewvc.jboss.org/cgi-bin/viewvc.cgi/jbosscache/experimental/jsr166/
+
+ This product contains a modified version of Robert Harder's Public Domain
+ Base64 Encoder and Decoder, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.base64.txt (Public Domain)
+ * HOMEPAGE:
+ * http://iharder.sourceforge.net/current/java/base64/
+
+ This product contains a modified portion of 'Webbit', an event based
+ WebSocket and HTTP server, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.webbit.txt (BSD License)
+ * HOMEPAGE:
+ * https://github.com/joewalnes/webbit
+
+ This product contains a modified portion of 'SLF4J', a simple logging
+ facade for Java, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.slf4j.txt (MIT License)
+ * HOMEPAGE:
+ * https://www.slf4j.org/
+
+ This product contains a modified portion of 'Apache Harmony', an open source
+ Java SE, which can be obtained at:
+
+ * NOTICE:
+ * license/NOTICE.harmony.txt
+ * LICENSE:
+ * license/LICENSE.harmony.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://archive.apache.org/dist/harmony/
+
+ This product contains a modified portion of 'jbzip2', a Java bzip2 compression
+ and decompression library written by Matthew J. Francis. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.jbzip2.txt (MIT License)
+ * HOMEPAGE:
+ * https://code.google.com/p/jbzip2/
+
+ This product contains a modified portion of 'libdivsufsort', a C API library to construct
+ the suffix array and the Burrows-Wheeler transformed string for any input string of
+ a constant-size alphabet written by Yuta Mori. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.libdivsufsort.txt (MIT License)
+ * HOMEPAGE:
+ * https://github.com/y-256/libdivsufsort
+
+ This product contains a modified portion of Nitsan Wakart's 'JCTools', Java Concurrency Tools for the JVM,
+ which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.jctools.txt (ASL2 License)
+ * HOMEPAGE:
+ * https://github.com/JCTools/JCTools
+
+ This product optionally depends on 'JZlib', a re-implementation of zlib in
+ pure Java, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.jzlib.txt (BSD style License)
+ * HOMEPAGE:
+ * http://www.jcraft.com/jzlib/
+
+ This product optionally depends on 'Compress-LZF', a Java library for encoding and
+ decoding data in LZF format, written by Tatu Saloranta. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.compress-lzf.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/ning/compress
+
+ This product optionally depends on 'lz4', a LZ4 Java compression
+ and decompression library written by Adrien Grand. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.lz4.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/jpountz/lz4-java
+
+ This product optionally depends on 'lzma-java', a LZMA Java compression
+ and decompression library, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.lzma-java.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/jponge/lzma-java
+
+ This product optionally depends on 'zstd-jni', a zstd-jni Java compression
+ and decompression library, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.zstd-jni.txt (BSD)
+ * HOMEPAGE:
+ * https://github.com/luben/zstd-jni
+
+ This product contains a modified portion of 'jfastlz', a Java port of FastLZ compression
+ and decompression library written by William Kinney. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.jfastlz.txt (MIT License)
+ * HOMEPAGE:
+ * https://code.google.com/p/jfastlz/
+
+ This product contains a modified portion of and optionally depends on 'Protocol Buffers', Google's data
+ interchange format, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.protobuf.txt (New BSD License)
+ * HOMEPAGE:
+ * https://github.com/google/protobuf
+
+ This product optionally depends on 'Bouncy Castle Crypto APIs' to generate
+ a temporary self-signed X.509 certificate when the JVM does not provide the
+ equivalent functionality. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.bouncycastle.txt (MIT License)
+ * HOMEPAGE:
+ * https://www.bouncycastle.org/
+
+ This product optionally depends on 'Snappy', a compression library produced
+ by Google Inc, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.snappy.txt (New BSD License)
+ * HOMEPAGE:
+ * https://github.com/google/snappy
+
+ This product optionally depends on 'JBoss Marshalling', an alternative Java
+ serialization API, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.jboss-marshalling.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/jboss-remoting/jboss-marshalling
+
+ This product optionally depends on 'Caliper', Google's micro-
+ benchmarking framework, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.caliper.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/google/caliper
+
+ This product optionally depends on 'Apache Commons Logging', a logging
+ framework, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.commons-logging.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://commons.apache.org/logging/
+
+ This product optionally depends on 'Apache Log4J', a logging framework, which
+ can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.log4j.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://logging.apache.org/log4j/
+
+ This product optionally depends on 'Aalto XML', an ultra-high performance
+ non-blocking XML processor, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.aalto-xml.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://wiki.fasterxml.com/AaltoHome
+
+ This product contains a modified version of 'HPACK', a Java implementation of
+ the HTTP/2 HPACK algorithm written by Twitter. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.hpack.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/twitter/hpack
+
+ This product contains a modified version of 'HPACK', a Java implementation of
+ the HTTP/2 HPACK algorithm written by Cory Benfield. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.hyper-hpack.txt (MIT License)
+ * HOMEPAGE:
+ * https://github.com/python-hyper/hpack/
+
+ This product contains a modified version of 'HPACK', a Java implementation of
+ the HTTP/2 HPACK algorithm written by Tatsuhiro Tsujikawa. It can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.nghttp2-hpack.txt (MIT License)
+ * HOMEPAGE:
+ * https://github.com/nghttp2/nghttp2/
+
+ This product contains a modified portion of 'Apache Commons Lang', a Java library
+ provides utilities for the java.lang API, which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.commons-lang.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://commons.apache.org/proper/commons-lang/
+
+
+ This product contains the Maven wrapper scripts from 'Maven Wrapper', that provides an easy way to ensure a user has everything necessary to run the Maven build.
+
+ * LICENSE:
+ * license/LICENSE.mvn-wrapper.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/takari/maven-wrapper
+
+ This product contains the dnsinfo.h header file, that provides a way to retrieve the system DNS configuration on MacOS.
+ This private header is also used by Apple's open source
+ mDNSResponder (https://opensource.apple.com/tarballs/mDNSResponder/).
+
+ * LICENSE:
+ * license/LICENSE.dnsinfo.txt (Apple Public Source License 2.0)
+ * HOMEPAGE:
+ * https://www.opensource.apple.com/source/configd/configd-453.19/dnsinfo/dnsinfo.h
+
+ This product optionally depends on 'Brotli4j', Brotli compression and
+ decompression for Java., which can be obtained at:
+
+ * LICENSE:
+ * license/LICENSE.brotli4j.txt (Apache License 2.0)
+ * HOMEPAGE:
+ * https://github.com/hyperxpro/Brotli4j
+
+ This product is statically linked against Quiche.
+
+ * LICENSE:
+ * license/LICENSE.quiche.txt (BSD2)
+ * HOMEPAGE:
+ * https://github.com/cloudflare/quiche
+
+
+ This product is statically linked against boringssl.
+
+ * LICENSE (Combination ISC and OpenSSL license)
+ * license/LICENSE.boringssl.txt (Combination ISC and OpenSSL license)
+ * HOMEPAGE:
+ * https://boringssl.googlesource.com/boringssl/
+
+************************
+Berkeley Software Distribution License 3-Clause
+************************
+
+The following binary components are provided under the BSD License. See project link for details.
+
+ (BSD License) Protocol Buffers
+
+******************
+MIT No Attribution License
+******************
+
+The following binary components are provided under the MIT No Attribution License
+
+ (MIT-0) Reactive Streams
+ The following NOTICE information applies:
+ Reactive Streams
+ Copyright 2014 Reactive Streams
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/pom.xml
new file mode 100644
index 0000000000..30173e880b
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/pom.xml
@@ -0,0 +1,115 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ 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.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+
+ <parent>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-bundle</artifactId>
+ <version>2.7.0-SNAPSHOT</version>
+ </parent>
+
+ <artifactId>nifi-aws-kinesis</artifactId>
+ <packaging>jar</packaging>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-regions</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-service-api</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-proxy-configuration-api</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-record-serialization-service-api</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>software.amazon.kinesis</groupId>
+ <artifactId>amazon-kinesis-client</artifactId>
+ <version>3.1.3</version>
+ <exclusions>
+ <exclusion>
+ <groupId>com.google.code.findbugs</groupId>
+ <artifactId>jsr305</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>commons-logging</groupId>
+ <artifactId>commons-logging</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.json</groupId>
+ <artifactId>json</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>software.amazon.glue</groupId>
+ <artifactId>schema-registry-common</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>software.amazon.glue</groupId>
+ <artifactId>schema-registry-serde</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>netty-nio-client</artifactId>
+ </dependency>
+
+ <!-- TEST DEPENDENCIES -->
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-aws-processors</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-mock</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-mock-record-utils</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-record-serialization-services</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-ssl-context-service-api</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.testcontainers</groupId>
+ <artifactId>testcontainers-localstack</artifactId>
+ <scope>test</scope>
+ </dependency>
+ </dependencies>
+</project>
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
new file mode 100644
index 0000000000..8700ac343f
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesis.java
@@ -0,0 +1,749 @@
+/*
+ * 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;
+
+import jakarta.annotation.Nullable;
+import org.apache.nifi.annotation.behavior.InputRequirement;
+import org.apache.nifi.annotation.behavior.SystemResource;
+import org.apache.nifi.annotation.behavior.SystemResourceConsideration;
+import org.apache.nifi.annotation.behavior.WritesAttribute;
+import org.apache.nifi.annotation.behavior.WritesAttributes;
+import org.apache.nifi.annotation.configuration.DefaultSettings;
+import org.apache.nifi.annotation.documentation.CapabilityDescription;
+import org.apache.nifi.annotation.documentation.Tags;
+import org.apache.nifi.annotation.lifecycle.OnScheduled;
+import org.apache.nifi.annotation.lifecycle.OnStopped;
+import org.apache.nifi.components.DescribedValue;
+import org.apache.nifi.components.PropertyDescriptor;
+import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.processor.AbstractProcessor;
+import org.apache.nifi.processor.DataUnit;
+import org.apache.nifi.processor.ProcessContext;
+import org.apache.nifi.processor.ProcessSession;
+import org.apache.nifi.processor.Relationship;
+import org.apache.nifi.processor.exception.ProcessException;
+import org.apache.nifi.processor.util.StandardValidators;
+import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService;
+import org.apache.nifi.processors.aws.kinesis.MemoryBoundRecordBuffer.Lease;
+import org.apache.nifi.processors.aws.kinesis.ReaderRecordProcessor.ProcessingResult;
+import org.apache.nifi.processors.aws.kinesis.RecordBuffer.ShardBufferId;
+import org.apache.nifi.processors.aws.region.RegionUtilV2;
+import org.apache.nifi.proxy.ProxyConfiguration;
+import org.apache.nifi.proxy.ProxyConfigurationService;
+import org.apache.nifi.proxy.ProxySpec;
+import org.apache.nifi.serialization.RecordReaderFactory;
+import org.apache.nifi.serialization.RecordSetWriterFactory;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.http.Protocol;
+import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
+import software.amazon.awssdk.http.nio.netty.Http2Configuration;
+import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.cloudwatch.CloudWatchAsyncClient;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
+import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
+import software.amazon.awssdk.services.kinesis.KinesisAsyncClientBuilder;
+import software.amazon.kinesis.common.ConfigsBuilder;
+import software.amazon.kinesis.common.InitialPositionInStream;
+import software.amazon.kinesis.common.InitialPositionInStreamExtended;
+import software.amazon.kinesis.coordinator.Scheduler;
+import software.amazon.kinesis.lifecycle.events.InitializationInput;
+import software.amazon.kinesis.lifecycle.events.LeaseLostInput;
+import software.amazon.kinesis.lifecycle.events.ProcessRecordsInput;
+import software.amazon.kinesis.lifecycle.events.ShardEndedInput;
+import software.amazon.kinesis.lifecycle.events.ShutdownRequestedInput;
+import software.amazon.kinesis.metrics.LogMetricsFactory;
+import software.amazon.kinesis.metrics.MetricsFactory;
+import software.amazon.kinesis.metrics.NullMetricsFactory;
+import software.amazon.kinesis.processor.ShardRecordProcessor;
+import software.amazon.kinesis.processor.ShardRecordProcessorFactory;
+import software.amazon.kinesis.processor.SingleStreamTracker;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.net.URI;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Date;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
+
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.APPROXIMATE_ARRIVAL_TIMESTAMP;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.FIRST_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.FIRST_SUB_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.LAST_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.LAST_SUB_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.MIME_TYPE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.PARTITION_KEY;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_COUNT;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_ERROR_MESSAGE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.SHARD_ID;
+
+@InputRequirement(InputRequirement.Requirement.INPUT_FORBIDDEN)
+@Tags({"amazon", "aws", "kinesis", "consume", "stream", "record"})
+@CapabilityDescription("""
+ Consumes data from the specified AWS Kinesis stream and outputs a FlowFile for every processed Record (raw)
+ or a FlowFile for a batch of processed records if a Record Reader and Record Writer are configured.
+ The processor may take a few minutes on the first start and several seconds on subsequent starts
+ to initialise before starting to fetch data.
+ Uses DynamoDB for check pointing and coordination, and (optional) CloudWatch for metrics.
+ """)
+@WritesAttributes({
+ @WritesAttribute(attribute = ConsumeKinesisAttributes.STREAM_NAME,
+ description = "The name of the Kinesis Stream from which all Kinesis Records in the Flow File were read"),
+ @WritesAttribute(attribute = SHARD_ID,
+ description = "Shard ID from which all Kinesis Records in the Flow File were read"),
+ @WritesAttribute(attribute = PARTITION_KEY,
+ description = "Partition key of the last Kinesis Record in the Flow File"),
+ @WritesAttribute(attribute = FIRST_SEQUENCE_NUMBER,
+ description = "A Sequence Number of the first Kinesis Record in the Flow File"),
+ @WritesAttribute(attribute = FIRST_SUB_SEQUENCE_NUMBER,
+ description = "A SubSequence Number of the first Kinesis Record in the Flow File. Generated by KPL when aggregating records into a single Kinesis Record"),
+ @WritesAttribute(attribute = LAST_SEQUENCE_NUMBER,
+ description = "A Sequence Number of the last Kinesis Record in the Flow File"),
+ @WritesAttribute(attribute = LAST_SUB_SEQUENCE_NUMBER,
+ description = "A SubSequence Number of the last Kinesis Record in the Flow File. Generated by KPL when aggregating records into a single Kinesis Record"),
+ @WritesAttribute(attribute = APPROXIMATE_ARRIVAL_TIMESTAMP,
+ description = "Approximate arrival timestamp of the last Kinesis Record in the Flow File"),
+ @WritesAttribute(attribute = MIME_TYPE,
+ description = "Sets the mime.type attribute to the MIME Type specified by the Record Writer (if configured)"),
+ @WritesAttribute(attribute = RECORD_COUNT,
+ description = "Number of records written to the FlowFiles by the Record Writer (if configured)"),
+ @WritesAttribute(attribute = RECORD_ERROR_MESSAGE,
+ description = "This attribute provides on failure the error message encountered by the Record Reader or Record Writer (if configured)")
+})
+@DefaultSettings(yieldDuration = "100 millis")
+@SystemResourceConsideration(resource = SystemResource.CPU, description = """
+ The processor uses additional CPU resources when consuming data from Kinesis.
+ The consumption is started immediately after this Processor is scheduled. The consumption ends only when the Processor is stopped.""")
+@SystemResourceConsideration(resource = SystemResource.NETWORK, description = """
+ The processor will continually poll for new Records,
+ requesting up to a maximum number of Records/bytes per call. This can result in sustained network usage.""")
+@SystemResourceConsideration(resource = SystemResource.MEMORY, description = """
+ ConsumeKinesis buffers Kinesis Records in memory until they can be processed.
+ The maximum size of the buffer is controlled by the 'Max Bytes to Buffer' property.
+ In addition, the processor may cache some amount of data for each shard when the processor's buffer is full.""")
+public class ConsumeKinesis extends AbstractProcessor {
+
+ private static final Duration HTTP_CLIENTS_CONNECTION_TIMEOUT = Duration.ofSeconds(30);
+ private static final Duration HTTP_CLIENTS_READ_TIMEOUT = Duration.ofMinutes(3);
+
+ /**
+ * Best balance between throughput and CPU usage by KCL.
+ */
+ private static final int KINESIS_HTTP_CLIENT_CONCURRENCY_PER_TASK = 16;
+ private static final int KINESIS_HTTP_CLIENT_WINDOW_SIZE_BYTES = 512 * 1024; // 512 KiB
+ private static final Duration KINESIS_HTTP_HEALTH_CHECK_PERIOD = Duration.ofMinutes(1);
+
+ private static final Duration KINESIS_SCHEDULER_GRACEFUL_SHUTDOWN_TIMEOUT = Duration.ofMinutes(3);
+
+ static final PropertyDescriptor STREAM_NAME = new PropertyDescriptor.Builder()
+ .name("Stream Name")
+ .description("The name of the Kinesis stream to consume from.")
+ .required(true)
+ .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
+ .build();
+
+ static final PropertyDescriptor APPLICATION_NAME = new PropertyDescriptor.Builder()
+ .name("Application Name")
+ .description("The name of the Kinesis application. This is used for DynamoDB table naming and worker coordination.")
+ .required(true)
+ .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
+ .build();
+
+ static final PropertyDescriptor AWS_CREDENTIALS_PROVIDER_SERVICE = new PropertyDescriptor.Builder()
+ .name("AWS Credentials Provider Service")
+ .description("""
+ The Controller Service that is used to obtain AWS credentials provider.
+ Ensure that the credentials provided have access to Kinesis, DynamoDB and (optional) CloudWatch.
+ (See processor's additional details for more information.)
+ """)
+ .required(true)
+ .identifiesControllerService(AWSCredentialsProviderService.class)
+ .build();
+
+ static final PropertyDescriptor REGION = new PropertyDescriptor.Builder()
+ .name("Region")
+ .description("AWS Region in which the Kinesis stream is located.")
+ .required(true)
+ .allowableValues(RegionUtilV2.getAvailableRegions())
+ .defaultValue(RegionUtilV2.createAllowableValue(Region.US_WEST_2).getValue())
+ .build();
+
+ static final PropertyDescriptor PROCESSING_STRATEGY = new PropertyDescriptor.Builder()
+ .name("Processing Strategy")
+ .description("Strategy for processing Kinesis Records and writing serialized output to FlowFiles.")
+ .required(true)
+ .allowableValues(ProcessingStrategy.class)
+ .defaultValue(ProcessingStrategy.FLOW_FILE)
+ .build();
+
+ static final PropertyDescriptor RECORD_READER = new PropertyDescriptor.Builder()
+ .name("Record Reader")
+ .description("""
+ The Record Reader to use for parsing the data received from Kinesis.
+
+ The Record Reader is responsible for providing schemas for the records. If the schemas change frequently,
+ it might hinder performance of the processor. (See processor's additional details for more information.)
+ """)
+ .required(true)
+ .dependsOn(PROCESSING_STRATEGY, ProcessingStrategy.RECORD)
+ .identifiesControllerService(RecordReaderFactory.class)
+ .build();
+
+ static final PropertyDescriptor RECORD_WRITER = new PropertyDescriptor.Builder()
+ .name("Record Writer")
+ .description("The Record Writer to use for serializing records.")
+ .required(true)
+ .dependsOn(PROCESSING_STRATEGY, ProcessingStrategy.RECORD)
+ .identifiesControllerService(RecordSetWriterFactory.class)
+ .build();
+
+ static final PropertyDescriptor INITIAL_STREAM_POSITION = new PropertyDescriptor.Builder()
+ .name("Initial Stream Position")
+ .description("The position in the stream where the processor should start reading.")
+ .required(true)
+ .allowableValues(InitialPosition.class)
+ .defaultValue(InitialPosition.TRIM_HORIZON)
+ .build();
+
+ static final PropertyDescriptor STREAM_POSITION_TIMESTAMP = new PropertyDescriptor.Builder()
+ .name("Stream Position Timestamp")
+ .description("Timestamp position in stream from which to start reading Kinesis Records. The timestamp must be in ISO 8601 format.")
+ .required(true)
+ .addValidator(StandardValidators.ISO8601_INSTANT_VALIDATOR)
+ .dependsOn(INITIAL_STREAM_POSITION, InitialPosition.AT_TIMESTAMP)
+ .build();
+
+ static final PropertyDescriptor MAX_BYTES_TO_BUFFER = new PropertyDescriptor.Builder()
+ .name("Max Bytes to Buffer")
+ .description("""
+ The maximum size of Kinesis Records that can be buffered in memory before being processed by NiFi.
+ If the buffer size exceeds the limit, the processor will stop consuming new records until free space is available.
+
+ Using a larger value may increase the throughput, but will do so at the expense of using more memory.
+ """)
+ .required(true)
+ .addValidator(StandardValidators.DATA_SIZE_VALIDATOR)
+ .defaultValue("100 MB")
+ .build();
+
+ static final PropertyDescriptor CHECKPOINT_INTERVAL = new PropertyDescriptor.Builder()
+ .name("Checkpoint Interval")
+ .description("""
+ Interval between checkpointing consumed Kinesis records. To checkpoint records each time the Processor is run, set this value to 0 seconds.
+
+ More frequent checkpoint may reduce performance and increase DynamoDB costs,
+ but less frequent checkpointing may result in duplicates when a Shard lease is lost or NiFi is restarted.
+ """)
+ .required(true)
+ .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR)
+ .defaultValue("5 sec")
+ .build();
+
+ static final PropertyDescriptor METRICS_PUBLISHING = new PropertyDescriptor.Builder()
+ .name("Metrics Publishing")
+ .description("Specifies where Kinesis usage metrics are published to.")
+ .required(true)
+ .allowableValues(MetricsPublishing.class)
+ .defaultValue(MetricsPublishing.DISABLED)
+ .build();
+
+ static final PropertyDescriptor PROXY_CONFIGURATION_SERVICE = ProxyConfiguration.createProxyConfigPropertyDescriptor(ProxySpec.HTTP, ProxySpec.HTTP_AUTH);
+
+ private static final List<PropertyDescriptor> PROPERTY_DESCRIPTORS = List.of(
+ STREAM_NAME,
+ APPLICATION_NAME,
+ AWS_CREDENTIALS_PROVIDER_SERVICE,
+ REGION,
+ PROCESSING_STRATEGY,
+ RECORD_READER,
+ RECORD_WRITER,
+ INITIAL_STREAM_POSITION,
+ STREAM_POSITION_TIMESTAMP,
+ MAX_BYTES_TO_BUFFER,
+ CHECKPOINT_INTERVAL,
+ PROXY_CONFIGURATION_SERVICE,
+ METRICS_PUBLISHING
+ );
+
+ static final Relationship REL_SUCCESS = new Relationship.Builder()
+ .name("success")
+ .description("FlowFiles that are created when records are successfully read from Kinesis and parsed")
+ .build();
+
+ static final Relationship REL_PARSE_FAILURE = new Relationship.Builder()
+ .name("parse.failure")
+ .description("FlowFiles that failed to parse using the configured Record Reader")
+ .build();
+
+ private static final Set<Relationship> RAW_FILE_RELATIONSHIPS = Set.of(REL_SUCCESS);
+ private static final Set<Relationship> RECORD_FILE_RELATIONSHIPS = Set.of(REL_SUCCESS, REL_PARSE_FAILURE);
+
+ private volatile DynamoDbAsyncClient dynamoDbClient;
+ private volatile CloudWatchAsyncClient cloudWatchClient;
+ private volatile KinesisAsyncClient kinesisClient;
+ private volatile Scheduler kinesisScheduler;
+
+ private volatile String streamName;
+ private volatile RecordBuffer.ForProcessor<Lease> recordBuffer;
+
+ private volatile @Nullable ReaderRecordProcessor readerRecordProcessor;
+
+ // An instance filed, so that it can be read in getRelationships.
+ private volatile ProcessingStrategy processingStrategy = ProcessingStrategy.from(
+ PROCESSING_STRATEGY.getDefaultValue());
+
+ @Override
+ protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
+ return PROPERTY_DESCRIPTORS;
+ }
+
+ @Override
+ public Set<Relationship> getRelationships() {
+ return switch (processingStrategy) {
+ case FLOW_FILE -> RAW_FILE_RELATIONSHIPS;
+ case RECORD -> RECORD_FILE_RELATIONSHIPS;
+ };
+ }
+
+ @Override
+ public void onPropertyModified(final PropertyDescriptor descriptor, final String oldValue, final String newValue) {
+ if (descriptor.equals(PROCESSING_STRATEGY)) {
+ processingStrategy = ProcessingStrategy.from(newValue);
+ }
+ }
+
+ @OnScheduled
+ public void setup(final ProcessContext context) {
+ readerRecordProcessor = switch (processingStrategy) {
+ case FLOW_FILE -> null;
+ case RECORD -> createReaderRecordProcessor(context);
+ };
+
+ final Region region = Region.of(context.getProperty(REGION).getValue());
+ final AwsCredentialsProvider credentialsProvider = context.getProperty(AWS_CREDENTIALS_PROVIDER_SERVICE)
+ .asControllerService(AWSCredentialsProviderService.class).getAwsCredentialsProvider();
+
+ kinesisClient = KinesisAsyncClient.builder()
+ .region(region)
+ .credentialsProvider(credentialsProvider)
+ .endpointOverride(getKinesisEndpointOverride())
+ .httpClient(createKinesisHttpClient(context))
+ .build();
+
+ dynamoDbClient = DynamoDbAsyncClient.builder()
+ .region(region)
+ .credentialsProvider(credentialsProvider)
+ .endpointOverride(getDynamoDbEndpointOverride())
+ .httpClient(createHttpClientBuilder(context).build())
+ .build();
+
+ cloudWatchClient = CloudWatchAsyncClient.builder()
+ .region(region)
+ .credentialsProvider(credentialsProvider)
+ .endpointOverride(getCloudwatchEndpointOverride())
+ .httpClient(createHttpClientBuilder(context).build())
+ .build();
+
+ streamName = context.getProperty(STREAM_NAME).getValue();
+ final InitialPositionInStreamExtended initialPositionExtended = getInitialPosition(context);
+ final SingleStreamTracker streamTracker = new SingleStreamTracker(streamName, initialPositionExtended);
+
+ final long maxBytesToBuffer = context.getProperty(MAX_BYTES_TO_BUFFER).asDataSize(DataUnit.B).longValue();
+ final Duration checkpointInterval = context.getProperty(CHECKPOINT_INTERVAL).asDuration();
+ final MemoryBoundRecordBuffer memoryBoundRecordBuffer = new MemoryBoundRecordBuffer(getLogger(), maxBytesToBuffer, checkpointInterval);
+ recordBuffer = memoryBoundRecordBuffer;
+ final ShardRecordProcessorFactory recordProcessorFactory = () -> new ConsumeKinesisRecordProcessor(memoryBoundRecordBuffer);
+
+ final String applicationName = context.getProperty(APPLICATION_NAME).getValue();
+ final String workerId = getIdentifier();
+ final ConfigsBuilder configsBuilder = new ConfigsBuilder(streamTracker, applicationName, kinesisClient, dynamoDbClient, cloudWatchClient, workerId, recordProcessorFactory);
+
+ final MetricsFactory metricsFactory = configureMetricsFactory(context);
+
+ kinesisScheduler = new Scheduler(
+ configsBuilder.checkpointConfig(),
+ configsBuilder.coordinatorConfig(),
+ configsBuilder.leaseManagementConfig(),
+ configsBuilder.lifecycleConfig(),
+ configsBuilder.metricsConfig().metricsFactory(metricsFactory),
+ configsBuilder.processorConfig(),
+ configsBuilder.retrievalConfig()
+ );
+
+ final String schedulerThreadName = "%s-Scheduler-%s".formatted(getClass().getSimpleName(), getIdentifier());
+ final Thread schedulerThread = new Thread(kinesisScheduler, schedulerThreadName);
+ schedulerThread.setDaemon(true);
+ schedulerThread.start();
+ // The thread is stopped when kinesisScheduler is shutdown in the onStopped method.
+ }
+
+ /**
+ * Creating Kinesis HTTP client, as per
+ * {@link software.amazon.kinesis.common.KinesisClientUtil#adjustKinesisClientBuilder(KinesisAsyncClientBuilder)}.
+ */
+ private static SdkAsyncHttpClient createKinesisHttpClient(final ProcessContext context) {
+ final int maxConcurrency = KINESIS_HTTP_CLIENT_CONCURRENCY_PER_TASK * context.getMaxConcurrentTasks();
+
+ return createHttpClientBuilder(context)
+ .protocol(Protocol.HTTP2)
+ .maxConcurrency(maxConcurrency)
+ .http2Configuration(Http2Configuration.builder()
+ .initialWindowSize(KINESIS_HTTP_CLIENT_WINDOW_SIZE_BYTES)
+ .healthCheckPingPeriod(KINESIS_HTTP_HEALTH_CHECK_PERIOD)
+ .build())
+ .build();
+ }
+
+ private static NettyNioAsyncHttpClient.Builder createHttpClientBuilder(final ProcessContext context) {
+ final NettyNioAsyncHttpClient.Builder builder = NettyNioAsyncHttpClient.builder()
+ .connectionTimeout(HTTP_CLIENTS_CONNECTION_TIMEOUT)
+ .readTimeout(HTTP_CLIENTS_READ_TIMEOUT);
+
+ final ProxyConfigurationService proxyConfigService = context.getProperty(PROXY_CONFIGURATION_SERVICE).asControllerService(ProxyConfigurationService.class);
+ if (proxyConfigService != null) {
+ final ProxyConfiguration proxyConfig = proxyConfigService.getConfiguration();
+
+ final software.amazon.awssdk.http.nio.netty.ProxyConfiguration.Builder proxyConfigBuilder = software.amazon.awssdk.http.nio.netty.ProxyConfiguration.builder()
+ .host(proxyConfig.getProxyServerHost())
+ .port(proxyConfig.getProxyServerPort());
+
+ if (proxyConfig.hasCredential()) {
+ proxyConfigBuilder.username(proxyConfig.getProxyUserName());
+ proxyConfigBuilder.password(proxyConfig.getProxyUserPassword());
+ }
+
+ builder.proxyConfiguration(proxyConfigBuilder.build());
+ }
+
+ return builder;
+ }
+
+ private ReaderRecordProcessor createReaderRecordProcessor(final ProcessContext context) {
+ final RecordReaderFactory recordReaderFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
+ final RecordSetWriterFactory recordWriterFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
+ return new ReaderRecordProcessor(recordReaderFactory, recordWriterFactory, getLogger());
+ }
+
+ private static InitialPositionInStreamExtended getInitialPosition(final ProcessContext context) {
+ final InitialPosition initialPosition = context.getProperty(INITIAL_STREAM_POSITION).asAllowableValue(InitialPosition.class);
+ return switch (initialPosition) {
+ case TRIM_HORIZON ->
+ InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.TRIM_HORIZON);
+ case LATEST -> InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.LATEST);
+ case AT_TIMESTAMP -> {
+ final String timestampValue = context.getProperty(STREAM_POSITION_TIMESTAMP).getValue();
+ final Instant timestamp = Instant.parse(timestampValue);
+ yield InitialPositionInStreamExtended.newInitialPositionAtTimestamp(Date.from(timestamp));
+ }
+ };
+ }
+
+ private static @Nullable MetricsFactory configureMetricsFactory(final ProcessContext context) {
+ final MetricsPublishing metricsPublishing = context.getProperty(METRICS_PUBLISHING).asAllowableValue(MetricsPublishing.class);
+ return switch (metricsPublishing) {
+ case DISABLED -> new NullMetricsFactory();
+ case LOGS -> new LogMetricsFactory();
+ case CLOUDWATCH -> null; // If no metrics factory was provided, CloudWatch metrics factory is used by default.
+ };
+ }
+
+ @OnStopped
+ public void onStopped() {
+ if (kinesisScheduler != null) {
+ shutdownScheduler();
+ kinesisScheduler = null;
+ }
+ if (kinesisClient != null) {
+ kinesisClient.close();
+ kinesisClient = null;
+ }
+ if (dynamoDbClient != null) {
+ dynamoDbClient.close();
+ dynamoDbClient = null;
+ }
+ if (cloudWatchClient != null) {
+ cloudWatchClient.close();
+ cloudWatchClient = null;
+ }
+
+ recordBuffer = null;
+ readerRecordProcessor = null;
+ }
+
+ private void shutdownScheduler() {
+ final long start = System.nanoTime();
+ getLogger().debug("Shutting down Kinesis Scheduler");
+
+ boolean gracefulShutdownSucceeded;
+ try {
+ gracefulShutdownSucceeded = kinesisScheduler.startGracefulShutdown().get(KINESIS_SCHEDULER_GRACEFUL_SHUTDOWN_TIMEOUT.getSeconds(), SECONDS);
+ if (!gracefulShutdownSucceeded) {
+ getLogger().warn("Failed to shutdown Kinesis Scheduler gracefully. See the logs for more details");
+ }
+ } catch (final RuntimeException | InterruptedException | ExecutionException | TimeoutException e) {
+ if (e instanceof TimeoutException) {
+ getLogger().warn("Failed to shutdown Kinesis Scheduler gracefully after {} seconds", KINESIS_SCHEDULER_GRACEFUL_SHUTDOWN_TIMEOUT.getSeconds(), e);
+ } else {
+ getLogger().warn("Failed to shutdown Kinesis Scheduler gracefully", e);
+ }
+ gracefulShutdownSucceeded = false;
+ }
+
+ if (!gracefulShutdownSucceeded) {
+ kinesisScheduler.shutdown();
+ }
+
+ final long finish = System.nanoTime();
+ getLogger().debug("Kinesis Scheduler shutdown finished after {} seconds", NANOSECONDS.toSeconds(finish - start));
+ }
+
+ @Override
+ public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
+ final Optional<Lease> leaseAcquired = recordBuffer.acquireBufferLease();
+
+ leaseAcquired.ifPresentOrElse(
+ lease -> processRecordsFromBuffer(session, lease),
+ context::yield
+ );
+ }
+
+ private void processRecordsFromBuffer(final ProcessSession session, final Lease lease) {
+ try {
+ final List<KinesisClientRecord> records = recordBuffer.consumeRecords(lease);
+
+ if (records.isEmpty()) {
+ recordBuffer.returnBufferLease(lease);
+ return;
+ }
+
+ final String shardId = lease.shardId();
+ final ReaderRecordProcessor processor = readerRecordProcessor;
+ if (processor != null) {
+ processRecordsWithReader(processor, session, shardId, records);
+ } else {
+ processRecordsAsRaw(session, shardId, records);
+ }
+
+ session.adjustCounter("Records Processed", records.size(), false);
+
+ session.commitAsync(
+ () -> commitRecords(lease),
+ __ -> rollbackRecords(lease)
+ );
+ } catch (final RuntimeException e) {
+ rollbackRecords(lease);
+ throw e;
+ }
+ }
+
+ private void commitRecords(final Lease lease) {
+ try {
+ recordBuffer.commitConsumedRecords(lease);
+ } finally {
+ recordBuffer.returnBufferLease(lease);
+ }
+ }
+
+ private void rollbackRecords(final Lease lease) {
+ try {
+ recordBuffer.rollbackConsumedRecords(lease);
+ } finally {
+ recordBuffer.returnBufferLease(lease);
+ }
+ }
+
+ private void processRecordsAsRaw(final ProcessSession session, final String shardId, final List<KinesisClientRecord> records) {
+ for (final KinesisClientRecord record : records) {
+ FlowFile flowFile = session.create();
+ flowFile = session.putAllAttributes(flowFile, ConsumeKinesisAttributes.fromKinesisRecords(streamName, shardId, record, record));
+
+ flowFile = session.write(flowFile, out -> {
+ try (final WritableByteChannel channel = Channels.newChannel(out)) {
+ channel.write(record.data());
+ }
+ });
+
+ session.transfer(flowFile, REL_SUCCESS);
+ }
+ }
+
+ private void processRecordsWithReader(final ReaderRecordProcessor recordProcessor, final ProcessSession session, final String shardId, final List<KinesisClientRecord> records) {
+ final ProcessingResult result = recordProcessor.processRecords(session, streamName, shardId, records);
+
+ session.transfer(result.successFlowFiles(), REL_SUCCESS);
+ session.transfer(result.parseFailureFlowFiles(), REL_PARSE_FAILURE);
+ }
+
+ /**
+ * An adapter between Kinesis Consumer Library and {@link RecordBuffer}.
+ */
+ private static class ConsumeKinesisRecordProcessor implements ShardRecordProcessor {
+
+ private final RecordBuffer.ForKinesisClientLibrary recordBuffer;
+ private volatile @Nullable ShardBufferId bufferId;
+
+ ConsumeKinesisRecordProcessor(final MemoryBoundRecordBuffer recordBuffer) {
+ this.recordBuffer = recordBuffer;
+ }
+
+ @Override
+ public void initialize(final InitializationInput initializationInput) {
+ bufferId = recordBuffer.createBuffer(initializationInput.shardId());
+ }
+
+ @Override
+ public void processRecords(final ProcessRecordsInput processRecordsInput) {
+ if (bufferId == null) {
+ throw new IllegalStateException("Buffer ID not found: Record Processor not initialized");
+ }
+ recordBuffer.addRecords(bufferId, processRecordsInput.records(), processRecordsInput.checkpointer());
+ }
+
+ @Override
+ public void leaseLost(final LeaseLostInput leaseLostInput) {
+ if (bufferId != null) {
+ recordBuffer.consumerLeaseLost(bufferId);
+ }
+ }
+
+ @Override
+ public void shardEnded(final ShardEndedInput shardEndedInput) {
+ if (bufferId != null) {
+ recordBuffer.checkpointEndedShard(bufferId, shardEndedInput.checkpointer());
+ }
+ }
+
+ @Override
+ public void shutdownRequested(final ShutdownRequestedInput shutdownRequestedInput) {
+ if (bufferId != null) {
+ recordBuffer.shutdownShardConsumption(bufferId, shutdownRequestedInput.checkpointer());
+ }
+ }
+ }
+
+ enum ProcessingStrategy implements DescribedValue {
+ FLOW_FILE("Write one FlowFile for each consumed Kinesis Record"),
+ RECORD("Write one FlowFile containing multiple consumed Kinesis Records processed with Record Reader and Record Writer");
+
+ private final String description;
+
+ ProcessingStrategy(final String description) {
+ this.description = description;
+ }
+
+ @Override
+ public String getValue() {
+ return name();
+ }
+
+ @Override
+ public String getDisplayName() {
+ return name();
+ }
+
+ @Override
+ public String getDescription() {
+ return description;
+ }
+
+ private static ProcessingStrategy from(final String name) {
+ // As long as getValue() returns name(), using valueOf is fine.
+ return ProcessingStrategy.valueOf(name);
+ }
+ }
+
+ enum InitialPosition implements DescribedValue {
+ TRIM_HORIZON("Trim Horizon", "Start reading at the last untrimmed record in the shard in the system, which is the oldest data record in the shard."),
+ LATEST("Latest", "Start reading just after the most recent record in the shard, so that you always read the most recent data in the shard."),
+ AT_TIMESTAMP("At Timestamp", "Start reading at the record with the specified timestamp.");
+
+ private final String displayName;
+ private final String description;
+
+ InitialPosition(final String displayName, final String description) {
+ this.displayName = displayName;
+ this.description = description;
+ }
+
+ @Override
+ public String getValue() {
+ return name();
+ }
+
+ @Override
+ public String getDisplayName() {
+ return displayName;
+ }
+
+ @Override
+ public String getDescription() {
+ return description;
+ }
+ }
+
+ enum MetricsPublishing implements DescribedValue {
+ DISABLED("Disabled", "No metrics are published"),
+ LOGS("Logs", "Metrics are published to application logs"),
+ CLOUDWATCH("CloudWatch", "Metrics are published to Amazon CloudWatch");
+
+ private final String displayName;
+ private final String description;
+
+ MetricsPublishing(final String displayName, final String description) {
+ this.displayName = displayName;
+ this.description = description;
+ }
+
+ @Override
+ public String getValue() {
+ return name();
+ }
+
+ @Override
+ public String getDisplayName() {
+ return displayName;
+ }
+
+ @Override
+ public String getDescription() {
+ return description;
+ }
+ }
+
+ // Visible for tests only.
+ @Nullable URI getKinesisEndpointOverride() {
+ return null;
+ }
+
+ // Visible for tests only.
+ @Nullable URI getDynamoDbEndpointOverride() {
+ return null;
+ }
+
+ // Visible for tests only.
+ @Nullable URI getCloudwatchEndpointOverride() {
+ return null;
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisAttributes.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisAttributes.java
new file mode 100644
index 0000000000..9011c2a367
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisAttributes.java
@@ -0,0 +1,80 @@
+/*
+ * 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;
+
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.util.HashMap;
+import java.util.Map;
+
+final class ConsumeKinesisAttributes {
+
+ private static final String PREFIX = "aws.kinesis.";
+
+ // AWS Kinesis attributes.
+ static final String STREAM_NAME = PREFIX + "stream.name";
+ static final String SHARD_ID = PREFIX + "shard.id";
+ static final String FIRST_SEQUENCE_NUMBER = PREFIX + "first.sequence.number";
+ static final String FIRST_SUB_SEQUENCE_NUMBER = PREFIX + "first.subsequence.number";
+ static final String LAST_SEQUENCE_NUMBER = PREFIX + "last.sequence.number";
+ static final String LAST_SUB_SEQUENCE_NUMBER = PREFIX + "last.subsequence.number";
+
+ static final String PARTITION_KEY = PREFIX + "partition.key";
+ static final String APPROXIMATE_ARRIVAL_TIMESTAMP = PREFIX + "approximate.arrival.timestamp.ms";
+
+ // Record attributes.
+ static final String MIME_TYPE = "mime.type";
+ static final String RECORD_COUNT = "record.count";
+ static final String RECORD_ERROR_MESSAGE = "record.error.message";
+
+ /**
+ * Creates a map of FlowFile attributes from the provided Kinesis records.
+ *
+ * @param streamName the name of the Kinesis stream the FileFile records came from.
+ * @param shardId the shard ID the FlowFile records came from.
+ * @param firstRecord the first Kinesis record in the FlowFile.
+ * @param lastRecord the last Kinesis record in the FlowFile.
+ * @return a <b>mutable</b> map with kinesis attributes.
+ */
+ static Map<String, String> fromKinesisRecords(
+ final String streamName,
+ final String shardId,
+ final KinesisClientRecord firstRecord,
+ final KinesisClientRecord lastRecord) {
+ final Map<String, String> attributes = new HashMap<>(8);
+
+ attributes.put(STREAM_NAME, streamName);
+ attributes.put(SHARD_ID, shardId);
+
+ attributes.put(FIRST_SEQUENCE_NUMBER, firstRecord.sequenceNumber());
+ attributes.put(FIRST_SUB_SEQUENCE_NUMBER, String.valueOf(firstRecord.subSequenceNumber()));
+
+ attributes.put(LAST_SEQUENCE_NUMBER, lastRecord.sequenceNumber());
+ attributes.put(LAST_SUB_SEQUENCE_NUMBER, String.valueOf(lastRecord.subSequenceNumber()));
+
+ attributes.put(PARTITION_KEY, lastRecord.partitionKey());
+
+ if (lastRecord.approximateArrivalTimestamp() != null) {
+ attributes.put(APPROXIMATE_ARRIVAL_TIMESTAMP, String.valueOf(lastRecord.approximateArrivalTimestamp().toEpochMilli()));
+ }
+
+ return attributes;
+ }
+
+ private ConsumeKinesisAttributes() {
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/MemoryBoundRecordBuffer.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/MemoryBoundRecordBuffer.java
new file mode 100644
index 0000000000..c20d695263
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/MemoryBoundRecordBuffer.java
@@ -0,0 +1,684 @@
+/*
+ * 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;
+
+import jakarta.annotation.Nullable;
+import org.apache.nifi.logging.ComponentLog;
+import org.apache.nifi.processors.aws.kinesis.RecordBuffer.ShardBufferId;
+import org.apache.nifi.processors.aws.kinesis.RecordBuffer.ShardBufferLease;
+import software.amazon.kinesis.exceptions.InvalidStateException;
+import software.amazon.kinesis.exceptions.KinesisClientLibDependencyException;
+import software.amazon.kinesis.exceptions.ShutdownException;
+import software.amazon.kinesis.exceptions.ThrottlingException;
+import software.amazon.kinesis.processor.RecordProcessorCheckpointer;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.Random;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static java.util.Collections.emptyList;
+
+/**
+ * A record buffer which limits the maximum memory usage across all shard buffers.
+ * If the memory limit is reached, adding new records will block until enough memory is freed.
+ */
+final class MemoryBoundRecordBuffer implements RecordBuffer.ForKinesisClientLibrary, RecordBuffer.ForProcessor<MemoryBoundRecordBuffer.Lease> {
+
+ private final ComponentLog logger;
+
+ private final long checkpointIntervalMillis;
+ private final BlockingMemoryTracker memoryTracker;
+
+ private final AtomicLong bufferIdCounter = new AtomicLong(0);
+
+ /**
+ * All shard buffers stored by their ids.
+ * <p>
+ * When a buffer is invalidated, it is removed from this map, but its id may still be present in the buffersToLease queue.
+ * Since the buffer can be invalidated concurrently, it's possible for some buffer operations to be called
+ * after the buffer was removed from this map. In that case the operations should take no effect.
+ */
+ private final ConcurrentMap<ShardBufferId, ShardBuffer> shardBuffers = new ConcurrentHashMap<>();
+
+ /**
+ * A queue with ids shard buffers available for leasing.
+ * <p>
+ * Note: when a buffer is invalidated its id is NOT removed from the queue immediately.
+ */
+ private final Queue<ShardBufferId> buffersToLease = new ConcurrentLinkedQueue<>();
+
+ MemoryBoundRecordBuffer(final ComponentLog logger, final long maxMemoryBytes, final Duration checkpointInterval) {
+ this.logger = logger;
+ this.memoryTracker = new BlockingMemoryTracker(logger, maxMemoryBytes);
+ this.checkpointIntervalMillis = checkpointInterval.toMillis();
+ }
+
+ @Override
+ public ShardBufferId createBuffer(final String shardId) {
+ final ShardBufferId id = new ShardBufferId(shardId, bufferIdCounter.getAndIncrement());
+
+ logger.debug("Creating new buffer for shard {} with id {}", shardId, id);
+
+ shardBuffers.put(id, new ShardBuffer(id, logger, checkpointIntervalMillis));
+ buffersToLease.add(id);
+ return id;
+ }
+
+ @Override
+ public void addRecords(final ShardBufferId bufferId, final List<KinesisClientRecord> records, final RecordProcessorCheckpointer checkpointer) {
+ if (records.isEmpty()) {
+ return;
+ }
+
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+ if (buffer == null) {
+ logger.debug("Buffer with id {} not found. Cannot add records with sequence and subsequence numbers: {}.{} - {}.{}",
+ bufferId,
+ records.getFirst().sequenceNumber(),
+ records.getFirst().subSequenceNumber(),
+ records.getLast().sequenceNumber(),
+ records.getLast().subSequenceNumber());
+ return;
+ }
+
+ final RecordBatch recordBatch = new RecordBatch(records, checkpointer, calculateMemoryUsage(records));
+ memoryTracker.reserveMemory(recordBatch);
+ final boolean addedRecords = buffer.offer(recordBatch);
+
+ if (addedRecords) {
+ logger.debug("Successfully added records with sequence and subsequence numbers: {}.{} - {}.{} to buffer with id {}",
+ records.getFirst().sequenceNumber(),
+ records.getFirst().subSequenceNumber(),
+ records.getLast().sequenceNumber(),
+ records.getLast().subSequenceNumber(),
+ bufferId);
+ } else {
+ logger.debug("Buffer with id {} was invalidated. Cannot add records with sequence and subsequence numbers: {}.{} - {}.{}",
+ bufferId,
+ records.getFirst().sequenceNumber(),
+ records.getFirst().subSequenceNumber(),
+ records.getLast().sequenceNumber(),
+ records.getLast().subSequenceNumber());
+ // If the buffer was invalidated, we should free memory reserved for these records.
+ memoryTracker.freeMemory(List.of(recordBatch));
+ }
+ }
+
+ @Override
+ public void checkpointEndedShard(final ShardBufferId bufferId, final RecordProcessorCheckpointer checkpointer) {
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+ if (buffer == null) {
+ logger.debug("Buffer with id {} not found. Cannot checkpoint the ended shard", bufferId);
+ return;
+ }
+
+ logger.debug("Finishing consumption for buffer {}. Checkpointing the ended shard", bufferId);
+ buffer.checkpointEndedShard(checkpointer);
+
+ logger.debug("Removing buffer with id {} after successful ended shard checkpoint", bufferId);
+ shardBuffers.remove(bufferId);
+ }
+
+ @Override
+ public void shutdownShardConsumption(final ShardBufferId bufferId, final RecordProcessorCheckpointer checkpointer) {
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+ if (buffer == null) {
+ logger.debug("Buffer with id {} not found. Cannot shutdown shard consumption", bufferId);
+ return;
+ }
+
+ logger.debug("Shutting down the buffer {}. Checkpointing last consumed record", bufferId);
+ buffer.shutdownBuffer(checkpointer);
+
+ logger.debug("Removing buffer with id {} after successful last consumed record checkpoint", bufferId);
+ shardBuffers.remove(bufferId);
+ }
+
+ @Override
+ public void consumerLeaseLost(final ShardBufferId bufferId) {
+ final ShardBuffer buffer = shardBuffers.remove(bufferId);
+
+ logger.debug("Lease lost for buffer {}: Invalidating", bufferId);
+
+ if (buffer != null) {
+ final Collection<RecordBatch> invalidatedBatches = buffer.invalidate();
+ memoryTracker.freeMemory(invalidatedBatches);
+ }
+ }
+
+ @Override
+ public Optional<Lease> acquireBufferLease() {
+ final Set<ShardBufferId> seenBuffers = new HashSet<>();
+
+ while (true) {
+ final ShardBufferId bufferId = buffersToLease.poll();
+ if (bufferId == null) {
+ // The queue is empty or all buffers were seen already. Nothing to consume.
+ return Optional.empty();
+ }
+
+ if (seenBuffers.contains(bufferId)) {
+ // If the same buffer is seen again, there is a high chance we iterated through most of the buffers and didn't find any that isn't empty.
+ // To avoid burning CPU we return empty here, even if some buffer received records in the meantime. It will be picked up in the next iteration.
+ buffersToLease.add(bufferId);
+ return Optional.empty();
+ }
+
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+
+ if (buffer == null) {
+ // By the time the bufferId is polled, it might have been invalidated. No need to return it to the queue.
+ logger.debug("Buffer with id {} was removed while polling for lease. Continuing to poll", bufferId);
+ } else if (buffer.isEmpty()) {
+ seenBuffers.add(bufferId);
+ buffersToLease.add(bufferId);
+ logger.debug("Buffer with id {} is empty. Continuing to poll", bufferId);
+ } else {
+ logger.debug("Acquired lease for buffer {}", bufferId);
+ return Optional.of(new Lease(bufferId));
+ }
+ }
+ }
+
+ @Override
+ public List<KinesisClientRecord> consumeRecords(final Lease lease) {
+ if (lease.isReturnedToPool()) {
+ logger.warn("Attempting to consume records from a buffer that was already returned to the pool. Ignoring");
+ return emptyList();
+ }
+
+ final ShardBufferId bufferId = lease.bufferId();
+
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+ if (buffer == null) {
+ logger.debug("Buffer with id {} not found. Cannot consume records", bufferId);
+ return emptyList();
+ }
+
+ return buffer.consumeRecords();
+ }
+
+ @Override
+ public void commitConsumedRecords(final Lease lease) {
+ if (lease.isReturnedToPool()) {
+ logger.warn("Attempting to commit records from a buffer that was already returned to the pool. Ignoring");
+ return;
+ }
+
+ final ShardBufferId bufferId = lease.bufferId();
+
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+ if (buffer == null) {
+ logger.debug("Buffer with id {} not found. Cannot commit consumed records", bufferId);
+ return;
+ }
+
+ final List<RecordBatch> consumedBatches = buffer.commitConsumedRecords();
+ memoryTracker.freeMemory(consumedBatches);
+ }
+
+ @Override
+ public void rollbackConsumedRecords(final Lease lease) {
+ if (lease.isReturnedToPool()) {
+ logger.warn("Attempting to rollback records from a buffer that was already returned to the pool. Ignoring");
+ return;
+ }
+
+ final ShardBufferId bufferId = lease.bufferId();
+ final ShardBuffer buffer = shardBuffers.get(bufferId);
+
+ if (buffer != null) {
+ buffer.rollbackConsumedRecords();
+ }
+ }
+
+ @Override
+ public void returnBufferLease(final Lease lease) {
+ if (lease.returnToPool()) {
+ final ShardBufferId bufferId = lease.bufferId();
+ buffersToLease.add(bufferId);
+ logger.debug("The buffer {} is available for lease again", bufferId);
+ } else {
+ logger.warn("Attempting to return a buffer that was already returned to the pool. Ignoring");
+ }
+ }
+
+ static final class Lease implements ShardBufferLease {
+
+ private final ShardBufferId bufferId;
+ private final AtomicBoolean returnedToPool = new AtomicBoolean(false);
+
+ private Lease(final ShardBufferId bufferId) {
+ this.bufferId = bufferId;
+ }
+
+ @Override
+ public String shardId() {
+ return bufferId.shardId();
+ }
+
+ private ShardBufferId bufferId() {
+ return bufferId;
+ }
+
+ private boolean isReturnedToPool() {
+ return returnedToPool.get();
+ }
+
+ /**
+ * Marks the lease as returned to the pool.
+ * @return true if the lease was not returned before, false otherwise.
+ */
+ private boolean returnToPool() {
+ final boolean wasReturned = returnedToPool.getAndSet(true);
+ return !wasReturned;
+ }
+ }
+
+ /**
+ * A memory tracker which blocks a thread when the memory usage exceeds the allowed maximum.
+ * <p>
+ * In order to make progress, the memory consumption may exceed the limit, but any new records will not be accepted.
+ * This is done to support the case when a single record batch is larger than the allowed memory limit.
+ */
+ private static class BlockingMemoryTracker {
+
+ private static final long AWAIT_MILLIS = 100;
+
+ private final ComponentLog logger;
+
+ private final long maxMemoryBytes;
+
+ private final AtomicLong consumedMemoryBytes = new AtomicLong(0);
+ /**
+ * Whenever memory is freed a latch opens. Then replaced with a new one.
+ */
+ private final AtomicReference<CountDownLatch> memoryAvailableLatch = new AtomicReference<>(new CountDownLatch(1));
+
+ BlockingMemoryTracker(final ComponentLog logger, final long maxMemoryBytes) {
+ this.logger = logger;
+ this.maxMemoryBytes = maxMemoryBytes;
+ }
+
+ void reserveMemory(final RecordBatch recordBatch) {
+ final long consumedBytes = recordBatch.batchSizeBytes();
+
+ if (consumedBytes == 0) {
+ return;
+ }
+
+ while (true) {
+ final long currentlyConsumedBytes = consumedMemoryBytes.get();
+
+ if (currentlyConsumedBytes >= maxMemoryBytes) {
+ // Not enough memory available, need to wait.
+ try {
+ memoryAvailableLatch.get().await(AWAIT_MILLIS, TimeUnit.MILLISECONDS);
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Thread interrupted while waiting for available memory in RecordBuffer", e);
+ }
+ } else {
+ final long newConsumedBytes = currentlyConsumedBytes + consumedBytes;
+ if (consumedMemoryBytes.compareAndSet(currentlyConsumedBytes, newConsumedBytes)) {
+ logger.debug("Reserved {} bytes for {} records. Total consumed memory: {} bytes",
+ consumedBytes, recordBatch.size(), newConsumedBytes);
+ break;
+ }
+ // If we're here, the compare and set operation failed, as another thread has modified the gauge in meantime.
+ // Retrying the operation.
+ }
+ }
+ }
+
+ void freeMemory(final Collection<RecordBatch> consumedBatches) {
+ if (consumedBatches.isEmpty()) {
+ return;
+ }
+
+ long freedBytes = 0;
+ for (final RecordBatch batch : consumedBatches) {
+ freedBytes += batch.batchSizeBytes();
+ }
+
+ while (true) {
+ final long currentlyConsumedBytes = consumedMemoryBytes.get();
+ if (currentlyConsumedBytes < freedBytes) {
+ throw new IllegalStateException("Attempting to free more memory than currently used");
+ }
+
+ final long newTotal = currentlyConsumedBytes - freedBytes;
+ if (consumedMemoryBytes.compareAndSet(currentlyConsumedBytes, newTotal)) {
+ logger.debug("Freed {} bytes for {} batches. Total consumed memory: {} bytes",
+ freedBytes, consumedBatches.size(), newTotal);
+
+ final CountDownLatch oldLatch = memoryAvailableLatch.getAndSet(new CountDownLatch(1));
+ oldLatch.countDown(); // Release any waiting threads for free memory.
+ break;
+ }
+ // If we're here, the compare and set operation failed, as another thread has modified the gauge in meantime.
+ // Retrying the operation.
+ }
+ }
+ }
+
+ private record RecordBatch(List<KinesisClientRecord> records,
+ @Nullable RecordProcessorCheckpointer checkpointer,
+ long batchSizeBytes) {
+ int size() {
+ return records.size();
+ }
+ }
+
+ private long calculateMemoryUsage(final Collection<KinesisClientRecord> records) {
+ long totalBytes = 0;
+ for (final KinesisClientRecord record : records) {
+ final ByteBuffer data = record.data();
+ if (data != null) {
+ totalBytes += data.capacity();
+ }
+ }
+ return totalBytes;
+ }
+
+ /**
+ * ShardBuffer stores all record batches for a single shard in two queues:
+ * - IN_PROGRESS: record batches that have been consumed but not yet checkpointed.
+ * - PENDING: record batches that have been added but not yet consumed.
+ * <p>
+ * When consuming records all PENDING batches are moved to IN_PROGRESS.
+ * After a successful checkpoint all IN_PROGRESS batches are cleared.
+ * After a rollback all IN_PROGRESS batches are kept, allowing to retry consumption.
+ * <p>
+ * Each batch preserves the original grouping of records as provided by Kinesis
+ * along with their associated checkpointer, ensuring atomicity.
+ */
+ private static class ShardBuffer {
+
+ private static final long AWAIT_MILLIS = 100;
+
+ // Retry configuration.
+ private static final int MAX_RETRY_ATTEMPTS = 5;
+ private static final long BASE_RETRY_DELAY_MILLIS = 100;
+ private static final long MAX_RETRY_DELAY_MILLIS = 10_000;
+ private static final Random RANDOM = new Random();
+
+ private final ShardBufferId bufferId;
+ private final ComponentLog logger;
+
+ private final long checkpointIntervalMillis;
+ private volatile long nextCheckpointTimeMillis;
+ /**
+ * A last records checkpointer that was ignored due to the checkpoint interval.
+ * If null, the last checkpoint was successful or no checkpoint was attempted yet.
+ */
+ private volatile @Nullable RecordProcessorCheckpointer lastIgnoredCheckpointer = null;
+
+ /**
+ * Queues for managing record batches with their checkpointers in different states.
+ */
+ private final Queue<RecordBatch> inProgressBatches = new ConcurrentLinkedQueue<>();
+ private final Queue<RecordBatch> pendingBatches = new ConcurrentLinkedQueue<>();
+ /**
+ * Counter for tracking the number of batches in the buffer. Can be larger than the number of batches in the queues.
+ */
+ private final AtomicInteger batchesCount = new AtomicInteger(0);
+
+ /**
+ * A countdown latch that is used to signal when the buffer becomes empty. Used when ShardBuffer should be closed.
+ */
+ private volatile @Nullable CountDownLatch emptyBufferLatch = null;
+ private final AtomicBoolean invalidated = new AtomicBoolean(false);
+
+ ShardBuffer(final ShardBufferId bufferId, final ComponentLog logger, final long checkpointIntervalMillis) {
+ this.bufferId = bufferId;
+ this.logger = logger;
+ this.checkpointIntervalMillis = checkpointIntervalMillis;
+ this.nextCheckpointTimeMillis = System.currentTimeMillis() + checkpointIntervalMillis;
+ }
+
+ /**
+ * @param recordBatch record batch with records to add.
+ * @return true if the records were added successfully, false if a buffer was invalidated.
+ */
+ boolean offer(final RecordBatch recordBatch) {
+ if (invalidated.get()) {
+ return false;
+ }
+
+ // Batches count must be always equal to or larger than the number of batches in the queues.
+ // Thus, the ordering of the operations.
+ batchesCount.incrementAndGet();
+ pendingBatches.offer(recordBatch);
+
+ return true;
+ }
+
+ List<KinesisClientRecord> consumeRecords() {
+ if (invalidated.get()) {
+ return emptyList();
+ }
+
+ RecordBatch pendingBatch;
+ while ((pendingBatch = pendingBatches.poll()) != null) {
+ inProgressBatches.offer(pendingBatch);
+ }
+
+ final List<KinesisClientRecord> recordsToConsume = new ArrayList<>();
+ for (final RecordBatch batch : inProgressBatches) {
+ recordsToConsume.addAll(batch.records());
+ }
+
+ return recordsToConsume;
+ }
+
+ List<RecordBatch> commitConsumedRecords() {
+ if (invalidated.get()) {
+ return emptyList();
+ }
+
+ final List<RecordBatch> checkpointedBatches = new ArrayList<>();
+ RecordBatch batch;
+ while ((batch = inProgressBatches.poll()) != null) {
+ checkpointedBatches.add(batch);
+ }
+
+ if (checkpointedBatches.isEmpty()) {
+ // The buffer could be invalidated in the meantime, or no records were consumed.
+ return emptyList();
+ }
+
+ // Batches count must always be equal to or larger than the number of batches in the queues.
+ // To achieve so, the count is decreased only after the queue has been emptied.
+ batchesCount.addAndGet(-checkpointedBatches.size());
+
+ final RecordProcessorCheckpointer lastBatchCheckpointer = checkpointedBatches.getLast().checkpointer();
+ if (System.currentTimeMillis() >= nextCheckpointTimeMillis) {
+ checkpointSafely(lastBatchCheckpointer);
+ nextCheckpointTimeMillis = System.currentTimeMillis() + checkpointIntervalMillis;
+ lastIgnoredCheckpointer = null;
+ } else {
+ // Saving the checkpointer for later, in case shutdown happens before the next checkpoint.
+ lastIgnoredCheckpointer = lastBatchCheckpointer;
+ }
+
+ final CountDownLatch localEmptyBufferLatch = this.emptyBufferLatch;
+ if (localEmptyBufferLatch != null && isEmpty()) {
+ // If the latch is not null, it means we are waiting for the buffer to become empty.
+ localEmptyBufferLatch.countDown();
+ }
+
+ return checkpointedBatches;
+ }
+
+ void rollbackConsumedRecords() {
+ if (invalidated.get()) {
+ return;
+ }
+
+ for (final RecordBatch recordBatch : inProgressBatches) {
+ for (final KinesisClientRecord record : recordBatch.records()) {
+ record.data().rewind();
+ }
+ }
+ }
+
+ void checkpointEndedShard(final RecordProcessorCheckpointer checkpointer) {
+ while (true) {
+ if (invalidated.get()) {
+ return;
+ }
+
+ // Setting the latch first, so that if the buffer is being emptied concurrently in commitConsumedRecords
+ // the latch is guaranteed to be visible to the commitConsumedRecords, and, therefore, opened.
+ // This will eliminate unnecessary downtimes when waiting for a latch to be opened.
+ final CountDownLatch localEmptyBufferLatch = new CountDownLatch(1);
+ this.emptyBufferLatch = localEmptyBufferLatch;
+
+ if (batchesCount.get() == 0) {
+ // Buffer is empty, perform final checkpoint.
+ checkpointSafely(checkpointer);
+ return;
+ }
+
+ // Wait for the records to be consumed first.
+ try {
+ localEmptyBufferLatch.await(AWAIT_MILLIS, TimeUnit.MILLISECONDS);
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Thread interrupted while waiting for records to be consumed", e);
+ }
+ }
+ }
+
+ void shutdownBuffer(final RecordProcessorCheckpointer checkpointer) {
+ if (invalidated.get()) {
+ return;
+ }
+
+ if (batchesCount.get() == 0) {
+ checkpointSafely(checkpointer);
+ } else {
+ // If there are still records in the buffer, checkpointing with the latest provided checkpointer is not safe.
+ // But, if the records were committed without checkpointing in the past, we can checkpoint them now.
+ final RecordProcessorCheckpointer lastCheckpointer = this.lastIgnoredCheckpointer;
+ if (lastCheckpointer != null) {
+ checkpointSafely(lastCheckpointer);
+ }
+ }
+ }
+
+ Collection<RecordBatch> invalidate() {
+ if (invalidated.getAndSet(true)) {
+ return emptyList();
+ }
+
+ final List<RecordBatch> batches = new ArrayList<>();
+ RecordBatch batch;
+ // If both consumeRecords and invalidate are called concurrently, invalidation must always consume all batches.
+ // Since consumeRecords moves batches from pending to in_progress, during invalidation pending batches should be drained first.
+ while ((batch = pendingBatches.poll()) != null) {
+ batches.add(batch);
+ }
+ while ((batch = inProgressBatches.poll()) != null) {
+ batches.add(batch);
+ }
+
+ // No need to adjust batchesCount after invalidation.
+
+ return batches;
+ }
+
+ /**
+ * Checks if the buffer has any records. <b>Can produce false negatives.</b>
+ *
+ * @return whether there are any records in the buffer.
+ */
+ boolean isEmpty() {
+ return invalidated.get() || batchesCount.get() == 0;
+ }
+
+ /**
+ * Performs checkpointing using exponential backoff and jitter, if needed.
+ *
+ * @param checkpointer the checkpointer to use.
+ */
+ private void checkpointSafely(final @Nullable RecordProcessorCheckpointer checkpointer) {
+ if (checkpointer == null) {
+ logger.warn("Attempting to checkpoint records with a null checkpointer. Ignoring checkpoint");
+ return;
+ }
+
+ logger.debug("Performing checkpoint for buffer with id {}", bufferId);
+
+ for (int attempt = 1; attempt <= MAX_RETRY_ATTEMPTS; attempt++) {
+ try {
+ checkpointer.checkpoint();
+ if (attempt > 1) {
+ logger.debug("Checkpoint succeeded on attempt {}", attempt);
+ }
+ return;
+ } catch (final ThrottlingException | InvalidStateException | KinesisClientLibDependencyException e) {
+ if (attempt == MAX_RETRY_ATTEMPTS) {
+ logger.error("Failed to checkpoint after {} attempts, giving up", MAX_RETRY_ATTEMPTS, e);
+ return;
+ }
+
+ final long delayMillis = calculateRetryDelay(attempt);
+
+ logger.debug("Checkpoint failed on attempt {} with {}, retrying in {} ms",
+ attempt, e.getMessage(), delayMillis);
+
+ try {
+ Thread.sleep(delayMillis);
+ } catch (final InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ logger.warn("Thread interrupted while waiting to retry checkpoint. Exiting retry loop", ie);
+ return;
+ }
+ } catch (final ShutdownException e) {
+ logger.warn("Failed to checkpoint records due to shutdown. Ignoring checkpoint", e);
+ return;
+ }
+ }
+ }
+
+ private long calculateRetryDelay(final int attempt) {
+ final long desiredBaseDelayMillis = BASE_RETRY_DELAY_MILLIS * (1L << (attempt - 1));
+ final long baseDelayMillis = Math.min(desiredBaseDelayMillis, MAX_RETRY_DELAY_MILLIS);
+ final long jitterMillis = RANDOM.nextLong(baseDelayMillis / 4); // Up to 25% jitter.
+ return baseDelayMillis + jitterMillis;
+ }
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ReaderRecordProcessor.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ReaderRecordProcessor.java
new file mode 100644
index 0000000000..3e36475804
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/ReaderRecordProcessor.java
@@ -0,0 +1,265 @@
+/*
+ * 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;
+
+import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.logging.ComponentLog;
+import org.apache.nifi.processor.ProcessSession;
+import org.apache.nifi.processor.exception.ProcessException;
+import org.apache.nifi.schema.access.SchemaNotFoundException;
+import org.apache.nifi.serialization.MalformedRecordException;
+import org.apache.nifi.serialization.RecordReader;
+import org.apache.nifi.serialization.RecordReaderFactory;
+import org.apache.nifi.serialization.RecordSetWriter;
+import org.apache.nifi.serialization.RecordSetWriterFactory;
+import org.apache.nifi.serialization.WriteResult;
+import org.apache.nifi.serialization.record.Record;
+import org.apache.nifi.serialization.record.RecordSchema;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import static java.util.Collections.emptyMap;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.MIME_TYPE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_COUNT;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_ERROR_MESSAGE;
+
+final class ReaderRecordProcessor {
+
+ private final RecordReaderFactory recordReaderFactory;
+ private final RecordSetWriterFactory recordWriterFactory;
+ private final ComponentLog logger;
+
+ ReaderRecordProcessor(
+ final RecordReaderFactory recordReaderFactory,
+ final RecordSetWriterFactory recordWriterFactory,
+ final ComponentLog logger) {
+ this.recordReaderFactory = recordReaderFactory;
+ this.recordWriterFactory = recordWriterFactory;
+ this.logger = logger;
+ }
+
+ ProcessingResult processRecords(
+ final ProcessSession session,
+ final String streamName,
+ final String shardId,
+ final List<KinesisClientRecord> records) {
+ final List<FlowFile> successFlowFiles = new ArrayList<>();
+ final List<FlowFile> failureFlowFiles = new ArrayList<>();
+
+ ActiveFlowFile activeFlowFile = null;
+
+ for (final KinesisClientRecord kinesisRecord : records) {
+ final int dataSize = kinesisRecord.data().remaining();
+ final byte[] data = new byte[dataSize];
+ kinesisRecord.data().get(data);
+
+ try (final InputStream in = new ByteArrayInputStream(data);
+ final RecordReader reader = recordReaderFactory.createRecordReader(emptyMap(), in, data.length, logger)) {
+
+ Record record;
+ while ((record = reader.nextRecord()) != null) {
+ final RecordSchema writeSchema = recordWriterFactory.getSchema(emptyMap(), record.getSchema());
+
+ if (activeFlowFile == null) {
+ activeFlowFile = ActiveFlowFile.startNewFile(logger, session, recordWriterFactory, writeSchema, streamName, shardId);
+ } else if (!writeSchema.equals(activeFlowFile.schema())) {
+ // If the write schema has changed, we need to complete the current FlowFile and start a new one.
+ final FlowFile completedFlowFile = activeFlowFile.complete();
+ successFlowFiles.add(completedFlowFile);
+
+ activeFlowFile = ActiveFlowFile.startNewFile(logger, session, recordWriterFactory, writeSchema, streamName, shardId);
+ }
+
+ activeFlowFile.writeRecord(record, kinesisRecord);
+ }
+ } catch (final IOException | MalformedRecordException | SchemaNotFoundException e) {
+ logger.error("Reader or Writer failed to process Kinesis Record with Stream Name [{}] Shard Id [{}] Sequence Number [{}] SubSequence Number [{}]",
+ streamName, shardId, kinesisRecord.sequenceNumber(), kinesisRecord.subSequenceNumber(), e);
+ final FlowFile failureFlowFile = createParseFailureFlowFile(session, streamName, shardId, kinesisRecord, e);
+ failureFlowFiles.add(failureFlowFile);
+ }
+ }
+
+ if (activeFlowFile != null) {
+ final FlowFile completedFlowFile = activeFlowFile.complete();
+ successFlowFiles.add(completedFlowFile);
+ }
+
+ return new ProcessingResult(successFlowFiles, failureFlowFiles);
+ }
+
+ private static FlowFile createParseFailureFlowFile(
+ final ProcessSession session,
+ final String streamName,
+ final String shardId,
+ final KinesisClientRecord record,
+ final Exception e) {
+ FlowFile flowFile = session.create();
+
+ record.data().rewind();
+ flowFile = session.write(flowFile, out -> {
+ try (final WritableByteChannel channel = Channels.newChannel(out)) {
+ channel.write(record.data());
+ }
+ });
+
+ final Map<String, String> attributes = ConsumeKinesisAttributes.fromKinesisRecords(streamName, shardId, record, record);
+
+ final Throwable cause = e.getCause() != null ? e.getCause() : e;
+ attributes.put(RECORD_ERROR_MESSAGE, cause.toString());
+
+ return session.putAllAttributes(flowFile, attributes);
+ }
+
+ record ProcessingResult(List<FlowFile> successFlowFiles, List<FlowFile> parseFailureFlowFiles) {
+ }
+
+ /**
+ * A class that manages a single {@link FlowFile} with a static schema that is currently being written to.
+ * On a schema change the current {@link ActiveFlowFile} should be completed a new instance of this class
+ * with a new schema should be created.
+ *
+ * An {@link ActiveFlowFile} must have at least one record written to it before it can be completed.
+ */
+ private static final class ActiveFlowFile {
+
+ private final ComponentLog logger;
+
+ private final ProcessSession session;
+ private final FlowFile flowFile;
+ private final RecordSetWriter writer;
+ private final RecordSchema schema;
+
+ private final String streamName;
+ private final String shardId;
+
+ private KinesisClientRecord firstRecord;
+ private KinesisClientRecord lastRecord;
+
+ private ActiveFlowFile(
+ final ComponentLog logger,
+ final ProcessSession session,
+ final FlowFile flowFile,
+ final RecordSetWriter writer,
+ final RecordSchema schema,
+ final String streamName,
+ final String shardId) {
+ this.logger = logger;
+ this.session = session;
+ this.flowFile = flowFile;
+ this.writer = writer;
+ this.schema = schema;
+ this.streamName = streamName;
+ this.shardId = shardId;
+ }
+
+ static ActiveFlowFile startNewFile(
+ final ComponentLog logger,
+ final ProcessSession session,
+ final RecordSetWriterFactory recordWriterFactory,
+ final RecordSchema writeSchema,
+ final String streamName,
+ final String shardId) throws SchemaNotFoundException {
+ final FlowFile flowFile = session.create();
+ final OutputStream outputStream = session.write(flowFile);
+
+ try {
+ final RecordSetWriter writer = recordWriterFactory.createWriter(logger, writeSchema, outputStream, flowFile);
+ writer.beginRecordSet();
+
+ return new ActiveFlowFile(logger, session, flowFile, writer, writeSchema, streamName, shardId);
+
+ } catch (final SchemaNotFoundException e) {
+ logger.debug("Failed to find writeSchema for Kinesis stream record: {}", e.getMessage());
+ try {
+ outputStream.close();
+ } catch (final IOException ioe) {
+ e.addSuppressed(ioe);
+ }
+ throw e;
+
+ } catch (final IOException e) {
+ final ProcessException processException = new ProcessException("Failed to create a writer for a FlowFile", e);
+
+ logger.debug("Stopping Kinesis records processing. Failed to create a writer for a FlowFile: {}", e.getMessage());
+ try {
+ outputStream.close();
+ } catch (final IOException ioe) {
+ processException.addSuppressed(ioe);
+ }
+ throw processException;
+ }
+ }
+
+ RecordSchema schema() {
+ return schema;
+ }
+
+ void writeRecord(final Record record, final KinesisClientRecord kinesisRecord) {
+ try {
+ writer.write(record);
+ } catch (final IOException e) {
+ logger.debug("Stopping Kinesis records processing. Failed to write to a FlowFile: {}", e.getMessage());
+ throw new ProcessException("Failed to write a record into a FlowFile", e);
+ }
+
+ if (firstRecord == null) {
+ firstRecord = kinesisRecord;
+ }
+ lastRecord = kinesisRecord;
+ }
+
+ FlowFile complete() {
+ if (firstRecord == null || lastRecord == null) {
+ throw new IllegalStateException("Cannot complete an ActiveFlowFile that has no records");
+ }
+
+ try {
+ final WriteResult finalResult = writer.finishRecordSet();
+ writer.close();
+
+ final Map<String, String> attributes = ConsumeKinesisAttributes.fromKinesisRecords(streamName, shardId, firstRecord, lastRecord);
+ attributes.putAll(finalResult.getAttributes());
+ attributes.put(RECORD_COUNT, String.valueOf(finalResult.getRecordCount()));
+ attributes.put(MIME_TYPE, writer.getMimeType());
+
+ return session.putAllAttributes(flowFile, attributes);
+
+ } catch (final IOException e) {
+ final ProcessException processException = new ProcessException("Failed to complete a FlowFile", e);
+
+ logger.debug("Stopping Kinesis records processing. Failed to complete a FlowFile: {}", e.getMessage());
+ try {
+ writer.close();
+ } catch (final IOException ioe) {
+ processException.addSuppressed(ioe);
+ }
+
+ throw processException;
+ }
+ }
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/RecordBuffer.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/RecordBuffer.java
new file mode 100644
index 0000000000..12fa52c6fe
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/RecordBuffer.java
@@ -0,0 +1,96 @@
+/*
+ * 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;
+
+import software.amazon.kinesis.processor.RecordProcessorCheckpointer;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * RecordBuffer keeps track of all created Shard buffers, including exclusive read access via leasing.
+ * It acts as the main interface between KCL callbacks and the {@link ConsumeKinesis} processor,
+ * routing events to appropriate Shard buffers and ensuring thread-safe operations.
+ */
+interface RecordBuffer {
+
+ /**
+ * Interface for interactions from the Kinesis Client Library to the Record Buffer.
+ * Reflects the methods called by {@link software.amazon.kinesis.processor.ShardRecordProcessor}.
+ */
+ interface ForKinesisClientLibrary {
+
+ ShardBufferId createBuffer(String shardId);
+
+ void addRecords(ShardBufferId bufferId, List<KinesisClientRecord> records, RecordProcessorCheckpointer checkpointer);
+
+ /**
+ * Called when a shard ends - waits until the buffer is flushed then performs the final checkpoint.
+ */
+ void checkpointEndedShard(ShardBufferId bufferId, RecordProcessorCheckpointer checkpointer);
+
+ /**
+ * Called when a consumer is shut down. Performs the checkpoint and returns
+ * without waiting for the buffer to be flushed.
+ */
+ void shutdownShardConsumption(ShardBufferId bufferId, RecordProcessorCheckpointer checkpointer);
+
+ /**
+ * Called when lease is lost - immediately invalidates the buffer to prevent further operations.
+ */
+ void consumerLeaseLost(ShardBufferId bufferId);
+ }
+
+ /**
+ * Interface for interactions from {@link ConsumeKinesis} processor to the Record Buffer.
+ */
+ interface ForProcessor<LEASE extends ShardBufferLease> {
+
+ /**
+ * Acquires an exclusive lease for a buffer that has data available for consumption.
+ * If no data is available in the buffers, returns an empty Optional.
+ * <p>
+ * After acquiring a lease, the processor can consume records from the buffer.
+ * After consuming the records the processor must always {@link #returnBufferLease(LEASE)}.
+ */
+ Optional<LEASE> acquireBufferLease();
+
+ /**
+ * Consumes records from the buffer associated with the given lease.
+ * The records have to be committed or rolled back later.
+ */
+ List<KinesisClientRecord> consumeRecords(LEASE lease);
+
+ void commitConsumedRecords(LEASE lease);
+
+ void rollbackConsumedRecords(LEASE lease);
+
+ /**
+ * Returns the lease for a buffer back to the pool making it available for consumption again.
+ * The method can be called multiple times with the same lease, but only the first call will actually take an effect.
+ */
+ void returnBufferLease(LEASE lease);
+ }
+
+ record ShardBufferId(String shardId, long bufferId) {
+ }
+
+ interface ShardBufferLease {
+ String shardId();
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor
new file mode 100644
index 0000000000..ea384e93f2
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor
@@ -0,0 +1,16 @@
+# 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.
+
+org.apache.nifi.processors.aws.kinesis.ConsumeKinesis
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/resources/docs/org.apache.nifi.processors.aws.kinesis.ConsumeKinesis/additionalDetails.md b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/resources/docs/org.apache.nifi.processors.aws.kinesis.ConsumeKinesis/additionalDetails.md
new file mode 100644
index 0000000000..5fb223ab89
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/resources/docs/org.apache.nifi.processors.aws.kinesis.ConsumeKinesis/additionalDetails.md
@@ -0,0 +1,136 @@
+<!--
+ 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.
+-->
+
+# ConsumeKinesis
+
+## IAM permissions required for ConsumeKinesis
+
+You must add the following permissions to the IAM role or user configured in _AWS Credentials Provider Service_.
+
+Cloudwatch permissions are needed only when _Metrics Publishing_ is set to _CloudWatch_.
+
+| Service | Actions | Resources (ARNs) | Purpose [...]
+|-----------------------------|-------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------- [...]
+| Amazon Kinesis Data Streams | DescribeStream<br>DescribeStreamSummary<br>RegisterStreamConsumer | Kinesis data stream from which `ConsumeKinesis` will process the data.<br>`arn:aws:kinesis:{Region}:{Account}:stream/{Stream Name}` | Before attempting to read records, the consumer checks if the data stream exists, if [...]
+| Amazon Kinesis Data Streams | GetRecords<br>GetShardIterator<br>ListShards | Kinesis data stream from which `ConsumeKinesis` will process the data.<br>`arn:aws:kinesis:{Region}:{Account}:stream/{Stream Name}` | Reads records from a shard. [...]
+| Amazon Kinesis Data Streams | SubscribeToShard<br>DescribeStreamConsumer | Kinesis data stream from which `ConsumeKinesis` will process the data. Add this action only if you use enhanced fan-out (EFO) consumers.<br>`arn:aws:kinesis:{Region}:{Account}:stream/{Stream Name}/consumer/*` | Subscribes to a shard for enhanced fan-out (EFO) consumers. [...]
+| Amazon DynamoDB | CreateTable<br>DescribeTable<br>UpdateTable<br>Scan<br>GetItem<br>PutItem<br>UpdateItem<br>DeleteItem | Lease table (metadata table in DynamoDB created by `ConsumeKinesis`.<br>`arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}` | These actions are required for `ConsumeKinesis` to manage the lease table created in [...]
+| Amazon DynamoDB | CreateTable<br>DescribeTable<br>Scan<br>GetItem<br>PutItem<br>UpdateItem<br>DeleteItem | Worker metrics and coordinator state table (metadata tables in DynamoDB) created by `ConsumeKinesis`.<br>`arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}-WorkerMetricStats`<br>`arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}-CoordinatorState` | These actions are required for `ConsumeKinesis` to manage the worker metrics and coor [...]
+| Amazon DynamoDB | Query | Global secondary index on the lease table.<br>`arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}/index/*` | This action is required for `ConsumeKinesis` to read the global secondary index of th [...]
+| Amazon CloudWatch | PutMetricData | `*` | Upload metrics to CloudWatch that are useful for monitoring the application. The aste [...]
+
+The following is an example policy document for `ConsumeKinesis`.
+
+```json
+{
+ "Version":"2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Action": [
+ "kinesis:DescribeStream",
+ "kinesis:DescribeStreamSummary",
+ "kinesis:RegisterStreamConsumer",
+ "kinesis:GetRecords",
+ "kinesis:GetShardIterator",
+ "kinesis:ListShards"
+ ],
+ "Resource": [
+ "arn:aws:kinesis:{Region}:{Account}:stream/{Stream Name}"
+ ]
+ },
+ {
+ "Effect": "Allow",
+ "Action": [
+ "kinesis:SubscribeToShard",
+ "kinesis:DescribeStreamConsumer"
+ ],
+ "Resource": [
+ "arn:aws:kinesis:{Region}:{Account}:stream/{Stream Name}/consumer/*"
+ ]
+ },
+ {
+ "Effect": "Allow",
+ "Action": [
+ "dynamodb:CreateTable",
+ "dynamodb:DescribeTable",
+ "dynamodb:UpdateTable",
+ "dynamodb:GetItem",
+ "dynamodb:UpdateItem",
+ "dynamodb:PutItem",
+ "dynamodb:DeleteItem",
+ "dynamodb:Scan"
+ ],
+ "Resource": [
+ "arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}"
+ ]
+ },
+ {
+ "Effect": "Allow",
+ "Action": [
+ "dynamodb:CreateTable",
+ "dynamodb:DescribeTable",
+ "dynamodb:GetItem",
+ "dynamodb:UpdateItem",
+ "dynamodb:PutItem",
+ "dynamodb:DeleteItem",
+ "dynamodb:Scan"
+ ],
+ "Resource": [
+ "arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}-WorkerMetricStats",
+ "arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}-CoordinatorState"
+ ]
+ },
+ {
+ "Effect": "Allow",
+ "Action": [
+ "dynamodb:Query"
+ ],
+ "Resource": [
+ "arn:aws:dynamodb:{Region}:{Account}:table/{Application Name}/index/*"
+ ]
+ },
+ {
+ "Effect": "Allow",
+ "Action": [
+ "cloudwatch:PutMetricData"
+ ],
+ "Resource": "*"
+ }
+ ]
+}
+```
+
+**Note:** Replace "{Region}", "{Account}", "{Stream Name}", and "{Application Name}" in the ARNs with your own AWS region,
+AWS account number, Kinesis data stream name, and `ConsumeKinesis` _Application Name_ property respectively.
+
+## Record processing
+
+When _Processing Strategy_ property is set to _RECORD_, _ConsumeKinesis_ operates in Record mode.
+In this mode, the processor reads records from Kinesis streams using the configured _Record Reader_,
+and writes them to FlowFiles using the configured _Record Writer_.
+
+The processor tries to optimize the number of FlowFiles created by batching multiple records with the same schema
+into a single FlowFile.
+
+### Schema changes
+
+_ConsumeKinesis_ supports dynamically changing Kinesis record schemas. When a record with new schema is encountered,
+the currently open FlowFile is closed and a new FlowFile is created for the new schema. Thanks to this, the processor
+preserves record ordering for a single Shard, even with record schema changes.
+
+**Please note**, the processor relies on _Record Reader_ to provide correct schema for each record.
+Using a Reader with schema inference may produce a lot of different schemas, which may lead to excessive FlowFile creation.
+If performance is a concern, it is recommended to use a Reader with a predefined schema or schema registry.
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisIT.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisIT.java
new file mode 100644
index 0000000000..ecb658e0d3
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisIT.java
@@ -0,0 +1,642 @@
+/*
+ * 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;
+
+import org.apache.nifi.json.JsonRecordSetWriter;
+import org.apache.nifi.json.JsonTreeReader;
+import org.apache.nifi.processor.Processor;
+import org.apache.nifi.processor.Relationship;
+import org.apache.nifi.processor.exception.FlowFileHandlingException;
+import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderControllerService;
+import org.apache.nifi.reporting.InitializationException;
+import org.apache.nifi.schema.access.SchemaAccessUtils;
+import org.apache.nifi.schema.inference.SchemaInferenceUtil;
+import org.apache.nifi.util.MockFlowFile;
+import org.apache.nifi.util.MockProcessSession;
+import org.apache.nifi.util.SharedSessionState;
+import org.apache.nifi.util.TestRunner;
+import org.apache.nifi.util.TestRunners;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.api.parallel.Execution;
+import org.junit.jupiter.api.parallel.ExecutionMode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testcontainers.containers.localstack.LocalStackContainer;
+import org.testcontainers.containers.localstack.LocalStackContainer.Service;
+import org.testcontainers.utility.DockerImageName;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.kinesis.KinesisClient;
+import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsRequestEntry;
+import software.amazon.awssdk.services.kinesis.model.ScalingType;
+import software.amazon.awssdk.services.kinesis.model.StreamStatus;
+
+import java.net.URI;
+import java.time.Duration;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.IntStream;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
+import static java.util.stream.Collectors.toList;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.REL_PARSE_FAILURE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.REL_SUCCESS;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_COUNT;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_ERROR_MESSAGE;
+import static org.apache.nifi.processors.aws.kinesis.JsonRecordAssert.assertFlowFileRecordPayloads;
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Timeout.ThreadMode.SEPARATE_THREAD;
+
+/**
+ * Tests run in parallel to optimize execution time as Kinesis consumer coordination takes a lot.
+ */
+@Execution(ExecutionMode.CONCURRENT)
+@Timeout(value = 5, unit = MINUTES, threadMode = SEPARATE_THREAD)
+class ConsumeKinesisIT {
+
+ private static final Logger logger = LoggerFactory.getLogger(ConsumeKinesisIT.class);
+ private static final DockerImageName LOCALSTACK_IMAGE = DockerImageName.parse("localstack/localstack:4.6.0");
+
+ private static final LocalStackContainer localstack = new LocalStackContainer(LOCALSTACK_IMAGE)
+ .withServices(Service.KINESIS, Service.DYNAMODB, Service.CLOUDWATCH);
+
+ private static KinesisClient kinesisClient;
+ private static DynamoDbClient dynamoDbClient;
+
+ private String streamName;
+ private String applicationName;
+ private TestRunner runner;
+ private TestKinesisStreamClient streamClient;
+
+ @BeforeAll
+ static void oneTimeSetup() {
+ localstack.start();
+
+ final AwsCredentialsProvider credentialsProvider = StaticCredentialsProvider.create(
+ AwsBasicCredentials.create(localstack.getAccessKey(), localstack.getSecretKey())
+ );
+
+ kinesisClient = KinesisClient.builder()
+ .endpointOverride(localstack.getEndpointOverride(Service.KINESIS))
+ .credentialsProvider(credentialsProvider)
+ .region(Region.of(localstack.getRegion()))
+ .build();
+
+ dynamoDbClient = DynamoDbClient.builder()
+ .endpointOverride(localstack.getEndpointOverride(Service.DYNAMODB))
+ .credentialsProvider(credentialsProvider)
+ .region(Region.of(localstack.getRegion()))
+ .build();
+ }
+
+ @AfterAll
+ static void tearDown() {
+ if (kinesisClient != null) {
+ kinesisClient.close();
+ }
+ if (dynamoDbClient != null) {
+ dynamoDbClient.close();
+ }
+ localstack.stop();
+ }
+
+ @BeforeEach
+ void setUp() throws InitializationException {
+ final UUID testId = UUID.randomUUID();
+ streamName = "%s-kinesis-stream-%s".formatted(getClass().getSimpleName(), testId);
+ streamClient = new TestKinesisStreamClient(kinesisClient, streamName);
+ applicationName = "%s-test-kinesis-app-%s".formatted(getClass().getSimpleName(), testId);
+ runner = createTestRunner(streamName, applicationName);
+ }
+
+ @AfterEach
+ void tearDownEach() {
+ runner.stop();
+
+ if (streamClient != null) {
+ try {
+ streamClient.deleteStream();
+ } catch (final Exception e) {
+ logger.warn("Failed to delete stream {}: {}", streamName, e.getMessage());
+ }
+ }
+
+ // Removing tables generated by KCL.
+ deleteTable(applicationName);
+ deleteTable(applicationName + "-CoordinatorState");
+ deleteTable(applicationName + "-WorkerMetricStats");
+
+ }
+
+ private void deleteTable(final String tableName) {
+ try {
+ dynamoDbClient.deleteTable(req -> req.tableName(tableName));
+ } catch (final Exception e) {
+ logger.warn("Failed to delete DynamoDB table {}: {}", tableName, e.getMessage());
+ }
+ }
+
+ @Test
+ void testConsumeSingleMessageFromSingleShard() {
+ streamClient.createStream(1);
+
+ final String testMessage = "Hello, Kinesis!";
+ streamClient.putRecord("test-partition-key", testMessage);
+
+ runProcessorWithInitAndWaitForFiles(runner, 1);
+
+ runner.assertTransferCount(REL_SUCCESS, 1);
+ final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(REL_SUCCESS);
+ final MockFlowFile flowFile = flowFiles.getFirst();
+
+ flowFile.assertContentEquals(testMessage);
+ flowFile.assertAttributeEquals("aws.kinesis.partition.key", "test-partition-key");
+ assertNotNull(flowFile.getAttribute("aws.kinesis.first.sequence.number"));
+ assertNotNull(flowFile.getAttribute("aws.kinesis.last.sequence.number"));
+ assertNotNull(flowFile.getAttribute("aws.kinesis.shard.id"));
+ }
+
+ @Test
+ void testConsumeManyMessagesFromSingleShardWithOrdering() {
+ final int messageCount = 10;
+
+ streamClient.createStream(1);
+
+ final List<String> messages = IntStream.range(0, messageCount).mapToObj(i -> "Message-" + i).toList();
+ streamClient.putRecords("partition-key", messages);
+
+ runProcessorWithInitAndWaitForFiles(runner, messageCount);
+
+ runner.assertTransferCount(REL_SUCCESS, messageCount);
+ final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(REL_SUCCESS);
+ final List<String> flowFilesContent = flowFiles.stream().map(MockFlowFile::getContent).toList();
+
+ assertEquals(messages, flowFilesContent);
+ }
+
+ @Test
+ void testConsumeMessagesFromMultipleShardsStream() {
+ final int shardCount = 5;
+ final int messagesPerPartitionKey = 5;
+
+ streamClient.createStream(shardCount);
+
+ // Every shard has message with the same payload.
+ final List<String> shardMessages = IntStream.range(0, messagesPerPartitionKey).mapToObj(String::valueOf).toList();
+
+ IntStream.range(0, shardCount).forEach(shard -> streamClient.putRecords(String.valueOf(shard), shardMessages));
+
+ // Run processor and wait for all records
+ final int totalMessages = shardCount * messagesPerPartitionKey;
+ runProcessorWithInitAndWaitForFiles(runner, totalMessages);
+
+ // Verify results
+ runner.assertTransferCount(REL_SUCCESS, totalMessages);
+ final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(REL_SUCCESS);
+
+ final Map<String, List<String>> partitionKey2Messages = flowFiles.stream()
+ .collect(groupingBy(
+ f -> f.getAttribute("aws.kinesis.partition.key"),
+ mapping(MockFlowFile::getContent, toList())
+ ));
+
+ assertEquals(shardCount, partitionKey2Messages.size());
+
+ assertAll(
+ partitionKey2Messages.values().stream()
+ .map(actual -> () -> assertEquals(shardMessages, actual))
+ );
+ }
+
+ @Test
+ @Disabled("Does not work with LocalStack: https://github.com/localstack/localstack/issues/12833. Enable only when using real Kinesis.")
+ void testResharding_inParallelWithConsumption() {
+ // Using partition keys with uniformally distributed hashes to ensure the data is distributed across split shards.
+ final List<String> partitionKeys = List.of(
+ "pk-0-14", // 035517ff4ca68849589f43842c07362f
+ "pk-1-14", // 5f045ae51eea9bd124d76041a6a27073
+ "pk-2-2", // 85fb9a2b01b009904eb8a6fa13a21d6c
+ "pk-3-2" // dbf24a6e26910143c60188e2fcb53b4f
+ );
+
+ // Data to be produced at each stage
+ final int partitionRecordsPerStage = 3;
+ final int totalStages = 5; // initial + 4 resharding stages with data
+ final int totalRecordsPerPartition = partitionRecordsPerStage * totalStages;
+ final int totalRecords = partitionKeys.size() * totalRecordsPerPartition;
+
+ // Start resharding and data production operations in background thread.
+ final Thread reshardingThread = new Thread(() -> {
+ int messageSeq = 0; // For each partition key the message content are sequential numbers.
+
+ streamClient.createStream(1);
+ putRecords(partitionKeys, partitionRecordsPerStage, messageSeq);
+ messageSeq += partitionRecordsPerStage;
+
+ streamClient.reshardStream(2);
+ putRecords(partitionKeys, partitionRecordsPerStage, messageSeq);
+ messageSeq += partitionRecordsPerStage;
+
+ streamClient.reshardStream(4);
+ putRecords(partitionKeys, partitionRecordsPerStage, messageSeq);
+ messageSeq += partitionRecordsPerStage;
+
+ streamClient.reshardStream(3);
+ putRecords(partitionKeys, partitionRecordsPerStage, messageSeq);
+ messageSeq += partitionRecordsPerStage;
+
+ streamClient.reshardStream(2);
+ putRecords(partitionKeys, partitionRecordsPerStage, messageSeq);
+ });
+
+ reshardingThread.start();
+
+ runProcessorWithInitAndWaitForFiles(runner, totalRecords);
+
+ runner.assertTransferCount(REL_SUCCESS, totalRecords);
+ final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(REL_SUCCESS);
+
+ final Map<String, List<String>> partitionKeyToMessages = flowFiles.stream()
+ .collect(groupingBy(
+ f -> f.getAttribute("aws.kinesis.partition.key"),
+ mapping(MockFlowFile::getContent, toList())
+ ));
+
+ final List<String> expectedPartitionMessages = IntStream.range(0, totalRecordsPerPartition).mapToObj(Integer::toString).toList();
+ assertAll(
+ partitionKeyToMessages.entrySet().stream()
+ .map(actual -> () -> assertEquals(
+ expectedPartitionMessages,
+ actual.getValue(),
+ "Partition messages do not match expected for partition key: " + actual.getKey()))
+ );
+ }
+
+ private void putRecords(final Collection<String> partitionKeys, final int count, final int startIndex) {
+ IntStream.range(startIndex, startIndex + count).forEach(i -> {
+ final String message = Integer.toString(i);
+ partitionKeys.forEach(partitionKey -> streamClient.putRecord(partitionKey, message));
+ });
+ }
+
+ @Test
+ void testSessionRollback() throws InterruptedException {
+ streamClient.createStream(1);
+
+ // Initialize the processor.
+ runner.run(1, false, true);
+
+ final ConsumeKinesis processor = (ConsumeKinesis) runner.getProcessor();
+
+ final String firstMessage = "Initial-Rollback-Message";
+ streamClient.putRecord("key", firstMessage);
+
+ // First attempt with a failing session - should rollback.
+ while (true) {
+ final MockProcessSession failingSession = createFailingSession(processor);
+ try {
+ processor.onTrigger(runner.getProcessContext(), failingSession);
+ } catch (final FlowFileHandlingException __) {
+ failingSession.assertAllFlowFilesTransferred(REL_SUCCESS, 0);
+ break; // Expected rollback occurred
+ }
+ Thread.sleep(1000);
+ }
+
+ // Write another message.
+ final String secondMessage = "Another-Test-Message";
+ streamClient.putRecord("key", secondMessage);
+
+ runProcessorAndWaitForFiles(runner, 2);
+
+ // Verify the messages are written in the correct order.
+ runner.assertTransferCount(REL_SUCCESS, 2);
+ final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(REL_SUCCESS);
+ flowFiles.getFirst().assertContentEquals(firstMessage);
+ flowFiles.getLast().assertContentEquals(secondMessage);
+ }
+
+ @Test
+ void testRecordProcessingWithSchemaChangesAndInvalidRecords() throws InitializationException {
+ streamClient.createStream(1);
+
+ final TestRunner recordRunner = createRecordTestRunner(streamName, applicationName);
+
+ final List<String> testRecords = List.of(
+ "{\"name\":\"John\",\"age\":30}", // Schema A
+ "{\"name\":\"Jane\",\"age\":25}", // Schema A
+ "{invalid json}",
+ "{\"id\":\"123\",\"value\":\"test\"}" // Schema B
+ );
+
+ testRecords.forEach(record -> streamClient.putRecord("key", record));
+
+ runProcessorWithInitAndWaitForFiles(recordRunner, 3);
+
+ recordRunner.assertTransferCount(REL_SUCCESS, 2);
+ final List<MockFlowFile> successFlowFiles = recordRunner.getFlowFilesForRelationship(REL_SUCCESS);
+
+ final MockFlowFile firstFlowFile = successFlowFiles.getFirst();
+ assertEquals("2", firstFlowFile.getAttribute(RECORD_COUNT));
+ assertFlowFileRecordPayloads(firstFlowFile, testRecords.getFirst(), testRecords.get(1));
+
+ final MockFlowFile secondFlowFile = successFlowFiles.get(1);
+ assertEquals("1", secondFlowFile.getAttribute(RECORD_COUNT));
+ assertFlowFileRecordPayloads(secondFlowFile, testRecords.getLast());
+
+ // Verify failure record
+ recordRunner.assertTransferCount(REL_PARSE_FAILURE, 1);
+ final List<MockFlowFile> parseFailureFlowFiles = recordRunner.getFlowFilesForRelationship(REL_PARSE_FAILURE);
+ final MockFlowFile parseFailureFlowFile = parseFailureFlowFiles.getFirst();
+
+ parseFailureFlowFile.assertContentEquals(testRecords.get(2));
+ assertNotNull(parseFailureFlowFile.getAttribute(RECORD_ERROR_MESSAGE));
+ }
+
+ private TestRunner createTestRunner(final String streamName, final String applicationName) throws InitializationException {
+ final TestRunner runner = TestRunners.newTestRunner(TestConsumeKinesis.class);
+
+ final AWSCredentialsProviderControllerService credentialsService = new AWSCredentialsProviderControllerService();
+ runner.addControllerService("credentials", credentialsService);
+ runner.setProperty(credentialsService, AWSCredentialsProviderControllerService.ACCESS_KEY_ID, localstack.getAccessKey());
+ runner.setProperty(credentialsService, AWSCredentialsProviderControllerService.SECRET_KEY, localstack.getSecretKey());
+ runner.enableControllerService(credentialsService);
+
+ runner.setProperty(ConsumeKinesis.AWS_CREDENTIALS_PROVIDER_SERVICE, "credentials");
+ runner.setProperty(ConsumeKinesis.STREAM_NAME, streamName);
+ runner.setProperty(ConsumeKinesis.APPLICATION_NAME, applicationName);
+ runner.setProperty(ConsumeKinesis.REGION, localstack.getRegion());
+ runner.setProperty(ConsumeKinesis.INITIAL_STREAM_POSITION, ConsumeKinesis.InitialPosition.TRIM_HORIZON);
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, ConsumeKinesis.ProcessingStrategy.FLOW_FILE);
+
+ runner.setProperty(ConsumeKinesis.METRICS_PUBLISHING, ConsumeKinesis.MetricsPublishing.CLOUDWATCH);
+
+ runner.setProperty(ConsumeKinesis.MAX_BYTES_TO_BUFFER, "10 MB");
+
+ runner.assertValid();
+ return runner;
+ }
+
+ private TestRunner createRecordTestRunner(final String streamName, final String applicationName) throws InitializationException {
+ final TestRunner runner = createTestRunner(streamName, applicationName);
+
+ final JsonTreeReader jsonReader = new JsonTreeReader();
+ runner.addControllerService("json-reader", jsonReader);
+ runner.setProperty(jsonReader, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaInferenceUtil.INFER_SCHEMA.getValue());
+ runner.enableControllerService(jsonReader);
+
+ final JsonRecordSetWriter jsonWriter = new JsonRecordSetWriter();
+ runner.addControllerService("json-writer", jsonWriter);
+ runner.setProperty(jsonWriter, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.INHERIT_RECORD_SCHEMA.getValue());
+ runner.enableControllerService(jsonWriter);
+
+ runner.setProperty(ConsumeKinesis.PROCESSING_STRATEGY, ConsumeKinesis.ProcessingStrategy.RECORD);
+ runner.setProperty(ConsumeKinesis.RECORD_READER, "json-reader");
+ runner.setProperty(ConsumeKinesis.RECORD_WRITER, "json-writer");
+
+ runner.assertValid();
+ return runner;
+ }
+
+ private void runProcessorWithInitAndWaitForFiles(final TestRunner runner, final int expectedFlowFileCount) {
+ runProcessorAndWaitForFiles(runner, expectedFlowFileCount, true);
+ }
+
+ private void runProcessorAndWaitForFiles(final TestRunner runner, final int expectedFlowFileCount) {
+ runProcessorAndWaitForFiles(runner, expectedFlowFileCount, false);
+ }
+
+ private void runProcessorAndWaitForFiles(final TestRunner runner, final int expectedFlowFileCount, final boolean withInit) {
+ logger.info("Running processor and waiting for {} files", expectedFlowFileCount);
+
+ if (withInit) {
+ runner.run(1, false, true);
+ }
+
+ final Set<Relationship> relationships = runner.getProcessor().getRelationships();
+
+ while (true) {
+ runner.run(1, false, false);
+
+ final int currentCount = relationships.stream()
+ .map(runner::getFlowFilesForRelationship)
+ .mapToInt(Collection::size)
+ .sum();
+ logger.info("Current files count: {}, expected: {}", currentCount, expectedFlowFileCount);
+
+ if (currentCount >= expectedFlowFileCount) {
+ return;
+ }
+
+ try {
+ Thread.sleep(5_000);
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Thread interrupted while waiting for files", e);
+ }
+ }
+ }
+
+ private MockProcessSession createFailingSession(final Processor processor) {
+ final SharedSessionState sharedState = new SharedSessionState(processor, new AtomicLong());
+
+ return MockProcessSession.builder(sharedState, processor)
+ .failCommit()
+ .build();
+ }
+
+ public static class TestConsumeKinesis extends ConsumeKinesis {
+
+ @Override
+ URI getKinesisEndpointOverride() {
+ return localstack.getEndpointOverride(Service.KINESIS);
+ }
+
+ @Override
+ URI getDynamoDbEndpointOverride() {
+ return localstack.getEndpointOverride(Service.DYNAMODB);
+ }
+
+ @Override
+ URI getCloudwatchEndpointOverride() {
+ return localstack.getEndpointOverride(Service.CLOUDWATCH);
+ }
+ }
+
+ /**
+ * Test client wrapper for Kinesis operations with built-in retry logic and error handling.
+ */
+ private static class TestKinesisStreamClient {
+ private static final Logger logger = LoggerFactory.getLogger(TestKinesisStreamClient.class);
+
+ private static final int MAX_RETRIES = 10;
+ private static final long INITIAL_RETRY_DELAY_MILLIS = 1_000;
+ private static final long MAX_RETRY_DELAY_MILLIS = 60 * 1_000;
+
+ private static final Duration STREAM_WAIT_TIMEOUT = Duration.ofMinutes(2);
+
+ private final KinesisClient kinesisClient;
+ private final String streamName;
+
+ TestKinesisStreamClient(KinesisClient kinesisClient, String streamName) {
+ this.kinesisClient = kinesisClient;
+ this.streamName = streamName;
+ }
+
+ void createStream(final int shardCount) {
+ logger.info("Creating stream: {} with {} shards", streamName, shardCount);
+
+ executeWithRetry(
+ "createStream",
+ () -> kinesisClient.createStream(req -> req.streamName(streamName).shardCount(shardCount)));
+
+ waitForStreamActive();
+ logger.info("Stream {} is now active", streamName);
+ }
+
+ void deleteStream() {
+ logger.info("Deleting stream: {}", streamName);
+
+ executeWithRetry(
+ "deleteStream",
+ () -> kinesisClient.deleteStream(req -> req.streamName(streamName).enforceConsumerDeletion(true)));
+ }
+
+ void putRecord(final String partitionKey, final String data) {
+ final SdkBytes bytes = SdkBytes.fromString(data, UTF_8);
+
+ executeWithRetry(
+ "putRecord",
+ () -> kinesisClient.putRecord(req -> req.streamName(streamName).partitionKey(partitionKey).data(bytes)));
+ }
+
+ void putRecords(final String partitionKey, final List<String> data) {
+ final List<PutRecordsRequestEntry> records = data.stream()
+ .map(it -> PutRecordsRequestEntry.builder()
+ .data(SdkBytes.fromString(it, UTF_8))
+ .partitionKey(partitionKey)
+ .build())
+ .toList();
+
+ executeWithRetry(
+ "putRecords",
+ () -> kinesisClient.putRecords(req -> req.streamName(streamName).records(records)));
+ }
+
+ /**
+ * Adjusts a number of shards for the stream.
+ * <b>Note: in order to ensure new shards become active, the method waits for 30 seconds.</b>
+ */
+ void reshardStream(final int targetShardCount) {
+ logger.info("Resharding stream {} to {} shards", streamName, targetShardCount);
+
+ executeWithRetry(
+ "reshardStream",
+ () -> kinesisClient.updateShardCount(req -> req.streamName(streamName).targetShardCount(targetShardCount).scalingType(ScalingType.UNIFORM_SCALING)));
+
+ waitForStreamActive();
+
+ try {
+ // After resharding new messages can still be put into the older shards for some time, so we wait a bit.
+ Thread.sleep(30_000);
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+
+ logger.info("Stream {} resharding completed", streamName);
+ }
+
+ private void waitForStreamActive() {
+ final long timeoutMillis = System.currentTimeMillis() + STREAM_WAIT_TIMEOUT.toMillis();
+
+ while (System.currentTimeMillis() < timeoutMillis) {
+ try {
+ final DescribeStreamResponse response = kinesisClient.describeStream(req -> req.streamName(streamName));
+
+ final StreamStatus status = response.streamDescription().streamStatus();
+ if (status == StreamStatus.ACTIVE) {
+ return;
+ }
+
+ logger.info("Stream {} status: {}, waiting...", streamName, status);
+ Thread.sleep(1000);
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Thread interrupted while waiting for stream to be active", e);
+ } catch (final RuntimeException e) {
+ logger.warn("Error checking stream status for {}: {}", streamName, e.getMessage());
+ try {
+ Thread.sleep(1000);
+ } catch (final InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Thread interrupted while waiting for stream to be active", ie);
+ }
+ }
+ }
+
+ throw new IllegalStateException("Stream " + streamName + " did not become active within timeout");
+ }
+
+ private void executeWithRetry(final String operation, final Runnable op) {
+ RuntimeException lastException = null;
+
+ for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) {
+ try {
+ op.run();
+ return;
+ } catch (final RuntimeException e) {
+ lastException = e;
+ logger.warn("Attempt {} of {} failed for operation {}: {}",
+ attempt, MAX_RETRIES, operation, e.getMessage());
+
+ if (attempt < MAX_RETRIES) {
+ try {
+ final long delayMillis = INITIAL_RETRY_DELAY_MILLIS * (1 << (attempt - 1));
+ Thread.sleep(Math.min(delayMillis, MAX_RETRY_DELAY_MILLIS));
+ } catch (final InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Thread interrupted during retry delay", ie);
+ }
+ }
+ }
+ }
+
+ throw new IllegalStateException("Operation " + operation + " failed after " + MAX_RETRIES + " attempts", lastException);
+ }
+ }
+}
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
new file mode 100644
index 0000000000..63216b8073
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ConsumeKinesisTest.java
@@ -0,0 +1,60 @@
+/*
+ * 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;
+
+import org.apache.nifi.processor.Relationship;
+import org.apache.nifi.util.TestRunner;
+import org.apache.nifi.util.TestRunners;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.Set;
+
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.PROCESSING_STRATEGY;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.ProcessingStrategy.FLOW_FILE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.ProcessingStrategy.RECORD;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.REL_PARSE_FAILURE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesis.REL_SUCCESS;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class ConsumeKinesisTest {
+
+ private TestRunner testRunner;
+
+ @BeforeEach
+ void setUp() {
+ testRunner = TestRunners.newTestRunner(ConsumeKinesis.class);
+ }
+
+ @Test
+ void getRelationshipsForFlowFileProcessingStrategy() {
+ testRunner.setProperty(PROCESSING_STRATEGY, FLOW_FILE);
+
+ final Set<Relationship> relationships = testRunner.getProcessor().getRelationships();
+
+ assertEquals(Set.of(REL_SUCCESS), relationships);
+ }
+
+ @Test
+ void getRelationshipsForRecordProcessingStrategy() {
+ testRunner.setProperty(PROCESSING_STRATEGY, RECORD);
+
+ final Set<Relationship> relationships = testRunner.getProcessor().getRelationships();
+
+ assertEquals(Set.of(REL_SUCCESS, REL_PARSE_FAILURE), relationships);
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/JsonRecordAssert.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/JsonRecordAssert.java
new file mode 100644
index 0000000000..e0f3a40e47
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/JsonRecordAssert.java
@@ -0,0 +1,73 @@
+/*
+ * 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;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.util.MockFlowFile;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+
+final class JsonRecordAssert {
+
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ static void assertFlowFileRecords(final FlowFile flowFile, final KinesisClientRecord... expectedRecords) {
+ assertFlowFileRecords(flowFile, List.of(expectedRecords));
+ }
+
+ static void assertFlowFileRecords(final FlowFile flowFile, final List<KinesisClientRecord> expectedRecords) {
+ final List<String> expectedPayloads = expectedRecords.stream()
+ .map(KinesisRecordPayload::extract)
+ .toList();
+ assertFlowFileRecordPayloads(flowFile, expectedPayloads);
+ }
+
+ static void assertFlowFileRecordPayloads(final FlowFile flowFile, final String... expectedPayloads) {
+ assertFlowFileRecordPayloads(flowFile, List.of(expectedPayloads));
+ }
+
+ static void assertFlowFileRecordPayloads(final FlowFile flowFile, final List<String> expectedPayloads) {
+ try {
+ final MockFlowFile mockFlowFile = assertInstanceOf(MockFlowFile.class, flowFile, "A passed FlowFile should be an instance of MockFlowFile");
+
+ final JsonNode node = MAPPER.readTree(mockFlowFile.getContent());
+ final ArrayNode array = assertInstanceOf(ArrayNode.class, node, "FlowFile content is expected to be an array");
+
+ assertEquals(expectedPayloads.size(), array.size(), "Array size mismatch");
+
+ for (int i = 0; i < expectedPayloads.size(); i++) {
+ final JsonNode recordNode = array.get(i);
+ final JsonNode expectedNode = MAPPER.readTree(expectedPayloads.get(i));
+ assertEquals(recordNode, expectedNode, "Record at index " + i + " does not match expected JSON");
+ }
+
+ } catch (final JsonProcessingException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private JsonRecordAssert() {
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisRecordPayload.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisRecordPayload.java
new file mode 100644
index 0000000000..8a03d71d24
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisRecordPayload.java
@@ -0,0 +1,36 @@
+/*
+ * 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;
+
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+final class KinesisRecordPayload {
+
+ static String extract(final KinesisClientRecord record) {
+ record.data().rewind();
+
+ final byte[] buffer = new byte[record.data().remaining()];
+ record.data().get(buffer);
+
+ return new String(buffer, UTF_8);
+ }
+
+ private KinesisRecordPayload() {
+ }
+}
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/MemoryBoundRecordBufferTest.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/MemoryBoundRecordBufferTest.java
new file mode 100644
index 0000000000..608d6f6cc4
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/MemoryBoundRecordBufferTest.java
@@ -0,0 +1,734 @@
+/*
+ * 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;
+
+import org.apache.nifi.documentation.init.NopComponentLog;
+import org.apache.nifi.processors.aws.kinesis.MemoryBoundRecordBuffer.Lease;
+import org.apache.nifi.processors.aws.kinesis.RecordBuffer.ShardBufferId;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import software.amazon.awssdk.services.kinesis.model.Record;
+import software.amazon.kinesis.exceptions.InvalidStateException;
+import software.amazon.kinesis.exceptions.KinesisClientLibDependencyException;
+import software.amazon.kinesis.exceptions.ShutdownException;
+import software.amazon.kinesis.exceptions.ThrottlingException;
+import software.amazon.kinesis.processor.Checkpointer;
+import software.amazon.kinesis.processor.PreparedCheckpointer;
+import software.amazon.kinesis.processor.RecordProcessorCheckpointer;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
+
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class MemoryBoundRecordBufferTest {
+
+ private static final long MAX_MEMORY_BYTES = 1024L;
+ private static final String SHARD_ID_1 = "shard-1";
+ private static final String SHARD_ID_2 = "shard-2";
+ private static final Duration CHECKPOINT_INTERVAL = Duration.ZERO;
+
+ private MemoryBoundRecordBuffer recordBuffer;
+ private TestCheckpointer checkpointer1;
+ private TestCheckpointer checkpointer2;
+
+ @BeforeEach
+ void setUp() {
+ recordBuffer = new MemoryBoundRecordBuffer(new NopComponentLog(), MAX_MEMORY_BYTES, CHECKPOINT_INTERVAL);
+ checkpointer1 = new TestCheckpointer();
+ checkpointer2 = new TestCheckpointer();
+ }
+
+ @Test
+ void testCreateBuffer() {
+ final ShardBufferId bufferId1 = recordBuffer.createBuffer(SHARD_ID_1);
+ assertEquals(SHARD_ID_1, bufferId1.shardId());
+
+ final ShardBufferId bufferId2 = recordBuffer.createBuffer(SHARD_ID_2);
+ assertEquals(SHARD_ID_2, bufferId2.shardId());
+
+ final ShardBufferId newBufferId1 = recordBuffer.createBuffer(SHARD_ID_1);
+ assertEquals(SHARD_ID_1, newBufferId1.shardId());
+
+ assertNotEquals(bufferId1, bufferId2);
+ assertNotEquals(bufferId1, newBufferId1);
+ }
+
+ @Test
+ void testAddRecordsToBuffer() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ // Buffer without records is not available for leasing.
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty());
+
+ final List<KinesisClientRecord> records = createTestRecords(2);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ // Should be able to get buffer ID from pool since buffer has records.
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(SHARD_ID_1, lease.shardId());
+ }
+
+ @Test
+ void testAddEmptyRecordsList() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> emptyRecords = Collections.emptyList();
+
+ recordBuffer.addRecords(bufferId, emptyRecords, checkpointer1);
+
+ // Should not be able to get buffer ID from pool since no records were added.
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty());
+ }
+
+ @Test
+ void testConsumeRecords() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(3);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+ assertEquals(records, consumedRecords);
+ // Just consuming record should not checkpoint them.
+ assertFalse(checkpointer1.isCheckpointed());
+ }
+
+ @Test
+ void testCommitConsumedRecords() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(2);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ recordBuffer.consumeRecords(lease);
+ recordBuffer.commitConsumedRecords(lease);
+
+ assertTrue(checkpointer1.isCheckpointed());
+ }
+
+ @Test
+ void testRollbackConsumedRecords() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(3);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+ final List<String> messages = consumedRecords.stream()
+ .map(this::readContent)
+ .toList();
+
+ recordBuffer.rollbackConsumedRecords(lease);
+
+ // Checkpointer should not be called during rollback.
+ assertFalse(checkpointer1.isCheckpointed());
+
+ final List<String> rolledBackMessages = recordBuffer.consumeRecords(lease).stream()
+ .map(this::readContent)
+ .toList();
+ assertEquals(messages, rolledBackMessages);
+ }
+
+ private String readContent(final KinesisClientRecord record) {
+ final ByteBuffer data = record.data();
+ final byte[] buffer = new byte[data.remaining()];
+ data.get(buffer);
+ return new String(buffer, StandardCharsets.UTF_8);
+ }
+
+ @Test
+ void testReturnBufferIdToPool() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(2);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ final Lease lease1 = recordBuffer.acquireBufferLease().orElseThrow();
+
+ // Consume some records but don't return buffer id to the pool.
+ recordBuffer.consumeRecords(lease1);
+
+ recordBuffer.addRecords(bufferId, createTestRecords(1), checkpointer2);
+
+ // The buffer is still unavailable.
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty());
+
+ // After returning id to the pool it's possible to get the buffer from pool again.
+ recordBuffer.returnBufferLease(lease1);
+ final Lease lease2 = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(SHARD_ID_1, lease2.shardId());
+ }
+
+ @Test
+ void testReturnBufferIdToPool_multipleReturns() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(2);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ final Lease lease1 = recordBuffer.acquireBufferLease().orElseThrow();
+
+ recordBuffer.returnBufferLease(lease1);
+ recordBuffer.returnBufferLease(lease1);
+
+ // Can retrieve the id only once.
+ final Lease lease2 = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(SHARD_ID_1, lease2.shardId());
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty());
+ }
+
+ @Test
+ void testReturnBufferIdToPool_withUncommittedRecords() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(2);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ final Lease lease1 = recordBuffer.acquireBufferLease().orElseThrow();
+
+ // Consume some records, but don't commit them.
+ final List<KinesisClientRecord> lease1Records = recordBuffer.consumeRecords(lease1);
+ recordBuffer.returnBufferLease(lease1);
+
+ final Lease lease2 = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(SHARD_ID_1, lease2.shardId());
+ final List<KinesisClientRecord> lease2Records = recordBuffer.consumeRecords(lease2);
+
+ // Until committed, the records stay in the buffer.
+ assertEquals(lease1Records, lease2Records);
+ }
+
+ @Test
+ void testReturnBufferIdToPoolWithNoPendingRecords() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ // Get buffer from pool and consume all records.
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ recordBuffer.consumeRecords(lease);
+ recordBuffer.commitConsumedRecords(lease);
+
+ // Return buffer ID to pool - should not make buffer available since no pending records.
+ recordBuffer.returnBufferLease(lease);
+
+ // Should not be able to get buffer from pool.
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty());
+ }
+
+ @Test
+ void testConsumerLeaseLost() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(2);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+ // Before lease lost buffer should be available in the pool.
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(SHARD_ID_1, lease.shardId());
+
+ // Simulate lease lost.
+ recordBuffer.consumerLeaseLost(bufferId);
+
+ // Should not be able to consume records from invalidated buffer.
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+ assertTrue(consumedRecords.isEmpty());
+
+ // Should not be able to commit records for invalidated buffer.
+ recordBuffer.commitConsumedRecords(lease);
+ assertFalse(checkpointer1.isCheckpointed());
+
+ // Buffer should not be available in pool.
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty());
+ }
+
+ @Test
+ void testCheckpointEndedShard() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ recordBuffer.consumeRecords(lease);
+ recordBuffer.commitConsumedRecords(lease);
+
+ recordBuffer.checkpointEndedShard(bufferId, checkpointer2);
+ assertTrue(checkpointer2.isCheckpointed());
+
+ // Buffer should be removed and not available for operations.
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+ assertTrue(consumedRecords.isEmpty());
+ }
+
+ @Test
+ void testShutdownShardConsumption_forEmptyBuffer() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ recordBuffer.consumeRecords(lease);
+ recordBuffer.commitConsumedRecords(lease);
+
+ recordBuffer.shutdownShardConsumption(bufferId, checkpointer2);
+ assertTrue(checkpointer2.isCheckpointed());
+
+ // Buffer should be removed and not available for operations.
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+ assertTrue(consumedRecords.isEmpty());
+ }
+
+ @Test
+ void testShutdownShardConsumption_forBufferWithRecords() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+
+ recordBuffer.shutdownShardConsumption(bufferId, checkpointer2);
+ assertFalse(checkpointer1.isCheckpointed());
+ assertFalse(checkpointer2.isCheckpointed());
+
+ assertTrue(recordBuffer.acquireBufferLease().isEmpty(), "Buffer should not be available after shutdown");
+ }
+
+ @Test
+ @Timeout(value = 5, unit = SECONDS)
+ void testMemoryBackpressure() throws InterruptedException {
+ // Create buffer with small memory limit.
+ final MemoryBoundRecordBuffer recordBuffer = new MemoryBoundRecordBuffer(new NopComponentLog(), 100L, CHECKPOINT_INTERVAL);
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+
+ // Still fits into the buffer.
+ final List<KinesisClientRecord> initialRecords = List.of(createRecordWithSize(80), createRecordWithSize(20));
+ recordBuffer.addRecords(bufferId, initialRecords, checkpointer1);
+
+ final CountDownLatch startLatch = new CountDownLatch(1);
+
+ final List<KinesisClientRecord> notFittingRecords = List.of(createRecordWithSize(50));
+ // Thread that tries to add records (should block due to memory limit).
+ final Thread addRecordsThread = new Thread(() -> {
+ startLatch.countDown();
+ // Doesn't fit into the buffer.
+ recordBuffer.addRecords(bufferId, notFittingRecords, checkpointer1);
+ });
+
+ addRecordsThread.start();
+ startLatch.await();
+
+ // Wait for thread to try to add records and get blocked.
+ Thread.sleep(200);
+ assertTrue(addRecordsThread.isAlive(), "Thread should be blocked waiting for memory");
+
+ // Commit records in the buffer to free memory.
+ final Lease lease1 = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(initialRecords, recordBuffer.consumeRecords(lease1));
+ recordBuffer.commitConsumedRecords(lease1);
+ recordBuffer.returnBufferLease(lease1);
+
+ // Thread should get unblocked and add the message.
+ addRecordsThread.join();
+ final Lease lease2 = recordBuffer.acquireBufferLease().orElseThrow();
+ assertEquals(notFittingRecords, recordBuffer.consumeRecords(lease2));
+ }
+
+ @Test
+ void testMemoryBackpressure_forRecordsLargerThanMaxBuffer() {
+ // Create buffer with small memory limit.
+ final int bufferSize = 10;
+ final MemoryBoundRecordBuffer recordBuffer = new MemoryBoundRecordBuffer(new NopComponentLog(), bufferSize, CHECKPOINT_INTERVAL);
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+
+ // A single batch with size double the buffer size.
+ final List<KinesisClientRecord> reallyLargeBatch = List.of(createRecordWithSize(bufferSize), createRecordWithSize(bufferSize));
+
+ // It's possible to insert a batch that exceeds the buffer size.
+ assertDoesNotThrow(() -> recordBuffer.addRecords(bufferId, reallyLargeBatch, checkpointer1));
+ }
+
+ @Test
+ @Timeout(value = 5, unit = SECONDS)
+ void testConcurrentBufferAccess() throws InterruptedException {
+ final int numberOfShards = 10;
+ final int numberOfRecordConsumers = 10;
+ final int totalThreads = numberOfShards + numberOfRecordConsumers;
+
+ final int recordsPerShard = 5;
+
+ final ExecutorService executor = Executors.newFixedThreadPool(totalThreads);
+ // All threads start at the same moment.
+ final CountDownLatch startLatch = new CountDownLatch(totalThreads);
+ final CountDownLatch finishLatch = new CountDownLatch(totalThreads);
+
+ final List<Integer> shardIds = IntStream.range(0, numberOfShards).boxed().toList();
+ final List<ShardBufferId> bufferIds = shardIds.stream()
+ .map(shardId -> recordBuffer.createBuffer("shard-" + shardId))
+ .toList();
+
+ // Every producer returns a checkpointer for the records it produced.
+ final List<TestCheckpointer> checkpointers = shardIds.stream()
+ .map(shardId -> {
+ final ShardBufferId bufferId = bufferIds.get(shardId);
+ final List<KinesisClientRecord> records = createTestRecords(recordsPerShard);
+ final TestCheckpointer threadCheckpointer = new TestCheckpointer();
+
+ executor.submit(() -> {
+ try {
+ startLatch.countDown();
+ startLatch.await();
+
+ recordBuffer.addRecords(bufferId, records, threadCheckpointer);
+ } catch (final InterruptedException e) {
+ throw new RuntimeException(e);
+ } finally {
+ finishLatch.countDown();
+ }
+ });
+
+ return threadCheckpointer;
+ })
+ .toList();
+
+ // Every consumer returns a list of shardId:sequenceNumber strings for the records it consumed.
+ final List<Future<List<String>>> processedRecordsFutures = IntStream.range(0, numberOfRecordConsumers)
+ .mapToObj(__ -> executor.submit(() -> {
+ try {
+ startLatch.countDown();
+ startLatch.await();
+
+ Optional<Lease> maybeLease = Optional.empty();
+ while (maybeLease.isEmpty()) {
+ maybeLease = recordBuffer.acquireBufferLease();
+ Thread.sleep(100); // Wait for records to be added by producers.
+ }
+
+ final Lease lease = maybeLease.orElseThrow();
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+
+ recordBuffer.commitConsumedRecords(lease);
+
+ return consumedRecords.stream()
+ .map(record -> lease.shardId() + ":" + record.sequenceNumber())
+ .toList();
+ } catch (final InterruptedException e) {
+ throw new RuntimeException(e);
+ } finally {
+ finishLatch.countDown();
+ }
+ }))
+ .toList();
+
+ finishLatch.await();
+ executor.shutdown();
+
+ assertAll(
+ checkpointers.stream()
+ .map(it -> () -> assertTrue(it.isCheckpointed(), "Every checkpointer should have been called"))
+ );
+
+ final long uniqueRecordsProcessed = processedRecordsFutures.stream()
+ .map(future -> {
+ try {
+ return future.get();
+ } catch (final Exception e) {
+ throw new RuntimeException(e);
+ }
+ })
+ .flatMap(Collection::stream)
+ .distinct()
+ .count();
+ assertEquals(numberOfShards * recordsPerShard, uniqueRecordsProcessed);
+ }
+
+ @Test
+ void testGetMultipleBuffersFromPool() {
+ final ShardBufferId bufferId1 = recordBuffer.createBuffer(SHARD_ID_1);
+ final ShardBufferId bufferId2 = recordBuffer.createBuffer(SHARD_ID_2);
+
+ // Add records to both buffers.
+ recordBuffer.addRecords(bufferId1, createTestRecords(2), checkpointer1);
+ recordBuffer.addRecords(bufferId2, createTestRecords(3), checkpointer2);
+
+ // Should be able to get both buffer IDs from pool.
+ final Lease lease1 = recordBuffer.acquireBufferLease().orElseThrow();
+ final Lease lease2 = recordBuffer.acquireBufferLease().orElseThrow();
+
+ assertEquals(SHARD_ID_1, lease1.shardId());
+ assertEquals(SHARD_ID_2, lease2.shardId());
+
+ // Should get different buffer IDs.
+ assertNotEquals(lease1.shardId(), lease2.shardId());
+ }
+
+ @Test
+ void testCommitRecordsWhileNewRecordsArrive() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+
+ // Add records to pending queue.
+ final List<KinesisClientRecord> batch1 = createTestRecords(2);
+ recordBuffer.addRecords(bufferId, batch1, checkpointer1);
+
+ // Consume batch1 records (moves from pending to in-progress).
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+
+ final List<KinesisClientRecord> consumedRecords = recordBuffer.consumeRecords(lease);
+ assertEquals(batch1, consumedRecords);
+
+ // Add more records while others are in-progress.
+ final List<KinesisClientRecord> batch2 = createTestRecords(1);
+ recordBuffer.addRecords(bufferId, batch2, checkpointer2);
+
+ // Commit in-progress records.
+ recordBuffer.commitConsumedRecords(lease);
+
+ // Consume batch2 records.
+ final List<KinesisClientRecord> remainingRecords = recordBuffer.consumeRecords(lease);
+ assertEquals(batch2, remainingRecords);
+ }
+
+ @ParameterizedTest
+ @ValueSource(classes = {
+ KinesisClientLibDependencyException.class,
+ InvalidStateException.class,
+ ThrottlingException.class,
+ })
+ void testRetriableCheckpointExceptions(final Class<? extends Exception> exceptionClass) throws Exception {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ final Exception exception = exceptionClass.getDeclaredConstructor(String.class).newInstance("Thrown from test");
+ final TestCheckpointer failingCheckpointer = new TestCheckpointer(exception, 2);
+
+ recordBuffer.addRecords(bufferId, records, failingCheckpointer);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+ recordBuffer.consumeRecords(lease);
+
+ // Should handle all exception types gracefully.
+ recordBuffer.commitConsumedRecords(lease);
+
+ assertTrue(failingCheckpointer.isCheckpointed());
+ }
+
+ @Test
+ void testShutdownDuringCheckpoint() {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ final ShutdownException exception = new ShutdownException("Test shutdown exception");
+ final TestCheckpointer failingCheckpointer = new TestCheckpointer(exception, 1);
+
+ recordBuffer.addRecords(bufferId, records, failingCheckpointer);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+ recordBuffer.consumeRecords(lease);
+
+ recordBuffer.commitConsumedRecords(lease);
+ assertFalse(failingCheckpointer.isCheckpointed());
+ }
+
+ @Test
+ @Timeout(value = 3, unit = SECONDS)
+ void testCheckpointEndedShardWaitsForPendingRecords() throws InterruptedException {
+ final ShardBufferId bufferId = recordBuffer.createBuffer(SHARD_ID_1);
+ final List<KinesisClientRecord> records = createTestRecords(1);
+
+ recordBuffer.addRecords(bufferId, records, checkpointer1);
+ final Lease lease = recordBuffer.acquireBufferLease().orElseThrow();
+ recordBuffer.consumeRecords(lease);
+
+ final CountDownLatch finishStarted = new CountDownLatch(1);
+
+ // Start finishConsumption in another thread.
+ final Thread finishThread = new Thread(() -> {
+ finishStarted.countDown();
+ recordBuffer.checkpointEndedShard(bufferId, checkpointer2);
+ });
+
+ finishThread.start();
+ finishStarted.await();
+
+ // Give finishThread time to get blocked.
+ Thread.sleep(200);
+ assertTrue(finishThread.isAlive(), "finishConsumption should block until records are committed");
+
+ // Commit records to unblock finishConsumption.
+ recordBuffer.commitConsumedRecords(lease);
+ assertTrue(checkpointer1.isCheckpointed());
+
+ finishThread.join();
+ assertTrue(checkpointer2.isCheckpointed(), "Checkpointer should be called after finishConsumption unblocks");
+ }
+
+ private List<KinesisClientRecord> createTestRecords(int count) {
+ return IntStream.range(0, count)
+ .mapToObj(i -> {
+ final String data = "test-record-" + i;
+ return KinesisClientRecord.builder()
+ .data(ByteBuffer.wrap(data.getBytes(StandardCharsets.UTF_8)).asReadOnlyBuffer())
+ .partitionKey("partition-" + i)
+ .sequenceNumber(String.valueOf(i))
+ .build();
+ })
+ .toList();
+ }
+
+ private KinesisClientRecord createRecordWithSize(int sizeBytes) {
+ final byte[] data = new byte[sizeBytes];
+ Arrays.fill(data, (byte) 'X');
+
+ return KinesisClientRecord.builder()
+ .data(ByteBuffer.wrap(data).asReadOnlyBuffer())
+ .partitionKey("partition")
+ .sequenceNumber("1")
+ .build();
+ }
+
+ /**
+ * Test implementation of RecordProcessorCheckpointer that tracks checkpoint calls
+ * and can simulate various exception scenarios.
+ */
+ private static class TestCheckpointer implements RecordProcessorCheckpointer {
+
+ private final AtomicBoolean checkpointed = new AtomicBoolean(false);
+ private final Exception exceptionToThrow;
+ private final AtomicInteger throwsLeft;
+
+ TestCheckpointer() {
+ this.exceptionToThrow = null;
+ this.throwsLeft = new AtomicInteger(0);
+ }
+
+ TestCheckpointer(final Exception exceptionToThrow, final int maxThrows) {
+ this.exceptionToThrow = exceptionToThrow;
+ this.throwsLeft = new AtomicInteger(maxThrows);
+ }
+
+ @Override
+ public void checkpoint() throws KinesisClientLibDependencyException, InvalidStateException, ThrottlingException, ShutdownException {
+ if (exceptionToThrow != null && throwsLeft.decrementAndGet() == 0) {
+ switch (exceptionToThrow) {
+ case KinesisClientLibDependencyException e -> throw e;
+ case InvalidStateException e -> throw e;
+ case ThrottlingException e -> throw e;
+ case ShutdownException e -> throw e;
+ default -> throw new RuntimeException(exceptionToThrow);
+ }
+ }
+
+ if (checkpointed.getAndSet(true)) {
+ throw new IllegalStateException("TestCheckpointer has already been checkpointed");
+ }
+ }
+
+ @Override
+ public void checkpoint(Record record) {
+ throw notImplemented();
+ }
+
+ @Override
+ public void checkpoint(String sequenceNumber) {
+ throw notImplemented();
+ }
+
+ @Override
+ public void checkpoint(String sequenceNumber, long subSequenceNumber) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint() {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(byte[] applicationState) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(Record record) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(Record record, byte[] applicationState) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(String sequenceNumber) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(String sequenceNumber, byte[] applicationState) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(String sequenceNumber, long subSequenceNumber) {
+ throw notImplemented();
+ }
+
+ @Override
+ public PreparedCheckpointer prepareCheckpoint(String sequenceNumber, long subSequenceNumber, byte[] applicationState) {
+ throw notImplemented();
+ }
+
+ @Override
+ public Checkpointer checkpointer() {
+ throw notImplemented();
+ }
+
+ boolean isCheckpointed() {
+ return checkpointed.get();
+ }
+
+ private static RuntimeException notImplemented() {
+ return new UnsupportedOperationException("Not implemented for test");
+ }
+ }
+}
\ No newline at end of file
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ReaderRecordProcessorTest.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ReaderRecordProcessorTest.java
new file mode 100644
index 0000000000..48a4abd687
--- /dev/null
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/ReaderRecordProcessorTest.java
@@ -0,0 +1,410 @@
+/*
+ * 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;
+
+import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.json.JsonRecordSetWriter;
+import org.apache.nifi.json.JsonTreeReader;
+import org.apache.nifi.logging.ComponentLog;
+import org.apache.nifi.processors.aws.kinesis.ReaderRecordProcessor.ProcessingResult;
+import org.apache.nifi.reporting.InitializationException;
+import org.apache.nifi.schema.access.SchemaAccessUtils;
+import org.apache.nifi.schema.inference.SchemaInferenceUtil;
+import org.apache.nifi.serialization.MalformedRecordException;
+import org.apache.nifi.serialization.RecordReader;
+import org.apache.nifi.serialization.RecordReaderFactory;
+import org.apache.nifi.serialization.record.MockRecordParser;
+import org.apache.nifi.serialization.record.Record;
+import org.apache.nifi.serialization.record.RecordSchema;
+import org.apache.nifi.util.MockFlowFile;
+import org.apache.nifi.util.MockProcessSession;
+import org.apache.nifi.util.SharedSessionState;
+import org.apache.nifi.util.TestRunner;
+import org.apache.nifi.util.TestRunners;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.time.Instant;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.APPROXIMATE_ARRIVAL_TIMESTAMP;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.LAST_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.LAST_SUB_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.MIME_TYPE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.PARTITION_KEY;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_COUNT;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.RECORD_ERROR_MESSAGE;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.FIRST_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.SHARD_ID;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.FIRST_SUB_SEQUENCE_NUMBER;
+import static org.apache.nifi.processors.aws.kinesis.ConsumeKinesisAttributes.STREAM_NAME;
+import static org.apache.nifi.processors.aws.kinesis.JsonRecordAssert.assertFlowFileRecords;
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class ReaderRecordProcessorTest {
+
+ private static final String TEST_STREAM_NAME = "stream-test";
+ private static final String TEST_SHARD_ID = "shardId-test";
+
+ private static final String USER_JSON_1 = "{\"name\":\"John\",\"age\":30}";
+ private static final String USER_JSON_2 = "{\"name\":\"Jane\",\"age\":25}";
+ private static final String USER_JSON_3 = "{\"name\":\"Bob\",\"age\":35}";
+
+ private static final String CITY_JSON_1 = "{\"name\":\"Seattle\",\"country\":\"US\"}";
+ private static final String CITY_JSON_2 = "{\"name\":\"Warsaw\",\"country\":\"PL\"}";
+
+ private static final String INVALID_JSON = "{invalid json}";
+
+ private MockProcessSession session;
+ private ComponentLog logger;
+
+ private JsonTreeReader jsonReader;
+ private JsonRecordSetWriter jsonWriter;
+
+ @BeforeEach
+ void setUp() throws InitializationException {
+ final TestRunner runner = TestRunners.newTestRunner(ConsumeKinesis.class);
+ final SharedSessionState sharedState = new SharedSessionState(runner.getProcessor(), new AtomicLong(0));
+ session = new MockProcessSession(sharedState, runner.getProcessor());
+ logger = runner.getLogger();
+
+ jsonReader = new JsonTreeReader();
+ runner.addControllerService("json-reader", jsonReader);
+ runner.setProperty(jsonReader, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaInferenceUtil.INFER_SCHEMA.getValue());
+ runner.enableControllerService(jsonReader);
+
+ jsonWriter = new JsonRecordSetWriter();
+ runner.addControllerService("json-writer", jsonWriter);
+ runner.setProperty(jsonWriter, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.INHERIT_RECORD_SCHEMA.getValue());
+ runner.enableControllerService(jsonWriter);
+ }
+
+ @Test
+ void testProcessSingleRecord() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final KinesisClientRecord record = KinesisClientRecord.builder()
+ .data(ByteBuffer.wrap(USER_JSON_1.getBytes(UTF_8)))
+ .sequenceNumber("1")
+ .subSequenceNumber(2)
+ .approximateArrivalTimestamp(Instant.now())
+ .partitionKey("key-123")
+ .build();
+ final List<KinesisClientRecord> records = List.of(record);
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(1, result.successFlowFiles().size());
+ assertEquals(0, result.parseFailureFlowFiles().size());
+
+ final FlowFile successFlowFile = result.successFlowFiles().getFirst();
+
+ assertEquals(TEST_STREAM_NAME, successFlowFile.getAttribute(STREAM_NAME));
+ assertEquals(TEST_SHARD_ID, successFlowFile.getAttribute(SHARD_ID));
+
+ assertEquals(record.sequenceNumber(), successFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals(String.valueOf(record.subSequenceNumber()), successFlowFile.getAttribute(FIRST_SUB_SEQUENCE_NUMBER));
+ assertEquals(record.sequenceNumber(), successFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertEquals(String.valueOf(record.subSequenceNumber()), successFlowFile.getAttribute(LAST_SUB_SEQUENCE_NUMBER));
+
+ assertEquals(record.partitionKey(), successFlowFile.getAttribute(PARTITION_KEY));
+ assertEquals(String.valueOf(record.approximateArrivalTimestamp().toEpochMilli()), successFlowFile.getAttribute(APPROXIMATE_ARRIVAL_TIMESTAMP));
+
+ assertEquals("application/json", successFlowFile.getAttribute(MIME_TYPE));
+ assertEquals("1", successFlowFile.getAttribute(RECORD_COUNT));
+
+ assertFlowFileRecords(successFlowFile, records);
+ }
+
+ @Test
+ void testProcessMultipleRecordsWithSameSchema() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final List<KinesisClientRecord> records = List.of(
+ createKinesisRecord(USER_JSON_1, "1"),
+ createKinesisRecord(USER_JSON_2, "2"),
+ createKinesisRecord(USER_JSON_3, "3")
+ );
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(1, result.successFlowFiles().size());
+ assertEquals(0, result.parseFailureFlowFiles().size());
+
+ final FlowFile successFlowFile = result.successFlowFiles().getFirst();
+ assertEquals(TEST_STREAM_NAME, successFlowFile.getAttribute(STREAM_NAME));
+ assertEquals(TEST_SHARD_ID, successFlowFile.getAttribute(SHARD_ID));
+ assertEquals("3", successFlowFile.getAttribute(RECORD_COUNT));
+
+ assertEquals(records.getFirst().sequenceNumber(), successFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals(String.valueOf(records.getFirst().subSequenceNumber()), successFlowFile.getAttribute(FIRST_SUB_SEQUENCE_NUMBER));
+ assertEquals(records.getLast().sequenceNumber(), successFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertEquals(String.valueOf(records.getLast().subSequenceNumber()), successFlowFile.getAttribute(LAST_SUB_SEQUENCE_NUMBER));
+
+ assertFlowFileRecords(successFlowFile, records);
+ }
+
+ @Test
+ void testEmptyRecordsList() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, Collections.emptyList());
+
+ assertEquals(0, result.successFlowFiles().size());
+ assertEquals(0, result.parseFailureFlowFiles().size());
+ }
+
+ @Test
+ void testSchemaChangeCreatesNewFlowFile() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final List<KinesisClientRecord> records = List.of(
+ createKinesisRecord(USER_JSON_1, "1"),
+ createKinesisRecord(CITY_JSON_1, "2")
+ );
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(2, result.successFlowFiles().size()); // Two different schemas = two FlowFiles
+ assertEquals(0, result.parseFailureFlowFiles().size());
+
+ final FlowFile firstFlowFile = result.successFlowFiles().getFirst();
+ assertEquals("1", firstFlowFile.getAttribute(RECORD_COUNT));
+ assertFlowFileRecords(firstFlowFile, records.getFirst());
+
+ final FlowFile secondFlowFile = result.successFlowFiles().get(1);
+ assertEquals("1", secondFlowFile.getAttribute(RECORD_COUNT));
+ assertFlowFileRecords(secondFlowFile, records.get(1));
+ }
+
+ @Test
+ void testSchemaChangeWithMultipleRecordsInBetween() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final List<KinesisClientRecord> records = List.of(
+ createKinesisRecord(USER_JSON_1, "1"),
+ createKinesisRecord(USER_JSON_2, "2"),
+ createKinesisRecord(CITY_JSON_1, "3"),
+ createKinesisRecord(CITY_JSON_2, "4")
+ );
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(2, result.successFlowFiles().size());
+ assertEquals(0, result.parseFailureFlowFiles().size());
+
+ final FlowFile firstFlowFile = result.successFlowFiles().getFirst();
+ assertEquals("2", firstFlowFile.getAttribute(RECORD_COUNT));
+ assertFlowFileRecords(firstFlowFile, records.subList(0, 2));
+
+ final FlowFile secondFlowFile = result.successFlowFiles().get(1);
+ assertEquals("2", secondFlowFile.getAttribute(RECORD_COUNT));
+ assertFlowFileRecords(secondFlowFile, records.subList(2, 4));
+ }
+
+ @Test
+ void testSingleMalformedRecord() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final List<KinesisClientRecord> records = List.of(
+ createKinesisRecord(INVALID_JSON, "1")
+ );
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(0, result.successFlowFiles().size());
+ assertEquals(1, result.parseFailureFlowFiles().size());
+
+ final MockFlowFile failureFlowFile = (MockFlowFile) result.parseFailureFlowFiles().getFirst();
+ assertEquals(TEST_SHARD_ID, failureFlowFile.getAttribute(SHARD_ID));
+ assertEquals("1", failureFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("1", failureFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertNotNull(failureFlowFile.getAttribute(RECORD_ERROR_MESSAGE));
+
+ failureFlowFile.assertContentEquals(INVALID_JSON, UTF_8);
+ }
+
+ @Test
+ void testMalformedRecordBetweenValid() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final List<KinesisClientRecord> records = List.of(
+ createKinesisRecord(USER_JSON_1, "1"),
+ createKinesisRecord(INVALID_JSON, "2"),
+ createKinesisRecord(USER_JSON_2, "3"),
+ createKinesisRecord(INVALID_JSON, "4"),
+ createKinesisRecord(USER_JSON_3, "5")
+ );
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(1, result.successFlowFiles().size());
+ assertEquals(2, result.parseFailureFlowFiles().size());
+
+ final FlowFile successFlowFile = result.successFlowFiles().getFirst();
+ assertEquals(TEST_SHARD_ID, successFlowFile.getAttribute(SHARD_ID));
+ assertEquals("3", successFlowFile.getAttribute(RECORD_COUNT));
+ assertEquals("1", successFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("5", successFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertFlowFileRecords(successFlowFile, records.get(0), records.get(2), records.get(4));
+
+ assertAll(result.parseFailureFlowFiles().stream().map(
+ failureFlowFile -> () -> {
+ assertNotNull(failureFlowFile.getAttribute(RECORD_ERROR_MESSAGE));
+ assertEquals(TEST_SHARD_ID, failureFlowFile.getAttribute(SHARD_ID));
+ }
+ ));
+ }
+
+ @Test
+ void testIOExceptionDuringReaderCreation() {
+ final RecordReaderFactory failingReaderFactory = new MockRecordParser() {
+ @Override
+ public RecordReader createRecordReader(Map<String, String> variables, InputStream in, long inputLength, ComponentLog logger) throws IOException {
+ throw new IOException("Failed to create reader");
+ }
+ };
+
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(failingReaderFactory, jsonWriter, logger);
+
+ final KinesisClientRecord record = createKinesisRecord(USER_JSON_1, "1");
+ final List<KinesisClientRecord> records = List.of(record);
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(0, result.successFlowFiles().size());
+ assertEquals(1, result.parseFailureFlowFiles().size());
+
+ final MockFlowFile failureFlowFile = (MockFlowFile) result.parseFailureFlowFiles().getFirst();
+ assertTrue(failureFlowFile.getAttribute(RECORD_ERROR_MESSAGE).contains("Failed to create reader"));
+ failureFlowFile.assertContentEquals(KinesisRecordPayload.extract(record), UTF_8);
+ }
+
+ @Test
+ void testMalformedRecordExceptionDuringReading() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(getMalformedRecordExceptionReader(), jsonWriter, logger);
+
+ final KinesisClientRecord record = createKinesisRecord(USER_JSON_1, "1");
+ final List<KinesisClientRecord> records = Collections.singletonList(record);
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(0, result.successFlowFiles().size());
+ assertEquals(1, result.parseFailureFlowFiles().size());
+
+ final MockFlowFile failureFlowFile = (MockFlowFile) result.parseFailureFlowFiles().getFirst();
+ assertTrue(failureFlowFile.getAttribute(RECORD_ERROR_MESSAGE).contains("Test exception"));
+ failureFlowFile.assertContentEquals(KinesisRecordPayload.extract(record), UTF_8);
+ }
+
+ @Test
+ void testInvalidRecordsWithSchemaEvolution() {
+ final ReaderRecordProcessor processor = new ReaderRecordProcessor(jsonReader, jsonWriter, logger);
+
+ final List<KinesisClientRecord> records = List.of(
+ createKinesisRecord(USER_JSON_1, "1"), // Schema A
+ createKinesisRecord(USER_JSON_2, "2"), // Schema A
+ createKinesisRecord(INVALID_JSON, "3"),
+ createKinesisRecord(CITY_JSON_1, "4"), // Schema B
+ createKinesisRecord(INVALID_JSON, "5"),
+ createKinesisRecord(CITY_JSON_2, "6"), // Schema B
+ createKinesisRecord("{\"category\":\"electronics\",\"price\":99.99}", "7") // Schema C
+ );
+
+ final ProcessingResult result = processor.processRecords(session, TEST_STREAM_NAME, TEST_SHARD_ID, records);
+
+ assertEquals(3, result.successFlowFiles().size());
+ assertEquals(2, result.parseFailureFlowFiles().size());
+
+ final FlowFile firstFlowFile = result.successFlowFiles().getFirst();
+ assertEquals("2", firstFlowFile.getAttribute(RECORD_COUNT));
+ assertEquals("1", firstFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("2", firstFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertFlowFileRecords(firstFlowFile, records.subList(0, 2));
+
+ final FlowFile secondFlowFile = result.successFlowFiles().get(1);
+ assertEquals("2", secondFlowFile.getAttribute(RECORD_COUNT));
+ assertEquals("4", secondFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("6", secondFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertFlowFileRecords(secondFlowFile, records.get(3), records.get(5));
+
+ final FlowFile thirdFlowFile = result.successFlowFiles().get(2);
+ assertEquals("1", thirdFlowFile.getAttribute(RECORD_COUNT));
+ assertEquals("7", thirdFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("7", thirdFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertFlowFileRecords(thirdFlowFile, records.get(6));
+
+ final List<FlowFile> failureFlowFiles = result.parseFailureFlowFiles();
+
+ final MockFlowFile firstFailureFlowFile = (MockFlowFile) failureFlowFiles.getFirst();
+ assertEquals("3", firstFailureFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("3", firstFailureFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertNotNull(firstFailureFlowFile.getAttribute(RECORD_ERROR_MESSAGE));
+ assertEquals(TEST_SHARD_ID, firstFailureFlowFile.getAttribute(SHARD_ID));
+ firstFailureFlowFile.assertContentEquals(KinesisRecordPayload.extract(records.get(2)), UTF_8);
+
+ final MockFlowFile secondFailureFlowFile = (MockFlowFile) failureFlowFiles.get(1);
+ assertEquals("5", secondFailureFlowFile.getAttribute(FIRST_SEQUENCE_NUMBER));
+ assertEquals("5", secondFailureFlowFile.getAttribute(LAST_SEQUENCE_NUMBER));
+ assertNotNull(secondFailureFlowFile.getAttribute(RECORD_ERROR_MESSAGE));
+ assertEquals(TEST_SHARD_ID, secondFailureFlowFile.getAttribute(SHARD_ID));
+ secondFailureFlowFile.assertContentEquals(KinesisRecordPayload.extract(records.get(4)), UTF_8);
+ }
+
+ private static KinesisClientRecord createKinesisRecord(final String data, final String sequenceNumber) {
+ return KinesisClientRecord.builder()
+ .data(ByteBuffer.wrap(data.getBytes(UTF_8)))
+ .sequenceNumber(sequenceNumber)
+ .partitionKey("key")
+ .approximateArrivalTimestamp(Instant.now())
+ .build();
+ }
+
+ private static RecordReaderFactory getMalformedRecordExceptionReader() {
+ return new MockRecordParser() {
+ @Override
+ public RecordReader createRecordReader(Map<String, String> variables, InputStream in, long inputLength, ComponentLog logger) {
+ return new RecordReader() {
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public Record nextRecord(boolean coerceTypes, boolean dropUnknownFields) throws MalformedRecordException {
+ throw new MalformedRecordException("Test exception");
+ }
+
+ @Override
+ public RecordSchema getSchema() {
+ return null;
+ }
+ };
+ }
+ };
+ }
+}
\ No newline at end of file
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-nar/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-nar/pom.xml
index 1b4fb08bab..f1b9849771 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-nar/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-nar/pom.xml
@@ -40,12 +40,6 @@
<artifactId>sts</artifactId>
<version>${software.amazon.awssdk.version}</version>
<scope>provided</scope>
- <exclusions>
- <exclusion>
- <groupId>software.amazon.awssdk</groupId>
- <artifactId>netty-nio-client</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml
index ce2f47de84..b26d7719f7 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml
@@ -98,7 +98,7 @@
<dependency>
<groupId>software.amazon.kinesis</groupId>
<artifactId>amazon-kinesis-client</artifactId>
- <version>${aws-kinesis-client-library-version}</version>
+ <version>2.7.2</version>
<exclusions>
<exclusion>
<groupId>commons-logging</groupId>
@@ -188,11 +188,6 @@
<version>2.7.0-SNAPSHOT</version>
<scope>test</scope>
</dependency>
- <dependency>
- <groupId>org.testcontainers</groupId>
- <artifactId>testcontainers</artifactId>
- <scope>test</scope>
- </dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>dynamodb</artifactId>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-regions/pom.xml
similarity index 62%
copy from nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml
copy to nifi-extension-bundles/nifi-aws-bundle/nifi-aws-regions/pom.xml
index 89d27793d4..7a5a291836 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-regions/pom.xml
@@ -13,7 +13,9 @@
See the License for the specific language governing permissions and
limitations under the License.
-->
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
@@ -22,20 +24,14 @@
<version>2.7.0-SNAPSHOT</version>
</parent>
- <artifactId>nifi-aws-service-api-nar</artifactId>
- <packaging>nar</packaging>
+ <artifactId>nifi-aws-regions</artifactId>
+ <packaging>jar</packaging>
<dependencies>
<dependency>
- <groupId>org.apache.nifi</groupId>
- <artifactId>nifi-standard-shared-nar</artifactId>
- <version>2.7.0-SNAPSHOT</version>
- <type>nar</type>
- </dependency>
- <dependency>
- <groupId>org.apache.nifi</groupId>
- <artifactId>nifi-aws-service-api</artifactId>
- <scope>compile</scope>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>regions</artifactId>
</dependency>
</dependencies>
+
</project>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/RegionUtilV2.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-regions/src/main/java/org/apache/nifi/processors/aws/region/RegionUtilV2.java
similarity index 97%
rename from nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/RegionUtilV2.java
rename to nifi-extension-bundles/nifi-aws-bundle/nifi-aws-regions/src/main/java/org/apache/nifi/processors/aws/region/RegionUtilV2.java
index 40be091198..e198736d13 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-abstract-processors/src/main/java/org/apache/nifi/processors/aws/v2/RegionUtilV2.java
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-regions/src/main/java/org/apache/nifi/processors/aws/region/RegionUtilV2.java
@@ -14,7 +14,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.nifi.processors.aws.v2;
+package org.apache.nifi.processors.aws.region;
import org.apache.nifi.components.AllowableValue;
import software.amazon.awssdk.regions.Region;
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml
index 89d27793d4..6176af1c1a 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api-nar/pom.xml
@@ -37,5 +37,37 @@
<artifactId>nifi-aws-service-api</artifactId>
<scope>compile</scope>
</dependency>
+ <!-- Shared dependency on Netty NIO Client for Asynchronous HTTP requests -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>netty-nio-client</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>compile</scope>
+ </dependency>
+ <!-- Shared dependencies for Kinesis -->
+ <dependency>
+ <groupId>com.google.guava</groupId>
+ <artifactId>guava</artifactId>
+ <version>${guava.version}</version>
+ <scope>compile</scope>
+ </dependency>
+ <dependency>
+ <groupId>com.google.protobuf</groupId>
+ <artifactId>protobuf-java</artifactId>
+ <version>${protobuf.version}</version>
+ <scope>compile</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>dynamodb</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>compile</scope>
+ </dependency>
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>cloudwatch</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>compile</scope>
+ </dependency>
</dependencies>
</project>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api/pom.xml
index 9cbd38ba73..9cefd636b3 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-service-api/pom.xml
@@ -35,12 +35,6 @@
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>sts</artifactId>
- <exclusions>
- <exclusion>
- <groupId>software.amazon.awssdk</groupId>
- <artifactId>netty-nio-client</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
diff --git a/nifi-extension-bundles/nifi-aws-bundle/pom.xml b/nifi-extension-bundles/nifi-aws-bundle/pom.xml
index c89c58d907..23a3192217 100644
--- a/nifi-extension-bundles/nifi-aws-bundle/pom.xml
+++ b/nifi-extension-bundles/nifi-aws-bundle/pom.xml
@@ -27,11 +27,15 @@
<packaging>pom</packaging>
<properties>
- <aws-kinesis-client-library-version>2.7.2</aws-kinesis-client-library-version>
+ <guava.version>33.5.0-jre</guava.version>
+ <protobuf.version>4.33.0</protobuf.version>
</properties>
<modules>
+ <module>nifi-aws-kinesis</module>
+ <module>nifi-aws-kinesis-nar</module>
<module>nifi-aws-processors</module>
+ <module>nifi-aws-regions</module>
<module>nifi-aws-nar</module>
<module>nifi-aws-service-api</module>
<module>nifi-aws-service-api-nar</module>
@@ -43,12 +47,6 @@
<dependencyManagement>
<dependencies>
- <!-- Override Guava 31.0.1 from amazon-kinesis-client -->
- <dependency>
- <groupId>com.google.guava</groupId>
- <artifactId>guava</artifactId>
- <version>33.5.0-jre</version>
- </dependency>
<dependency>
<groupId>org.apache.nifi</groupId>
<artifactId>nifi-aws-service-api</artifactId>
@@ -59,12 +57,41 @@
<groupId>software.amazon.awssdk</groupId>
<artifactId>sts</artifactId>
<version>${software.amazon.awssdk.version}</version>
- <exclusions>
- <exclusion>
- <groupId>software.amazon.awssdk</groupId>
- <artifactId>netty-nio-client</artifactId>
- </exclusion>
- </exclusions>
+ </dependency>
+ <!-- Shared Guava provided in nifi-aws-service-api-nar -->
+ <dependency>
+ <groupId>com.google.guava</groupId>
+ <artifactId>guava</artifactId>
+ <version>${guava.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Shared Protobuf provided in nifi-aws-service-api-nar -->
+ <dependency>
+ <groupId>com.google.protobuf</groupId>
+ <artifactId>protobuf-java</artifactId>
+ <version>${protobuf.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Netty NIO Client provided in nifi-aws-service-api-nar -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>netty-nio-client</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- DynamoDB provided in nifi-aws-service-api-nar -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>dynamodb</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Cloudwatch provided in nifi-aws-service-api-nar -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>cloudwatch</artifactId>
+ <version>${software.amazon.awssdk.version}</version>
+ <scope>provided</scope>
</dependency>
<!-- Override kotlin-stdlib-common from amazon-kinesis-client -->
<!-- can be removed when not relying on kotlin-stdlib-common:jar anymore -->