Skip to content

Commit

Permalink
Fixup rate limiter
Browse files Browse the repository at this point in the history
  • Loading branch information
John Luo committed May 21, 2021
1 parent 8e48fbb commit 52e0b9c
Showing 1 changed file with 7 additions and 9 deletions.
16 changes: 7 additions & 9 deletions src/Middleware/RequestLimiter/src/RateLimiter.cs
Original file line number Diff line number Diff line change
Expand Up @@ -49,13 +49,10 @@ public override ValueTask<Resource> AcquireAsync(long requestedCount, Cancellati
Interlocked.Add(ref _resourceCount, -requestedCount);

var registration = new RateLimitRequest(requestedCount);
_queue.Enqueue(registration);

if (WaitHandle.WaitAny(new[] { registration.MRE.WaitHandle, cancellationToken.WaitHandle }) == 0)
{
return ValueTask.FromResult(Resource.NoopResource);
}

throw new InvalidOperationException("Limit exceeded");
// handle cancellation
return new ValueTask<Resource>(registration.TCS.Task);
}

private static void Replenish(object? state)
Expand Down Expand Up @@ -86,7 +83,7 @@ private static void Replenish(object? state)
if (requestToFulfill == request)
{
// If requestToFulfill == request, the fulfillment is successful.
requestToFulfill.MRE.Set();
requestToFulfill.TCS.SetResult(Resource.NoopResource);
}
else
{
Expand Down Expand Up @@ -114,12 +111,13 @@ private class RateLimitRequest
public RateLimitRequest(long count)
{
Count = count;
MRE = new ManualResetEventSlim();
// Use VoidAsyncOperationWithData<T> instead
TCS = new TaskCompletionSource<Resource>();
}

public long Count { get; }

public ManualResetEventSlim MRE { get; }
public TaskCompletionSource<Resource> TCS { get; }
}
}
}

0 comments on commit 52e0b9c

Please sign in to comment.