Skip to content

NIFI-15535 - Fixed AWS Connection Pool Shutdown on EKS with STS Credential Refresh #11020

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pvillard31 wants to merge 1 commit into from
Closed

NIFI-15535 - Fixed AWS Connection Pool Shutdown on EKS with STS Credential Refresh #11020

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions .../processors/aws/credentials/provider/service/AWSCredentialsProviderControllerService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.apache.nifi.annotation.behavior.Restriction;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnDisabled;
import org.apache.nifi.annotation.lifecycle.OnEnabled;
import org.apache.nifi.components.AllowableValue;
import org.apache.nifi.components.PropertyDescriptor;
Expand Down Expand Up @@ -244,7 +245,7 @@ public class AWSCredentialsProviderControllerService extends AbstractControllerS
OAUTH2_ACCESS_TOKEN_PROVIDER
);

private volatile AwsCredentialsProvider credentialsProvider;
private volatile ConfigurationContext configurationContext;

private final List<CredentialsStrategy> strategies = List.of(
// Primary Credential Strategies
Expand Down Expand Up @@ -292,7 +293,9 @@ public void migrateProperties(PropertyConfiguration config) {

@Override
public AwsCredentialsProvider getAwsCredentialsProvider() {
return credentialsProvider;
final AwsCredentialsProvider provider = createCredentialsProvider(configurationContext);
getLogger().debug("Created AwsCredentialsProvider [{}] instance [{}]", provider.getClass().getSimpleName(), System.identityHashCode(provider));
return provider;
}

private AwsCredentialsProvider createCredentialsProvider(final PropertyContext context) {
Expand Down Expand Up @@ -361,8 +364,12 @@ protected Collection<ValidationResult> customValidate(final ValidationContext va

@OnEnabled
public void onConfigured(final ConfigurationContext context) {
credentialsProvider = createCredentialsProvider(context);
getLogger().debug("Using credentials provider: {}", credentialsProvider.getClass());
this.configurationContext = context;
}

@OnDisabled
public void onDisabled() {
configurationContext = null;
}

public static AllowableValue[] getAvailableRegions() {
Expand Down
68 changes: 68 additions & 0 deletions ...cessors/aws/credentials/provider/service/AWSCredentialsProviderControllerServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNotSame;

public class AWSCredentialsProviderControllerServiceTest {

Expand Down Expand Up @@ -409,6 +410,73 @@ public void testWebIdentityDoesNotChainAssumeRoleDerivedProvider() throws Throwa
"Derived AssumeRole should not be chained when OAuth2 (Web Identity) is configured");
}

@Test
public void testGetAwsCredentialsProviderReturnsIndependentInstances() throws Throwable {
final TestRunner runner = TestRunners.newTestRunner(FetchS3Object.class);
final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();
runner.addControllerService("awsCredentialsProvider", serviceImpl);
runner.enableControllerService(serviceImpl);

final AwsCredentialsProviderService service = (AwsCredentialsProviderService) runner.getProcessContext()
.getControllerServiceLookup().getControllerService("awsCredentialsProvider");

final AwsCredentialsProvider first = service.getAwsCredentialsProvider();
final AwsCredentialsProvider second = service.getAwsCredentialsProvider();

assertNotNull(first);
assertNotNull(second);
assertNotSame(first, second, "Each call should return an independent AwsCredentialsProvider to avoid shared connection pool issues");
assertEquals(DefaultCredentialsProvider.class, first.getClass());
assertEquals(DefaultCredentialsProvider.class, second.getClass());
}

@Test
public void testGetAwsCredentialsProviderReturnsIndependentInstancesWithAccessKeys() throws Throwable {
final TestRunner runner = TestRunners.newTestRunner(FetchS3Object.class);
final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();
runner.addControllerService("awsCredentialsProvider", serviceImpl);
runner.setProperty(serviceImpl, ACCESS_KEY_ID, "awsAccessKey");
runner.setProperty(serviceImpl, SECRET_KEY, "awsSecretKey");
runner.enableControllerService(serviceImpl);

final AwsCredentialsProviderService service = (AwsCredentialsProviderService) runner.getProcessContext()
.getControllerServiceLookup().getControllerService("awsCredentialsProvider");

final AwsCredentialsProvider first = service.getAwsCredentialsProvider();
final AwsCredentialsProvider second = service.getAwsCredentialsProvider();

assertNotNull(first);
assertNotNull(second);
assertNotSame(first, second, "Each call should return an independent AwsCredentialsProvider to avoid shared connection pool issues");
assertEquals("awsAccessKey", first.resolveCredentials().accessKeyId());
assertEquals("awsAccessKey", second.resolveCredentials().accessKeyId());
}

@Test
public void testGetAwsCredentialsProviderReturnsIndependentInstancesWithAssumeRole() throws Throwable {
final TestRunner runner = TestRunners.newTestRunner(FetchS3Object.class);
final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();
runner.addControllerService("awsCredentialsProvider", serviceImpl);
runner.setProperty(serviceImpl, ACCESS_KEY_ID, "awsAccessKey");
runner.setProperty(serviceImpl, SECRET_KEY, "awsSecretKey");
runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_STS_REGION, Region.US_WEST_1.id());
runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_ARN, "Role");
runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_NAME, "RoleName");
runner.enableControllerService(serviceImpl);

final AwsCredentialsProviderService service = (AwsCredentialsProviderService) runner.getProcessContext()
.getControllerServiceLookup().getControllerService("awsCredentialsProvider");

final AwsCredentialsProvider first = service.getAwsCredentialsProvider();
final AwsCredentialsProvider second = service.getAwsCredentialsProvider();

assertNotNull(first);
assertNotNull(second);
assertNotSame(first, second, "Each call should return an independent AwsCredentialsProvider to avoid shared connection pool issues");
assertEquals(StsAssumeRoleCredentialsProvider.class, first.getClass());
assertEquals(StsAssumeRoleCredentialsProvider.class, second.getClass());
}

@Test
void testMigration() {
final Map<String, String> propertyValues = Map.of(
Expand Down
Loading