Rewire LoadBalancerFactory to allow injecting custom load balancers, by introducing LoadBalancerCreators

src/Ocelot/DependencyInjection/OcelotBuilder.cs

@@ -87,6 +87,10 @@ namespace Ocelot.DependencyInjection
             Services.TryAddSingleton<IFileConfigurationRepository, DiskFileConfigurationRepository>();
             Services.TryAddSingleton<IFileConfigurationSetter, FileAndInternalConfigurationSetter>();
             Services.TryAddSingleton<IServiceDiscoveryProviderFactory, ServiceDiscoveryProviderFactory>();
+            Services.AddSingleton<ILoadBalancerCreator, NoLoadBalancerCreator>();
+            Services.AddSingleton<ILoadBalancerCreator, RoundRobinCreator>();
+            Services.AddSingleton<ILoadBalancerCreator, CookieStickySessionsCreator>();
+            Services.AddSingleton<ILoadBalancerCreator, LeastConnectionCreator>();
             Services.TryAddSingleton<ILoadBalancerFactory, LoadBalancerFactory>();
             Services.TryAddSingleton<ILoadBalancerHouse, LoadBalancerHouse>();
             Services.TryAddSingleton<IOcelotLoggerFactory, AspDotNetLoggerFactory>();

src/Ocelot/LoadBalancer/LoadBalancers/CookieStickySessionsCreator.cs

@@ -0,0 +1,20 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System.Threading.Tasks;
+    using Ocelot.Configuration;
+    using Ocelot.Infrastructure;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class CookieStickySessionsCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            var loadBalancer = new RoundRobin(async () => await serviceProvider.Get());
+            var bus = new InMemoryBus<StickySession>();
+            return new CookieStickySessions(loadBalancer, reRoute.LoadBalancerOptions.Key,
+                reRoute.LoadBalancerOptions.ExpiryInMs, bus);
+        }
+
+        public string Type => nameof(CookieStickySessions);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/ILoadBalancerCreator.cs

@@ -0,0 +1,11 @@
+using Ocelot.Configuration;
+using Ocelot.ServiceDiscovery.Providers;
+
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    public interface ILoadBalancerCreator
+    {
+        ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider);
+        string Type { get; }
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/LeastConnectionCreator.cs

@@ -0,0 +1,15 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class LeastConnectionCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new LeastConnection(async () => await serviceProvider.Get(), reRoute.ServiceName);
+        }
+
+        public string Type => nameof(LeastConnection);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/LoadBalancerFactory.cs

@@ -1,47 +1,42 @@
-using Ocelot.Configuration;
-using Ocelot.Infrastructure;
-using Ocelot.Responses;
-using Ocelot.ServiceDiscovery;
-using System.Threading.Tasks;
-
-namespace Ocelot.LoadBalancer.LoadBalancers
+namespace Ocelot.LoadBalancer.LoadBalancers
 {
+    using System.Collections.Generic;
+    using System.Linq;
+    using Ocelot.Configuration;
+    using Ocelot.Responses;
+    using System.Threading.Tasks;
+    using Ocelot.ServiceDiscovery;
+
     public class LoadBalancerFactory : ILoadBalancerFactory
     {
         private readonly IServiceDiscoveryProviderFactory _serviceProviderFactory;
+        private readonly IEnumerable<ILoadBalancerCreator> _loadBalancerCreators;
 
-        public LoadBalancerFactory(IServiceDiscoveryProviderFactory serviceProviderFactory)
+        public LoadBalancerFactory(IServiceDiscoveryProviderFactory serviceProviderFactory, IEnumerable<ILoadBalancerCreator> loadBalancerCreators)
         {
             _serviceProviderFactory = serviceProviderFactory;
+            _loadBalancerCreators = loadBalancerCreators;
         }
 
-        public async Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
+        public Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
         {
-            var response = _serviceProviderFactory.Get(config, reRoute);
+            var serviceProviderFactoryResponse = _serviceProviderFactory.Get(config, reRoute);
 
-            if (response.IsError)
+            Response<ILoadBalancer> response;
+            if (serviceProviderFactoryResponse.IsError)
             {
-                return new ErrorResponse<ILoadBalancer>(response.Errors);
+                response = new ErrorResponse<ILoadBalancer>(serviceProviderFactoryResponse.Errors);
+            }
+            else
+            {
+                var serviceProvider = serviceProviderFactoryResponse.Data;
+                var requestedType = reRoute.LoadBalancerOptions?.Type ?? nameof(NoLoadBalancer);
+                var applicableCreator = _loadBalancerCreators.Single(c => c.Type == requestedType);
+                var createdLoadBalancer = applicableCreator.Create(reRoute, serviceProvider);
+                response = new OkResponse<ILoadBalancer>(createdLoadBalancer);
             }
 
-            var serviceProvider = response.Data;
-
-            switch (reRoute.LoadBalancerOptions?.Type)
-            {
-                case nameof(RoundRobin):
-                    return new OkResponse<ILoadBalancer>(new RoundRobin(async () => await serviceProvider.Get()));
-
-                case nameof(LeastConnection):
-                    return new OkResponse<ILoadBalancer>(new LeastConnection(async () => await serviceProvider.Get(), reRoute.ServiceName));
-
-                case nameof(CookieStickySessions):
-                    var loadBalancer = new RoundRobin(async () => await serviceProvider.Get());
-                    var bus = new InMemoryBus<StickySession>();
-                    return new OkResponse<ILoadBalancer>(new CookieStickySessions(loadBalancer, reRoute.LoadBalancerOptions.Key, reRoute.LoadBalancerOptions.ExpiryInMs, bus));
-
-                default:
-                    return new OkResponse<ILoadBalancer>(new NoLoadBalancer(async () => await serviceProvider.Get()));
-            }
+            return Task.FromResult(response);
         }
     }
 }

src/Ocelot/LoadBalancer/LoadBalancers/NoLoadBalancerCreator.cs

@@ -0,0 +1,15 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{ 
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class NoLoadBalancerCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new NoLoadBalancer(async () => await serviceProvider.Get());
+        }
+
+        public string Type => nameof(NoLoadBalancer);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/RoundRobin.cs

@@ -1,12 +1,12 @@
-using Ocelot.Middleware;
-using Ocelot.Responses;
-using Ocelot.Values;
-using System;
-using System.Collections.Generic;
-using System.Threading.Tasks;
-
-namespace Ocelot.LoadBalancer.LoadBalancers
+namespace Ocelot.LoadBalancer.LoadBalancers
 {
+    using Ocelot.Middleware;
+    using Ocelot.Responses;
+    using Ocelot.Values;
+    using System;
+    using System.Collections.Generic;
+    using System.Threading.Tasks;
+
     public class RoundRobin : ILoadBalancer
     {
         private readonly Func<Task<List<Service>>> _services;

src/Ocelot/LoadBalancer/LoadBalancers/RoundRobinCreator.cs

@@ -0,0 +1,15 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+    
+    public class RoundRobinCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new RoundRobin(async () => await serviceProvider.Get());
+        }
+
+        public string Type => nameof(RoundRobin);
+    }
+}