Added IClusterDestinationResolver

src/ReverseProxy/LoadBalancing/ILoadBalancingDestinationSelector.cs

@@ -0,0 +1,17 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+using Microsoft.AspNetCore.Http;
+using Yarp.ReverseProxy.Model;
+
+namespace Yarp.ReverseProxy.LoadBalancing;
+
+internal interface ILoadBalancingDestinationSelector
+{
+    DestinationState? PickDestination(
+        HttpContext? context,
+        ClusterState cluster,
+        IReadOnlyList<DestinationState> availableDestinations,
+        string? loadBalancingPolicy = null);
+}

src/ReverseProxy/LoadBalancing/LoadBalancingDestinationSelector.cs

@@ -0,0 +1,49 @@
+// 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.Collections.Frozen;
+using System.Collections.Generic;
+using Microsoft.AspNetCore.Http;
+using Yarp.ReverseProxy.Model;
+using Yarp.ReverseProxy.Utilities;
+
+namespace Yarp.ReverseProxy.LoadBalancing;
+
+internal sealed class LoadBalancingDestinationSelector : ILoadBalancingDestinationSelector
+{
+    private readonly FrozenDictionary<string, ILoadBalancingPolicy> _loadBalancingPolicies;
+
+    public LoadBalancingDestinationSelector(IEnumerable<ILoadBalancingPolicy> loadBalancingPolicies)
+    {
+        ArgumentNullException.ThrowIfNull(loadBalancingPolicies);
+        _loadBalancingPolicies = loadBalancingPolicies.ToDictionaryByUniqueId(p => p.Name);
+    }
+
+    public DestinationState? PickDestination(
+        HttpContext? context,
+        ClusterState cluster,
+        IReadOnlyList<DestinationState> availableDestinations,
+        string? loadBalancingPolicy = null)
+    {
+        ArgumentNullException.ThrowIfNull(cluster);
+        ArgumentNullException.ThrowIfNull(availableDestinations);
+
+        var destinationCount = availableDestinations.Count;
+
+        if (destinationCount == 0)
+        {
+            return null;
+        }
+
+        if (destinationCount == 1)
+        {
+            return availableDestinations[0];
+        }
+
+        var currentPolicy = _loadBalancingPolicies.GetRequiredServiceById(
+            loadBalancingPolicy ?? cluster.Model.Config.LoadBalancingPolicy,
+            LoadBalancingPolicies.PowerOfTwoChoices);
+        return currentPolicy.PickDestination(context ?? new DefaultHttpContext(), cluster, availableDestinations);
+    }
+}

src/ReverseProxy/LoadBalancing/LoadBalancingMiddleware.cs

@@ -2,13 +2,10 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
-using System.Collections.Frozen;
-using System.Collections.Generic;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Logging;
 using Yarp.ReverseProxy.Model;
-using Yarp.ReverseProxy.Utilities;
 
 namespace Yarp.ReverseProxy.LoadBalancing;
 
@@ -18,44 +15,31 @@ namespace Yarp.ReverseProxy.LoadBalancing;
 internal sealed class LoadBalancingMiddleware
 {
     private readonly ILogger _logger;
-    private readonly FrozenDictionary<string, ILoadBalancingPolicy> _loadBalancingPolicies;
+    private readonly ILoadBalancingDestinationSelector _destinationSelector;
     private readonly RequestDelegate _next;
 
     public LoadBalancingMiddleware(
         RequestDelegate next,
         ILogger<LoadBalancingMiddleware> logger,
-        IEnumerable<ILoadBalancingPolicy> loadBalancingPolicies)
+        ILoadBalancingDestinationSelector destinationSelector)
     {
         ArgumentNullException.ThrowIfNull(next);
         ArgumentNullException.ThrowIfNull(logger);
-        ArgumentNullException.ThrowIfNull(loadBalancingPolicies);
+        ArgumentNullException.ThrowIfNull(destinationSelector);
         _next = next;
         _logger = logger;
-        _loadBalancingPolicies = loadBalancingPolicies.ToDictionaryByUniqueId(p => p.Name);
+        _destinationSelector = destinationSelector;
     }
 
     public Task Invoke(HttpContext context)
     {
         var proxyFeature = context.GetReverseProxyFeature();
 
-        var destinations = proxyFeature.AvailableDestinations;
-        var destinationCount = destinations.Count;
-
-        DestinationState? destination;
-
-        if (destinationCount == 0)
-        {
-            destination = null;
-        }
-        else if (destinationCount == 1)
-        {
-            destination = destinations[0];
-        }
-        else
-        {
-            var currentPolicy = _loadBalancingPolicies.GetRequiredServiceById(proxyFeature.Cluster.Config.LoadBalancingPolicy, LoadBalancingPolicies.PowerOfTwoChoices);
-            destination = currentPolicy.PickDestination(context, proxyFeature.Route.Cluster!, destinations);
-        }
+        var destination = _destinationSelector.PickDestination(
+            context,
+            proxyFeature.Route.Cluster!,
+            proxyFeature.AvailableDestinations,
+            proxyFeature.Cluster.Config.LoadBalancingPolicy);
 
         if (destination is null)
         {

src/ReverseProxy/Management/ClusterDestinationResolver.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.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Yarp.ReverseProxy.LoadBalancing;
+using Yarp.ReverseProxy.Model;
+
+namespace Yarp.ReverseProxy.Management;
+
+internal sealed class ClusterDestinationResolver : IClusterDestinationResolver
+{
+    private readonly IProxyStateLookup _proxyStateLookup;
+    private readonly ILoadBalancingDestinationSelector _destinationSelector;
+
+    public ClusterDestinationResolver(
+        IProxyStateLookup proxyStateLookup,
+        ILoadBalancingDestinationSelector destinationSelector)
+    {
+        ArgumentNullException.ThrowIfNull(proxyStateLookup);
+        ArgumentNullException.ThrowIfNull(destinationSelector);
+
+        _proxyStateLookup = proxyStateLookup;
+        _destinationSelector = destinationSelector;
+    }
+
+    public ValueTask<DestinationState?> GetDestinationAsync(
+        string clusterId,
+        HttpContext? context = null,
+        CancellationToken cancellationToken = default)
+    {
+        ArgumentException.ThrowIfNullOrEmpty(clusterId);
+        cancellationToken.ThrowIfCancellationRequested();
+
+        if (!_proxyStateLookup.TryGetCluster(clusterId, out var cluster))
+        {
+            throw new KeyNotFoundException($"No cluster was found for the id '{clusterId}'.");
+        }
+
+        return ValueTask.FromResult(
+            _destinationSelector.PickDestination(context, cluster, cluster.DestinationsState.AvailableDestinations));
+    }
+}

src/ReverseProxy/Management/IClusterDestinationResolver.cs

@@ -0,0 +1,27 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Yarp.ReverseProxy.Model;
+
+namespace Yarp.ReverseProxy;
+
+/// <summary>
+/// Resolves a destination for a cluster using the current runtime state and configured load balancing policy.
+/// </summary>
+public interface IClusterDestinationResolver
+{
+    /// <summary>
+    /// Resolves a destination for the given cluster.
+    /// </summary>
+    /// <param name="clusterId">The cluster id.</param>
+    /// <param name="context">Optional request context used by load balancing policies.</param>
+    /// <param name="cancellationToken">Cancellation token.</param>
+    /// <returns>The selected destination, or <see langword="null"/> if no destinations are currently available.</returns>
+    ValueTask<DestinationState?> GetDestinationAsync(
+        string clusterId,
+        HttpContext? context = null,
+        CancellationToken cancellationToken = default);
+}

src/ReverseProxy/Management/IClusterDestinationResolverExtensions.cs

@@ -0,0 +1,40 @@
+// 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.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+
+namespace Yarp.ReverseProxy;
+
+/// <summary>
+/// Extension methods for <see cref="IClusterDestinationResolver"/>.
+/// </summary>
+public static class IClusterDestinationResolverExtensions
+{
+    /// <summary>
+    /// Resolves the destination URI for the given cluster.
+    /// </summary>
+    /// <param name="resolver">The destination resolver.</param>
+    /// <param name="clusterId">The cluster id.</param>
+    /// <param name="context">Optional request context used by load balancing policies.</param>
+    /// <param name="cancellationToken">Cancellation token.</param>
+    /// <returns>The selected destination URI, or <see langword="null"/> if no destinations are currently available.</returns>
+    public static async ValueTask<Uri?> GetDestinationUriAsync(
+        this IClusterDestinationResolver resolver,
+        string clusterId,
+        HttpContext? context = null,
+        CancellationToken cancellationToken = default)
+    {
+        ArgumentNullException.ThrowIfNull(resolver);
+
+        var destination = await resolver.GetDestinationAsync(clusterId, context, cancellationToken);
+        if (destination is null)
+        {
+            return null;
+        }
+
+        return new Uri(destination.Model.Config.Address, UriKind.Absolute);
+    }
+}

src/ReverseProxy/Management/IReverseProxyBuilderExtensions.cs

@@ -72,6 +72,7 @@ public static IReverseProxyBuilder AddConfigManager(this IReverseProxyBuilder bu
     {
         builder.Services.TryAddSingleton<ProxyConfigManager>();
         builder.Services.TryAddSingleton<IProxyStateLookup>(sp => sp.GetRequiredService<ProxyConfigManager>());
+        builder.Services.TryAddSingleton<IClusterDestinationResolver, ClusterDestinationResolver>();
         return builder;
     }
 
@@ -86,6 +87,7 @@ public static IReverseProxyBuilder AddProxy(this IReverseProxyBuilder builder)
     public static IReverseProxyBuilder AddLoadBalancingPolicies(this IReverseProxyBuilder builder)
     {
         builder.Services.TryAddSingleton<IRandomFactory, RandomFactory>();
+        builder.Services.TryAddSingleton<ILoadBalancingDestinationSelector, LoadBalancingDestinationSelector>();
 
         builder.Services.TryAddEnumerable(new[] {
             ServiceDescriptor.Singleton<ILoadBalancingPolicy, FirstLoadBalancingPolicy>(),

test/ReverseProxy.Tests/LoadBalancing/LoadBalancerMiddlewareTests.cs

@@ -24,10 +24,12 @@ private static LoadBalancingMiddleware CreateMiddleware(RequestDelegate next, pa
             .Setup(l => l.IsEnabled(It.IsAny<LogLevel>()))
             .Returns(true);
 
+        var destinationSelector = new LoadBalancingDestinationSelector(loadBalancingPolicies);
+
         return new LoadBalancingMiddleware(
             next,
             logger.Object,
-            loadBalancingPolicies);
+            destinationSelector);
     }
 
     [Fact]