Add async ISessionAffinityPolicy APIs (#1990)

* Add async ISessionAffinityPolicy APIs

* Consume CT changes

Author: MihaZupan

src/ReverseProxy/SessionAffinity/AffinitizeTransform.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Diagnostics;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Yarp.ReverseProxy.Model;
@@ -31,8 +32,15 @@ public override ValueTask ApplyAsync(ResponseTransformContext context)
         {
             return default;
         }
-        var selectedDestination = proxyFeature.ProxiedDestination!;
-        _sessionAffinityPolicy.AffinitizeResponse(context.HttpContext, proxyFeature.Route.Cluster!, options!, selectedDestination);
-        return default;
+
+        Debug.Assert(proxyFeature.Route.Cluster is not null);
+        Debug.Assert(proxyFeature.ProxiedDestination is not null);
+
+        return _sessionAffinityPolicy.AffinitizeResponseAsync(
+            context.HttpContext,
+            proxyFeature.Route.Cluster,
+            options,
+            proxyFeature.ProxiedDestination,
+            context.CancellationToken);
     }
 }

src/ReverseProxy/SessionAffinity/ISessionAffinityPolicy.cs

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Yarp.ReverseProxy.Configuration;
 using Yarp.ReverseProxy.Model;
@@ -28,6 +30,22 @@ public interface ISessionAffinityPolicy
     /// <returns><see cref="AffinityResult"/> carrying the found affinitized destinations if any and the <see cref="AffinityStatus"/>.</returns>
     AffinityResult FindAffinitizedDestinations(HttpContext context, ClusterState cluster, SessionAffinityConfig config, IReadOnlyList<DestinationState> destinations);
 
+    /// <summary>
+    /// Finds <see cref="DestinationState"/> to which the current request is affinitized by the affinity key.
+    /// </summary>
+    /// <param name="context">Current request's context.</param>
+    /// <param name="cluster">Current request's cluster.</param>
+    /// <param name="config">Affinity config.</param>
+    /// <param name="destinations"><see cref="DestinationState"/>s available for the request.</param>
+    /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
+    /// <returns><see cref="AffinityResult"/> carrying the found affinitized destinations if any and the <see cref="AffinityStatus"/>.</returns>
+    ValueTask<AffinityResult> FindAffinitizedDestinationsAsync(HttpContext context, ClusterState cluster, SessionAffinityConfig config, IReadOnlyList<DestinationState> destinations, CancellationToken cancellationToken)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+
+        return new ValueTask<AffinityResult>(FindAffinitizedDestinations(context, cluster, config, destinations));
+    }
+
     /// <summary>
     /// Affinitize the current response to the given <see cref="DestinationState"/> by setting the affinity key extracted from <see cref="DestinationState"/>.
     /// </summary>
@@ -36,4 +54,20 @@ public interface ISessionAffinityPolicy
     /// <param name="config">Affinity config.</param>
     /// <param name="destination"><see cref="DestinationState"/> to which request is to be affinitized.</param>
     void AffinitizeResponse(HttpContext context, ClusterState cluster, SessionAffinityConfig config, DestinationState destination);
+
+    /// <summary>
+    /// Affinitize the current response to the given <see cref="DestinationState"/> by setting the affinity key extracted from <see cref="DestinationState"/>.
+    /// </summary>
+    /// <param name="context">Current request's context.</param>
+    /// <param name="cluster">Current request's cluster.</param>
+    /// <param name="config">Affinity config.</param>
+    /// <param name="destination"><see cref="DestinationState"/> to which request is to be affinitized.</param>
+    /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
+    ValueTask AffinitizeResponseAsync(HttpContext context, ClusterState cluster, SessionAffinityConfig config, DestinationState destination, CancellationToken cancellationToken)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+
+        AffinitizeResponse(context, cluster, config, destination);
+        return default;
+    }
 }

src/ReverseProxy/SessionAffinity/SessionAffinityMiddleware.cs

@@ -54,7 +54,7 @@ private async Task InvokeInternal(HttpContext context, IReverseProxyFeature prox
         var cluster = proxyFeature.Route.Cluster!;
 
         var policy = _sessionAffinityPolicies.GetRequiredServiceById(config.Policy, SessionAffinityConstants.Policies.HashCookie);
-        var affinityResult = policy.FindAffinitizedDestinations(context, cluster, config, destinations);
+        var affinityResult = await policy.FindAffinitizedDestinationsAsync(context, cluster, config, destinations, context.RequestAborted);
 
         switch (affinityResult.Status)
         {

test/ReverseProxy.Tests/SessionAffinity/AffinitizeTransformTests.cs

@@ -3,6 +3,7 @@
 
 using System.Linq;
 using System.Net.Http;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Moq;
@@ -22,7 +23,14 @@ public async Task ApplyAsync_InvokeAffinitizeRequest()
         var cluster = GetCluster();
         var destination = cluster.Destinations.Values.First();
         var provider = new Mock<ISessionAffinityPolicy>(MockBehavior.Strict);
-        provider.Setup(p => p.AffinitizeResponse(It.IsAny<HttpContext>(), It.IsAny<ClusterState>(), It.IsNotNull<SessionAffinityConfig>(), It.IsAny<DestinationState>()));
+        provider
+            .Setup(p => p.AffinitizeResponseAsync(
+                It.IsAny<HttpContext>(),
+                It.IsAny<ClusterState>(),
+                It.IsNotNull<SessionAffinityConfig>(),
+                It.IsAny<DestinationState>(),
+                It.IsAny<CancellationToken>()))
+            .Returns(new ValueTask());
 
         var transform = new AffinitizeTransform(provider.Object);
 

test/ReverseProxy.Tests/SessionAffinity/SessionAffinityMiddlewareTests.cs

@@ -13,6 +13,7 @@
 using Yarp.ReverseProxy.Configuration;
 using Yarp.ReverseProxy.Model;
 using Yarp.ReverseProxy.Forwarder;
+using System.Threading;
 
 namespace Yarp.ReverseProxy.SessionAffinity.Tests;
 
@@ -153,21 +154,24 @@ internal IReadOnlyList<Mock<ISessionAffinityPolicy>> RegisterAffinityPolicies(
             policy.SetupGet(p => p.Name).Returns(mode);
             if (lookupMiddlewareTest)
             {
-                policy.Setup(p => p.FindAffinitizedDestinations(
+                policy.Setup(p => p.FindAffinitizedDestinationsAsync(
                     It.IsAny<HttpContext>(),
                     It.IsAny<ClusterState>(),
                     ClusterConfig.Config.SessionAffinity,
-                    expectedDestinations))
-                .Returns(new AffinityResult(destinations, status.Value))
+                    expectedDestinations,
+                    It.IsAny<CancellationToken>()))
+                .Returns(new ValueTask<AffinityResult>(new AffinityResult(destinations, status.Value)))
                 .Callback(() => callback(policy.Object));
             }
             else
             {
-                policy.Setup(p => p.AffinitizeResponse(
+                policy.Setup(p => p.AffinitizeResponseAsync(
                     It.IsAny<HttpContext>(),
                     It.IsAny<ClusterState>(),
                     ClusterConfig.Config.SessionAffinity,
-                    expectedDestinations[0]))
+                    expectedDestinations[0],
+                    It.IsAny<CancellationToken>()))
+                .Returns(new ValueTask())
                 .Callback(() => callback(policy.Object));
             }
             result.Add(policy);