├── .gitignore ├── DistributedRateLimiting.Orleans.sln ├── DistributedRateLimiting.Orleans ├── DistributedRateLimiter.cs ├── DistributedRateLimiterCoordinator.cs ├── DistributedRateLimiterOptions.cs ├── DistributedRateLimiting.Orleans.csproj ├── IDistributedRateLimiterCoordinator.cs ├── IRateLimiterClient.cs └── ServiceCollectionExtensions.cs ├── LICENSE └── TestApp ├── Program.cs └── TestApp.csproj /.gitignore: -------------------------------------------------------------------------------- 1 | ## Ignore Visual Studio temporary files, build results, and 2 | ## files generated by popular Visual Studio add-ons. 3 | ## 4 | ## Get latest from https://github.com/github/gitignore/blob/master/VisualStudio.gitignore 5 | 6 | # User-specific files 7 | *.rsuser 8 | *.suo 9 | *.user 10 | *.userosscache 11 | *.sln.docstates 12 | 13 | # User-specific files (MonoDevelop/Xamarin Studio) 14 | *.userprefs 15 | 16 | # Mono auto generated files 17 | mono_crash.* 18 | 19 | # Build results 20 | [Dd]ebug/ 21 | [Dd]ebugPublic/ 22 | [Rr]elease/ 23 | [Rr]eleases/ 24 | x64/ 25 | x86/ 26 | [Aa][Rr][Mm]/ 27 | [Aa][Rr][Mm]64/ 28 | bld/ 29 | [Bb]in/ 30 | [Oo]bj/ 31 | [Ll]og/ 32 | [Ll]ogs/ 33 | 34 | # Visual Studio 2015/2017 cache/options directory 35 | .vs/ 36 | # Uncomment if you have tasks that create the project's static files in wwwroot 37 | #wwwroot/ 38 | 39 | # Visual Studio 2017 auto generated files 40 | Generated\ Files/ 41 | 42 | # MSTest test Results 43 | [Tt]est[Rr]esult*/ 44 | [Bb]uild[Ll]og.* 45 | 46 | # NUnit 47 | *.VisualState.xml 48 | TestResult.xml 49 | nunit-*.xml 50 | 51 | # Build Results of an ATL Project 52 | [Dd]ebugPS/ 53 | [Rr]eleasePS/ 54 | dlldata.c 55 | 56 | # Benchmark Results 57 | BenchmarkDotNet.Artifacts/ 58 | 59 | # .NET Core 60 | project.lock.json 61 | project.fragment.lock.json 62 | artifacts/ 63 | 64 | # StyleCop 65 | StyleCopReport.xml 66 | 67 | # Files built by Visual Studio 68 | *_i.c 69 | *_p.c 70 | *_h.h 71 | *.ilk 72 | *.meta 73 | *.obj 74 | *.iobj 75 | *.pch 76 | *.pdb 77 | *.ipdb 78 | *.pgc 79 | *.pgd 80 | *.rsp 81 | *.sbr 82 | *.tlb 83 | *.tli 84 | *.tlh 85 | *.tmp 86 | *.tmp_proj 87 | *_wpftmp.csproj 88 | *.log 89 | *.vspscc 90 | *.vssscc 91 | .builds 92 | *.pidb 93 | *.svclog 94 | *.scc 95 | 96 | # Chutzpah Test files 97 | _Chutzpah* 98 | 99 | # Visual C++ cache files 100 | ipch/ 101 | *.aps 102 | *.ncb 103 | *.opendb 104 | *.opensdf 105 | *.sdf 106 | *.cachefile 107 | *.VC.db 108 | *.VC.VC.opendb 109 | 110 | # Visual Studio profiler 111 | *.psess 112 | *.vsp 113 | *.vspx 114 | *.sap 115 | 116 | # Visual Studio Trace Files 117 | *.e2e 118 | 119 | # TFS 2012 Local Workspace 120 | $tf/ 121 | 122 | # Guidance Automation Toolkit 123 | *.gpState 124 | 125 | # ReSharper is a .NET coding add-in 126 | _ReSharper*/ 127 | *.[Rr]e[Ss]harper 128 | *.DotSettings.user 129 | 130 | # TeamCity is a build add-in 131 | _TeamCity* 132 | 133 | # DotCover is a Code Coverage Tool 134 | *.dotCover 135 | 136 | # AxoCover is a Code Coverage Tool 137 | .axoCover/* 138 | !.axoCover/settings.json 139 | 140 | # Visual Studio code coverage results 141 | *.coverage 142 | *.coveragexml 143 | 144 | # NCrunch 145 | _NCrunch_* 146 | .*crunch*.local.xml 147 | nCrunchTemp_* 148 | 149 | # MightyMoose 150 | *.mm.* 151 | AutoTest.Net/ 152 | 153 | # Web workbench (sass) 154 | .sass-cache/ 155 | 156 | # Installshield output folder 157 | [Ee]xpress/ 158 | 159 | # DocProject is a documentation generator add-in 160 | DocProject/buildhelp/ 161 | DocProject/Help/*.HxT 162 | DocProject/Help/*.HxC 163 | DocProject/Help/*.hhc 164 | DocProject/Help/*.hhk 165 | DocProject/Help/*.hhp 166 | DocProject/Help/Html2 167 | DocProject/Help/html 168 | 169 | # Click-Once directory 170 | publish/ 171 | 172 | # Publish Web Output 173 | *.[Pp]ublish.xml 174 | *.azurePubxml 175 | # Note: Comment the next line if you want to checkin your web deploy settings, 176 | # but database connection strings (with potential passwords) will be unencrypted 177 | *.pubxml 178 | *.publishproj 179 | 180 | # Microsoft Azure Web App publish settings. Comment the next line if you want to 181 | # checkin your Azure Web App publish settings, but sensitive information contained 182 | # in these scripts will be unencrypted 183 | PublishScripts/ 184 | 185 | # NuGet Packages 186 | *.nupkg 187 | # NuGet Symbol Packages 188 | *.snupkg 189 | # The packages folder can be ignored because of Package Restore 190 | **/[Pp]ackages/* 191 | # except build/, which is used as an MSBuild target. 192 | !**/[Pp]ackages/build/ 193 | # Uncomment if necessary however generally it will be regenerated when needed 194 | #!**/[Pp]ackages/repositories.config 195 | # NuGet v3's project.json files produces more ignorable files 196 | *.nuget.props 197 | *.nuget.targets 198 | 199 | # Microsoft Azure Build Output 200 | csx/ 201 | *.build.csdef 202 | 203 | # Microsoft Azure Emulator 204 | ecf/ 205 | rcf/ 206 | 207 | # Windows Store app package directories and files 208 | AppPackages/ 209 | BundleArtifacts/ 210 | Package.StoreAssociation.xml 211 | _pkginfo.txt 212 | *.appx 213 | *.appxbundle 214 | *.appxupload 215 | 216 | # Visual Studio cache files 217 | # files ending in .cache can be ignored 218 | *.[Cc]ache 219 | # but keep track of directories ending in .cache 220 | !?*.[Cc]ache/ 221 | 222 | # Others 223 | ClientBin/ 224 | ~$* 225 | *~ 226 | *.dbmdl 227 | *.dbproj.schemaview 228 | *.jfm 229 | *.pfx 230 | *.publishsettings 231 | orleans.codegen.cs 232 | 233 | # Including strong name files can present a security risk 234 | # (https://github.com/github/gitignore/pull/2483#issue-259490424) 235 | #*.snk 236 | 237 | # Since there are multiple workflows, uncomment next line to ignore bower_components 238 | # (https://github.com/github/gitignore/pull/1529#issuecomment-104372622) 239 | #bower_components/ 240 | 241 | # RIA/Silverlight projects 242 | Generated_Code/ 243 | 244 | # Backup & report files from converting an old project file 245 | # to a newer Visual Studio version. Backup files are not needed, 246 | # because we have git ;-) 247 | _UpgradeReport_Files/ 248 | Backup*/ 249 | UpgradeLog*.XML 250 | UpgradeLog*.htm 251 | ServiceFabricBackup/ 252 | *.rptproj.bak 253 | 254 | # SQL Server files 255 | *.mdf 256 | *.ldf 257 | *.ndf 258 | 259 | # Business Intelligence projects 260 | *.rdl.data 261 | *.bim.layout 262 | *.bim_*.settings 263 | *.rptproj.rsuser 264 | *- [Bb]ackup.rdl 265 | *- [Bb]ackup ([0-9]).rdl 266 | *- [Bb]ackup ([0-9][0-9]).rdl 267 | 268 | # Microsoft Fakes 269 | FakesAssemblies/ 270 | 271 | # GhostDoc plugin setting file 272 | *.GhostDoc.xml 273 | 274 | # Node.js Tools for Visual Studio 275 | .ntvs_analysis.dat 276 | node_modules/ 277 | 278 | # Visual Studio 6 build log 279 | *.plg 280 | 281 | # Visual Studio 6 workspace options file 282 | *.opt 283 | 284 | # Visual Studio 6 auto-generated workspace file (contains which files were open etc.) 285 | *.vbw 286 | 287 | # Visual Studio LightSwitch build output 288 | **/*.HTMLClient/GeneratedArtifacts 289 | **/*.DesktopClient/GeneratedArtifacts 290 | **/*.DesktopClient/ModelManifest.xml 291 | **/*.Server/GeneratedArtifacts 292 | **/*.Server/ModelManifest.xml 293 | _Pvt_Extensions 294 | 295 | # Paket dependency manager 296 | .paket/paket.exe 297 | paket-files/ 298 | 299 | # FAKE - F# Make 300 | .fake/ 301 | 302 | # CodeRush personal settings 303 | .cr/personal 304 | 305 | # Python Tools for Visual Studio (PTVS) 306 | __pycache__/ 307 | *.pyc 308 | 309 | # Cake - Uncomment if you are using it 310 | # tools/** 311 | # !tools/packages.config 312 | 313 | # Tabs Studio 314 | *.tss 315 | 316 | # Telerik's JustMock configuration file 317 | *.jmconfig 318 | 319 | # BizTalk build output 320 | *.btp.cs 321 | *.btm.cs 322 | *.odx.cs 323 | *.xsd.cs 324 | 325 | # OpenCover UI analysis results 326 | OpenCover/ 327 | 328 | # Azure Stream Analytics local run output 329 | ASALocalRun/ 330 | 331 | # MSBuild Binary and Structured Log 332 | *.binlog 333 | 334 | # NVidia Nsight GPU debugger configuration file 335 | *.nvuser 336 | 337 | # MFractors (Xamarin productivity tool) working folder 338 | .mfractor/ 339 | 340 | # Local History for Visual Studio 341 | .localhistory/ 342 | 343 | # BeatPulse healthcheck temp database 344 | healthchecksdb 345 | 346 | # Backup folder for Package Reference Convert tool in Visual Studio 2017 347 | MigrationBackup/ 348 | 349 | # Ionide (cross platform F# VS Code tools) working folder 350 | .ionide/ 351 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans.sln: -------------------------------------------------------------------------------- 1 |  2 | Microsoft Visual Studio Solution File, Format Version 12.00 3 | # Visual Studio Version 17 4 | VisualStudioVersion = 17.3.32519.111 5 | MinimumVisualStudioVersion = 10.0.40219.1 6 | Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "DistributedRateLimiting.Orleans", "DistributedRateLimiting.Orleans\DistributedRateLimiting.Orleans.csproj", "{0CD94AB2-353E-49EA-9504-91C65934904A}" 7 | EndProject 8 | Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "TestApp", "TestApp\TestApp.csproj", "{393583DF-42F3-4ECF-BD29-0B9EE6164ED0}" 9 | EndProject 10 | Global 11 | GlobalSection(SolutionConfigurationPlatforms) = preSolution 12 | Debug|Any CPU = Debug|Any CPU 13 | Release|Any CPU = Release|Any CPU 14 | EndGlobalSection 15 | GlobalSection(ProjectConfigurationPlatforms) = postSolution 16 | {0CD94AB2-353E-49EA-9504-91C65934904A}.Debug|Any CPU.ActiveCfg = Debug|Any CPU 17 | {0CD94AB2-353E-49EA-9504-91C65934904A}.Debug|Any CPU.Build.0 = Debug|Any CPU 18 | {0CD94AB2-353E-49EA-9504-91C65934904A}.Release|Any CPU.ActiveCfg = Release|Any CPU 19 | {0CD94AB2-353E-49EA-9504-91C65934904A}.Release|Any CPU.Build.0 = Release|Any CPU 20 | {393583DF-42F3-4ECF-BD29-0B9EE6164ED0}.Debug|Any CPU.ActiveCfg = Debug|Any CPU 21 | {393583DF-42F3-4ECF-BD29-0B9EE6164ED0}.Debug|Any CPU.Build.0 = Debug|Any CPU 22 | {393583DF-42F3-4ECF-BD29-0B9EE6164ED0}.Release|Any CPU.ActiveCfg = Release|Any CPU 23 | {393583DF-42F3-4ECF-BD29-0B9EE6164ED0}.Release|Any CPU.Build.0 = Release|Any CPU 24 | EndGlobalSection 25 | GlobalSection(SolutionProperties) = preSolution 26 | HideSolutionNode = FALSE 27 | EndGlobalSection 28 | GlobalSection(ExtensibilityGlobals) = postSolution 29 | SolutionGuid = {E59AFF41-EE72-4754-8FE7-44A70F1A0A9D} 30 | EndGlobalSection 31 | EndGlobal 32 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/DistributedRateLimiter.cs: -------------------------------------------------------------------------------- 1 | using Microsoft.Extensions.Logging; 2 | using Microsoft.Extensions.Options; 3 | using Orleans; 4 | using System.Diagnostics; 5 | using System.Diagnostics.CodeAnalysis; 6 | using System.Threading.RateLimiting; 7 | 8 | namespace DistributedRateLimiting.Orleans; 9 | 10 | /// 11 | /// implementation that helps manage concurrent access to a resource. 12 | /// 13 | public sealed class DistributedRateLimiter : RateLimiter, IRateLimiterClient 14 | { 15 | private static readonly Lease SuccessfulLease = new Lease(true, null, 0); 16 | private static readonly Lease FailedLease = new Lease(false, null, 0); 17 | private static readonly Lease QueueLimitLease = new Lease(false, null, 0, "Queue limit reached"); 18 | private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency; 19 | private readonly IGrainFactory _grainFactory; 20 | private readonly ILogger _logger; 21 | private readonly DistributedRateLimiterOptions _options; 22 | private readonly Queue _queue = new Queue(); 23 | private readonly IDistributedRateLimiterCoordinator _coordinator; 24 | private readonly Task _processTask; 25 | private readonly SemaphoreSlim _processSignal = new(1); 26 | 27 | private int _localAvailablePermitCount; 28 | private int _outstandingRequestedPermitCount; 29 | private long? _idleSince = Stopwatch.GetTimestamp(); 30 | private bool _disposed; 31 | 32 | // Use the queue as the lock field so we don't need to allocate another object for a lock and have another field in the object 33 | private object Lock => _queue; 34 | 35 | /// 36 | public override TimeSpan? IdleDuration => _idleSince is null ? null : new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency)); 37 | 38 | /// 39 | /// Initializes the . 40 | /// 41 | /// Options to specify the behavior of the . 42 | public DistributedRateLimiter( 43 | ILogger logger, 44 | IOptions options, 45 | IGrainFactory grainFactory) 46 | { 47 | _grainFactory = grainFactory; 48 | _logger = logger; 49 | _options = options?.Value ?? throw new ArgumentNullException(nameof(options)); 50 | _coordinator = grainFactory.GetGrain("default"); 51 | _processTask = Task.Run(ProcessRequests); 52 | } 53 | 54 | /// 55 | public override int GetAvailablePermits() => _localAvailablePermitCount; 56 | 57 | /// 58 | protected override RateLimitLease AcquireCore(int permitCount) 59 | { 60 | // These amounts of resources can never be acquired 61 | if (permitCount > _options.GlobalPermitCount) 62 | { 63 | throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, "Permit limit exceeded"); 64 | } 65 | 66 | ThrowIfDisposed(); 67 | 68 | // Return SuccessfulLease or FailedLease to indicate limiter state 69 | if (permitCount == 0) 70 | { 71 | return _localAvailablePermitCount > 0 ? SuccessfulLease : FailedLease; 72 | } 73 | 74 | // Perf: Check SemaphoreSlim implementation instead of locking 75 | if (_localAvailablePermitCount >= permitCount) 76 | { 77 | lock (Lock) 78 | { 79 | if (TryLeaseUnsynchronized(permitCount, out RateLimitLease? lease)) 80 | { 81 | return lease; 82 | } 83 | } 84 | } 85 | 86 | return FailedLease; 87 | } 88 | 89 | /// 90 | protected override ValueTask WaitAsyncCore(int permitCount, CancellationToken cancellationToken = default) 91 | { 92 | // These amounts of resources can never be acquired 93 | if (permitCount > _options.GlobalPermitCount) 94 | { 95 | throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, "Permit limit exceeded"); 96 | } 97 | 98 | // Return SuccessfulLease if requestedCount is 0 and resources are available 99 | if (permitCount == 0 && _localAvailablePermitCount > 0 && !_disposed) 100 | { 101 | return new ValueTask(SuccessfulLease); 102 | } 103 | 104 | // Perf: Check SemaphoreSlim implementation instead of locking 105 | lock (Lock) 106 | { 107 | if (TryLeaseUnsynchronized(permitCount, out RateLimitLease? lease)) 108 | { 109 | return new ValueTask(lease); 110 | } 111 | 112 | // Avoid integer overflow by using subtraction instead of addition 113 | Debug.Assert(_options.QueueLimit >= _outstandingRequestedPermitCount); 114 | if (_options.QueueLimit - _outstandingRequestedPermitCount < permitCount) 115 | { 116 | return new ValueTask(QueueLimitLease); 117 | } 118 | 119 | CancelQueueState tcs = new CancelQueueState(permitCount, this, cancellationToken); 120 | CancellationTokenRegistration ctr = default; 121 | if (cancellationToken.CanBeCanceled) 122 | { 123 | ctr = cancellationToken.Register(static obj => 124 | { 125 | ((CancelQueueState)obj!).TrySetCanceled(); 126 | }, tcs); 127 | } 128 | 129 | RequestRegistration request = new RequestRegistration(permitCount, tcs, ctr); 130 | _queue.Enqueue(request); 131 | _outstandingRequestedPermitCount += permitCount; 132 | Debug.Assert(_outstandingRequestedPermitCount <= _options.QueueLimit); 133 | 134 | _processSignal.Release(); 135 | return new ValueTask(request.Tcs.Task); 136 | } 137 | } 138 | 139 | private bool TryLeaseUnsynchronized(int permitCount, [NotNullWhen(true)] out RateLimitLease? lease) 140 | { 141 | ThrowIfDisposed(); 142 | 143 | // if permitCount is 0 we want to queue it if there are no available permits 144 | if (_localAvailablePermitCount >= permitCount && _localAvailablePermitCount != 0) 145 | { 146 | if (permitCount == 0) 147 | { 148 | // Edge case where the check before the lock showed 0 available permits but when we got the lock some permits were now available 149 | lease = SuccessfulLease; 150 | return true; 151 | } 152 | 153 | // If there are no items queued we can lease 154 | if (_outstandingRequestedPermitCount == 0) 155 | { 156 | _idleSince = null; 157 | _localAvailablePermitCount -= permitCount; 158 | Debug.Assert(_localAvailablePermitCount >= 0); 159 | lease = new Lease(true, this, permitCount); 160 | return true; 161 | } 162 | } 163 | 164 | lease = null; 165 | return false; 166 | } 167 | 168 | private void Release(int releaseCount) 169 | { 170 | lock (Lock) 171 | { 172 | if (_disposed) 173 | { 174 | return; 175 | } 176 | 177 | _localAvailablePermitCount += releaseCount; 178 | Debug.Assert(_localAvailablePermitCount <= _options.GlobalPermitCount); 179 | 180 | while (_queue.Count > 0) 181 | { 182 | RequestRegistration nextPendingRequest = _queue.Peek(); 183 | 184 | if (_localAvailablePermitCount >= nextPendingRequest.Count) 185 | { 186 | nextPendingRequest = _queue.Dequeue(); 187 | 188 | _localAvailablePermitCount -= nextPendingRequest.Count; 189 | _outstandingRequestedPermitCount -= nextPendingRequest.Count; 190 | Debug.Assert(_localAvailablePermitCount >= 0); 191 | 192 | Lease lease = nextPendingRequest.Count == 0 ? SuccessfulLease : new Lease(true, this, nextPendingRequest.Count); 193 | // Check if request was canceled 194 | if (!nextPendingRequest.Tcs.TrySetResult(lease)) 195 | { 196 | // Queued item was canceled so add count back 197 | _localAvailablePermitCount += nextPendingRequest.Count; 198 | // Updating queue count is handled by the cancellation code 199 | _outstandingRequestedPermitCount += nextPendingRequest.Count; 200 | } 201 | nextPendingRequest.CancellationTokenRegistration.Dispose(); 202 | Debug.Assert(_outstandingRequestedPermitCount >= 0); 203 | } 204 | else 205 | { 206 | break; 207 | } 208 | } 209 | 210 | if (_outstandingRequestedPermitCount == 0) 211 | { 212 | //Debug.Assert(_idleSince is null); 213 | //Debug.Assert(_outstandingRequestedPermitCount == 0); 214 | _idleSince = Stopwatch.GetTimestamp(); 215 | } 216 | 217 | _processSignal.Release(); 218 | } 219 | } 220 | 221 | private async Task ProcessRequests() 222 | { 223 | // RequestId allows for idempotency. 224 | // This is employed to prevent the remote service from process a given request twice. 225 | // For example, in the case of a retry after a communication failure. 226 | int requestId = 1; 227 | IRateLimiterClient? selfReference = null; 228 | var pendingToRelease = 0; 229 | var pendingToAcquire = 0; 230 | Stopwatch lastRefresh = Stopwatch.StartNew(); 231 | while (true) 232 | { 233 | try 234 | { 235 | selfReference ??= await _grainFactory.CreateObjectReference(this); 236 | 237 | // If there are no pending requests, wait until a signal arrives to indicate that there may be some acquire/release 238 | // work to process. 239 | if (pendingToRelease == 0 && pendingToAcquire == 0) 240 | { 241 | if (lastRefresh.Elapsed > _options.ClientLeaseRefreshInterval) 242 | { 243 | await _coordinator.RefreshLeases(selfReference); 244 | lastRefresh.Restart(); 245 | } 246 | else 247 | { 248 | var periodUntilNextRefresh = _options.ClientLeaseRefreshInterval.TotalMilliseconds - lastRefresh.Elapsed.TotalMilliseconds; 249 | await _processSignal.WaitAsync((int)periodUntilNextRefresh); 250 | } 251 | } 252 | 253 | lock (Lock) 254 | { 255 | // We first check that we don't have to retry a previous request. 256 | if (pendingToAcquire == 0) 257 | { 258 | pendingToAcquire = Math.Clamp(0, _options.TargetPermitsPerClient - _localAvailablePermitCount, _options.TargetPermitsPerClient); 259 | 260 | // If a large request comes in, make sure that we are servicing it. 261 | if (_queue.TryPeek(out var request)) 262 | { 263 | pendingToAcquire = Math.Max(pendingToAcquire, request.Count - _localAvailablePermitCount); 264 | } 265 | } 266 | } 267 | 268 | // If there is a pending request, try to fulfill it. 269 | if (pendingToAcquire > 0) 270 | { 271 | var acquired = await _coordinator.TryAcquire(selfReference, requestId, pendingToAcquire); 272 | ++requestId; 273 | pendingToAcquire = 0; 274 | Release(acquired); 275 | } 276 | 277 | lock (Lock) 278 | { 279 | // We first check that we don't have to retry a previous request. 280 | if (pendingToRelease == 0) 281 | { 282 | var surplus = _localAvailablePermitCount - _options.TargetPermitsPerClient; 283 | if (surplus > 0) 284 | { 285 | pendingToRelease = surplus; 286 | _localAvailablePermitCount -= pendingToRelease; 287 | } 288 | } 289 | } 290 | 291 | if (pendingToRelease > 0) 292 | { 293 | await _coordinator.Release(selfReference, requestId, pendingToRelease); 294 | ++requestId; 295 | pendingToRelease = 0; 296 | } 297 | } 298 | catch (Exception exception) 299 | { 300 | _logger.LogError(exception, "Exception satisfying RateLimiter request"); 301 | } 302 | } 303 | } 304 | 305 | protected override void Dispose(bool disposing) 306 | { 307 | if (!disposing) 308 | { 309 | return; 310 | } 311 | 312 | lock (Lock) 313 | { 314 | if (_disposed) 315 | { 316 | return; 317 | } 318 | _disposed = true; 319 | while (_queue.Count > 0) 320 | { 321 | RequestRegistration next = _queue.Dequeue(); 322 | next.CancellationTokenRegistration.Dispose(); 323 | next.Tcs.TrySetResult(FailedLease); 324 | } 325 | } 326 | } 327 | 328 | protected override ValueTask DisposeAsyncCore() 329 | { 330 | Dispose(true); 331 | 332 | return default; 333 | } 334 | 335 | private void ThrowIfDisposed() 336 | { 337 | if (_disposed) 338 | { 339 | throw new ObjectDisposedException(nameof(DistributedRateLimiter)); 340 | } 341 | } 342 | 343 | public void OnPermitsAvailable(int availablePermits) 344 | { 345 | _processSignal.Release(); 346 | } 347 | 348 | private sealed class Lease : RateLimitLease 349 | { 350 | private static readonly string[] s_allMetadataNames = new[] { MetadataName.ReasonPhrase.Name }; 351 | 352 | private bool _disposed; 353 | private readonly DistributedRateLimiter? _limiter; 354 | private readonly int _count; 355 | private readonly string? _reason; 356 | 357 | public Lease(bool isAcquired, DistributedRateLimiter? limiter, int count, string? reason = null) 358 | { 359 | IsAcquired = isAcquired; 360 | _limiter = limiter; 361 | _count = count; 362 | _reason = reason; 363 | 364 | // No need to set the limiter if count is 0, Dispose will noop 365 | Debug.Assert(count == 0 ? limiter is null : true); 366 | } 367 | 368 | public override bool IsAcquired { get; } 369 | 370 | public override IEnumerable MetadataNames => s_allMetadataNames; 371 | 372 | public override bool TryGetMetadata(string metadataName, out object? metadata) 373 | { 374 | if (_reason is not null && metadataName == MetadataName.ReasonPhrase.Name) 375 | { 376 | metadata = _reason; 377 | return true; 378 | } 379 | metadata = default; 380 | return false; 381 | } 382 | 383 | protected override void Dispose(bool disposing) 384 | { 385 | if (_disposed) 386 | { 387 | return; 388 | } 389 | 390 | _disposed = true; 391 | 392 | _limiter?.Release(_count); 393 | } 394 | } 395 | 396 | private readonly struct RequestRegistration 397 | { 398 | public RequestRegistration(int requestedCount, TaskCompletionSource tcs, 399 | CancellationTokenRegistration cancellationTokenRegistration) 400 | { 401 | Count = requestedCount; 402 | // Perf: Use AsyncOperation instead 403 | Tcs = tcs; 404 | CancellationTokenRegistration = cancellationTokenRegistration; 405 | } 406 | 407 | public int Count { get; } 408 | 409 | public TaskCompletionSource Tcs { get; } 410 | 411 | public CancellationTokenRegistration CancellationTokenRegistration { get; } 412 | } 413 | 414 | private sealed class CancelQueueState : TaskCompletionSource 415 | { 416 | private readonly int _permitCount; 417 | private readonly DistributedRateLimiter _limiter; 418 | private readonly CancellationToken _cancellationToken; 419 | 420 | public CancelQueueState(int permitCount, DistributedRateLimiter limiter, CancellationToken cancellationToken) 421 | : base(TaskCreationOptions.RunContinuationsAsynchronously) 422 | { 423 | _permitCount = permitCount; 424 | _limiter = limiter; 425 | _cancellationToken = cancellationToken; 426 | } 427 | 428 | public new bool TrySetCanceled() 429 | { 430 | if (TrySetCanceled(_cancellationToken)) 431 | { 432 | lock (_limiter.Lock) 433 | { 434 | _limiter._outstandingRequestedPermitCount -= _permitCount; 435 | } 436 | return true; 437 | } 438 | return false; 439 | } 440 | } 441 | } 442 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/DistributedRateLimiterCoordinator.cs: -------------------------------------------------------------------------------- 1 | #nullable enable 2 | 3 | using Microsoft.Extensions.Options; 4 | using Orleans; 5 | using Orleans.Concurrency; 6 | using System.Diagnostics; 7 | 8 | namespace DistributedRateLimiting.Orleans; 9 | 10 | [Reentrant] 11 | internal class DistributedRateLimiterCoordinator : Grain, IDistributedRateLimiterCoordinator, IDisposable 12 | { 13 | private readonly DistributedRateLimiterOptions _options; 14 | private readonly Dictionary _clients = new(); 15 | private readonly Queue _requests = new(); 16 | private IDisposable? _clientPurgeTimer; 17 | private int _availablePermits; 18 | 19 | public DistributedRateLimiterCoordinator(IOptions options) 20 | { 21 | _options = options.Value; 22 | _availablePermits = options.Value.GlobalPermitCount; 23 | } 24 | 25 | public override Task OnActivateAsync(CancellationToken cancellationToken) 26 | { 27 | _clientPurgeTimer = this.RegisterTimer(static state => ((DistributedRateLimiterCoordinator)state).OnPurgeClients(), this, TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(5)); 28 | return Task.CompletedTask; 29 | } 30 | 31 | public ValueTask TryAcquire(IRateLimiterClient client, int sequenceNumber, int permitCount) 32 | { 33 | if (!_clients.TryGetValue(client, out var state)) 34 | { 35 | state = _clients[client] = new ClientState(); 36 | } 37 | 38 | state.LastSeen.Restart(); 39 | var existingSequenceNumber = state.SequenceNumber; 40 | if (existingSequenceNumber >= sequenceNumber) 41 | { 42 | // Respond to duplicate requests with the result of the last request. 43 | // This allows for acquire requests to be safely retried. 44 | // The assumption/requirement here is that the client is awaiting those requests in sequence. 45 | return new(state.LastAcquiredPermitCount); 46 | } 47 | 48 | DropIdleClients(); 49 | 50 | // Acquire additional permits 51 | int acquiredPermits; 52 | if (_availablePermits >= permitCount) 53 | { 54 | state.InUsePermitCount += permitCount; 55 | _availablePermits -= permitCount; 56 | 57 | if (state.HasPendingRequest) 58 | { 59 | state.HasPendingRequest = false; 60 | } 61 | 62 | acquiredPermits = permitCount; 63 | } 64 | else 65 | { 66 | // Enqueue the client so that its request can be serviced later. 67 | if (!state.HasPendingRequest) 68 | { 69 | state.PermitsToAcquire = permitCount; 70 | state.HasPendingRequest = true; 71 | _requests.Enqueue(client); 72 | } 73 | 74 | acquiredPermits = 0; 75 | } 76 | 77 | ServicePendingRequests(); 78 | 79 | state.SequenceNumber = sequenceNumber; 80 | state.LastAcquiredPermitCount = acquiredPermits; 81 | return new(acquiredPermits); 82 | } 83 | 84 | public ValueTask Release(IRateLimiterClient client, int sequenceNumber, int permitCount) 85 | { 86 | if (!_clients.TryGetValue(client, out var state)) 87 | { 88 | state = _clients[client] = new ClientState(); 89 | } 90 | 91 | state.LastSeen.Restart(); 92 | var existingSequenceNumber = state.SequenceNumber; 93 | if (existingSequenceNumber >= sequenceNumber) 94 | { 95 | // Ignore duplicate requests: the result has already been incorporated into this instance's state. 96 | return default; 97 | } 98 | 99 | DropIdleClients(); 100 | 101 | state.InUsePermitCount -= permitCount; 102 | _availablePermits = Math.Min(_availablePermits + permitCount, _options.GlobalPermitCount); 103 | 104 | ServicePendingRequests(); 105 | 106 | state.SequenceNumber = sequenceNumber; 107 | state.LastAcquiredPermitCount = 0; 108 | return default; 109 | } 110 | 111 | public ValueTask RefreshLeases(IRateLimiterClient client) 112 | { 113 | if (_clients.TryGetValue(client, out var state)) 114 | { 115 | state.LastSeen.Restart(); 116 | } 117 | 118 | return default; 119 | } 120 | 121 | public ValueTask Unregister(IRateLimiterClient client) 122 | { 123 | if (_clients.TryGetValue(client, out var state)) 124 | { 125 | RemoveClient(client, state); 126 | } 127 | 128 | return default; 129 | } 130 | 131 | private void ServicePendingRequests() 132 | { 133 | while (_availablePermits > 0 && _requests.TryPeek(out var client)) 134 | { 135 | if (!_clients.TryGetValue(client, out var state)) 136 | { 137 | // The client must have been dropped 138 | _ = _requests.Dequeue(); 139 | continue; 140 | } 141 | 142 | if (_availablePermits > state.PermitsToAcquire) 143 | { 144 | // Asynchronously signal the client to tell it that there are some permits available. 145 | client.OnPermitsAvailable(_availablePermits); 146 | state.HasPendingRequest = false; 147 | _requests.Dequeue(); 148 | } 149 | else 150 | { 151 | break; 152 | } 153 | } 154 | } 155 | 156 | private Task OnPurgeClients() 157 | { 158 | DropIdleClients(); 159 | ServicePendingRequests(); 160 | return Task.CompletedTask; 161 | } 162 | 163 | private void DropIdleClients() 164 | { 165 | var dropped = default(List>); 166 | foreach (var (client, state) in _clients) 167 | { 168 | if (state.LastSeen.Elapsed > _options.IdleClientTimeout) 169 | { 170 | dropped ??= new(); 171 | dropped.Add(new KeyValuePair(client, state)); 172 | } 173 | } 174 | 175 | if (dropped is { Count: > 0 }) 176 | { 177 | foreach (var (client, state) in dropped) 178 | { 179 | var removed = RemoveClient(client, state); 180 | Debug.Assert(removed); 181 | } 182 | } 183 | } 184 | 185 | private bool RemoveClient(IRateLimiterClient client, ClientState? state) 186 | { 187 | if (state is not null) 188 | { 189 | // Return all in-use permits to the pool. 190 | _availablePermits += state.InUsePermitCount; 191 | } 192 | 193 | return _clients.Remove(client); 194 | } 195 | 196 | public void Dispose() 197 | { 198 | _clientPurgeTimer?.Dispose(); 199 | } 200 | 201 | private class ClientState 202 | { 203 | public ClientState() 204 | { 205 | LastSeen = Stopwatch.StartNew(); 206 | } 207 | 208 | // Monotinically increasing sequence number for this client 209 | public int SequenceNumber { get; set; } 210 | 211 | // The permits which are currently in-use by this client 212 | public int InUsePermitCount { get; set; } 213 | 214 | // The unfullfilled permit count of the next request for the client 215 | public int PermitsToAcquire { get; set; } 216 | 217 | public bool HasPendingRequest { get; set; } 218 | 219 | // The number of permits which were acquired in the previous request. 220 | public int LastAcquiredPermitCount { get; set; } 221 | 222 | public Stopwatch LastSeen { get; } 223 | } 224 | } 225 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/DistributedRateLimiterOptions.cs: -------------------------------------------------------------------------------- 1 | #nullable enable 2 | 3 | namespace DistributedRateLimiting.Orleans; 4 | 5 | /// 6 | /// Options to specify the behavior of a . 7 | /// 8 | public sealed class DistributedRateLimiterOptions 9 | { 10 | /// 11 | /// Maximum number of permits that can be leased by all rate limiters combined. 12 | /// 13 | public int GlobalPermitCount { get; set; } 14 | 15 | /// 16 | /// The number of permits to try to maintain leased on all clients. 17 | /// 18 | public int TargetPermitsPerClient { get; set; } 19 | 20 | /// 21 | /// Maximum number of permits that can be queued concurrently. 22 | /// 23 | public int QueueLimit { get; set; } 24 | 25 | /// 26 | /// The period of time after which to drop clients. 27 | /// 28 | public TimeSpan IdleClientTimeout { get; set; } = TimeSpan.FromMinutes(1); 29 | 30 | /// 31 | /// The period of time after which clients refresh their current leases. 32 | /// 33 | public TimeSpan ClientLeaseRefreshInterval { get; set; } = TimeSpan.FromSeconds(30); 34 | } 35 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/DistributedRateLimiting.Orleans.csproj: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | net7.0 5 | enable 6 | enable 7 | true 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | all 17 | runtime; build; native; contentfiles; analyzers; buildtransitive 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/IDistributedRateLimiterCoordinator.cs: -------------------------------------------------------------------------------- 1 | using Orleans; 2 | 3 | namespace DistributedRateLimiting.Orleans; 4 | 5 | internal interface IDistributedRateLimiterCoordinator : IGrainWithStringKey 6 | { 7 | ValueTask TryAcquire(IRateLimiterClient client, int sequenceNumber, int permitCount); 8 | ValueTask Release(IRateLimiterClient client, int sequenceNumber, int permitCount); 9 | ValueTask RefreshLeases(IRateLimiterClient client); 10 | ValueTask Unregister(IRateLimiterClient client); 11 | } 12 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/IRateLimiterClient.cs: -------------------------------------------------------------------------------- 1 | #nullable enable 2 | 3 | using Orleans; 4 | 5 | namespace DistributedRateLimiting.Orleans; 6 | 7 | internal interface IRateLimiterClient : IGrainObserver 8 | { 9 | void OnPermitsAvailable(int availablePermits); 10 | } 11 | -------------------------------------------------------------------------------- /DistributedRateLimiting.Orleans/ServiceCollectionExtensions.cs: -------------------------------------------------------------------------------- 1 | using Microsoft.Extensions.DependencyInjection; 2 | using System.Threading.RateLimiting; 3 | 4 | namespace DistributedRateLimiting.Orleans 5 | { 6 | public static class ServiceCollectionExtensions 7 | { 8 | public static IServiceCollection AddDistributedRateLimiter(this IServiceCollection services, Action configureOptions) 9 | { 10 | services.AddOptions().Configure(configureOptions); 11 | services.AddSingleton(); 12 | return services; 13 | } 14 | } 15 | } 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Reuben Bond 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /TestApp/Program.cs: -------------------------------------------------------------------------------- 1 | // See https://aka.ms/new-console-template for more information 2 | 3 | using DistributedRateLimiting.Orleans; 4 | using Microsoft.Extensions.DependencyInjection; 5 | using Microsoft.Extensions.Hosting; 6 | using Orleans.Hosting; 7 | using System.Net; 8 | using System.Threading.RateLimiting; 9 | 10 | var hostBuilder = Host.CreateDefaultBuilder(args); 11 | hostBuilder 12 | .UseOrleans(silo => 13 | { 14 | int instanceId = 0; 15 | if (args is { Length: > 0 }) 16 | { 17 | instanceId = int.Parse(args[0]); 18 | } 19 | 20 | silo.UseLocalhostClustering( 21 | siloPort: 11111 + instanceId, 22 | gatewayPort: 30000 + instanceId, 23 | primarySiloEndpoint: new IPEndPoint(IPAddress.Loopback, 11111)); 24 | }) 25 | .ConfigureServices(services => 26 | { 27 | services.AddDistributedRateLimiter(options => 28 | { 29 | options.QueueLimit = 200; 30 | options.GlobalPermitCount = 100; 31 | options.TargetPermitsPerClient = 20; 32 | }); 33 | }) 34 | .UseConsoleLifetime(); 35 | var host = await hostBuilder.StartAsync(); 36 | 37 | var rateLimiter = host.Services.GetRequiredService(); 38 | long numLeaseHolders = 0; 39 | var cancellationToken = new CancellationTokenSource(); 40 | var tasks = new List(); 41 | for (var i = 0; i < 5; i++) 42 | { 43 | var name = $"worker-{i}"; 44 | tasks.Add(Task.Run(() => RunWorker(rateLimiter, name, cancellationToken.Token))); 45 | } 46 | 47 | Console.CancelKeyPress += (_, __) => cancellationToken.Cancel(); 48 | await Task.WhenAll(tasks); 49 | await host.StopAsync(); 50 | 51 | host.Dispose(); 52 | 53 | async Task RunWorker(RateLimiter rateLimiter, string name, CancellationToken cancellationToken) 54 | { 55 | const int PermitCount = 25; 56 | while (!cancellationToken.IsCancellationRequested) 57 | { 58 | RateLimitLease lease; 59 | do 60 | { 61 | lease = rateLimiter.Acquire(PermitCount); 62 | if (lease.IsAcquired) break; 63 | Console.WriteLine($"{name}: Waiting for lease (holders: {numLeaseHolders})"); 64 | lease = await rateLimiter.WaitAsync(PermitCount); 65 | } while (!lease.IsAcquired); 66 | 67 | var holders = Interlocked.Increment(ref numLeaseHolders); 68 | Console.WriteLine($"{name}: Acquired lease (holders: {holders})"); 69 | await Task.Delay(500); 70 | lease.Dispose(); 71 | holders = Interlocked.Decrement(ref numLeaseHolders); 72 | Console.WriteLine($"{name}: Disposed lease (holders: {holders})"); 73 | await Task.Delay(500); 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /TestApp/TestApp.csproj: -------------------------------------------------------------------------------- 1 |  2 | 3 | 4 | Exe 5 | net7.0 6 | enable 7 | enable 8 | 9 | 10 | 11 | 12 | 13 | 14 | all 15 | runtime; build; native; contentfiles; analyzers; buildtransitive 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | --------------------------------------------------------------------------------