Add methods allowing retrieval of HttpRequestMessage from ResilienceContext (#5460)

Fixes #4957

- Adds extension methods allowing retrieval of the request message from resilience context
- Replaces usage of ResilienceKeys.RequestMessage variable with corresponding Get/Set methods

Author: iliar-turdushev

src/Libraries/Microsoft.Extensions.Http.Resilience/Hedging/Internals/RequestMessageSnapshotStrategy.cs

@@ -20,7 +20,9 @@ protected override async ValueTask<Outcome<TResult>> ExecuteCore<TResult, TState
         ResilienceContext context,
         TState state)
     {
-        if (!context.Properties.TryGetValue(ResilienceKeys.RequestMessage, out var request) || request is null)
+        HttpRequestMessage? request = context.GetRequestMessage();
+
+        if (request is null)
         {
             Throw.InvalidOperationException("The HTTP request message was not found in the resilience context.");
         }

src/Libraries/Microsoft.Extensions.Http.Resilience/Hedging/ResilienceHttpClientBuilderExtensions.Hedging.cs

@@ -94,7 +94,7 @@ public static IStandardHedgingHandlerBuilder AddStandardHedgingHandler(this IHtt
                 requestMessage.SetResilienceContext(args.ActionContext);
 
                 // replace the request message
-                args.ActionContext.Properties.Set(ResilienceKeys.RequestMessage, requestMessage);
+                args.ActionContext.SetRequestMessage(requestMessage);
 
                 if (args.PrimaryContext.Properties.TryGetValue(ResilienceKeys.RoutingStrategy, out var routingPipeline))
                 {

src/Libraries/Microsoft.Extensions.Http.Resilience/Internal/ResilienceKeys.cs

@@ -10,7 +10,7 @@ namespace Microsoft.Extensions.Http.Resilience.Internal;
 
 internal static class ResilienceKeys
 {
-    public static readonly ResiliencePropertyKey<HttpRequestMessage> RequestMessage = new("Resilience.Http.RequestMessage");
+    public static readonly ResiliencePropertyKey<HttpRequestMessage?> RequestMessage = new("Resilience.Http.RequestMessage");
 
     public static readonly ResiliencePropertyKey<RequestRoutingStrategy> RoutingStrategy = new("Resilience.Http.RequestRoutingStrategy");
 

src/Libraries/Microsoft.Extensions.Http.Resilience/Resilience/HttpResilienceContextExtensions.cs

@@ -0,0 +1,46 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+using System.Net.Http;
+using Microsoft.Extensions.Http.Resilience.Internal;
+using Microsoft.Shared.DiagnosticIds;
+using Microsoft.Shared.Diagnostics;
+using Polly;
+
+namespace Polly;
+
+/// <summary>
+/// Provides utility methods for working with <see cref="ResilienceContext"/>.
+/// </summary>
+[Experimental(diagnosticId: DiagnosticIds.Experiments.Resilience, UrlFormat = DiagnosticIds.UrlFormat)]
+public static class HttpResilienceContextExtensions
+{
+    /// <summary>
+    /// Gets the request message from the <see cref="ResilienceContext"/>.
+    /// </summary>
+    /// <param name="context">The resilience context.</param>
+    /// <returns>
+    /// The request message.
+    /// If the request message is not present in the <see cref="ResilienceContext"/> the method returns <see langword="null"/>.
+    /// </returns>
+    /// <exception cref="ArgumentNullException"><paramref name="context"/> is <see langword="null"/>.</exception>
+    public static HttpRequestMessage? GetRequestMessage(this ResilienceContext context)
+    {
+        _ = Throw.IfNull(context);
+        return context.Properties.GetValue(ResilienceKeys.RequestMessage, default);
+    }
+
+    /// <summary>
+    /// Sets the request message on the <see cref="ResilienceContext"/>.
+    /// </summary>
+    /// <param name="context">The resilience context.</param>
+    /// <param name="requestMessage">The request message.</param>
+    /// <exception cref="ArgumentNullException"><paramref name="context"/> is <see langword="null"/>.</exception>
+    public static void SetRequestMessage(this ResilienceContext context, HttpRequestMessage? requestMessage)
+    {
+        _ = Throw.IfNull(context);
+        context.Properties.Set(ResilienceKeys.RequestMessage, requestMessage);
+    }
+}

src/Libraries/Microsoft.Extensions.Http.Resilience/Resilience/ResilienceHandler.cs

@@ -58,7 +58,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
 
         ResilienceContext context = GetOrSetResilienceContext(request, cancellationToken, out bool created);
         TrySetRequestMetadata(context, request);
-        SetRequestMessage(context, request);
+        context.SetRequestMessage(request);
 
         try
         {
@@ -117,7 +117,7 @@ protected override HttpResponseMessage Send(HttpRequestMessage request, Cancella
 
         ResilienceContext context = GetOrSetResilienceContext(request, cancellationToken, out bool created);
         TrySetRequestMetadata(context, request);
-        SetRequestMessage(context, request);
+        context.SetRequestMessage(request);
 
         try
         {
@@ -165,11 +165,8 @@ private static void TrySetRequestMetadata(ResilienceContext context, HttpRequest
         }
     }
 
-    private static void SetRequestMessage(ResilienceContext context, HttpRequestMessage request)
-        => context.Properties.Set(ResilienceKeys.RequestMessage, request);
-
     private static HttpRequestMessage GetRequestMessage(ResilienceContext context, HttpRequestMessage request)
-        => context.Properties.GetValue(ResilienceKeys.RequestMessage, request);
+        => context.GetRequestMessage() ?? request;
 
     private static void RestoreResilienceContext(ResilienceContext context, HttpRequestMessage request, bool created)
     {

src/Libraries/Microsoft.Extensions.Http.Resilience/Routing/Internal/RoutingResilienceStrategy.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Net.Http;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Http.Resilience.Internal;
 using Microsoft.Shared.Diagnostics;
@@ -26,7 +27,9 @@ protected override async ValueTask<Outcome<TResult>> ExecuteCore<TResult, TState
         ResilienceContext context,
         TState state)
     {
-        if (!context.Properties.TryGetValue(ResilienceKeys.RequestMessage, out var request))
+        HttpRequestMessage? request = context.GetRequestMessage();
+
+        if (request is null)
         {
             Throw.InvalidOperationException("The HTTP request message was not found in the resilience context.");
         }

test/Libraries/Microsoft.Extensions.Http.Resilience.Tests/Internal/RequestMessageSnapshotStrategyTests.cs

@@ -20,7 +20,7 @@ public async Task ExecuteAsync_EnsureSnapshotAttached()
         var strategy = Create();
         var context = ResilienceContextPool.Shared.Get();
         using var request = new HttpRequestMessage();
-        context.Properties.Set(ResilienceKeys.RequestMessage, request);
+        context.SetRequestMessage(request);
 
         using var response = await strategy.ExecuteAsync(
             context =>
@@ -39,5 +39,15 @@ public void ExecuteAsync_RequestMessageNotFound_Throws()
         strategy.Invoking(s => s.Execute(() => { })).Should().Throw<InvalidOperationException>();
     }
 
+    [Fact]
+    public void ExecuteAsync_RequestMessageIsNull_Throws()
+    {
+        var strategy = Create();
+        var context = ResilienceContextPool.Shared.Get();
+        context.SetRequestMessage(null);
+
+        strategy.Invoking(s => s.Execute(_ => { }, context)).Should().Throw<InvalidOperationException>();
+    }
+
     private static ResiliencePipeline Create() => new ResiliencePipelineBuilder().AddStrategy(_ => new RequestMessageSnapshotStrategy(), Mock.Of<ResilienceStrategyOptions>()).Build();
 }

test/Libraries/Microsoft.Extensions.Http.Resilience.Tests/Resilience/HttpResilienceContextExtensionsTests.cs

@@ -0,0 +1,77 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Net.Http;
+using Microsoft.Extensions.Http.Resilience.Internal;
+using Polly;
+using Xunit;
+
+namespace Microsoft.Extensions.Http.Resilience.Test.Resilience;
+
+public class HttpResilienceContextExtensionsTests
+{
+    [Fact]
+    public void GetRequestMessage_ResilienceContextIsNull_Throws()
+    {
+        ResilienceContext context = null!;
+        Assert.Throws<ArgumentNullException>(context.GetRequestMessage);
+    }
+
+    [Fact]
+    public void GetRequestMessage_RequestMessageIsMissing_ReturnsNull()
+    {
+        var context = ResilienceContextPool.Shared.Get();
+
+        Assert.Null(context.GetRequestMessage());
+    }
+
+    [Fact]
+    public void GetRequestMessage_RequestMessageIsNull_ReturnsNull()
+    {
+        var context = ResilienceContextPool.Shared.Get();
+        context.Properties.Set(ResilienceKeys.RequestMessage, null);
+
+        Assert.Null(context.GetRequestMessage());
+    }
+
+    [Fact]
+    public void GetRequestMessage_RequestMessageIsPresent_ReturnsRequestMessage()
+    {
+        var context = ResilienceContextPool.Shared.Get();
+        using var request = new HttpRequestMessage();
+        context.Properties.Set(ResilienceKeys.RequestMessage, request);
+
+        Assert.Same(request, context.GetRequestMessage());
+    }
+
+    [Fact]
+    public void SetRequestMessage_ResilienceContextIsNull_Throws()
+    {
+        ResilienceContext context = null!;
+        using var request = new HttpRequestMessage();
+
+        Assert.Throws<ArgumentNullException>(() => context.SetRequestMessage(request));
+    }
+
+    [Fact]
+    public void SetRequestMessage_RequestMessageIsNull_SetsNullRequestMessage()
+    {
+        var context = ResilienceContextPool.Shared.Get();
+        context.SetRequestMessage(null);
+
+        Assert.True(context.Properties.TryGetValue(ResilienceKeys.RequestMessage, out HttpRequestMessage? request));
+        Assert.Null(request);
+    }
+
+    [Fact]
+    public void SetRequestMessage_RequestMessageIsNotNull_SetsRequestMessage()
+    {
+        var context = ResilienceContextPool.Shared.Get();
+        using var request = new HttpRequestMessage();
+        context.SetRequestMessage(request);
+
+        Assert.True(context.Properties.TryGetValue(ResilienceKeys.RequestMessage, out HttpRequestMessage? actualRequest));
+        Assert.Same(request, actualRequest);
+    }
+}

test/Libraries/Microsoft.Extensions.Http.Resilience.Tests/Resilience/ResilienceHandlerTest.cs

@@ -8,7 +8,6 @@
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.Extensions.Http.Diagnostics;
-using Microsoft.Extensions.Http.Resilience.Internal;
 using Microsoft.Extensions.Http.Resilience.Test.Helpers;
 using Polly;
 using Xunit;
@@ -108,7 +107,7 @@ public async Task Send_EnsureInvoker(bool executionContextSet, bool asynchronous
         handler.InnerHandler = new TestHandlerStub((r, _) =>
         {
             r.GetResilienceContext().Should().NotBeNull();
-            r.GetResilienceContext()!.Properties.GetValue(ResilienceKeys.RequestMessage, null!).Should().BeSameAs(r);
+            r.GetResilienceContext()!.GetRequestMessage().Should().BeSameAs(r);
 
             return Task.FromResult(new HttpResponseMessage { StatusCode = HttpStatusCode.Created });
         });

test/Libraries/Microsoft.Extensions.Http.Resilience.Tests/Routing/RoutingResilienceStrategyTests.cs

@@ -4,7 +4,6 @@
 using System;
 using System.Net.Http;
 using FluentAssertions;
-using Microsoft.Extensions.Http.Resilience.Internal;
 using Microsoft.Extensions.Http.Resilience.Routing.Internal;
 using Moq;
 using Polly;
@@ -22,14 +21,24 @@ public void NoRequestMessage_Throws()
         strategy.Invoking(s => s.Execute(() => { })).Should().Throw<InvalidOperationException>().WithMessage("The HTTP request message was not found in the resilience context.");
     }
 
+    [Fact]
+    public void RequestMessageIsNull_Throws()
+    {
+        var strategy = Create(() => Mock.Of<RequestRoutingStrategy>());
+        var context = ResilienceContextPool.Shared.Get();
+        context.SetRequestMessage(null);
+
+        strategy.Invoking(s => s.Execute(_ => { }, context)).Should().Throw<InvalidOperationException>().WithMessage("The HTTP request message was not found in the resilience context.");
+    }
+
     [Fact]
     public void NoRoutingProvider_Ok()
     {
         using var request = new HttpRequestMessage();
 
         var strategy = Create(null);
         var context = ResilienceContextPool.Shared.Get();
-        context.Properties.Set(ResilienceKeys.RequestMessage, request);
+        context.SetRequestMessage(request);
 
         strategy.Invoking(s => s.Execute(_ => { }, context)).Should().NotThrow();
     }