liz/PluralKit
liz/PluralKit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add basic support for multi-node clustering

Browse Source
This commit is contained in:
Ske 2021-06-09 14:49:12 +02:00
parent b09a34d9cd
commit d61f61fead
3 changed files with 53 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
Myriad/Gateway/Cluster.cs
View File

@ -32,17 +32,19 @@ namespace Myriad.Gateway
public async Task Start(GatewayInfo.Bot info) public async Task Start(GatewayInfo.Bot info)
{ {
var concurrency = GetActualShardConcurrency(info.SessionStartLimit.MaxConcurrency); await Start(info.Url, 0, info.Shards - 1, info.Shards, info.SessionStartLimit.MaxConcurrency);
_ratelimiter = new(_logger, concurrency);
await Start(info.Url, info.Shards);
} }
public async Task Start(string url, int shardCount) public async Task Start(string url, int shardMin, int shardMax, int shardTotal, int concurrency)
{ {
_logger.Information("Starting {ShardCount} shards at {Url}", shardCount, url); concurrency = GetActualShardConcurrency(concurrency);
for (var i = 0; i < shardCount; i++) _ratelimiter = new(_logger, concurrency);
CreateAndAddShard(url, new ShardInfo(i, shardCount));
var shardCount = shardMax - shardMin + 1;
_logger.Information("Starting {ShardCount} of {ShardTotal} shards (#{ShardMin}-#{ShardMax}) at {Url}",
shardCount, shardTotal, shardMin, shardMax, url);
for (var i = shardMin; i <= shardMax; i++)
CreateAndAddShard(url, new ShardInfo(i, shardTotal));
await StartShards(); await StartShards();
} }

9
PluralKit.Bot/BotConfig.cs
View File

@ -15,5 +15,14 @@ namespace PluralKit.Bot
public int? MaxShardConcurrency { get; set; } public int? MaxShardConcurrency { get; set; }
public ulong? AdminRole { get; set; } public ulong? AdminRole { get; set; }
public ClusterSettings? Cluster { get; set; }
public record ClusterSettings
{
public string NodeName { get; set; }
public int TotalShards { get; set; }
public int TotalNodes { get; set; }
}
} }
} }

36
PluralKit.Bot/Init.cs
View File

@ -1,4 +1,5 @@
using System; using System;
using System.Linq;
using System.Threading; using System.Threading;
using System.Threading.Tasks; using System.Threading.Tasks;
@ -48,8 +49,7 @@ namespace PluralKit.Bot
// Start the Discord shards themselves (handlers already set up) // Start the Discord shards themselves (handlers already set up)
logger.Information("Connecting to Discord"); logger.Information("Connecting to Discord");
var info = await services.Resolve<DiscordApiClient>().GetGatewayBot(); await StartCluster(services);
await services.Resolve<Cluster>().Start(info);
logger.Information("Connected! All is good (probably)."); logger.Information("Connected! All is good (probably).");
// Lastly, we just... wait. Everything else is handled in the DiscordClient event loop // Lastly, we just... wait. Everything else is handled in the DiscordClient event loop
@ -128,5 +128,37 @@ namespace PluralKit.Bot
builder.RegisterModule<BotModule>(); builder.RegisterModule<BotModule>();
return builder.Build(); return builder.Build();
} }
private static async Task StartCluster(IComponentContext services)
{
var info = await services.Resolve<DiscordApiClient>().GetGatewayBot();
var cluster = services.Resolve<Cluster>();
var config = services.Resolve<BotConfig>();
if (config.Cluster != null)
{
// For multi-instance deployments, calculate the "span" of shards this node is responsible for
var totalNodes = config.Cluster.TotalNodes;
var totalShards = config.Cluster.TotalShards;
var nodeIndex = ExtractNodeIndex(config.Cluster.NodeName);
// Should evenly distribute shards even with an uneven amount of nodes
var shardMin = (int) Math.Round(totalShards * (float) nodeIndex / totalNodes);
var shardMax = (int) Math.Round(totalShards * (float) (nodeIndex + 1) / totalNodes) - 1;
await cluster.Start(info.Url, shardMin, shardMax, totalShards, info.SessionStartLimit.MaxConcurrency);
}
else
{
await cluster.Start(info);
}
}
private static int ExtractNodeIndex(string nodeName)
{
// Node name eg. "pluralkit-3", want to extract the 3. blame k8s :p
return int.Parse(nodeName.Split("-").Last());
}
} }
} }