[RateLimiting] Add statistics API (#72306)

dotnet · Aug 12, 2022 · e055449 · e055449
1 parent 30bec96
commit e055449
Show file tree

Hide file tree

Showing 27 changed files with 1,074 additions and 141 deletions.
diff --git a/...ystem.Threading.RateLimiting/ref/System.Threading.RateLimiting.Typeforwards.netcoreapp.cs b/...ystem.Threading.RateLimiting/ref/System.Threading.RateLimiting.Typeforwards.netcoreapp.cs
@@ -0,0 +1,6 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+// The compiler emits a reference to the internal copy of this type in our non-NETCoreApp assembly
+// so we must include a forward to be compatible with libraries compiled against non-NETCoreApp System.Threading.RateLimiting
+[assembly: System.Runtime.CompilerServices.TypeForwardedTo(typeof(System.Runtime.CompilerServices.IsExternalInit))]
diff --git a/src/libraries/System.Threading.RateLimiting/ref/System.Threading.RateLimiting.cs b/src/libraries/System.Threading.RateLimiting/ref/System.Threading.RateLimiting.cs
@@ -14,7 +14,7 @@ public ConcurrencyLimiter(System.Threading.RateLimiting.ConcurrencyLimiterOption
         protected override System.Threading.RateLimiting.RateLimitLease AttemptAcquireCore(int permitCount) { throw null; }
         protected override void Dispose(bool disposing) { }
         protected override System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
-        public override int GetAvailablePermits() { throw null; }
+        public override System.Threading.RateLimiting.RateLimiterStatistics? GetStatistics() { throw null; }
     }
     public sealed partial class ConcurrencyLimiterOptions
     {
@@ -33,7 +33,7 @@ public FixedWindowRateLimiter(System.Threading.RateLimiting.FixedWindowRateLimit
         protected override System.Threading.RateLimiting.RateLimitLease AttemptAcquireCore(int requestCount) { throw null; }
         protected override void Dispose(bool disposing) { }
         protected override System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
-        public override int GetAvailablePermits() { throw null; }
+        public override System.Threading.RateLimiting.RateLimiterStatistics? GetStatistics() { throw null; }
         public override bool TryReplenish() { throw null; }
     }
     public sealed partial class FixedWindowRateLimiterOptions
@@ -78,7 +78,7 @@ public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
         public System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
         protected virtual System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
-        public abstract int GetAvailablePermits(TResource resource);
+        public abstract System.Threading.RateLimiting.RateLimiterStatistics? GetStatistics(TResource resource);
         public System.Threading.RateLimiting.PartitionedRateLimiter<TOuter> WithTranslatedKey<TOuter>(System.Func<TOuter, TResource> keyAdapter, bool leaveOpen) { throw null; }
     }
     public enum QueueProcessingOrder
@@ -98,7 +98,15 @@ public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
         public System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
         protected virtual System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
-        public abstract int GetAvailablePermits();
+        public abstract System.Threading.RateLimiting.RateLimiterStatistics? GetStatistics();
+    }
+    public partial class RateLimiterStatistics
+    {
+        public RateLimiterStatistics() { }
+        public long CurrentAvailablePermits { get { throw null; } set { } }
+        public long CurrentQueuedCount { get { throw null; } set { } }
+        public long TotalFailedLeases { get { throw null; } set { } }
+        public long TotalSuccessfulLeases { get { throw null; } set { } }
     }
     public abstract partial class RateLimitLease : System.IDisposable
     {
@@ -146,7 +154,7 @@ public SlidingWindowRateLimiter(System.Threading.RateLimiting.SlidingWindowRateL
         protected override System.Threading.RateLimiting.RateLimitLease AttemptAcquireCore(int requestCount) { throw null; }
         protected override void Dispose(bool disposing) { }
         protected override System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
-        public override int GetAvailablePermits() { throw null; }
+        public override System.Threading.RateLimiting.RateLimiterStatistics? GetStatistics() { throw null; }
         public override bool TryReplenish() { throw null; }
     }
     public sealed partial class SlidingWindowRateLimiterOptions
@@ -169,7 +177,7 @@ public TokenBucketRateLimiter(System.Threading.RateLimiting.TokenBucketRateLimit
         protected override System.Threading.RateLimiting.RateLimitLease AttemptAcquireCore(int tokenCount) { throw null; }
         protected override void Dispose(bool disposing) { }
         protected override System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
-        public override int GetAvailablePermits() { throw null; }
+        public override System.Threading.RateLimiting.RateLimiterStatistics? GetStatistics() { throw null; }
         public override bool TryReplenish() { throw null; }
     }
     public sealed partial class TokenBucketRateLimiterOptions

diff --git a/src/libraries/System.Threading.RateLimiting/ref/System.Threading.RateLimiting.csproj b/src/libraries/System.Threading.RateLimiting/ref/System.Threading.RateLimiting.csproj
@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <TargetFrameworks>$(NetCoreAppCurrent);$(NetCoreAppMinimum);netstandard2.0;$(NetFrameworkMinimum)</TargetFrameworks>
   </PropertyGroup>
@@ -7,7 +7,11 @@
     <PackageReference Include="System.Threading.Tasks.Extensions" Version="$(SystemThreadingTasksExtensionsVersion)" />
     <ProjectReference Include="$(LibrariesProjectRoot)Microsoft.Bcl.AsyncInterfaces\ref\Microsoft.Bcl.AsyncInterfaces.csproj" />
   </ItemGroup>
+
   <ItemGroup>
     <Compile Include="System.Threading.RateLimiting.cs" />
+    <Compile Include="System.Threading.RateLimiting.Typeforwards.netcoreapp.cs" Condition="'$(TargetFrameworkIdentifier)' == '.NETCoreApp'" />
+    <Compile Include="$(CoreLibSharedDir)System\Runtime\CompilerServices\IsExternalInit.cs" Condition="'$(TargetFrameworkIdentifier)' != '.NETCoreApp'"
+          Link="Common\System\Runtime\CompilerServices\IsExternalInit.cs" />
   </ItemGroup>
 </Project>
diff --git a/src/libraries/System.Threading.RateLimiting/src/System.Threading.RateLimiting.csproj b/src/libraries/System.Threading.RateLimiting/src/System.Threading.RateLimiting.csproj
@@ -27,12 +27,14 @@ System.Threading.RateLimiting.RateLimitLease</PackageDescription>
     <Compile Include="System\Threading\RateLimiting\PartitionedRateLimiter.T.cs" />
     <Compile Include="System\Threading\RateLimiting\QueueProcessingOrder.cs" />
     <Compile Include="System\Threading\RateLimiting\RateLimiter.cs" />
+    <Compile Include="System\Threading\RateLimiting\RateLimiterStatistics.cs" />
     <Compile Include="System\Threading\RateLimiting\RateLimitLease.cs" />
     <Compile Include="System\Threading\RateLimiting\RateLimitPartition.cs" />
     <Compile Include="System\Threading\RateLimiting\RateLimitPartition.T.cs" />
     <Compile Include="System\Threading\RateLimiting\ReplenishingRateLimiter.cs" />
     <Compile Include="System\Threading\RateLimiting\SlidingWindowRateLimiter.cs" />
     <Compile Include="System\Threading\RateLimiting\SlidingWindowRateLimiterOptions.cs" />
+    <Compile Include="System\Threading\RateLimiting\System.Threading.RateLimiting.Typeforwards.netcoreapp.cs" />
     <Compile Include="System\Threading\RateLimiting\TimerAwaitable.cs" />
     <Compile Include="System\Threading\RateLimiting\TokenBucketRateLimiter.cs" />
     <Compile Include="System\Threading\RateLimiting\TokenBucketRateLimiterOptions.cs" />
@@ -44,4 +46,9 @@ System.Threading.RateLimiting.RateLimitLease</PackageDescription>
     <PackageReference Include="System.Threading.Tasks.Extensions" Version="$(SystemThreadingTasksExtensionsVersion)" />
     <ProjectReference Include="$(LibrariesProjectRoot)Microsoft.Bcl.AsyncInterfaces\src\Microsoft.Bcl.AsyncInterfaces.csproj" />
   </ItemGroup>
+
+  <ItemGroup Condition="'$(TargetFrameworkIdentifier)' != '.NETCoreApp'">
+    <Compile Remove="System\Threading\RateLimiting\System.Threading.RateLimiting.Typeforwards.netcoreapp.cs" />
+    <Compile Include="$(CoreLibSharedDir)System\Runtime\CompilerServices\IsExternalInit.cs" Link="Common\System\Runtime\CompilerServices\IsExternalInit.cs" />
+  </ItemGroup>
 </Project>
diff --git a/...Threading.RateLimiting/src/System/Threading/RateLimiting/ChainedPartitionedRateLimiter.cs b/...Threading.RateLimiting/src/System/Threading/RateLimiting/ChainedPartitionedRateLimiter.cs
@@ -23,21 +23,33 @@ public ChainedPartitionedRateLimiter(PartitionedRateLimiter<TResource>[] limiter
             _limiters = limiters;
         }
 
-        public override int GetAvailablePermits(TResource resource)
+        public override RateLimiterStatistics? GetStatistics(TResource resource)
         {
             ThrowIfDisposed();
-            int lowestPermitCount = int.MaxValue;
+            long lowestAvailablePermits = long.MaxValue;
+            long currentQueuedCount = 0;
+            long totalFailedLeases = 0;
+            long innerMostSuccessfulLeases = 0;
             foreach (PartitionedRateLimiter<TResource> limiter in _limiters)
             {
-                int permitCount = limiter.GetAvailablePermits(resource);
-
-                if (permitCount < lowestPermitCount)
+                if (limiter.GetStatistics(resource) is { } statistics)
                 {
-                    lowestPermitCount = permitCount;
+                    if (statistics.CurrentAvailablePermits < lowestAvailablePermits)
+                    {
+                        lowestAvailablePermits = statistics.CurrentAvailablePermits;
+                    }
+                    currentQueuedCount += statistics.CurrentQueuedCount;
+                    totalFailedLeases += statistics.TotalFailedLeases;
+                    innerMostSuccessfulLeases = statistics.TotalSuccessfulLeases;
                 }
             }
-
-            return lowestPermitCount;
+            return new RateLimiterStatistics()
+            {
+                CurrentAvailablePermits = lowestAvailablePermits,
+                CurrentQueuedCount = currentQueuedCount,
+                TotalFailedLeases = totalFailedLeases,
+                TotalSuccessfulLeases = innerMostSuccessfulLeases,
+            };
         }
 
         protected override RateLimitLease AttemptAcquireCore(TResource resource, int permitCount)

diff --git a/...ies/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs b/...ies/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs
@@ -18,6 +18,9 @@ public sealed class ConcurrencyLimiter : RateLimiter
         private long? _idleSince = Stopwatch.GetTimestamp();
         private bool _disposed;
 
+        private long _failedLeasesCount;
+        private long _successfulLeasesCount;
+
         private readonly ConcurrencyLimiterOptions _options;
         private readonly Deque<RequestRegistration> _queue = new Deque<RequestRegistration>();
 
@@ -62,7 +65,17 @@ public ConcurrencyLimiter(ConcurrencyLimiterOptions options)
         }
 
         /// <inheritdoc/>
-        public override int GetAvailablePermits() => _permitCount;
+        public override RateLimiterStatistics? GetStatistics()
+        {
+            ThrowIfDisposed();
+            return new RateLimiterStatistics()
+            {
+                CurrentAvailablePermits = _permitCount,
+                CurrentQueuedCount = _queueCount,
+                TotalFailedLeases = Interlocked.Read(ref _failedLeasesCount),
+                TotalSuccessfulLeases = Interlocked.Read(ref _successfulLeasesCount),
+            };
+        }
 
         /// <inheritdoc/>
         protected override RateLimitLease AttemptAcquireCore(int permitCount)
@@ -78,7 +91,13 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
             // Return SuccessfulLease or FailedLease to indicate limiter state
             if (permitCount == 0)
             {
-                return _permitCount > 0 ? SuccessfulLease : FailedLease;
+                if (_permitCount > 0)
+                {
+                    Interlocked.Increment(ref _successfulLeasesCount);
+                    return SuccessfulLease;
+                }
+                Interlocked.Increment(ref _failedLeasesCount);
+                return FailedLease;
             }
 
             // Perf: Check SemaphoreSlim implementation instead of locking
@@ -93,6 +112,7 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
                 }
             }
 
+            Interlocked.Increment(ref _failedLeasesCount);
             return FailedLease;
         }
 
@@ -108,6 +128,7 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, C
             // Return SuccessfulLease if requestedCount is 0 and resources are available
             if (permitCount == 0 && _permitCount > 0 && !_disposed)
             {
+                Interlocked.Increment(ref _successfulLeasesCount);
                 return new ValueTask<RateLimitLease>(SuccessfulLease);
             }
 
@@ -136,11 +157,16 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, C
                                 // Updating queue count is handled by the cancellation code
                                 _queueCount += oldestRequest.Count;
                             }
+                            else
+                            {
+                                Interlocked.Increment(ref _failedLeasesCount);
+                            }
                         }
                         while (_options.QueueLimit - _queueCount < permitCount);
                     }
                     else
                     {
+                        Interlocked.Increment(ref _failedLeasesCount);
                         // Don't queue if queue limit reached and QueueProcessingOrder is OldestFirst
                         return new ValueTask<RateLimitLease>(QueueLimitLease);
                     }
@@ -174,6 +200,7 @@ private bool TryLeaseUnsynchronized(int permitCount, [NotNullWhen(true)] out Rat
             {
                 if (permitCount == 0)
                 {
+                    Interlocked.Increment(ref _successfulLeasesCount);
                     // Edge case where the check before the lock showed 0 available permits but when we got the lock some permits were now available
                     lease = SuccessfulLease;
                     return true;
@@ -186,6 +213,7 @@ private bool TryLeaseUnsynchronized(int permitCount, [NotNullWhen(true)] out Rat
                     _idleSince = null;
                     _permitCount -= permitCount;
                     Debug.Assert(_permitCount >= 0);
+                    Interlocked.Increment(ref _successfulLeasesCount);
                     lease = new ConcurrencyLease(true, this, permitCount);
                     return true;
                 }
@@ -234,6 +262,10 @@ private void Release(int releaseCount)
                             // Updating queue count is handled by the cancellation code
                             _queueCount += nextPendingRequest.Count;
                         }
+                        else
+                        {
+                            Interlocked.Increment(ref _successfulLeasesCount);
+                        }
                         nextPendingRequest.CancellationTokenRegistration.Dispose();
                         Debug.Assert(_queueCount >= 0);
                     }

diff --git a/...Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs b/...Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs
@@ -65,9 +65,9 @@ private async Task RunTimer()
             _timer.Dispose();
         }
 
-        public override int GetAvailablePermits(TResource resource)
+        public override RateLimiterStatistics? GetStatistics(TResource resource)
         {
-            return GetRateLimiter(resource).GetAvailablePermits();
+            return GetRateLimiter(resource).GetStatistics();
         }
 
         protected override RateLimitLease AttemptAcquireCore(TResource resource, int permitCount)

diff --git a/...System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs b/...System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs
@@ -19,6 +19,9 @@ public sealed class FixedWindowRateLimiter : ReplenishingRateLimiter
         private long? _idleSince;
         private bool _disposed;
 
+        private long _failedLeasesCount;
+        private long _successfulLeasesCount;
+
         private readonly Timer? _renewTimer;
         private readonly FixedWindowRateLimiterOptions _options;
         private readonly Deque<RequestRegistration> _queue = new Deque<RequestRegistration>();
@@ -81,7 +84,17 @@ public FixedWindowRateLimiter(FixedWindowRateLimiterOptions options)
         }
 
         /// <inheritdoc/>
-        public override int GetAvailablePermits() => _requestCount;
+        public override RateLimiterStatistics? GetStatistics()
+        {
+            ThrowIfDisposed();
+            return new RateLimiterStatistics()
+            {
+                CurrentAvailablePermits = _requestCount,
+                CurrentQueuedCount = _queueCount,
+                TotalFailedLeases = Interlocked.Read(ref _failedLeasesCount),
+                TotalSuccessfulLeases = Interlocked.Read(ref _successfulLeasesCount),
+            };
+        }
 
         /// <inheritdoc/>
         protected override RateLimitLease AttemptAcquireCore(int requestCount)
@@ -100,9 +113,11 @@ protected override RateLimitLease AttemptAcquireCore(int requestCount)
                 // Requests will be allowed if the total served request is less than the max allowed requests (permit limit).
                 if (_requestCount > 0)
                 {
+                    Interlocked.Increment(ref _successfulLeasesCount);
                     return SuccessfulLease;
                 }
 
+                Interlocked.Increment(ref _failedLeasesCount);
                 return CreateFailedWindowLease(requestCount);
             }
 
@@ -113,6 +128,7 @@ protected override RateLimitLease AttemptAcquireCore(int requestCount)
                     return lease;
                 }
 
+                Interlocked.Increment(ref _failedLeasesCount);
                 return CreateFailedWindowLease(requestCount);
             }
         }
@@ -131,6 +147,7 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int requestCount,
             // Return SuccessfulAcquisition if requestCount is 0 and resources are available
             if (requestCount == 0 && _requestCount > 0)
             {
+                Interlocked.Increment(ref _successfulLeasesCount);
                 return new ValueTask<RateLimitLease>(SuccessfulLease);
             }
 
@@ -157,11 +174,16 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int requestCount,
                             {
                                 _queueCount += oldestRequest.Count;
                             }
+                            else
+                            {
+                                Interlocked.Increment(ref _failedLeasesCount);
+                            }
                         }
                         while (_options.QueueLimit - _queueCount < requestCount);
                     }
                     else
                     {
+                        Interlocked.Increment(ref _failedLeasesCount);
                         // Don't queue if queue limit reached and QueueProcessingOrder is OldestFirst
                         return new ValueTask<RateLimitLease>(CreateFailedWindowLease(requestCount));
                     }
@@ -204,6 +226,7 @@ private bool TryLeaseUnsynchronized(int requestCount, [NotNullWhen(true)] out Ra
             {
                 if (requestCount == 0)
                 {
+                    Interlocked.Increment(ref _successfulLeasesCount);
                     // Edge case where the check before the lock showed 0 available permit counters but when we got the lock, some permits were now available
                     lease = SuccessfulLease;
                     return true;
@@ -216,6 +239,7 @@ private bool TryLeaseUnsynchronized(int requestCount, [NotNullWhen(true)] out Ra
                     _idleSince = null;
                     _requestCount -= requestCount;
                     Debug.Assert(_requestCount >= 0);
+                    Interlocked.Increment(ref _successfulLeasesCount);
                     lease = SuccessfulLease;
                     return true;
                 }
@@ -314,6 +338,10 @@ private void ReplenishInternal(long nowTicks)
                             // Updating queue count is handled by the cancellation code
                             _queueCount += nextPendingRequest.Count;
                         }
+                        else
+                        {
+                            Interlocked.Increment(ref _successfulLeasesCount);
+                        }
                         nextPendingRequest.CancellationTokenRegistration.Dispose();
                         Debug.Assert(_queueCount >= 0);
                     }