Configure request limits via DI (#6730)

dotnet · May 28, 2024 · 12cc023 · 12cc023
1 parent 37959c3
commit 12cc023
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 13 deletions.
diff --git a/src/Microsoft.Diagnostics.Monitoring.WebApi/RequestThrottling/RequestLimit.cs b/src/Microsoft.Diagnostics.Monitoring.WebApi/RequestThrottling/RequestLimit.cs
@@ -0,0 +1,7 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.Diagnostics.Monitoring.WebApi
+{
+    public sealed record RequestLimit(string Key, int Limit);
+}
diff --git a/src/Microsoft.Diagnostics.Monitoring.WebApi/RequestThrottling/RequestLimitTracker.cs b/src/Microsoft.Diagnostics.Monitoring.WebApi/RequestThrottling/RequestLimitTracker.cs
@@ -27,19 +27,13 @@ private sealed class RequestCount : IDisposable
         private readonly ConcurrentDictionary<string, RequestCount> _requestCounts = new();
         private readonly ILogger<RequestLimitTracker> _logger;
 
-        public RequestLimitTracker(ILogger<RequestLimitTracker> logger)
+        public RequestLimitTracker(ILogger<RequestLimitTracker> logger, IEnumerable<RequestLimit> limits)
         {
-            //CONSIDER Should we have configuration for these?
+            foreach (RequestLimit requestLimit in limits)
+            {
+                _requestLimitTable.Add(requestLimit.Key, requestLimit.Limit);
+            }
 
-            _requestLimitTable.Add(Utilities.ArtifactType_Dump, 1);
-            _requestLimitTable.Add(Utilities.ArtifactType_GCDump, 1);
-            _requestLimitTable.Add(Utilities.ArtifactType_Logs, 3);
-            _requestLimitTable.Add(Utilities.ArtifactType_Trace, 3);
-            _requestLimitTable.Add(Utilities.ArtifactType_Metrics, 3);
-            _requestLimitTable.Add(Utilities.ArtifactType_Stacks, 1);
-            _requestLimitTable.Add(Utilities.ArtifactType_Exceptions, 1);
-            _requestLimitTable.Add(Utilities.ArtifactType_Parameters, 1);
-            _requestLimitTable.Add(Unlimited, int.MaxValue);
 
             _logger = logger;
         }

diff --git a/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTestCommon/TestHostHelper.cs b/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTestCommon/TestHostHelper.cs
@@ -110,7 +110,7 @@ public static IHost CreateHost(
                     }
 
                     services.ConfigureEgress();
-                    services.AddSingleton<IRequestLimitTracker, RequestLimitTracker>();
+                    services.ConfigureRequestLimits();
                     services.ConfigureOperationStore();
 
                     services.ConfigureDiagnosticPort(context.Configuration);

diff --git a/src/Tools/dotnet-monitor/Commands/CollectCommandHandler.cs b/src/Tools/dotnet-monitor/Commands/CollectCommandHandler.cs
@@ -119,7 +119,7 @@ private static IHostBuilder Configure(this IHostBuilder builder, StartupAuthenti
                 services.AddSingleton<IDiagnosticServices, DiagnosticServices>();
                 services.AddSingleton<IDumpService, DumpService>();
                 services.AddSingleton<IEndpointInfoSourceCallbacks, OperationTrackerServiceEndpointInfoSourceCallback>();
-                services.AddSingleton<IRequestLimitTracker, RequestLimitTracker>();
+                services.ConfigureRequestLimits();
                 services.ConfigureOperationStore();
                 services.ConfigureExtensions();
                 services.ConfigureExtensionLocations(settings);

diff --git a/src/Tools/dotnet-monitor/ServiceCollectionExtensions.cs b/src/Tools/dotnet-monitor/ServiceCollectionExtensions.cs
@@ -41,6 +41,7 @@
 using Microsoft.Extensions.Options;
 using System;
 using System.IO;
+using Utils = Microsoft.Diagnostics.Monitoring.WebApi.Utilities;
 
 namespace Microsoft.Diagnostics.Tools.Monitor
 {
@@ -365,6 +366,22 @@ public static IServiceCollection ConfigureHostingStartup(this IServiceCollection
             return services;
         }
 
+        public static IServiceCollection ConfigureRequestLimits(this IServiceCollection services)
+        {
+            services.AddSingleton<IRequestLimitTracker, RequestLimitTracker>();
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Dump, 1); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_GCDump, 1); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Logs, 3); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Trace, 3); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Metrics, 3); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Stacks, 1); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Exceptions, 1); });
+            services.AddSingleton((_) => { return new RequestLimit(Utils.ArtifactType_Parameters, 1); });
+            services.AddSingleton((_) => { return new RequestLimit(RequestLimitTracker.Unlimited, int.MaxValue); });
+
+            return services;
+        }
+
         public static IServiceCollection ConfigureStartupHook(this IServiceCollection services)
         {
             services.AddScoped<StartupHookApplicator>();