mirror of
https://github.com/PluralKit/PluralKit.git
synced 2026-02-13 17:20:14 +00:00
Add support for Twilight gateway queue
This commit is contained in:
parent
333530d24d
commit
26dc69e5a4
8 changed files with 70 additions and 15 deletions
|
|
@ -1,73 +0,0 @@
|
|||
using System;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Threading.Tasks;
|
||||
|
||||
using Serilog;
|
||||
|
||||
namespace Myriad.Gateway
|
||||
{
|
||||
public class ShardIdentifyRatelimiter
|
||||
{
|
||||
// docs specify 5 seconds, but we're actually throttling connections, not identify, so we need a bit of leeway
|
||||
private static readonly TimeSpan BucketLength = TimeSpan.FromSeconds(6);
|
||||
|
||||
private readonly ConcurrentDictionary<int, ConcurrentQueue<TaskCompletionSource>> _buckets = new();
|
||||
private readonly int _maxConcurrency;
|
||||
|
||||
private Task? _refillTask;
|
||||
private readonly ILogger _logger;
|
||||
|
||||
public ShardIdentifyRatelimiter(ILogger logger, int maxConcurrency)
|
||||
{
|
||||
_logger = logger.ForContext<ShardIdentifyRatelimiter>();
|
||||
_maxConcurrency = maxConcurrency;
|
||||
}
|
||||
|
||||
public Task Acquire(int shard)
|
||||
{
|
||||
var bucket = shard % _maxConcurrency;
|
||||
var queue = _buckets.GetOrAdd(bucket, _ => new ConcurrentQueue<TaskCompletionSource>());
|
||||
var tcs = new TaskCompletionSource();
|
||||
queue.Enqueue(tcs);
|
||||
|
||||
ScheduleRefill();
|
||||
|
||||
return tcs.Task;
|
||||
}
|
||||
|
||||
private void ScheduleRefill()
|
||||
{
|
||||
if (_refillTask != null && !_refillTask.IsCompleted)
|
||||
return;
|
||||
|
||||
_refillTask?.Dispose();
|
||||
_refillTask = RefillTask();
|
||||
}
|
||||
|
||||
private async Task RefillTask()
|
||||
{
|
||||
await Task.Delay(TimeSpan.FromMilliseconds(250));
|
||||
|
||||
while (true)
|
||||
{
|
||||
var isClear = true;
|
||||
foreach (var (bucket, queue) in _buckets)
|
||||
{
|
||||
if (!queue.TryDequeue(out var tcs))
|
||||
continue;
|
||||
|
||||
_logger.Information(
|
||||
"Allowing identify for bucket {BucketId} through ({QueueLength} left in bucket queue)",
|
||||
bucket, queue.Count);
|
||||
tcs.SetResult();
|
||||
isClear = false;
|
||||
}
|
||||
|
||||
if (isClear)
|
||||
return;
|
||||
|
||||
await Task.Delay(BucketLength);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue