Add cluster access API, reassignment (#1538)

docs/docfx/articles/ab-testing.md

@@ -0,0 +1,51 @@
+# A/B Testing and Rolling Upgrades
+
+## Introduction
+
+A/B testing and rolling upgrades require procedures for dynamically assigning incoming traffic to evaluate changes in the destination application. YARP does not have a built-in model for this, but it does expose some infrastructure useful for building such a system. See [issue #126](https://github.com/microsoft/reverse-proxy/issues/126) for additional details about this scenario.
+
+## Example
+
+```
+    public void Configure(IApplicationBuilder app, IProxyStateLookup lookup)
+    {
+        app.UseRouting();
+        app.UseEndpoints(endpoints =>
+        {
+            endpoints.MapReverseProxy(proxyPipeline =>
+            {
+                // Custom cluster selection
+                proxyPipeline.Use((context, next) =>
+                {
+                    if (lookup.TryGetCluster(ChooseCluster(context), out var cluster))
+                    {
+                        context.ReassignProxyRequest(cluster);
+                    }
+
+                    return next();
+                });
+                proxyPipeline.UseSessionAffinity();
+                proxyPipeline.UseLoadBalancing();
+            });
+        });
+    }
+
+    private string ChooseCluster(HttpContext context)
+    {
+        // Decide which cluster to use. This could be random, weighted, based on headers, etc.
+        return Random.Shared.Next(2) == 1 ? "cluster1" : "cluster2";
+    }
+```
+
+## Usage
+
+This scenario makes use of two APIs, [IProxyStateLookup](xref:Yarp.ReverseProxy.IProxyStateLookup) and [ReassignProxyRequest](xref:Microsoft.AspNetCore.Http.HttpContextFeaturesExtensions.ReassignProxyRequest), called from a custom proxy middleware as shown in the sample above.
+
+`IProxyStateLookup` is a service available in the Dependency Injection container that can be used to look up or enumerate the current routes and clusters. Note this data may change if the configuration changes. An A/B orchestration algorithm can examine the request, decide which cluster to send it to, and then retrieve that cluster from `IProxyStateLookup.TryGetCluster`.
+
+Once the cluster is selected, `ReassignProxyRequest` can be called to assign the request to that cluster. This updates the [IReverseProxyFeature](xref:Yarp.ReverseProxy.Model.IReverseProxyFeature) with the new cluster and destination information needed for the rest of the proxy middleware pipeline to handle the request.
+
+## Session affinity
+
+Note that session affinity functionality is split between middleware, which reads it settings from the current cluster, and transforms, which are part of the original route. Clusters used for A/B testing should use the same session affinity configuration to avoid conflicts.
+

docs/docfx/articles/toc.yml

@@ -40,3 +40,5 @@
   href: packages-refs.md
 - name: Diagnosing proxy issues
   href: diagnosing-yarp-issues.md
+- name: A/B Testing
+  href: ab-testing.md

src/ReverseProxy/Management/IProxyStateLookup.cs

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using Yarp.ReverseProxy.Model;
+
+namespace Yarp.ReverseProxy;
+
+/// <summary>
+/// Allows access to the proxy's current set of routes and clusters.
+/// </summary>
+public interface IProxyStateLookup
+{
+    /// <summary>
+    /// Retrieves a specific route by id, if present.
+    /// </summary>
+    bool TryGetRoute(string id, [NotNullWhen(true)] out RouteModel? route);
+
+    /// <summary>
+    /// Enumerates all current routes. This is thread safe but the collection may change mid enumeration if the configuration is reloaded.
+    /// </summary>
+    IEnumerable<RouteModel> GetRoutes();
+
+    /// <summary>
+    /// Retrieves a specific cluster by id, if present.
+    /// </summary>
+    bool TryGetCluster(string id, [NotNullWhen(true)] out ClusterState? cluster);
+
+    /// <summary>
+    /// Enumerates all current clusters. This is thread safe but the collection may change mid enumeration if the configuration is reloaded.
+    /// </summary>
+    IEnumerable<ClusterState> GetClusters();
+}

src/ReverseProxy/Management/IReverseProxyBuilderExtensions.cs

@@ -50,6 +50,7 @@ public static IReverseProxyBuilder AddRuntimeStateManagers(this IReverseProxyBui
     public static IReverseProxyBuilder AddConfigManager(this IReverseProxyBuilder builder)
     {
         builder.Services.TryAddSingleton<ProxyConfigManager>();
+        builder.Services.TryAddSingleton<IProxyStateLookup>(sp => sp.GetRequiredService<ProxyConfigManager>());
         return builder;
     }
 

src/ReverseProxy/Management/ProxyConfigManager.cs

@@ -30,7 +30,7 @@ namespace Yarp.ReverseProxy.Management;
 /// in a thread-safe manner while avoiding locks on the hot path.
 /// </summary>
 // https://github.com/dotnet/aspnetcore/blob/cbe16474ce9db7ff588aed89596ff4df5c3f62e1/src/Mvc/Mvc.Core/src/Routing/ActionEndpointDataSourceBase.cs
-internal sealed class ProxyConfigManager : EndpointDataSource, IDisposable
+internal sealed class ProxyConfigManager : EndpointDataSource, IProxyStateLookup, IDisposable
 {
     private static readonly IReadOnlyDictionary<string, ClusterConfig> _emptyClusterDictionary = new ReadOnlyDictionary<string, ClusterConfig>(new Dictionary<string, ClusterConfig>());
 
@@ -680,6 +680,39 @@ private RouteModel BuildRouteModel(RouteConfig source, ClusterState? cluster)
         return new RouteModel(source, cluster, transforms);
     }
 
+    public bool TryGetRoute(string id, [NotNullWhen(true)] out RouteModel? route)
+    {
+        if (_routes.TryGetValue(id, out var routeState))
+        {
+            route = routeState.Model;
+            return true;
+        }
+
+        route = null;
+        return false;
+    }
+
+    public IEnumerable<RouteModel> GetRoutes()
+    {
+        foreach (var (_, route) in _routes)
+        {
+            yield return route.Model;
+        }
+    }
+
+    public bool TryGetCluster(string id, [NotNullWhen(true)] out ClusterState? cluster)
+    {
+        return _clusters.TryGetValue(id, out cluster!);
+    }
+
+    public IEnumerable<ClusterState> GetClusters()
+    {
+        foreach (var (_, cluster) in _clusters)
+        {
+            yield return cluster;
+        }
+    }
+
     public void Dispose()
     {
         _configChangeSource.Dispose();

src/ReverseProxy/Model/HttpContextFeaturesExtensions.cs

@@ -40,4 +40,24 @@ public static IReverseProxyFeature GetReverseProxyFeature(this HttpContext conte
     {
         return context.Features.Get<IForwarderErrorFeature>();
     }
+
+    // Compare to ProxyPipelineInitializerMiddleware
+    /// <summary>
+    /// Replaces the assigned cluster and destinations in <see cref="IReverseProxyFeature"/> with the new <see cref="ClusterState"/>,
+    /// causing the request to be sent to the new cluster instead.
+    /// </summary>
+    public static void ReassignProxyRequest(this HttpContext context, ClusterState cluster)
+    {
+        var oldFeature = context.GetReverseProxyFeature();
+        var destinations = cluster.DestinationsState;
+        var newFeature = new ReverseProxyFeature()
+        {
+            Route = oldFeature.Route,
+            Cluster = cluster.Model,
+            AllDestinations = destinations.AllDestinations,
+            AvailableDestinations = destinations.AvailableDestinations,
+            ProxiedDestination = oldFeature.ProxiedDestination,
+        };
+        context.Features.Set<IReverseProxyFeature>(newFeature);
+    }
 }

src/ReverseProxy/SessionAffinity/AffinitizeTransform.cs

@@ -27,7 +27,11 @@ public override ValueTask ApplyAsync(ResponseTransformContext context)
         var proxyFeature = context.HttpContext.GetReverseProxyFeature();
         var options = proxyFeature.Cluster.Config.SessionAffinity;
         // The transform should only be added to routes that have affinity enabled.
-        Debug.Assert(options?.Enabled ?? true, "Session affinity is not enabled");
+        // However, the cluster can be re-assigned dynamically.
+        if (options == null || !options.Enabled.GetValueOrDefault())
+        {
+            return default;
+        }
         var selectedDestination = proxyFeature.ProxiedDestination!;
         _sessionAffinityPolicy.AffinitizeResponse(context.HttpContext, proxyFeature.Route.Cluster!, options!, selectedDestination);
         return default;

src/ReverseProxy/Utilities/NullableAttributes.cs

@@ -69,25 +69,6 @@ sealed class MaybeNullWhenAttribute : Attribute
     public bool ReturnValue { get; }
 }
 
-/// <summary>Specifies that when a method returns <see cref="ReturnValue"/>, the parameter will not be null even if the corresponding type allows it.</summary>
-[AttributeUsage(AttributeTargets.Parameter, Inherited = false)]
-#if INTERNAL_NULLABLE_ATTRIBUTES
-internal
-#else
-    public
-#endif
-        sealed class NotNullWhenAttribute : Attribute
-{
-    /// <summary>Initializes the attribute with the specified return value condition.</summary>
-    /// <param name="returnValue">
-    /// The return value condition. If the method returns this value, the associated parameter will not be null.
-    /// </param>
-    public NotNullWhenAttribute(bool returnValue) => ReturnValue = returnValue;
-
-    /// <summary>Gets the return value condition.</summary>
-    public bool ReturnValue { get; }
-}
-
 /// <summary>Specifies that the output will be non-null if the named parameter is non-null.</summary>
 [AttributeUsage(AttributeTargets.Parameter | AttributeTargets.Property | AttributeTargets.ReturnValue, AllowMultiple = true, Inherited = false)]
 #if INTERNAL_NULLABLE_ATTRIBUTES

test/ReverseProxy.Tests/Management/ProxyConfigManagerTests.cs

@@ -106,6 +106,20 @@ public async Task Endpoints_StartsEmpty()
         Assert.Empty(endpoints);
     }
 
+    [Fact]
+    public async Task Lookup_StartsEmpty()
+    {
+        var services = CreateServices(new List<RouteConfig>(), new List<ClusterConfig>());
+        var manager = services.GetRequiredService<ProxyConfigManager>();
+        var lookup = services.GetRequiredService<IProxyStateLookup>();
+        await manager.InitialLoadAsync();
+
+        Assert.Empty(lookup.GetRoutes());
+        Assert.Empty(lookup.GetClusters());
+        Assert.False(lookup.TryGetRoute("route1", out var _));
+        Assert.False(lookup.TryGetCluster("cluster1", out var _));
+    }
+
     [Fact]
     public async Task GetChangeToken_InitialValue()
     {
@@ -142,6 +156,7 @@ public async Task BuildConfig_OneClusterOneDestinationOneRoute_Works()
         var services = CreateServices(new List<RouteConfig>() { route }, new List<ClusterConfig>() { cluster });
 
         var manager = services.GetRequiredService<ProxyConfigManager>();
+        var lookup = services.GetRequiredService<IProxyStateLookup>();
         var dataSource = await manager.InitialLoadAsync();
 
         Assert.NotNull(dataSource);
@@ -150,6 +165,10 @@ public async Task BuildConfig_OneClusterOneDestinationOneRoute_Works()
         var routeConfig = endpoint.Metadata.GetMetadata<RouteModel>();
         Assert.NotNull(routeConfig);
         Assert.Equal("route1", routeConfig.Config.RouteId);
+        Assert.True(lookup.TryGetRoute("route1", out var routeModel));
+        Assert.Equal(route, routeModel.Config);
+        routeModel = Assert.Single(lookup.GetRoutes());
+        Assert.Equal(route, routeModel.Config);
 
         var clusterState = routeConfig.Cluster;
         Assert.NotNull(clusterState);
@@ -159,6 +178,10 @@ public async Task BuildConfig_OneClusterOneDestinationOneRoute_Works()
         Assert.NotNull(clusterState.Model);
         Assert.NotNull(clusterState.Model.HttpClient);
         Assert.Same(clusterState, routeConfig.Cluster);
+        Assert.True(lookup.TryGetCluster("cluster1", out clusterState));
+        Assert.Equal(cluster, clusterState.Model.Config);
+        clusterState = Assert.Single(lookup.GetClusters());
+        Assert.Equal(cluster, clusterState.Model.Config);
 
         var actualDestinations = clusterState.Destinations.Values;
         var destination = Assert.Single(actualDestinations);

test/ReverseProxy.Tests/Model/HttpContextFeaturesExtensions.cs

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Net.Http;
+using Microsoft.AspNetCore.Http;
+using Xunit;
+using Yarp.ReverseProxy.Configuration;
+using Yarp.ReverseProxy.Forwarder;
+
+namespace Yarp.ReverseProxy.Model.Tests;
+
+public class HttpContextFeaturesExtensions
+{
+    [Fact]
+    public void ReassignProxyRequest_Success()
+    {
+        var client = new HttpMessageInvoker(new SocketsHttpHandler());
+        var context = new DefaultHttpContext();
+        var d1 = new DestinationState("d1");
+        var d2 = new DestinationState("d2");
+        var cc1 = new ClusterConfig() { ClusterId = "c1" };
+        var cm1 = new ClusterModel(cc1, client);
+        var cs1 = new ClusterState("c1") { Model = cm1 };
+        var r1 = new RouteModel(new RouteConfig() { RouteId = "r1" }, cs1, HttpTransformer.Empty);
+        var feature = new ReverseProxyFeature()
+        {
+            AllDestinations = d1,
+            AvailableDestinations = d1,
+            Cluster = cm1,
+            Route = r1,
+            ProxiedDestination = d1,
+        };
+
+        context.Features.Set<IReverseProxyFeature>(feature);
+
+        var cc2 = new ClusterConfig() { ClusterId = "cc2" };
+        var cm2 = new ClusterModel(cc2, client);
+        var cs2 = new ClusterState("cs2")
+        {
+            DestinationsState = new ClusterDestinationsState(d2, d2),
+            Model = cm2,
+        };
+        context.ReassignProxyRequest(cs2);
+
+        var newFeature = context.GetReverseProxyFeature();
+        Assert.NotSame(feature, newFeature);
+        Assert.Same(d2, newFeature.AllDestinations);
+        Assert.Same(d2, newFeature.AvailableDestinations);
+        Assert.Same(d1, newFeature.ProxiedDestination); // Copied unmodified.
+        Assert.Same(cm2, newFeature.Cluster);
+        Assert.Same(r1, newFeature.Route);
+    }
+}

testassets/ReverseProxy.Code/Startup.cs

@@ -6,11 +6,12 @@
 using System.Net.Http.Headers;
 using Microsoft.AspNetCore.Authentication;
 using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.DependencyInjection;
 using Yarp.ReverseProxy.Configuration;
 using Yarp.ReverseProxy.Model;
-using Yarp.Telemetry.Consumption;
 using Yarp.ReverseProxy.Transforms;
+using Yarp.Telemetry.Consumption;
 
 namespace Yarp.ReverseProxy.Sample;
 
@@ -47,6 +48,14 @@ public void ConfigureServices(IServiceCollection services)
                 {
                     { "destination1", new DestinationConfig() { Address = "https://localhost:10000" } }
                 }
+            },
+            new ClusterConfig()
+            {
+                ClusterId = "cluster2",
+                Destinations = new Dictionary<string, DestinationConfig>(StringComparer.OrdinalIgnoreCase)
+                {
+                    { "destination2", new DestinationConfig() { Address = "https://localhost:10001" } }
+                }
             }
         };
 
@@ -100,7 +109,7 @@ public void ConfigureServices(IServiceCollection services)
     /// <summary>
     /// This method gets called by the runtime. Use this method to configure the HTTP request pipeline.
     /// </summary>
-    public void Configure(IApplicationBuilder app)
+    public void Configure(IApplicationBuilder app, IProxyStateLookup lookup)
     {
         app.UseRouting();
         app.UseAuthorization();
@@ -112,6 +121,11 @@ public void Configure(IApplicationBuilder app)
                 // Custom endpoint selection
                 proxyPipeline.Use((context, next) =>
                 {
+                    if (lookup.TryGetCluster("cluster2", out var cluster))
+                    {
+                        context.ReassignProxyRequest(cluster);
+                    }
+
                     var someCriteria = false; // MeetsCriteria(context);
                     if (someCriteria)
                     {