Fix a race condition in aggregate ProfileFileSupplier

Author: alextwoods

.changes/next-release/bugfix-AWSSDKforJavav2-0caa9ea.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Fix a race condition in aggregate ProfileFileSupplier that could cause credential resolution failures with shared DefaultCredentialsProvider."
+}

core/profiles/src/main/java/software/amazon/awssdk/profiles/ProfileFileSupplier.java

@@ -26,6 +26,7 @@
 import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.profiles.internal.ProfileFileRefresher;
+import software.amazon.awssdk.utils.Pair;
 
 /**
  * Encapsulates the logic for supplying either a single or multiple ProfileFile instances.
@@ -127,43 +128,41 @@ static ProfileFileSupplier aggregate(ProfileFileSupplier... suppliers) {
 
         return new ProfileFileSupplier() {
 
-            final AtomicReference<ProfileFile> currentAggregateProfileFile = new AtomicReference<>();
-            final Map<Supplier<ProfileFile>, ProfileFile> currentValuesBySupplier
-                = Collections.synchronizedMap(new LinkedHashMap<>());
+            final AtomicReference<Pair<Map<Supplier<ProfileFile>, ProfileFile>, ProfileFile>> state =
+                new AtomicReference<>(Pair.of(Collections.emptyMap(), ProfileFile.empty()));
 
             @Override
             public ProfileFile get() {
-                boolean refreshAggregate = false;
-                for (ProfileFileSupplier supplier : suppliers) {
-                    if (didSuppliedValueChange(supplier)) {
-                        refreshAggregate = true;
+                while(true) {
+                    Pair<Map<Supplier<ProfileFile>, ProfileFile>, ProfileFile> currentState = state.get();
+                    Map<Supplier<ProfileFile>, ProfileFile> nextValues = new LinkedHashMap<>(currentState.left());
+
+                    boolean refreshAggregate = false;
+
+                    for (ProfileFileSupplier supplier : suppliers) {
+                        ProfileFile next = supplier.get();
+                        ProfileFile prev = nextValues.put(supplier, next);
+                        // we ONLY care about if the reference has changed, we don't care about object equality here
+                        if (prev != next) {
+                            refreshAggregate = true;
+                        }
                     }
-                }
-
-                if (refreshAggregate) {
-                    refreshCurrentAggregate();
-                }
-
-                return  currentAggregateProfileFile.get();
-            }
 
-            private boolean didSuppliedValueChange(Supplier<ProfileFile> supplier) {
-                ProfileFile next = supplier.get();
-                ProfileFile current = currentValuesBySupplier.put(supplier, next);
+                    if (!refreshAggregate) {
+                        return currentState.right();
+                    }
 
-                return !Objects.equals(next, current);
-            }
+                    ProfileFile.Aggregator aggregator = ProfileFile.aggregator();
+                    nextValues.values().forEach(aggregator::addFile);
+                    ProfileFile nextAggregate = aggregator.build();
 
-            private void refreshCurrentAggregate() {
-                ProfileFile.Aggregator aggregator = ProfileFile.aggregator();
-                currentValuesBySupplier.values().forEach(aggregator::addFile);
-                ProfileFile current = currentAggregateProfileFile.get();
-                ProfileFile next = aggregator.build();
-                if (!Objects.equals(current, next)) {
-                    currentAggregateProfileFile.compareAndSet(current, next);
+                    Pair<Map<Supplier<ProfileFile>, ProfileFile>, ProfileFile> nextState = Pair.of(nextValues, nextAggregate);
+                    if (state.compareAndSet(currentState, nextState)) {
+                        return nextAggregate;
+                    }
+                    // else: another thread has modified the state in between, retry with the fresh state
                 }
             }
-
         };
     }
 

core/profiles/src/test/java/software/amazon/awssdk/profiles/ProfileFileSupplierTest.java

@@ -31,12 +31,20 @@
 import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.time.temporal.TemporalAmount;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -503,6 +511,50 @@ void aggregate_duplicateOptionsGivenReloadingProfileFirst_preservesPrecedence()
         assertThat(accessKeyId).isEqualTo("defaultAccessKey2");
     }
 
+    @Test
+    void aggregate_concurrentGetAlwaysReturnsCorrectAggregate() throws ExecutionException, InterruptedException {
+        ProfileFile credentialFile = credentialProfileFile("test1", "key1", "secret1");
+        ProfileFile configFile = configProfileFile("profile test",
+                                                   Pair.of("region", "us-west-2"),
+                                                   Pair.of("aws_account_id", "012354678922"));
+
+
+        ProfileFile expectedAggregate = ProfileFile.aggregator().addFile(credentialFile).addFile(configFile).build();
+
+        ProfileFileSupplier supplier = ProfileFileSupplier.aggregate(() -> credentialFile, () -> configFile);
+
+        ExecutorService executor = Executors.newFixedThreadPool(24);
+        CountDownLatch startLatch = new CountDownLatch(1);
+        List<Future<Boolean>> tasks = new ArrayList<>();
+
+        for(int i = 0; i < 24; i++) {
+            tasks.add(executor.submit(() -> {
+                try {
+                    startLatch.await();
+                    ProfileFile resolved = supplier.get();
+                    return Objects.equals(expectedAggregate, resolved);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+            }));
+        }
+        // All tasks are now submitted — release them
+        startLatch.countDown();
+        executor.shutdown();
+        try {
+            assertThat(executor.awaitTermination(10, TimeUnit.SECONDS))
+                .as("executor did not terminate")
+                .isTrue();
+        } finally {
+            executor.shutdownNow();
+        }
+
+        // assert that all concurrent get's returned the same, expected aggregate
+        for(Future<Boolean> task : tasks) {
+            assertThat(task.get()).isTrue();
+        }
+    }
+
     @Test
     void fixedProfileFile_nullProfileFile_returnsNonNullSupplier() {
         ProfileFile file = null;