From 80ab5f6a912f2225e2a3caf24d679560a52247ae Mon Sep 17 00:00:00 2001
From: David Lievrouw <david.lievrouw@recomatics.com>
Date: Wed, 10 Jul 2019 19:36:01 +0200
Subject: [PATCH] Extend OcelotBuilder with overloads of the
 AddCustomLoadBalancer registration helper method

---
 .../DependencyInjection/IOcelotBuilder.cs     | 20 +++++++++
 .../DependencyInjection/OcelotBuilder.cs      | 43 +++++++++++++++++++
 .../DelegateInvokingLoadBalancerCreator.cs    | 25 +++++++++++
 .../DependencyInjection/OcelotBuilderTests.cs | 25 ++++++-----
 4 files changed, 103 insertions(+), 10 deletions(-)
 create mode 100644 src/Ocelot/LoadBalancer/LoadBalancers/DelegateInvokingLoadBalancerCreator.cs

diff --git a/src/Ocelot/DependencyInjection/IOcelotBuilder.cs b/src/Ocelot/DependencyInjection/IOcelotBuilder.cs
index 13ec5a6a..33b3424e 100644
--- a/src/Ocelot/DependencyInjection/IOcelotBuilder.cs
+++ b/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();
     }
 }
diff --git a/src/Ocelot/DependencyInjection/OcelotBuilder.cs b/src/Ocelot/DependencyInjection/OcelotBuilder.cs
index a9290ad7..9e348183 100644
--- a/src/Ocelot/DependencyInjection/OcelotBuilder.cs
+++ b/src/Ocelot/DependencyInjection/OcelotBuilder.cs
@@ -1,3 +1,5 @@
+using Ocelot.ServiceDiscovery.Providers;
+
 using Ocelot.Configuration.ChangeTracking;
 
 namespace Ocelot.DependencyInjection
@@ -173,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>();
diff --git a/src/Ocelot/LoadBalancer/LoadBalancers/DelegateInvokingLoadBalancerCreator.cs b/src/Ocelot/LoadBalancer/LoadBalancers/DelegateInvokingLoadBalancerCreator.cs
new file mode 100644
index 00000000..535550dc
--- /dev/null
+++ b/src/Ocelot/LoadBalancer/LoadBalancers/DelegateInvokingLoadBalancerCreator.cs
@@ -0,0 +1,25 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System;
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    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 ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return _creatorFunc(reRoute, serviceProvider);
+        }
+
+        public string Type => typeof(T).Name;
+    }
+}
diff --git a/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs b/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
index 499c2a75..f85bbb11 100644
--- a/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
+++ b/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
@@ -1,6 +1,8 @@
-using Ocelot.Configuration;
+using System.Threading.Tasks;
 using Ocelot.LoadBalancer.LoadBalancers;
-using Ocelot.ServiceDiscovery.Providers;
+using Ocelot.Middleware;
+using Ocelot.Responses;
+using Ocelot.Values;
 
 namespace Ocelot.UnitTests.DependencyInjection
 {
@@ -166,7 +168,7 @@ namespace Ocelot.UnitTests.DependencyInjection
         public void should_add_custom_load_balancer_creators()
         {
             this.Given(x => WhenISetUpOcelotServices())
-                .When(x => AddSingletonLoadBalancerCreator<FakeCustomLoadBalancerCreator>())
+                .When(x => AddCustomLoadBalancer<FakeCustomLoadBalancer>())
                 .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
                 .BDDfy();
         }
@@ -177,10 +179,10 @@ namespace Ocelot.UnitTests.DependencyInjection
             _ocelotBuilder.AddSingletonDefinedAggregator<T>();
         }
         
-        private void AddSingletonLoadBalancerCreator<T>()
-            where T : class, ILoadBalancerCreator
+        private void AddCustomLoadBalancer<T>()
+            where T : ILoadBalancer, new()
         {
-            _ocelotBuilder.Services.AddSingleton<ILoadBalancerCreator, T>();
+            _ocelotBuilder.AddCustomLoadBalancer<T>();
         }
 
         private void AddTransientDefinedAggregator<T>()
@@ -266,7 +268,7 @@ namespace Ocelot.UnitTests.DependencyInjection
             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(FakeCustomLoadBalancerCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(DelegateInvokingLoadBalancerCreator<FakeCustomLoadBalancer>)).ShouldBe(1);
         }
 
         private void ThenTheAggregatorsAreTransient<TOne, TWo>()
@@ -354,14 +356,17 @@ namespace Ocelot.UnitTests.DependencyInjection
             _ex.ShouldBeNull();
         }
 
-        private class FakeCustomLoadBalancerCreator : ILoadBalancerCreator
+        private class FakeCustomLoadBalancer : ILoadBalancer
         {
-            public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
             {
                 throw new NotImplementedException();
             }
 
-            public string Type { get; }
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new NotImplementedException();
+            }
         }
     }
 }