From c1cc348c7fa7329639c4bec235955b2d624fa410 Mon Sep 17 00:00:00 2001
From: Ian Botsford <83236726+ianbotsf@users.noreply.github.com>
Date: Thu, 14 Nov 2024 18:10:54 +0000
Subject: [PATCH] fix: correctly deduplicate S3 Express credential cache
 refreshes

---
 .changes/903848fb-77d4-4431-8a4e-abb52b25b7c8.json  |  5 +++++
 .../express/DefaultS3ExpressCredentialsProvider.kt  | 13 ++++++++++++-
 .../DefaultS3ExpressCredentialsProviderTest.kt      |  1 -
 3 files changed, 17 insertions(+), 2 deletions(-)
 create mode 100644 .changes/903848fb-77d4-4431-8a4e-abb52b25b7c8.json

diff --git a/.changes/903848fb-77d4-4431-8a4e-abb52b25b7c8.json b/.changes/903848fb-77d4-4431-8a4e-abb52b25b7c8.json
new file mode 100644
index 00000000000..c0cf6e8c5f1
--- /dev/null
+++ b/.changes/903848fb-77d4-4431-8a4e-abb52b25b7c8.json
@@ -0,0 +1,5 @@
+{
+    "id": "903848fb-77d4-4431-8a4e-abb52b25b7c8",
+    "type": "bugfix",
+    "description": "Fix concurrency issues with S3 Express credential caching"
+}
\ No newline at end of file
diff --git a/services/s3/common/src/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProvider.kt b/services/s3/common/src/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProvider.kt
index c5f07c85945..3bfb627eccf 100644
--- a/services/s3/common/src/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProvider.kt
+++ b/services/s3/common/src/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProvider.kt
@@ -57,7 +57,18 @@ internal class DefaultS3ExpressCredentialsProvider(
                     client.logger.trace { "Credentials for ${key.bucket} are expiring in ${it.expiringCredentials.expiresAt} and are within their refresh window, performing asynchronous refresh..." }
                     launch(coroutineContext) {
                         try {
-                            it.sfg.singleFlight { createSessionCredentials(key, client) }
+                            it.sfg.singleFlight {
+                                // This coroutine/SFG may have started _after_ prior instances(s) finished, replacing
+                                // the cached value already. To prevent re-refreshing it, we need to re-acquire the
+                                // current cached value and verify whether it's expiring soon.
+                                val currentCreds = credentialsCache.get(key)
+
+                                if (currentCreds?.expiringCredentials?.isExpiringWithin(refreshBuffer) == true) {
+                                    createSessionCredentials(key, client)
+                                } else {
+                                    it.expiringCredentials
+                                }
+                            }
                         } catch (e: Exception) {
                             client.logger.warn(e) { "Asynchronous refresh for ${key.bucket} failed." }
                         }
diff --git a/services/s3/common/test/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProviderTest.kt b/services/s3/common/test/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProviderTest.kt
index bbd9f03f3d7..4b3c436714d 100644
--- a/services/s3/common/test/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProviderTest.kt
+++ b/services/s3/common/test/aws/sdk/kotlin/services/s3/express/DefaultS3ExpressCredentialsProviderTest.kt
@@ -116,7 +116,6 @@ class DefaultS3ExpressCredentialsProviderTest {
     }
 
     @Test
-    @Ignore // FIXME flaky test temporarily disabled to unblock preview builds
     fun testAsyncRefreshDebounce() = runTest {
         val timeSource = TestTimeSource()
         val clock = ManualClock()