Merge branch 'DavidLievrouw-CustomLoadBalancers'

docs/features/loadbalancer.rst

@@ -108,3 +108,112 @@ subsequent requests. This means the sessions will be stuck across ReRoutes.
 Please note that if you give more than one DownstreamHostAndPort or you are using a Service Discovery provider such as Consul 
 and this returns more than one service then CookieStickySessions uses round robin to select the next server. This is hard coded at the 
 moment but could be changed.
+
+Custom Load Balancers
+^^^^^^^^^^^^^^^^^^^^
+
+`DavidLievrouw <https://github.com/DavidLievrouw`_ implemented a way to provide Ocelot with custom load balancer in `PR 1155 <https://github.com/ThreeMammals/Ocelot/pull/1155`_.
+
+In order to create and use a custom load balancer you can do the following. Below we setup a basic load balancing config and not the Type is CustomLoadBalancer this is the name of a class we will
+setup to do load balancing.
+
+.. code-block:: json
+
+    {
+        "DownstreamPathTemplate": "/api/posts/{postId}",
+        "DownstreamScheme": "https",
+        "DownstreamHostAndPorts": [
+                {
+                    "Host": "10.0.1.10",
+                    "Port": 5000,
+                },
+                {
+                    "Host": "10.0.1.11",
+                    "Port": 5000,
+                }
+            ],
+        "UpstreamPathTemplate": "/posts/{postId}",
+        "LoadBalancerOptions": {
+            "Type": "CustomLoadBalancer"
+        },
+        "UpstreamHttpMethod": [ "Put", "Delete" ]
+    }
+
+
+Then you need to create a class that implements the ILoadBalancer interface. Below is a simple round robin example.
+
+.. code-block:: csharp
+
+        private class CustomLoadBalancer : ILoadBalancer
+        {
+            private readonly Func<Task<List<Service>>> _services;
+            private readonly object _lock = new object();
+
+            private int _last;
+
+            public CustomLoadBalancer(Func<Task<List<Service>>> services)
+            {
+                _services = services;
+            }
+
+            public async Task<Response<ServiceHostAndPort>> Lease(DownstreamContext downstreamContext)
+            {
+                var services = await _services();
+                lock (_lock)
+                {
+                    if (_last >= services.Count)
+                    {
+                        _last = 0;
+                    }
+
+                    var next = services[_last];
+                    _last++;
+                    return new OkResponse<ServiceHostAndPort>(next.HostAndPort);
+                }
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+            }
+        }
+
+Finally you need to register this class with Ocelot. I have used the most complex example below to show all of the data / types that can be passed into the factory that creates load balancers.
+
+.. code-block:: csharp
+
+            Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, CustomLoadBalancer> loadBalancerFactoryFunc = (serviceProvider, reRoute, serviceDiscoveryProvider) => new CustomLoadBalancer(serviceDiscoveryProvider.Get);
+
+            s.AddOcelot()
+                .AddCustomLoadBalancer(loadBalancerFactoryFunc);
+
+However there is a much simpler example that will work the same.
+
+.. code-block:: csharp
+
+            s.AddOcelot()
+                .AddCustomLoadBalancer<CustomLoadBalancer>();
+
+There are numerous extension methods to add a custom load balancer and the interface is as follows.
+
+.. code-block:: csharp
+
+        IOcelotBuilder AddCustomLoadBalancer<T>()
+            where T : ILoadBalancer, new();
+
+         IOcelotBuilder AddCustomLoadBalancer<T>(Func<T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+         IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+         IOcelotBuilder AddCustomLoadBalancer<T>(
+            Func<DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+         IOcelotBuilder AddCustomLoadBalancer<T>(
+            Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+When you enable custom load balancers Ocelot looks up your load balancer by its class name when it decides if it should do load balancing. If it finds a match it will use your load balaner to load balance. If Ocelot cannot match the load balancer type in your configuration with the name of registered load balancer class then you will receive a HTTP 500 internal server error. If your load balancer factory throw an exception when Ocelot calls it you will receive a HTTP 500 internal server error.
+
+Remember if you specify no load balancer in your config Ocelot will not try and load balance.

src/Ocelot/DependencyInjection/IOcelotBuilder.cs

@@ -3,6 +3,9 @@ using Microsoft.Extensions.DependencyInjection;
 using Ocelot.Middleware.Multiplexer;
 using System;
 using System.Net.Http;
+using Ocelot.Configuration;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.ServiceDiscovery.Providers;
 
 namespace Ocelot.DependencyInjection
 {
@@ -25,6 +28,23 @@ namespace Ocelot.DependencyInjection
         IOcelotBuilder AddTransientDefinedAggregator<T>()
             where T : class, IDefinedAggregator;
 
+        IOcelotBuilder AddCustomLoadBalancer<T>()
+            where T : ILoadBalancer, new();
+        
+        IOcelotBuilder AddCustomLoadBalancer<T>(Func<T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+        IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+        IOcelotBuilder AddCustomLoadBalancer<T>(
+            Func<DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+        IOcelotBuilder AddCustomLoadBalancer<T>(
+            Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
         IOcelotBuilder AddConfigPlaceholders();
     }
 }

src/Ocelot/DependencyInjection/OcelotBuilder.cs

@@ -1,3 +1,5 @@
+using Ocelot.ServiceDiscovery.Providers;
+
 using Ocelot.Configuration.ChangeTracking;
 
 namespace Ocelot.DependencyInjection
@@ -87,6 +89,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>();
@@ -169,6 +175,47 @@ namespace Ocelot.DependencyInjection
             return this;
         }
 
+        public IOcelotBuilder AddCustomLoadBalancer<T>()
+            where T : ILoadBalancer, new()
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) => new T());
+            return this;
+        }
+        
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<T> loadBalancerFactoryFunc) 
+            where T : ILoadBalancer
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) =>
+                loadBalancerFactoryFunc());
+            return this;
+        }
+
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, T> loadBalancerFactoryFunc) 
+            where T : ILoadBalancer
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) =>
+                loadBalancerFactoryFunc(provider));
+            return this;
+        }
+
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) =>
+                loadBalancerFactoryFunc(reRoute, serviceDiscoveryProvider));
+            return this;
+        }
+
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer
+        {
+            Services.AddSingleton<ILoadBalancerCreator>(provider =>
+                new DelegateInvokingLoadBalancerCreator<T>(
+                    (reRoute, serviceDiscoveryProvider) => 
+                        loadBalancerFactoryFunc(provider, reRoute, serviceDiscoveryProvider)));
+            return this;
+        }
+
         private void AddSecurity()
         {
             Services.TryAddSingleton<ISecurityOptionsCreator, SecurityOptionsCreator>();

src/Ocelot/Errors/OcelotErrorCode.cs

@@ -41,5 +41,7 @@
         QuotaExceededError = 36,
         RequestCanceled = 37,
         ConnectionToDownstreamServiceError = 38,
+        CouldNotFindLoadBalancerCreator = 39,
+        ErrorInvokingLoadBalancerCreator = 40,
     }
 }

src/Ocelot/LoadBalancer/LoadBalancers/CookieStickySessionsCreator.cs

@@ -0,0 +1,21 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System.Threading.Tasks;
+    using Ocelot.Configuration;
+    using Ocelot.Infrastructure;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Ocelot.Responses;
+
+    public class CookieStickySessionsCreator : ILoadBalancerCreator
+    {
+        public Response<ILoadBalancer> Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            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));
+        }
+
+        public string Type => nameof(CookieStickySessions);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/CouldNotFindLoadBalancerCreator.cs

@@ -0,0 +1,12 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Errors;
+
+    public class CouldNotFindLoadBalancerCreator : Error
+    {
+        public CouldNotFindLoadBalancerCreator(string message) 
+            : base(message, OcelotErrorCode.CouldNotFindLoadBalancerCreator)
+        {
+        }
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/DelegateInvokingLoadBalancerCreator.cs

@@ -0,0 +1,34 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System;
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Ocelot.Responses;
+
+    public class DelegateInvokingLoadBalancerCreator<T> : ILoadBalancerCreator
+        where T : ILoadBalancer
+    {
+        private readonly Func<DownstreamReRoute, IServiceDiscoveryProvider, ILoadBalancer> _creatorFunc;
+
+        public DelegateInvokingLoadBalancerCreator(
+            Func<DownstreamReRoute, IServiceDiscoveryProvider, ILoadBalancer> creatorFunc)
+        {
+            _creatorFunc = creatorFunc;
+        }
+
+        public Response<ILoadBalancer> Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            try
+            {
+                return new OkResponse<ILoadBalancer>(_creatorFunc(reRoute, serviceProvider));
+
+            }
+            catch (Exception e)
+            {
+                return new ErrorResponse<ILoadBalancer>(new ErrorInvokingLoadBalancerCreator(e));
+            }
+        }
+
+        public string Type => typeof(T).Name;
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/ErrorInvokingLoadBalancerCreator.cs

@@ -0,0 +1,12 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System;
+    using Errors;
+
+    public class ErrorInvokingLoadBalancerCreator : Error
+    {
+        public ErrorInvokingLoadBalancerCreator(Exception e) : base($"Error when invoking user provided load balancer creator function, Message: {e.Message}, StackTrace: {e.StackTrace}", OcelotErrorCode.ErrorInvokingLoadBalancerCreator)
+        {
+        }
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/ILoadBalancerCreator.cs

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

src/Ocelot/LoadBalancer/LoadBalancers/ILoadBalancerFactory.cs

@@ -2,10 +2,9 @@
 {
     using Ocelot.Configuration;
     using Ocelot.Responses;
-    using System.Threading.Tasks;
 
     public interface ILoadBalancerFactory
     {
-        Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config);
+        Response<ILoadBalancer> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config);
     }
 }

src/Ocelot/LoadBalancer/LoadBalancers/ILoadBalancerHouse.cs

@@ -1,11 +1,10 @@
 using Ocelot.Configuration;
 using Ocelot.Responses;
-using System.Threading.Tasks;
 
 namespace Ocelot.LoadBalancer.LoadBalancers
 {
     public interface ILoadBalancerHouse
     {
-        Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config);
+        Response<ILoadBalancer> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config);
     }
 }

src/Ocelot/LoadBalancer/LoadBalancers/LeastConnectionCreator.cs

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

src/Ocelot/LoadBalancer/LoadBalancers/LoadBalancerFactory.cs

@@ -1,47 +1,48 @@
-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 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 Response<ILoadBalancer> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
         {
-            var response = _serviceProviderFactory.Get(config, reRoute);
+            var serviceProviderFactoryResponse = _serviceProviderFactory.Get(config, reRoute);
 
-            if (response.IsError)
+            if (serviceProviderFactoryResponse.IsError)
             {
-                return new ErrorResponse<ILoadBalancer>(response.Errors);
+                return new ErrorResponse<ILoadBalancer>(serviceProviderFactoryResponse.Errors);
             }
 
-            var serviceProvider = response.Data;
+            var serviceProvider = serviceProviderFactoryResponse.Data;
+            var requestedType = reRoute.LoadBalancerOptions?.Type ?? nameof(NoLoadBalancer);
+            var applicableCreator = _loadBalancerCreators.SingleOrDefault(c => c.Type == requestedType);
 
-            switch (reRoute.LoadBalancerOptions?.Type)
+            if (applicableCreator == null)
             {
-                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 new ErrorResponse<ILoadBalancer>(new CouldNotFindLoadBalancerCreator($"Could not find load balancer creator for Type: {requestedType}, please check your config specified the correct load balancer and that you have registered a class with the same name."));
             }
+
+            var createdLoadBalancerResponse = applicableCreator.Create(reRoute, serviceProvider);
+
+            if (createdLoadBalancerResponse.IsError)
+            {
+                return new ErrorResponse<ILoadBalancer>(createdLoadBalancerResponse.Errors);
+            }
+
+            return new OkResponse<ILoadBalancer>(createdLoadBalancerResponse.Data);
         }
     }
 }

src/Ocelot/LoadBalancer/LoadBalancers/LoadBalancerHouse.cs

@@ -3,7 +3,6 @@ using Ocelot.Responses;
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Threading.Tasks;
 
 namespace Ocelot.LoadBalancer.LoadBalancers
 {
@@ -18,7 +17,7 @@ namespace Ocelot.LoadBalancer.LoadBalancers
             _loadBalancers = new ConcurrentDictionary<string, ILoadBalancer>();
         }
 
-        public async Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
+        public Response<ILoadBalancer> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
         {
             try
             {
@@ -30,7 +29,7 @@ namespace Ocelot.LoadBalancer.LoadBalancers
 
                     if (reRoute.LoadBalancerOptions.Type != loadBalancer.GetType().Name)
                     {
-                        result = await _factory.Get(reRoute, config);
+                        result = _factory.Get(reRoute, config);
                         if (result.IsError)
                         {
                             return new ErrorResponse<ILoadBalancer>(result.Errors);
@@ -43,7 +42,7 @@ namespace Ocelot.LoadBalancer.LoadBalancers
                     return new OkResponse<ILoadBalancer>(loadBalancer);
                 }
 
-                result = await _factory.Get(reRoute, config);
+                result = _factory.Get(reRoute, config);
 
                 if (result.IsError)
                 {
@@ -58,7 +57,7 @@ namespace Ocelot.LoadBalancer.LoadBalancers
             {
                 return new ErrorResponse<ILoadBalancer>(new List<Ocelot.Errors.Error>()
                 {
-                    new UnableToFindLoadBalancerError($"unabe to find load balancer for {reRoute.LoadBalancerKey} exception is {ex}")
+                    new UnableToFindLoadBalancerError($"unabe to find load balancer for {reRoute.LoadBalancerKey} exception is {ex}"),
                 });
             }
         }

src/Ocelot/LoadBalancer/LoadBalancers/NoLoadBalancerCreator.cs

@@ -0,0 +1,16 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{ 
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Ocelot.Responses;
+
+    public class NoLoadBalancerCreator : ILoadBalancerCreator
+    {
+        public Response<ILoadBalancer> Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new OkResponse<ILoadBalancer>(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,16 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Ocelot.Responses;
+
+    public class RoundRobinCreator : ILoadBalancerCreator
+    {
+        public Response<ILoadBalancer> Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new OkResponse<ILoadBalancer>(new RoundRobin(async () => await serviceProvider.Get()));
+        }
+
+        public string Type => nameof(RoundRobin);
+    }
+}

src/Ocelot/LoadBalancer/Middleware/LoadBalancingMiddleware.cs

@@ -22,7 +22,7 @@ namespace Ocelot.LoadBalancer.Middleware
 
         public async Task Invoke(DownstreamContext context)
         {
-            var loadBalancer = await _loadBalancerHouse.Get(context.DownstreamReRoute, context.Configuration.ServiceProviderConfiguration);
+            var loadBalancer = _loadBalancerHouse.Get(context.DownstreamReRoute, context.Configuration.ServiceProviderConfiguration);
             if (loadBalancer.IsError)
             {
                 Logger.LogDebug("there was an error retriving the loadbalancer, setting pipeline error");

src/Ocelot/Responder/ErrorsToHttpStatusCodeMapper.cs

@@ -45,7 +45,9 @@ namespace Ocelot.Responder
                 return 502;
             }
 
-            if (errors.Any(e => e.Code == OcelotErrorCode.UnableToCompleteRequestError))
+            if (errors.Any(e => e.Code == OcelotErrorCode.UnableToCompleteRequestError
+                || e.Code == OcelotErrorCode.CouldNotFindLoadBalancerCreator
+                || e.Code == OcelotErrorCode.ErrorInvokingLoadBalancerCreator))
             {
                 return 500;
             }

test/Ocelot.AcceptanceTests/LoadBalancerTests.cs

@@ -6,7 +6,13 @@ namespace Ocelot.AcceptanceTests
     using Shouldly;
     using System;
     using System.Collections.Generic;
+    using System.Threading.Tasks;
+    using Configuration;
+    using Middleware;
+    using Responses;
+    using ServiceDiscovery.Providers;
     using TestStack.BDDfy;
+    using Values;
     using Xunit;
 
     public class LoadBalancerTests : IDisposable
@@ -122,6 +128,88 @@ namespace Ocelot.AcceptanceTests
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_load_balance_request_with_custom_load_balancer()
+        {
+            var downstreamPortOne = RandomPortFinder.GetRandomPort();
+            var downstreamPortTwo = RandomPortFinder.GetRandomPort();
+            var downstreamServiceOneUrl = $"http://localhost:{downstreamPortOne}";
+            var downstreamServiceTwoUrl = $"http://localhost:{downstreamPortTwo}";
+
+            var configuration = new FileConfiguration
+            {
+                ReRoutes = new List<FileReRoute>
+                    {
+                        new FileReRoute
+                        {
+                            DownstreamPathTemplate = "/",
+                            DownstreamScheme = "http",
+                            UpstreamPathTemplate = "/",
+                            UpstreamHttpMethod = new List<string> { "Get" },
+                            LoadBalancerOptions = new FileLoadBalancerOptions { Type = nameof(CustomLoadBalancer) },
+                            DownstreamHostAndPorts = new List<FileHostAndPort>
+                            {
+                                new FileHostAndPort
+                                {
+                                    Host = "localhost",
+                                    Port = downstreamPortOne,
+                                },
+                                new FileHostAndPort
+                                {
+                                    Host = "localhost",
+                                    Port = downstreamPortTwo,
+                                },
+                            },
+                        },
+                    },
+                GlobalConfiguration = new FileGlobalConfiguration(),
+            };
+
+            Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, CustomLoadBalancer> loadBalancerFactoryFunc = (serviceProvider, reRoute, serviceDiscoveryProvider) => new CustomLoadBalancer(serviceDiscoveryProvider.Get);
+
+            this.Given(x => x.GivenProductServiceOneIsRunning(downstreamServiceOneUrl, 200))
+                .And(x => x.GivenProductServiceTwoIsRunning(downstreamServiceTwoUrl, 200))
+                .And(x => _steps.GivenThereIsAConfiguration(configuration))
+                .And(x => _steps.GivenOcelotIsRunningWithCustomLoadBalancer(loadBalancerFactoryFunc))
+                .When(x => _steps.WhenIGetUrlOnTheApiGatewayMultipleTimes("/", 50))
+                .Then(x => x.ThenTheTwoServicesShouldHaveBeenCalledTimes(50))
+                .And(x => x.ThenBothServicesCalledRealisticAmountOfTimes(24, 26))
+                .BDDfy();
+        }
+
+        private class CustomLoadBalancer : ILoadBalancer
+        {
+            private readonly Func<Task<List<Service>>> _services;
+            private readonly object _lock = new object();
+
+            private int _last;
+
+            public CustomLoadBalancer(Func<Task<List<Service>>> services)
+            {
+                _services = services;
+            }
+
+            public async Task<Response<ServiceHostAndPort>> Lease(DownstreamContext downstreamContext)
+            {
+                var services = await _services();
+                lock (_lock)
+                {
+                    if (_last >= services.Count)
+                    {
+                        _last = 0;
+                    }
+
+                    var next = services[_last];
+                    _last++;
+                    return new OkResponse<ServiceHostAndPort>(next.HostAndPort);
+                }
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+            }
+        }
+
         private void ThenBothServicesCalledRealisticAmountOfTimes(int bottom, int top)
         {
             _counterOne.ShouldBeInRange(bottom, top);

test/Ocelot.AcceptanceTests/Steps.cs

@@ -39,6 +39,9 @@ namespace Ocelot.AcceptanceTests
     using System.Text;
     using System.Threading;
     using System.Threading.Tasks;
+    using Configuration;
+    using LoadBalancer.LoadBalancers;
+    using ServiceDiscovery.Providers;
     using static Ocelot.AcceptanceTests.HttpDelegatingHandlersTests;
     using ConfigurationBuilder = Microsoft.Extensions.Configuration.ConfigurationBuilder;
     using CookieHeaderValue = Microsoft.Net.Http.Headers.CookieHeaderValue;
@@ -255,6 +258,39 @@ namespace Ocelot.AcceptanceTests
             _ocelotClient = _ocelotServer.CreateClient();
         }
 
+        /// <summary>
+        /// This is annoying cos it should be in the constructor but we need to set up the file before calling startup so its a step.
+        /// </summary>
+        public void GivenOcelotIsRunningWithCustomLoadBalancer<T>(Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer
+        {
+            _webHostBuilder = new WebHostBuilder();
+
+            _webHostBuilder
+                .ConfigureAppConfiguration((hostingContext, config) =>
+                {
+                    config.SetBasePath(hostingContext.HostingEnvironment.ContentRootPath);
+                    var env = hostingContext.HostingEnvironment;
+                    config.AddJsonFile("appsettings.json", optional: true, reloadOnChange: false)
+                        .AddJsonFile($"appsettings.{env.EnvironmentName}.json", optional: true, reloadOnChange: false);
+                    config.AddJsonFile("ocelot.json", false, false);
+                    config.AddEnvironmentVariables();
+                })
+                .ConfigureServices(s =>
+                {
+                    s.AddOcelot()
+                        .AddCustomLoadBalancer(loadBalancerFactoryFunc);
+                })
+                .Configure(app =>
+                {
+                    app.UseOcelot().Wait();
+                });
+
+            _ocelotServer = new TestServer(_webHostBuilder);
+
+            _ocelotClient = _ocelotServer.CreateClient();
+        }
+
         public void GivenOcelotIsRunningWithConsul()
         {
             _webHostBuilder = new WebHostBuilder();

test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs

@@ -1,3 +1,9 @@
+using System.Threading.Tasks;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.Middleware;
+using Ocelot.Responses;
+using Ocelot.Values;
+
 namespace Ocelot.UnitTests.DependencyInjection
 {
     using Microsoft.AspNetCore.Hosting;
@@ -148,6 +154,42 @@ namespace Ocelot.UnitTests.DependencyInjection
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_default_ctor()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer<FakeCustomLoadBalancer>())
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_factory_method()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer(() => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_di_factory_method()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer(provider => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_factory_method_with_arguments()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer((reroute, discoveryProvider) => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+        
         [Fact]
         public void should_replace_iplaceholder()
         {
@@ -158,6 +200,15 @@ namespace Ocelot.UnitTests.DependencyInjection
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_add_custom_load_balancer_creators()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer((provider, reroute, discoveryProvider) => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
         private void AddSingletonDefinedAggregator<T>()
             where T : class, IDefinedAggregator
         {
@@ -239,6 +290,17 @@ namespace Ocelot.UnitTests.DependencyInjection
             handlers[1].ShouldBeOfType<TWo>();
         }
         
+        private void ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators()
+        {
+            _serviceProvider = _services.BuildServiceProvider();
+            var creators = _serviceProvider.GetServices<ILoadBalancerCreator>().ToList();
+            creators.Count(c => c.GetType() == typeof(NoLoadBalancerCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(RoundRobinCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(CookieStickySessionsCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(LeastConnectionCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(DelegateInvokingLoadBalancerCreator<FakeCustomLoadBalancer>)).ShouldBe(1);
+        }
+
         private void ThenTheAggregatorsAreTransient<TOne, TWo>()
         {
             var aggregators = _serviceProvider.GetServices<IDefinedAggregator>().ToList();
@@ -323,5 +385,20 @@ namespace Ocelot.UnitTests.DependencyInjection
         {
             _ex.ShouldBeNull();
         }
+
+        private class FakeCustomLoadBalancer : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                // Not relevant for these tests
+                throw new NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                // Not relevant for these tests
+                throw new NotImplementedException();
+            }
+        }
     }
 }

test/Ocelot.UnitTests/LoadBalancer/CookieStickySessionsCreatorTests.cs

@@ -0,0 +1,74 @@
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    using Moq;
+    using Ocelot.Configuration;
+    using Ocelot.Configuration.Builder;
+    using Ocelot.LoadBalancer.LoadBalancers;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Ocelot.Responses;
+    using Shouldly;
+    using TestStack.BDDfy;
+    using Xunit;
+
+    public class CookieStickySessionsCreatorTests
+    {
+        private readonly CookieStickySessionsCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private Response<ILoadBalancer> _loadBalancer;
+        private string _typeName;
+
+        public CookieStickySessionsCreatorTests()
+        {
+            _creator = new CookieStickySessionsCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .WithLoadBalancerOptions(new LoadBalancerOptions("myType", "myKey", 1000))
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<CookieStickySessions>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("CookieStickySessions"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.Data.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/DelegateInvokingLoadBalancerCreatorTests.cs

@@ -0,0 +1,127 @@
+using System;
+using System.Threading.Tasks;
+using Moq;
+using Ocelot.Configuration;
+using Ocelot.Configuration.Builder;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.Middleware;
+using Ocelot.Responses;
+using Ocelot.ServiceDiscovery.Providers;
+using Ocelot.Values;
+using Shouldly;
+using TestStack.BDDfy;
+using Xunit;
+
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    public class DelegateInvokingLoadBalancerCreatorTests
+    {
+        private DelegateInvokingLoadBalancerCreator<FakeLoadBalancer> _creator;
+        private Func<DownstreamReRoute, IServiceDiscoveryProvider, ILoadBalancer> _creatorFunc;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private Response<ILoadBalancer> _loadBalancer;
+        private string _typeName;
+
+        public DelegateInvokingLoadBalancerCreatorTests()
+        {
+            _creatorFunc = (reRoute, serviceDiscoveryProvider) =>
+                new FakeLoadBalancer(reRoute, serviceDiscoveryProvider);
+            _creator = new DelegateInvokingLoadBalancerCreator<FakeLoadBalancer>(_creatorFunc);
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("FakeLoadBalancer"))
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_return_result_of_specified_creator_func()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<FakeLoadBalancer>())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_return_error()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .And(x => x.GivenTheCreatorFuncThrows())
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenAnErrorIsReturned())
+                .BDDfy();
+        }
+
+        private void GivenTheCreatorFuncThrows()
+        {
+            _creatorFunc = (reRoute, serviceDiscoveryProvider) => throw new Exception();
+
+            _creator = new DelegateInvokingLoadBalancerCreator<FakeLoadBalancer>(_creatorFunc);
+        }
+
+        private void ThenAnErrorIsReturned()
+        {
+            _loadBalancer.IsError.ShouldBeTrue();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.Data.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+
+        private class FakeLoadBalancer : ILoadBalancer
+        {
+            public FakeLoadBalancer(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceDiscoveryProvider)
+            {
+                ReRoute = reRoute;
+                ServiceDiscoveryProvider = serviceDiscoveryProvider;
+            }
+
+            public DownstreamReRoute ReRoute { get; }
+            public IServiceDiscoveryProvider ServiceDiscoveryProvider { get; }
+
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new NotImplementedException();
+            }
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/LeastConnectionCreatorTests.cs

@@ -0,0 +1,74 @@
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    using Moq;
+    using Ocelot.Configuration;
+    using Ocelot.Configuration.Builder;
+    using Ocelot.LoadBalancer.LoadBalancers;
+    using Ocelot.Responses;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Shouldly;
+    using TestStack.BDDfy;
+    using Xunit;
+
+    public class LeastConnectionCreatorTests
+    {
+        private readonly LeastConnectionCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private Response<ILoadBalancer> _loadBalancer;
+        private string _typeName;
+
+        public LeastConnectionCreatorTests()
+        {
+            _creator = new LeastConnectionCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .WithServiceName("myService")
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<LeastConnection>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("LeastConnection"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.Data.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/LoadBalancerFactoryTests.cs

@@ -7,17 +7,24 @@ using Ocelot.ServiceDiscovery;
 using Ocelot.ServiceDiscovery.Providers;
 using Shouldly;
 using System.Collections.Generic;
+using System.Threading.Tasks;
+using Ocelot.Infrastructure.RequestData;
+using Ocelot.Middleware;
+using Ocelot.Values;
 using TestStack.BDDfy;
 using Xunit;
 
 namespace Ocelot.UnitTests.LoadBalancer
 {
+    using System;
+
     public class LoadBalancerFactoryTests
     {
         private DownstreamReRoute _reRoute;
         private readonly LoadBalancerFactory _factory;
         private Response<ILoadBalancer> _result;
         private readonly Mock<IServiceDiscoveryProviderFactory> _serviceProviderFactory;
+        private readonly IEnumerable<ILoadBalancerCreator> _loadBalancerCreators;
         private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
         private ServiceProviderConfiguration _serviceProviderConfig;
 
@@ -25,11 +32,18 @@ namespace Ocelot.UnitTests.LoadBalancer
         {
             _serviceProviderFactory = new Mock<IServiceDiscoveryProviderFactory>();
             _serviceProvider = new Mock<IServiceDiscoveryProvider>();
-            _factory = new LoadBalancerFactory(_serviceProviderFactory.Object);
+            _loadBalancerCreators = new ILoadBalancerCreator[]
+            {
+                new FakeLoadBalancerCreator<FakeLoadBalancerOne>(),
+                new FakeLoadBalancerCreator<FakeLoadBalancerTwo>(),
+                new FakeLoadBalancerCreator<FakeNoLoadBalancer>(nameof(NoLoadBalancer)),
+                new BrokenLoadBalancerCreator<BrokenLoadBalancer>(),
+            };
+            _factory = new LoadBalancerFactory(_serviceProviderFactory.Object, _loadBalancerCreators);
         }
 
         [Fact]
-        public void should_return_no_load_balancer()
+        public void should_return_no_load_balancer_by_default()
         {
             var reRoute = new DownstreamReRouteBuilder()
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
@@ -39,15 +53,15 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
                 .And(x => x.GivenTheServiceProviderFactoryReturns())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<NoLoadBalancer>())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<FakeNoLoadBalancer>())
                 .BDDfy();
         }
 
         [Fact]
-        public void should_return_round_robin_load_balancer()
+        public void should_return_matching_load_balancer()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("RoundRobin", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("FakeLoadBalancerTwo", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
@@ -55,15 +69,15 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
                 .And(x => x.GivenTheServiceProviderFactoryReturns())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<RoundRobin>())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<FakeLoadBalancerTwo>())
                 .BDDfy();
         }
 
         [Fact]
-        public void should_return_round_least_connection_balancer()
+        public void should_return_error_response_if_cannot_find_load_balancer_creator()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("LeastConnection", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("DoesntExistLoadBalancer", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
@@ -71,7 +85,24 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
                 .And(x => x.GivenTheServiceProviderFactoryReturns())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<LeastConnection>())
+                .Then(x => x.ThenAnErrorResponseIsReturned())
+                .And(x => x.ThenTheErrorMessageIsCorrect())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_return_error_response_if_creator_errors()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .WithLoadBalancerOptions(new LoadBalancerOptions("BrokenLoadBalancer", "", 0))
+                .WithUpstreamHttpMethod(new List<string> { "Get" })
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
+                .And(x => x.GivenTheServiceProviderFactoryReturns())
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenAnErrorResponseIsReturned())
                 .BDDfy();
         }
 
@@ -79,7 +110,7 @@ namespace Ocelot.UnitTests.LoadBalancer
         public void should_call_service_provider()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("RoundRobin", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("FakeLoadBalancerOne", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
@@ -92,18 +123,18 @@ namespace Ocelot.UnitTests.LoadBalancer
         }
 
         [Fact]
-        public void should_return_sticky_session()
+        public void should_return_error_response_when_call_to_service_provider_fails()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("CookieStickySessions", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("FakeLoadBalancerOne", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
             this.Given(x => x.GivenAReRoute(reRoute))
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
-                .And(x => x.GivenTheServiceProviderFactoryReturns())
+                .And(x => x.GivenTheServiceProviderFactoryFails())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<CookieStickySessions>())
+                .Then(x => x.ThenAnErrorResponseIsReturned())
                 .BDDfy();
         }
 
@@ -119,6 +150,13 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .Returns(new OkResponse<IServiceDiscoveryProvider>(_serviceProvider.Object));
         }
 
+        private void GivenTheServiceProviderFactoryFails()
+        {
+            _serviceProviderFactory
+                .Setup(x => x.Get(It.IsAny<ServiceProviderConfiguration>(), It.IsAny<DownstreamReRoute>()))
+                .Returns(new ErrorResponse<IServiceDiscoveryProvider>(new CannotFindDataError("For tests")));
+        }
+
         private void ThenTheServiceProviderIsCalledCorrectly()
         {
             _serviceProviderFactory
@@ -132,12 +170,113 @@ namespace Ocelot.UnitTests.LoadBalancer
 
         private void WhenIGetTheLoadBalancer()
         {
-            _result = _factory.Get(_reRoute, _serviceProviderConfig).Result;
+            _result = _factory.Get(_reRoute, _serviceProviderConfig);
         }
 
         private void ThenTheLoadBalancerIsReturned<T>()
         {
             _result.Data.ShouldBeOfType<T>();
         }
+
+        private void ThenAnErrorResponseIsReturned()
+        {
+            _result.IsError.ShouldBeTrue();
+        }
+
+        private void ThenTheErrorMessageIsCorrect()
+        {
+            _result.Errors[0].Message.ShouldBe("Could not find load balancer creator for Type: DoesntExistLoadBalancer, please check your config specified the correct load balancer and that you have registered a class with the same name.");
+        }
+
+        private class FakeLoadBalancerCreator<T> : ILoadBalancerCreator
+            where T : ILoadBalancer, new()
+        {
+
+            public FakeLoadBalancerCreator()
+            {
+                Type = typeof(T).Name;
+            }
+
+            public FakeLoadBalancerCreator(string type)
+            {
+                Type = type;
+            }
+
+            public Response<ILoadBalancer> Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+            {
+                return new OkResponse<ILoadBalancer>(new T());
+            }
+            
+            public string Type { get; }
+        }
+
+        private class BrokenLoadBalancerCreator<T> : ILoadBalancerCreator
+            where T : ILoadBalancer, new()
+        {
+            public BrokenLoadBalancerCreator()
+            {
+                Type = typeof(T).Name;
+            }
+
+            public Response<ILoadBalancer> Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+            {
+                return new ErrorResponse<ILoadBalancer>(new ErrorInvokingLoadBalancerCreator(new Exception()));
+            }
+
+            public string Type { get; }
+        }
+
+        private class FakeLoadBalancerOne : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
+        private class FakeLoadBalancerTwo : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
+        private class FakeNoLoadBalancer : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
+        private class BrokenLoadBalancer : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
     }
 }

test/Ocelot.UnitTests/LoadBalancer/LoadBalancerHouseTests.cs

@@ -111,8 +111,8 @@ namespace Ocelot.UnitTests.LoadBalancer
         private void WhenIGetTheReRouteWithTheSameKeyButDifferentLoadBalancer(DownstreamReRoute reRoute)
         {
             _reRoute = reRoute;
-            _factory.Setup(x => x.Get(_reRoute, _serviceProviderConfig)).ReturnsAsync(new OkResponse<ILoadBalancer>(new LeastConnection(null, null)));
-            _getResult = _loadBalancerHouse.Get(_reRoute, _serviceProviderConfig).Result;
+            _factory.Setup(x => x.Get(_reRoute, _serviceProviderConfig)).Returns(new OkResponse<ILoadBalancer>(new LeastConnection(null, null)));
+            _getResult = _loadBalancerHouse.Get(_reRoute, _serviceProviderConfig);
         }
 
         private void ThenAnErrorIsReturned()
@@ -138,13 +138,13 @@ namespace Ocelot.UnitTests.LoadBalancer
         {
             _reRoute = reRoute;
             _loadBalancer = loadBalancer;
-            _factory.Setup(x => x.Get(_reRoute, _serviceProviderConfig)).ReturnsAsync(new OkResponse<ILoadBalancer>(loadBalancer));
-            _getResult = _loadBalancerHouse.Get(reRoute, _serviceProviderConfig).Result;
+            _factory.Setup(x => x.Get(_reRoute, _serviceProviderConfig)).Returns(new OkResponse<ILoadBalancer>(loadBalancer));
+            _getResult = _loadBalancerHouse.Get(reRoute, _serviceProviderConfig);
         }
 
         private void WhenWeGetTheLoadBalancer(DownstreamReRoute reRoute)
         {
-            _getResult = _loadBalancerHouse.Get(reRoute, _serviceProviderConfig).Result;
+            _getResult = _loadBalancerHouse.Get(reRoute, _serviceProviderConfig);
         }
 
         private void ThenItIsReturned()

test/Ocelot.UnitTests/LoadBalancer/LoadBalancerMiddlewareTests.cs

@@ -182,7 +182,7 @@ namespace Ocelot.UnitTests.LoadBalancer
         {
             _loadBalancerHouse
                 .Setup(x => x.Get(It.IsAny<DownstreamReRoute>(), It.IsAny<ServiceProviderConfiguration>()))
-                .ReturnsAsync(new OkResponse<ILoadBalancer>(_loadBalancer.Object));
+                .Returns(new OkResponse<ILoadBalancer>(_loadBalancer.Object));
         }
 
         private void GivenTheLoadBalancerHouseReturnsAnError()
@@ -194,7 +194,7 @@ namespace Ocelot.UnitTests.LoadBalancer
 
             _loadBalancerHouse
                 .Setup(x => x.Get(It.IsAny<DownstreamReRoute>(), It.IsAny<ServiceProviderConfiguration>()))
-                .ReturnsAsync(_getLoadBalancerHouseError);
+                .Returns(_getLoadBalancerHouseError);
         }
 
         private void ThenAnErrorStatingLoadBalancerCouldNotBeFoundIsSetOnPipeline()

test/Ocelot.UnitTests/LoadBalancer/NoLoadBalancerCreatorTests.cs

@@ -0,0 +1,73 @@
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    using Moq;
+    using Ocelot.Configuration;
+    using Ocelot.Configuration.Builder;
+    using Ocelot.LoadBalancer.LoadBalancers;
+    using Ocelot.Responses;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Shouldly;
+    using TestStack.BDDfy;
+    using Xunit;
+
+    public class NoLoadBalancerCreatorTests
+    {
+        private readonly NoLoadBalancerCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private Response<ILoadBalancer> _loadBalancer;
+        private string _typeName;
+
+        public NoLoadBalancerCreatorTests()
+        {
+            _creator = new NoLoadBalancerCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<NoLoadBalancer>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("NoLoadBalancer"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.Data.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/RoundRobinCreatorTests.cs

@@ -0,0 +1,73 @@
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    using Moq;
+    using Ocelot.Configuration;
+    using Ocelot.Configuration.Builder;
+    using Ocelot.LoadBalancer.LoadBalancers;
+    using Ocelot.Responses;
+    using Ocelot.ServiceDiscovery.Providers;
+    using Shouldly;
+    using TestStack.BDDfy;
+    using Xunit;
+
+    public class RoundRobinCreatorTests
+    {
+        private readonly RoundRobinCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private Response<ILoadBalancer> _loadBalancer;
+        private string _typeName;
+
+        public RoundRobinCreatorTests()
+        {
+            _creator = new RoundRobinCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<RoundRobin>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("RoundRobin"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.Data.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/Responder/ErrorsToHttpStatusCodeMapperTests.cs

@@ -47,6 +47,8 @@ namespace Ocelot.UnitTests.Responder
 
         [Theory]
         [InlineData(OcelotErrorCode.UnableToCompleteRequestError)]
+        [InlineData(OcelotErrorCode.CouldNotFindLoadBalancerCreator)]
+        [InlineData(OcelotErrorCode.ErrorInvokingLoadBalancerCreator)]
         public void should_return_internal_server_error(OcelotErrorCode errorCode)
         {
             ShouldMapErrorToStatusCode(errorCode, HttpStatusCode.InternalServerError);
@@ -120,7 +122,7 @@ namespace Ocelot.UnitTests.Responder
             var errors = new List<OcelotErrorCode>
             {
                 OcelotErrorCode.CannotAddDataError,
-                OcelotErrorCode.RequestTimedOutError
+                OcelotErrorCode.RequestTimedOutError,
             };
 
             ShouldMapErrorsToStatusCode(errors, HttpStatusCode.ServiceUnavailable);
@@ -132,7 +134,7 @@ namespace Ocelot.UnitTests.Responder
             // If this test fails then it's because the number of error codes has changed.
             // You should make the appropriate changes to the test cases here to ensure
             // they cover all the error codes, and then modify this assertion.
-            Enum.GetNames(typeof(OcelotErrorCode)).Length.ShouldBe(39, "Looks like the number of error codes has changed. Do you need to modify ErrorsToHttpStatusCodeMapper?");
+            Enum.GetNames(typeof(OcelotErrorCode)).Length.ShouldBe(41, "Looks like the number of error codes has changed. Do you need to modify ErrorsToHttpStatusCodeMapper?");
         }
 
         private void ShouldMapErrorToStatusCode(OcelotErrorCode errorCode, HttpStatusCode expectedHttpStatusCode)