-
Notifications
You must be signed in to change notification settings - Fork 4.6k
/
PartitionedRateLimiter.T.cs
140 lines (124 loc) · 6.91 KB
/
PartitionedRateLimiter.T.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using System.Threading.Tasks;
namespace System.Threading.RateLimiting
{
/// <summary>
/// Represents a limiter type that users interact with to determine if an operation can proceed given a specific <typeparamref name="TResource"/>.
/// </summary>
/// <typeparam name="TResource">The resource type that is being limited.</typeparam>
public abstract class PartitionedRateLimiter<TResource> : IAsyncDisposable, IDisposable
{
/// <summary>
/// Gets a snapshot of the statistics for the <paramref name="resource"/> if available.
/// </summary>
/// <returns>An instance of <see cref="RateLimiterStatistics"/> containing a snapshot of the statistics for a <paramref name="resource"/>.</returns>
public abstract RateLimiterStatistics? GetStatistics(TResource resource);
/// <summary>
/// Fast synchronous attempt to acquire permits.
/// </summary>
/// <remarks>
/// Set <paramref name="permitCount"/> to 0 to get whether permits are exhausted.
/// </remarks>
/// <param name="resource">The resource to limit.</param>
/// <param name="permitCount">Number of permits to try and acquire.</param>
/// <returns>A successful or failed lease.</returns>
/// <exception cref="ArgumentOutOfRangeException"></exception>
public RateLimitLease AttemptAcquire(TResource resource, int permitCount = 1)
{
if (permitCount < 0)
{
throw new ArgumentOutOfRangeException(nameof(permitCount));
}
return AttemptAcquireCore(resource, permitCount);
}
/// <summary>
/// Method that <see cref="PartitionedRateLimiter{TResource}"/> implementations implement for <see cref="AttemptAcquire"/>.
/// </summary>
/// <param name="resource">The resource to limit.</param>
/// <param name="permitCount">Number of permits to try and acquire.</param>
/// <returns></returns>
protected abstract RateLimitLease AttemptAcquireCore(TResource resource, int permitCount);
/// <summary>
/// Wait until the requested permits are available or permits can no longer be acquired.
/// </summary>
/// <remarks>
/// Set <paramref name="permitCount"/> to 0 to wait until permits are replenished.
/// </remarks>
/// <param name="resource">The resource to limit.</param>
/// <param name="permitCount">Number of permits to try and acquire.</param>
/// <param name="cancellationToken">Optional token to allow canceling a queued request for permits.</param>
/// <returns>A task that completes when the requested permits are acquired or when the requested permits are denied.</returns>
/// <exception cref="ArgumentOutOfRangeException"></exception>
public ValueTask<RateLimitLease> AcquireAsync(TResource resource, int permitCount = 1, CancellationToken cancellationToken = default)
{
if (permitCount < 0)
{
throw new ArgumentOutOfRangeException(nameof(permitCount));
}
if (cancellationToken.IsCancellationRequested)
{
return new ValueTask<RateLimitLease>(Task.FromCanceled<RateLimitLease>(cancellationToken));
}
return AcquireAsyncCore(resource, permitCount, cancellationToken);
}
/// <summary>
/// Method that <see cref="PartitionedRateLimiter{TResource}"/> implementations implement for <see cref="AcquireAsync"/>.
/// </summary>
/// <param name="resource">The resource to limit.</param>
/// <param name="permitCount">Number of permits to try and acquire.</param>
/// <param name="cancellationToken">Optional token to allow canceling a queued request for permits.</param>
/// <returns>A task that completes when the requested permits are acquired or when the requested permits are denied.</returns>
protected abstract ValueTask<RateLimitLease> AcquireAsyncCore(TResource resource, int permitCount, CancellationToken cancellationToken);
/// <summary>
/// Dispose method for implementations to write.
/// </summary>
/// <param name="disposing"></param>
protected virtual void Dispose(bool disposing) { }
/// <summary>
/// Disposes the RateLimiter. This completes any queued acquires with a failed lease.
/// </summary>
public void Dispose()
{
// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
Dispose(disposing: true);
GC.SuppressFinalize(this);
}
/// <summary>
/// DisposeAsync method for implementations to write.
/// </summary>
protected virtual ValueTask DisposeAsyncCore()
{
return default;
}
/// <summary>
/// Disposes the RateLimiter asynchronously.
/// </summary>
/// <returns>ValueTask representing the completion of the disposal.</returns>
public async ValueTask DisposeAsync()
{
// Perform async cleanup.
await DisposeAsyncCore().ConfigureAwait(false);
// Dispose of unmanaged resources.
Dispose(false);
// Suppress finalization.
GC.SuppressFinalize(this);
}
/// <summary>
/// Translates PartitionedRateLimiter<TOuter> into the current <see cref="PartitionedRateLimiter{TResource}"/>
/// using the <paramref name="keyAdapter"/> to translate <typeparamref name="TOuter"/> to <typeparamref name="TResource"/>.
/// </summary>
/// <typeparam name="TOuter">The type to translate into <typeparamref name="TResource"/>.</typeparam>
/// <param name="keyAdapter">The function to be called every time a <typeparamref name="TOuter"/> is passed to
/// PartitionedRateLimiter<TOuter>.Acquire(TOuter, int) or PartitionedRateLimiter<TOuter>.WaitAsync(TOuter, int, CancellationToken).
/// <para />
/// <remarks><paramref name="keyAdapter"/> should be implemented in a thread-safe way.</remarks></param>
/// <param name="leaveOpen">Specifies whether the returned <see cref="PartitionedRateLimiter{TOuter}"/> will dispose the wrapped <see cref="PartitionedRateLimiter{TResource}"/>.</param>
/// <returns>A new PartitionedRateLimiter<TOuter> that translates <typeparamref name="TOuter"/>
/// to <typeparamref name="TResource"/> and calls the inner <see cref="PartitionedRateLimiter{TResource}"/>.</returns>
public PartitionedRateLimiter<TOuter> WithTranslatedKey<TOuter>(Func<TOuter, TResource> keyAdapter, bool leaveOpen)
{
return new TranslatingLimiter<TResource, TOuter>(this, keyAdapter, leaveOpen);
}
}
}