From 203a2dd04b4a87d6300311c48419b03edb0a8ad5 Mon Sep 17 00:00:00 2001
From: Garrett Beatty <gcbeatty@amazon.com>
Date: Thu, 2 Jul 2026 15:10:27 -0400
Subject: [PATCH] Retry IAM calls on throttling in RuntimeSupport integration
 tests

The RuntimeSupport integration tests now run in parallel, and each test
class provisions its own IAM role (GetRole/CreateRole/AttachRolePolicy at
startup, ListAttachedRolePolicies/DetachRolePolicy/DeleteRole at teardown).
The IAM control-plane APIs have low request-rate limits, so the parallel
burst intermittently throttles with 'Rate exceeded', failing
TestAllNET10HandlersAsync and TestThreadingLogging in setup rather than on
any assertion.

Wrap the IAM calls in a throttle-aware retry with exponential backoff so
the parallel startup/teardown burst is smoothed instead of failing.
---
 .../BaseCustomRuntimeTest.cs                  | 69 ++++++++++++++++---
 1 file changed, 61 insertions(+), 8 deletions(-)

diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/BaseCustomRuntimeTest.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/BaseCustomRuntimeTest.cs
index fd1e748f0..19316d5fd 100644
--- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/BaseCustomRuntimeTest.cs
+++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/BaseCustomRuntimeTest.cs
@@ -9,6 +9,7 @@
 using Amazon.IdentityManagement;
 using Amazon.IdentityManagement.Model;
 using Amazon.Lambda.Model;
+using Amazon.Runtime;
 using Amazon.S3;
 using Amazon.S3.Model;
 
@@ -81,22 +82,22 @@ public async Task CleanUpTestResources(AmazonS3Client s3Client, AmazonLambdaClie
                     {
                         RoleName = ExecutionRoleName
                     };
-                    var attachedPolicies = await iamClient.ListAttachedRolePoliciesAsync(listAttachedPoliciesRequest);
+                    var attachedPolicies = await RetryOnThrottleAsync(() => iamClient.ListAttachedRolePoliciesAsync(listAttachedPoliciesRequest));
 
                     foreach (var policy in attachedPolicies.AttachedPolicies)
                     {
-                        await iamClient.DetachRolePolicyAsync(new DetachRolePolicyRequest
+                        await RetryOnThrottleAsync(() => iamClient.DetachRolePolicyAsync(new DetachRolePolicyRequest
                         {
                             RoleName = ExecutionRoleName,
                             PolicyArn = policy.PolicyArn
-                        });
+                        }));
                     }
 
                     var deleteRoleRequest = new DeleteRoleRequest
                     {
                         RoleName = ExecutionRoleName
                     };
-                    await iamClient.DeleteRoleAsync(deleteRoleRequest);
+                    await RetryOnThrottleAsync(() => iamClient.DeleteRoleAsync(deleteRoleRequest));
                 }
                 catch (Exception)
                 {
@@ -130,7 +131,7 @@ private async Task<bool> ValidateAndSetIamRoleArn(IAmazonIdentityManagementServi
             };
             try
             {
-                ExecutionRoleArn = (await iamClient.GetRoleAsync(getRoleRequest)).Role.Arn;
+                ExecutionRoleArn = (await RetryOnThrottleAsync(() => iamClient.GetRoleAsync(getRoleRequest))).Role.Arn;
                 return true;
             }
             catch (NoSuchEntityException)
@@ -142,22 +143,74 @@ private async Task<bool> ValidateAndSetIamRoleArn(IAmazonIdentityManagementServi
                     Description = "Test role for CustomRuntimeTests.",
                     AssumeRolePolicyDocument = LAMBDA_ASSUME_ROLE_POLICY
                 };
-                ExecutionRoleArn = (await iamClient.CreateRoleAsync(createRoleRequest)).Role.Arn;
+                ExecutionRoleArn = (await RetryOnThrottleAsync(() => iamClient.CreateRoleAsync(createRoleRequest))).Role.Arn;
 
                 // Wait for role to propagate.
                 await Task.Delay(10000);
 
-                await iamClient.AttachRolePolicyAsync(new AttachRolePolicyRequest
+                await RetryOnThrottleAsync(() => iamClient.AttachRolePolicyAsync(new AttachRolePolicyRequest
                 {
                     PolicyArn = "arn:aws:iam::aws:policy/AWSLambdaExecute",
                     RoleName = ExecutionRoleName,
-                });
+                }));
 
 
                 return false;
             }
         }
 
+        /// <summary>
+        /// Invokes an AWS SDK call, retrying with exponential backoff when the service throttles the request.
+        /// Because every test class provisions its own IAM role and these tests now run in parallel, the
+        /// IAM control-plane APIs (which have low request-rate limits) are easily throttled at startup and
+        /// teardown, surfacing as "Rate exceeded". Retrying smooths the parallel burst instead of failing the test.
+        /// </summary>
+        private static async Task<T> RetryOnThrottleAsync<T>(Func<Task<T>> action)
+        {
+            const int maxAttempts = 8;
+            var delay = TimeSpan.FromSeconds(2);
+
+            for (var attempt = 1; ; attempt++)
+            {
+                try
+                {
+                    return await action();
+                }
+                catch (AmazonServiceException e) when (IsThrottlingException(e) && attempt < maxAttempts)
+                {
+                    await Task.Delay(delay);
+                    // Exponential backoff, capped at 30s, to let the IAM request rate recover.
+                    delay = TimeSpan.FromSeconds(Math.Min(delay.TotalSeconds * 2, 30));
+                }
+            }
+        }
+
+        /// <summary>
+        /// Overload of <see cref="RetryOnThrottleAsync{T}"/> for SDK calls whose response is not needed.
+        /// </summary>
+        private static async Task RetryOnThrottleAsync(Func<Task> action)
+        {
+            await RetryOnThrottleAsync<object>(async () =>
+            {
+                await action();
+                return null;
+            });
+        }
+
+        private static bool IsThrottlingException(AmazonServiceException e)
+        {
+            if (e.StatusCode == System.Net.HttpStatusCode.TooManyRequests)
+            {
+                return true;
+            }
+
+            var errorCode = e.ErrorCode ?? string.Empty;
+            return errorCode.IndexOf("Throttl", StringComparison.OrdinalIgnoreCase) >= 0
+                || errorCode.IndexOf("TooManyRequests", StringComparison.OrdinalIgnoreCase) >= 0
+                || errorCode.Equals("RequestLimitExceeded", StringComparison.OrdinalIgnoreCase)
+                || (e.Message?.IndexOf("Rate exceeded", StringComparison.OrdinalIgnoreCase) >= 0);
+        }
+
         private async Task CreateBucketWithDeploymentZipAsync(IAmazonS3 s3Client, string bucketName)
         {
             // create bucket if it doesn't exist