Fixes #6548 (#6618)

Adds an attempt to retrieve the request message from the resilience context in addition to retrieving the request from the response

Author: iliar-turdushev

src/Libraries/Microsoft.Extensions.Http.Resilience/Polly/HttpRetryStrategyOptionsExtensions.cs

@@ -8,6 +8,7 @@
 using Microsoft.Shared.DiagnosticIds;
 using Microsoft.Shared.Diagnostics;
 using Polly;
+using Polly.Retry;
 
 namespace Microsoft.Extensions.Http.Resilience;
 
@@ -52,15 +53,16 @@ public static void DisableFor(this HttpRetryStrategyOptions options, params Http
         {
             var result = await shouldHandle(args).ConfigureAwait(args.Context.ContinueOnCapturedContext);
 
-            if (result &&
-                args.Outcome.Result is HttpResponseMessage response &&
-                response.RequestMessage is HttpRequestMessage request)
+            if (result && GetRequestMessage(args) is HttpRequestMessage request)
             {
                 return !methods.Contains(request.Method);
             }
 
             return result;
         };
     }
+
+    private static HttpRequestMessage? GetRequestMessage(RetryPredicateArguments<HttpResponseMessage> args) =>
+        args.Outcome.Result?.RequestMessage ?? args.Context.GetRequestMessage();
 }
 

test/Libraries/Microsoft.Extensions.Http.Resilience.Tests/Polly/HttpRetryStrategyOptionsExtensionsTests.cs

@@ -65,12 +65,16 @@ public async Task DisableFor_RespectsOriginalShouldHandlePredicate()
     }
 
     [Fact]
-    public async Task DisableFor_ResponseMessageIsNull_DoesNotDisableRetries()
+    public async Task DisableFor_ResponseMessageIsNull_RetrievesRequestMessageFromContext()
     {
         var options = new HttpRetryStrategyOptions { ShouldHandle = _ => PredicateResult.True() };
         options.DisableFor(HttpMethod.Post);
 
-        Assert.True(await options.ShouldHandle(CreatePredicateArguments(null)));
+        using var request = new HttpRequestMessage { Method = HttpMethod.Post };
+        var context = ResilienceContextPool.Shared.Get();
+        context.SetRequestMessage(request);
+
+        Assert.False(await options.ShouldHandle(CreatePredicateArguments(null, context)));
     }
 
     [Fact]
@@ -80,8 +84,10 @@ public async Task DisableFor_RequestMessageIsNull_DoesNotDisableRetries()
         options.DisableFor(HttpMethod.Post);
 
         using var response = new HttpResponseMessage { RequestMessage = null };
+        var context = ResilienceContextPool.Shared.Get();
+        context.SetRequestMessage(null);
 
-        Assert.True(await options.ShouldHandle(CreatePredicateArguments(response)));
+        Assert.True(await options.ShouldHandle(CreatePredicateArguments(response, context)));
     }
 
     [Theory]
@@ -105,10 +111,10 @@ public async Task DisableForUnsafeHttpMethods_PositiveScenario(string httpMethod
         Assert.Equal(shouldHandle, await options.ShouldHandle(CreatePredicateArguments(response)));
     }
 
-    private static RetryPredicateArguments<HttpResponseMessage> CreatePredicateArguments(HttpResponseMessage? response)
+    private static RetryPredicateArguments<HttpResponseMessage> CreatePredicateArguments(HttpResponseMessage? response, ResilienceContext? context = null)
     {
         return new RetryPredicateArguments<HttpResponseMessage>(
-            ResilienceContextPool.Shared.Get(),
+            context ?? ResilienceContextPool.Shared.Get(),
             Outcome.FromResult(response),
             attemptNumber: 1);
     }