hudi-agent commented on code in PR #19383:
URL: https://github.com/apache/hudi/pull/19383#discussion_r3661017312
##########
hudi-utilities/src/main/java/org/apache/hudi/utilities/sources/helpers/KinesisOffsetGen.java:
##########
@@ -304,28 +308,58 @@ public KinesisOffsetGen(TypedProperties props) {
getStringWithAltKeys(props,
KinesisSourceConfig.KINESIS_STARTING_POSITION, true));
}
+ /**
+ * Assume-role credentials providers cached by {@code region|roleArn}.
Kinesis clients are built
+ * per {@code mapPartitions} task, so creating a fresh STS client each time
would leak one (with its
+ * own HTTP connection pool) per micro-batch on a long-lived streaming
executor — AWS SDK v2 does not
+ * close a user-supplied credentials provider (nor its injected StsClient)
when the KinesisClient
+ * closes. Caching means at most one provider/StsClient per distinct role
per executor JVM, reused
+ * across partitions and micro-batches. Providers auto-refresh and live for
the JVM lifetime, so they
+ * are intentionally never closed.
+ */
+ private static final Map<String, StsAssumeRoleCredentialsProvider>
ASSUME_ROLE_PROVIDERS =
+ new ConcurrentHashMap<>();
+
+ private static StsAssumeRoleCredentialsProvider assumeRoleProvider(String
region, String roleArn) {
+ return ASSUME_ROLE_PROVIDERS.computeIfAbsent(region + "|" + roleArn,
ignored ->
+ StsAssumeRoleCredentialsProvider.builder()
+ .stsClient(StsClient.builder().region(Region.of(region)).build())
+ .refreshRequest(AssumeRoleRequest.builder()
+ .roleArn(roleArn)
+ .roleSessionName("hudi-kinesis-source")
+ .build())
+ .build());
+ }
+
/**
* Builds a Kinesis client from explicit parameters. Used by both the
instance method
* {@link #createKinesisClient()} and by {@link
org.apache.hudi.utilities.sources.JsonKinesisSource}
* from serializable {@link KinesisReadConfig} in Spark closures.
*/
public static KinesisClient createKinesisClient(String region, String
endpointUrl,
- String accessKey, String secretKey) {
+ String accessKey, String secretKey, String roleArn) {
KinesisClientBuilder builder =
KinesisClient.builder().region(Region.of(region));
if (endpointUrl != null && !endpointUrl.isEmpty()) {
builder = builder.endpointOverride(URI.create(endpointUrl));
}
if (accessKey != null && !accessKey.isEmpty() && secretKey != null &&
!secretKey.isEmpty()) {
+ // Static credentials (e.g. LocalStack / custom endpoint) take
precedence.
Review Comment:
🤖 When both static access/secret keys and a role ARN are set, the static
branch silently wins and the role ARN is ignored — a cross-account read would
then fail with a hard-to-diagnose access-denied. Since KINESIS_ROLE_ARN's docs
don't mention this precedence, would it be worth logging a warning when both
are present (or noting the precedence in the ARN's documentation)?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
##########
hudi-utilities/src/main/java/org/apache/hudi/utilities/sources/helpers/KinesisOffsetGen.java:
##########
@@ -304,28 +308,58 @@ public KinesisOffsetGen(TypedProperties props) {
getStringWithAltKeys(props,
KinesisSourceConfig.KINESIS_STARTING_POSITION, true));
}
+ /**
+ * Assume-role credentials providers cached by {@code region|roleArn}.
Kinesis clients are built
+ * per {@code mapPartitions} task, so creating a fresh STS client each time
would leak one (with its
+ * own HTTP connection pool) per micro-batch on a long-lived streaming
executor — AWS SDK v2 does not
+ * close a user-supplied credentials provider (nor its injected StsClient)
when the KinesisClient
+ * closes. Caching means at most one provider/StsClient per distinct role
per executor JVM, reused
+ * across partitions and micro-batches. Providers auto-refresh and live for
the JVM lifetime, so they
+ * are intentionally never closed.
+ */
+ private static final Map<String, StsAssumeRoleCredentialsProvider>
ASSUME_ROLE_PROVIDERS =
+ new ConcurrentHashMap<>();
+
+ private static StsAssumeRoleCredentialsProvider assumeRoleProvider(String
region, String roleArn) {
+ return ASSUME_ROLE_PROVIDERS.computeIfAbsent(region + "|" + roleArn,
ignored ->
+ StsAssumeRoleCredentialsProvider.builder()
+ .stsClient(StsClient.builder().region(Region.of(region)).build())
+ .refreshRequest(AssumeRoleRequest.builder()
+ .roleArn(roleArn)
+ .roleSessionName("hudi-kinesis-source")
+ .build())
+ .build());
+ }
+
/**
Review Comment:
🤖 nit: `assumeRoleProvider` reads as a noun, so it's not immediately obvious
at the call site whether this creates a new provider or retrieves a cached one.
Could you rename it to something like `getOrCreateAssumeRoleProvider` to make
the lazy-cache semantics self-evident?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
