feat: Add AddLavalinkLoadBalancingStrategy helpers

angelobreuer · Sep 29, 2023 · 4bb7c2d · 4bb7c2d
1 parent f552b37
commit 4bb7c2d
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/src/Lavalink4NET.Cluster/Extensions/ServiceCollectionExtensions.cs b/src/Lavalink4NET.Cluster/Extensions/ServiceCollectionExtensions.cs
@@ -66,4 +66,34 @@ public static IServiceCollection ConfigureLavalinkCluster(this IServiceCollectio
     {
         return services.Configure(options);
     }
+
+    public static IServiceCollection AddLavalinkLoadBalancingStrategy<TLoadBalancingStrategy>(this IServiceCollection services)
+        where TLoadBalancingStrategy : class, INodeBalancingStrategy
+    {
+        ArgumentNullException.ThrowIfNull(services);
+
+        services.Replace(ServiceDescriptor.Singleton<INodeBalancingStrategy, TLoadBalancingStrategy>());
+
+        return services;
+    }
+
+    public static IServiceCollection AddLavalinkLoadBalancingStrategy<TLoadBalancingStrategy>(this IServiceCollection services, Func<IServiceProvider, TLoadBalancingStrategy> func)
+        where TLoadBalancingStrategy : class, INodeBalancingStrategy
+    {
+        ArgumentNullException.ThrowIfNull(services);
+
+        services.Replace(ServiceDescriptor.Singleton<INodeBalancingStrategy, TLoadBalancingStrategy>(func));
+
+        return services;
+    }
+
+    public static IServiceCollection AddLavalinkLoadBalancingStrategy<TLoadBalancingStrategy>(this IServiceCollection services, TLoadBalancingStrategy loadBalancingStrategy)
+        where TLoadBalancingStrategy : class, INodeBalancingStrategy
+    {
+        ArgumentNullException.ThrowIfNull(services);
+
+        services.Replace(ServiceDescriptor.Singleton<INodeBalancingStrategy>(loadBalancingStrategy));
+
+        return services;
+    }
 }