Support Rate Limiting Trace Sampling

src/libraries/System.Diagnostics.DiagnosticSource/src/System.Diagnostics.DiagnosticSource.csproj

@@ -55,6 +55,7 @@ System.Diagnostics.DiagnosticSource</PackageDescription>
     <Compile Include="System\Diagnostics\LegacyPropagator.cs" />
     <Compile Include="System\Diagnostics\NoOutputPropagator.cs" />
     <Compile Include="System\Diagnostics\PassThroughPropagator.cs" />
+    <Compile Include="System\Diagnostics\RateLimiter.cs" />
     <Compile Include="System\Diagnostics\RandomNumberGenerator.cs" />
     <Compile Include="System\Diagnostics\W3CPropagator.cs" />
     <Compile Include="System\Diagnostics\Metrics\AggregationManager.cs" />

src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/DsesFilterAndTransform.cs

@@ -27,6 +27,7 @@ internal sealed class DsesFilterAndTransform : IDisposable
 {
     private const string c_ActivitySourcePrefix = "[AS]";
     private const string c_ParentRatioSamplerPrefix = "ParentRatioSampler(";
+    private const string c_ParentRateLimitingSamplerPrefix = "ParentRateLimitingSampler(";
 
     /// <summary>
     /// Parses filterAndPayloadSpecs which is a list of lines each of which has the from
@@ -262,6 +263,29 @@ private static bool IsActivitySourceEntry(string filterAndPayloadSpec, int start
 
                         sampleFunc = DsesSamplerBuilder.CreateParentRatioSampler(ratio);
                     }
+                    else if (suffixPart.StartsWith(c_ParentRateLimitingSamplerPrefix.AsSpan(), StringComparison.OrdinalIgnoreCase))
+                    {
+                        int endingLocation = suffixPart.IndexOf(')');
+                        if (endingLocation < 0
+                            || !int.TryParse(
+#if NETFRAMEWORK || NETSTANDARD
+                                    suffixPart.Slice(c_ParentRateLimitingSamplerPrefix.Length, endingLocation - c_ParentRateLimitingSamplerPrefix.Length).ToString(),
+#else
+                                    suffixPart.Slice(c_ParentRateLimitingSamplerPrefix.Length, endingLocation - c_ParentRateLimitingSamplerPrefix.Length),
+#endif
+                                    NumberStyles.None, CultureInfo.InvariantCulture, out int maximumRatePerSecond)
+                            || maximumRatePerSecond <= 0)
+                        {
+                            if (eventSource.IsEnabled(EventLevel.Warning, DiagnosticSourceEventSource.Keywords.Messages))
+                            {
+                                eventSource.Message("DiagnosticSource: Ignoring filterAndPayloadSpec '[AS]" + entry.ToString() + "' because rate limiting sampling was invalid");
+                            }
+
+                            return next;
+                        }
+
+                        sampleFunc = DsesSamplerBuilder.CreateParentRateLimitingSampler(maximumRatePerSecond);
+                    }
                     else
                     {
                         if (eventSource.IsEnabled(EventLevel.Warning, DiagnosticSourceEventSource.Keywords.Messages))

src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/DsesSamplerBuilder.cs

@@ -66,4 +66,25 @@ static long GetLowerLong(ReadOnlySpan<byte> bytes)
             return result;
         }
     }
+
+    public static DsesSampleActivityFunc CreateParentRateLimitingSampler(int maximumRatePerSecond)
+    {
+        Debug.Assert(maximumRatePerSecond > 0, "maximumRatePerSecond must be greater than 0");
+
+        RateLimiter rateLimiter = new RateLimiter(maximumRatePerSecond);
+
+        return (bool hasActivityContext, ref ActivityCreationOptions<ActivityContext> options) =>
+        {
+            if (hasActivityContext && options.TraceId != default)
+            {
+                ActivityContext parentContext = options.Parent;
+
+                return parentContext == default || parentContext.IsRemote ? // root or remote activity
+                        (rateLimiter.TryAcquire() ? ActivitySamplingResult.AllDataAndRecorded : ActivitySamplingResult.None) :
+                        parentContext.TraceFlags.HasFlag(ActivityTraceFlags.Recorded) ? ActivitySamplingResult.AllDataAndRecorded : ActivitySamplingResult.None;
+            }
+
+            return ActivitySamplingResult.None;
+        };
+    }
 }

src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/RateLimiter.cs

@@ -0,0 +1,68 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Threading;
+
+namespace System.Diagnostics
+{
+    /// <summary>
+    /// A class that limits operations to a maximum count per second.
+    /// </summary>
+    internal sealed class RateLimiter
+    {
+        private readonly int _maxOperationsPerSecond;
+        private readonly Stopwatch _stopwatch;
+        private readonly long _ticksPerSecond;
+        private int _currentOperationCount;
+        private long _intervalStartTicks;
+
+        /// <summary>
+        /// Initializes a new instance of the RateLimiter class.
+        /// </summary>
+        /// <param name="maxOperationsPerSecond">Maximum number of operations allowed per second.</param>
+        internal RateLimiter(int maxOperationsPerSecond)
+        {
+            Debug.Assert(maxOperationsPerSecond > 0, "maxOperationsPerSecond must be greater than 0");
+
+            _maxOperationsPerSecond = maxOperationsPerSecond;
+            _stopwatch = new Stopwatch();
+            _stopwatch.Start();
+            _intervalStartTicks = _stopwatch.ElapsedTicks;
+            _currentOperationCount = 0;
+            _ticksPerSecond = Stopwatch.Frequency;
+        }
+
+        /// <summary>
+        /// Tries to perform an operation if the rate limit allows it.
+        /// </summary>
+        /// <returns>True if the operation is allowed within the rate limit, otherwise false.</returns>
+        internal bool TryAcquire()
+        {
+            do
+            {
+                long currentTicks = _stopwatch.ElapsedTicks;
+                long intervalStartTicks = Interlocked.Read(ref _intervalStartTicks);
+                int currentOperationCount = Volatile.Read(ref _currentOperationCount);
+                long elapsedTicks = currentTicks - intervalStartTicks;
+
+                // If a second has elapsed, reset the counter
+                if (elapsedTicks >= _ticksPerSecond)
+                {
+                    if (Interlocked.CompareExchange(ref _currentOperationCount, 1, currentOperationCount) != currentOperationCount)
+                    {
+                        // Another thread has already reset the counter, so we need to check again
+                        continue;
+                    }
+
+                    // Update the _intervalStartTicks if no-one else updated it in the meantime.
+                    Interlocked.CompareExchange(ref _intervalStartTicks, currentTicks, intervalStartTicks);
+
+                    return true; // Allow the operation
+                }
+
+                return Interlocked.Increment(ref _currentOperationCount) <= _maxOperationsPerSecond;
+            } while (true);
+        }
+    }
+}

src/libraries/System.Diagnostics.DiagnosticSource/tests/DiagnosticSourceEventSourceBridgeTests.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Generic;
 using System.Diagnostics.Tracing;
+using System.Globalization;
 using System.Reflection;
 using System.Text;
 using System.Threading;
@@ -1441,6 +1442,55 @@ public void TestMultipleRulesOnlyFirstTaken(string spec, string errorMessage)
             }, spec, errorMessage).Dispose();
         }
 
+        [ConditionalTheory(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))]
+        [InlineData(1)]
+        [InlineData(5)]
+        [InlineData(50)]
+        [InlineData(100)]
+        public void TestRateLimitingSampler(int maxOperationPerSecond)
+        {
+            RemoteExecutor.Invoke((maxOpPerSecond) =>
+            {
+                // Test RateLimitingSampler behavior
+                using TestDiagnosticSourceEventListener eventSourceListener = new TestDiagnosticSourceEventListener();
+                Activity a = new Activity(""); // we need this to ensure DiagnosticSourceEventSource.Logger creation.
+                Assert.Equal("", a.Source.Name);
+
+                Assert.Equal(0, eventSourceListener.EventCount);
+
+                Stopwatch sw = new Stopwatch();
+                sw.Start();
+
+                eventSourceListener.Enable($"[AS]*/-ParentRateLimitingSampler({maxOpPerSecond})");
+
+                while (sw.ElapsedMilliseconds < 3000)  // run for 3 seconds
+                {
+                    // ensure we are creating a root activity
+                    Activity.Current = null;
+                    using var nextRoot = a.Source.StartActivity("NextRoot");
+                }
+
+                int maxOps = int.Parse(maxOpPerSecond, CultureInfo.InvariantCulture);
+                // maxOperationPerSecond sampling allowed per second
+                // 2 events for every sampling (activity start and stop)
+                // 3 seconds of sampling
+                // tolerance of extra sample can be done if the second turn after the loop check sw.ElapsedMilliseconds. 2 extra events (start and stop).
+                Assert.True(maxOps * 2 * 3 + 2 >= eventSourceListener.EventCount, $"{eventSourceListener.EventCount} events were recorded, while maxOpPerSecond is {maxOpPerSecond}");
+
+                // Ensure the minimum number of events is recorded in 3 seconds. We are choosing reasonable small number as it depends on the machine speed.
+                Assert.True(eventSourceListener.EventCount >= 6, $"{eventSourceListener.EventCount} events were recorded, while maxOpPerSecond is {maxOpPerSecond}");
+
+                Thread.Sleep(1000); // ensure new allowance for root creation
+                Activity.Current = null;
+                using var root = a.Source.StartActivity("root");
+                Assert.NotNull(root);
+                Assert.True(root.Recorded);
+
+                using var child = a.Source.StartActivity("child");
+                Assert.NotNull(child); // Child should be created as the parent is recorded.
+            },  maxOperationPerSecond.ToString(CultureInfo.InvariantCulture)).Dispose();
+        }
+
         [ConditionalFact(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))]
         public void TestParentRatioSampler()
         {

src/libraries/System.Diagnostics.DiagnosticSource/tests/System.Diagnostics.DiagnosticSource.Tests.csproj

@@ -23,6 +23,7 @@
 
   <ItemGroup>
     <Compile Include="..\src\System\Diagnostics\DsesSamplerBuilder.cs" Link="DsesSamplerBuilder.cs" />
+    <Compile Include="..\src\System\Diagnostics\RateLimiter.cs" Link="RateLimiter.cs" />
     <Compile Include="..\src\System\Diagnostics\Helpers.cs" Link="Helpers.cs" />
     <Compile Include="..\src\System\Diagnostics\Metrics\Aggregator.cs" Link="Aggregator.cs" />
     <Compile Include="..\src\System\Diagnostics\Metrics\AggregatorStore.cs" Link="AggregatorStore.cs" />