tests to handle some error cases and docs

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 load balance your request. 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.
+
+Remember if you specify no load balancer in your config Ocelot will not try and load balance.

src/Ocelot/Errors/OcelotErrorCode.cs

@@ -41,5 +41,6 @@ public enum OcelotErrorCode
         QuotaExceededError = 36,
         RequestCanceled = 37,
         ConnectionToDownstreamServiceError = 38,
+        CouldNotFindLoadBalancerCreator = 39,
     }
 }

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/LoadBalancerFactory.cs

@@ -28,7 +28,13 @@ public Response<ILoadBalancer> Get(DownstreamReRoute reRoute, ServiceProviderCon
 
             var serviceProvider = serviceProviderFactoryResponse.Data;
             var requestedType = reRoute.LoadBalancerOptions?.Type ?? nameof(NoLoadBalancer);
-            var applicableCreator = _loadBalancerCreators.Single(c => c.Type == requestedType);
+            var applicableCreator = _loadBalancerCreators.SingleOrDefault(c => c.Type == requestedType);
+
+            if (applicableCreator == null)
+            {
+                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 createdLoadBalancer = applicableCreator.Create(reRoute, serviceProvider);
             return new OkResponse<ILoadBalancer>(createdLoadBalancer);
         }

src/Ocelot/Responder/ErrorsToHttpStatusCodeMapper.cs

@@ -45,7 +45,8 @@ public int Map(List<Error> errors)
                 return 502;
             }
 
-            if (errors.Any(e => e.Code == OcelotErrorCode.UnableToCompleteRequestError))
+            if (errors.Any(e => e.Code == OcelotErrorCode.UnableToCompleteRequestError
+                || e.Code == OcelotErrorCode.CouldNotFindLoadBalancerCreator))
             {
                 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 @@ public void should_load_balance_request_with_round_robin()
                 .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 @@ public void GivenOcelotIsRunning()
             _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/LoadBalancer/LoadBalancerFactoryTests.cs

@@ -69,7 +69,24 @@ public void should_return_matching_load_balancer()
                 .Then(x => x.ThenTheLoadBalancerIsReturned<FakeLoadBalancerTwo>())
                 .BDDfy();
         }
-
+
+        [Fact]
+        public void should_return_error_response_if_cannot_find_load_balancer_creator()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .WithLoadBalancerOptions(new LoadBalancerOptions("DoesntExistLoadBalancer", "", 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())
+                .And(x => x.ThenTheErrorMessageIsCorrect())
+                .BDDfy();
+        }
+
         [Fact]
         public void should_call_service_provider()
         {
@@ -147,6 +164,11 @@ 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()
         {

test/Ocelot.UnitTests/Responder/ErrorsToHttpStatusCodeMapperTests.cs

@@ -47,6 +47,7 @@ public void should_return_service_unavailable(OcelotErrorCode errorCode)
 
         [Theory]
         [InlineData(OcelotErrorCode.UnableToCompleteRequestError)]
+        [InlineData(OcelotErrorCode.CouldNotFindLoadBalancerCreator)]
         public void should_return_internal_server_error(OcelotErrorCode errorCode)
         {
             ShouldMapErrorToStatusCode(errorCode, HttpStatusCode.InternalServerError);
@@ -120,7 +121,7 @@ public void ServiceUnavailableErrorsHaveThirdHighestPriority()
             var errors = new List<OcelotErrorCode>
             {
                 OcelotErrorCode.CannotAddDataError,
-                OcelotErrorCode.RequestTimedOutError
+                OcelotErrorCode.RequestTimedOutError,
             };
 
             ShouldMapErrorsToStatusCode(errors, HttpStatusCode.ServiceUnavailable);
@@ -132,7 +133,7 @@ public void check_we_have_considered_all_errors_in_these_tests()
             // 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(40, "Looks like the number of error codes has changed. Do you need to modify ErrorsToHttpStatusCodeMapper?");
         }
 
         private void ShouldMapErrorToStatusCode(OcelotErrorCode errorCode, HttpStatusCode expectedHttpStatusCode)