Added tests for FlushAsync and Dispose

Author: jamescrosswell

src/Sentry/MetricAggregator.cs

@@ -7,9 +7,15 @@ namespace Sentry;
 
 internal class MetricAggregator : IMetricAggregator
 {
+    internal const string DisposingMessage = "Disposing MetricAggregator.";
+    internal const string AlreadyDisposedMessage = "Already disposed MetricAggregator.";
+    internal const string CancelledMessage = "Stopping the Metric Aggregator due to a cancellation.";
+    internal const string ShutdownScheduledMessage = "Shutdown scheduled. Stopping by: {0}.";
+    internal const string ShutdownImmediatelyMessage = "Exiting immediately due to 0 shutdown timeout.";
+    internal const string FlushShutdownMessage = "Shutdown token triggered. Exiting metric aggregator.";
+
     private readonly SentryOptions _options;
     private readonly IMetricHub _metricHub;
-    private readonly TimeSpan _flushInterval;
 
     private readonly SemaphoreSlim _codeLocationLock = new(1,1);
     private readonly ReaderWriterLockSlim _bucketsLock = new ReaderWriterLockSlim();
@@ -29,16 +35,14 @@ private readonly Lazy<Dictionary<long, ConcurrentDictionary<string, Metric>>> _b
     internal readonly ConcurrentDictionary<long, HashSet<MetricResourceIdentifier>> _seenLocations = new();
     internal Dictionary<long, Dictionary<MetricResourceIdentifier, SentryStackFrame>> _pendingLocations = new();
 
-    private readonly Task _loopTask;
+    internal readonly Task _loopTask;
 
     internal MetricAggregator(SentryOptions options, IMetricHub metricHub,
-        CancellationTokenSource? shutdownSource = null,
-        bool disableLoopTask = false, TimeSpan? flushInterval = null)
+        CancellationTokenSource? shutdownSource = null, bool disableLoopTask = false)
     {
         _options = options;
         _metricHub = metricHub;
         _shutdownSource = shutdownSource ?? new CancellationTokenSource();
-        _flushInterval = flushInterval ?? TimeSpan.FromSeconds(5);
 
         if (disableLoopTask)
         {
@@ -305,12 +309,12 @@ private async Task RunLoopAsync()
                 // If the cancellation was signaled, run until the end of the queue or shutdownTimeout
                 try
                 {
-                    await Task.Delay(_flushInterval, _shutdownSource.Token).ConfigureAwait(false);
+                    await Task.Delay(_options.ShutdownTimeout, _shutdownSource.Token).ConfigureAwait(false);
                 }
                 // Cancellation requested and no timeout allowed, so exit even if there are more items
                 catch (OperationCanceledException) when (_options.ShutdownTimeout == TimeSpan.Zero)
                 {
-                    _options.LogDebug("Exiting immediately due to 0 shutdown timeout.");
+                    _options.LogDebug(ShutdownImmediatelyMessage);
 
                     await shutdownTimeout.CancelAsync().ConfigureAwait(false);
 
@@ -319,9 +323,7 @@ private async Task RunLoopAsync()
                 // Cancellation requested, scheduled shutdown
                 catch (OperationCanceledException)
                 {
-                    _options.LogDebug(
-                        "Shutdown scheduled. Stopping by: {0}.",
-                        _options.ShutdownTimeout);
+                    _options.LogDebug(ShutdownScheduledMessage, _options.ShutdownTimeout);
 
                     shutdownTimeout.CancelAfterSafe(_options.ShutdownTimeout);
 
@@ -391,15 +393,20 @@ public async Task FlushAsync(bool force = true, CancellationToken cancellationTo
         }
         catch (OperationCanceledException)
         {
-            _options.LogInfo("Shutdown token triggered. Exiting metric aggregator.");
+            _options.LogInfo(FlushShutdownMessage);
         }
         catch (Exception exception)
         {
             _options.LogError(exception, "Error processing metrics.");
         }
         finally
         {
-            _flushLock.Release();
+            // If the shutdown token was cancelled before we start this method, we can get here
+            // without the _flushLock.CurrentCount (i.e. available threads) having been decremented
+            if (_flushLock.CurrentCount < 1)
+            {
+                _flushLock.Release();
+            }
         }
     }
 
@@ -495,11 +502,11 @@ private void ClearStaleLocations()
     /// <inheritdoc cref="IAsyncDisposable.DisposeAsync"/>
     public async ValueTask DisposeAsync()
     {
-        _options.LogDebug("Disposing MetricAggregator.");
+        _options.LogDebug(DisposingMessage);
 
         if (_disposed)
         {
-            _options.LogDebug("Already disposed MetricAggregator.");
+            _options.LogDebug(AlreadyDisposedMessage);
             return;
         }
 
@@ -518,7 +525,7 @@ public async ValueTask DisposeAsync()
         }
         catch (OperationCanceledException)
         {
-            _options.LogDebug("Stopping the Metric Aggregator due to a cancellation.");
+            _options.LogDebug(CancelledMessage);
         }
         catch (Exception exception)
         {

test/Sentry.Tests/MetricAggregatorTests.cs

@@ -4,19 +4,33 @@ namespace Sentry.Tests;
 
 public class MetricAggregatorTests
 {
-    class Fixture
+    private class Fixture
     {
-        public SentryOptions Options { get; set; } = new();
-        public IHub Hub { get; set; } = Substitute.For<IHub>();
-        public IMetricHub MetricHub { get; set; } = Substitute.For<IMetricHub>();
-        public bool DisableFlushLoop { get; set; } = true;
-        public TimeSpan? FlushInterval { get; set; }
-        public MetricAggregator GetSut()
-            => new(Options, MetricHub, disableLoopTask: DisableFlushLoop, flushInterval: FlushInterval);
+        public readonly IDiagnosticLogger Logger;
+        public readonly SentryOptions Options;
+
+        public readonly IMetricHub MetricHub;
+        public bool DisableFlushLoop;
+        public readonly CancellationTokenSource CancellationTokenSource;
+
+        public Fixture()
+        {
+            Logger = Substitute.For<IDiagnosticLogger>();
+            Options = new SentryOptions
+            {
+                Debug = true,
+                DiagnosticLogger = Logger
+            };
+            MetricHub = Substitute.For<IMetricHub>();
+            DisableFlushLoop = true;
+            CancellationTokenSource = new CancellationTokenSource();
+        }
+
+        public MetricAggregator GetSut() => new(Options, MetricHub, CancellationTokenSource, DisableFlushLoop);
     }
 
     // private readonly Fixture _fixture = new();
-    private static readonly Fixture _fixture = new();
+    private readonly Fixture _fixture = new();
 
     [Fact]
     public void GetMetricBucketKey_GeneratesExpectedKey()
@@ -176,7 +190,8 @@ public async Task GetFlushableBuckets_IsThreadsafe()
         var sent = 0;
         MetricHelper.FlushShift = 0.0;
         _fixture.DisableFlushLoop = false;
-        _fixture.FlushInterval = TimeSpan.FromMilliseconds(100);
+        // TODO: Remove
+        // _fixture.FlushInterval = TimeSpan.FromMilliseconds(100);
         _fixture.MetricHub.CaptureMetrics(Arg.Do<IEnumerable<Metric>>(metrics =>
             {
                 foreach (var metric in metrics)
@@ -339,4 +354,126 @@ public void RecordCodeLocation_BadStackLevel_AddsToSeenButNotPending()
 
         sut._pendingLocations.SelectMany(x => x.Value).Should().BeEmpty();
     }
+
+    [Fact]
+    public void Dispose_OnlyExecutesOnce()
+    {
+        // Arrange
+        _fixture.Logger.IsEnabled(Arg.Any<SentryLevel>()).Returns(true);
+        var sut = _fixture.GetSut();
+
+        // Act
+        sut.Dispose();
+        sut.Dispose();
+        sut.Dispose();
+
+        // Assert
+        _fixture.Logger.Received(2).Log(SentryLevel.Debug, MetricAggregator.AlreadyDisposedMessage, null);
+    }
+
+    [Fact]
+    public void Dispose_StopsLoopTask()
+    {
+        // Arrange
+        _fixture.Logger.IsEnabled(Arg.Any<SentryLevel>()).Returns(true);
+        _fixture.DisableFlushLoop = false;
+        _fixture.Options.ShutdownTimeout = TimeSpan.Zero;
+        var sut = _fixture.GetSut();
+
+        // Act
+        sut.Dispose();
+
+        // Assert
+        _fixture.Logger.Received(1).Log(SentryLevel.Debug, MetricAggregator.DisposingMessage, null);
+        sut._loopTask.Status.Should().BeOneOf(TaskStatus.RanToCompletion, TaskStatus.Faulted);
+    }
+
+    [Fact]
+    public async Task Dispose_SwallowsException()
+    {
+        // Arrange
+        _fixture.CancellationTokenSource.Dispose();
+        _fixture.DisableFlushLoop = false;
+        var sut = _fixture.GetSut();
+
+        // We expect an exception here, because we disposed the cancellation token source
+        await Assert.ThrowsAsync<ObjectDisposedException>(() => sut._loopTask);
+
+        // Act
+        await sut.DisposeAsync();
+
+        // Assert
+        sut._loopTask.Status.Should().Be(TaskStatus.Faulted);
+    }
+
+    [Fact]
+    public async Task Cancel_NonZeroTimeout_SchedulesShutdown()
+    {
+        // Arrange
+        _fixture.Logger.IsEnabled(Arg.Any<SentryLevel>()).Returns(true);
+        _fixture.DisableFlushLoop = false;
+        _fixture.Options.ShutdownTimeout = TimeSpan.FromSeconds(1);
+        var sut = _fixture.GetSut();
+
+        // Act
+        await _fixture.CancellationTokenSource.CancelAsync();
+        await Task.Delay(1000);
+
+        // Assert
+        _fixture.Logger.Received(1).Log(SentryLevel.Debug, MetricAggregator.ShutdownScheduledMessage, null, Arg.Any<TimeSpan>());
+    }
+
+    [Fact]
+    public async Task Cancel_ZeroTimeout_ShutdownImmediately()
+    {
+        // Arrange
+        _fixture.Logger.IsEnabled(Arg.Any<SentryLevel>()).Returns(true);
+        _fixture.DisableFlushLoop = false;
+        _fixture.Options.ShutdownTimeout = TimeSpan.Zero;
+        var sut = _fixture.GetSut();
+
+        // Act
+        await _fixture.CancellationTokenSource.CancelAsync();
+        await Task.Delay(1000);
+
+        // Assert
+        _fixture.Logger.Received(1).Log(SentryLevel.Debug, MetricAggregator.ShutdownImmediatelyMessage, null);
+    }
+
+    [Fact]
+    public async Task FlushAsync_FlushesPendingLocations()
+    {
+        // Arrange
+        var type = MetricType.Counter;
+        var key = "counter_key";
+        var unit = MeasurementUnit.None;
+        var stackLevel = 1;
+        var timestamp = DateTimeOffset.Now.Subtract(TimeSpan.FromSeconds(20));
+        var sut = _fixture.GetSut();
+        sut.RecordCodeLocation(type, key, unit, stackLevel, timestamp);
+
+        // Act
+        await sut.FlushAsync();
+
+        // Assert
+        _fixture.MetricHub.Received(1).CaptureCodeLocations(Arg.Any<CodeLocations>());
+    }
+
+    [Fact]
+    public async Task FlushAsync_Cancel_Exists()
+    {
+        // Arrange
+        _fixture.DisableFlushLoop = false;
+        _fixture.Logger.IsEnabled(Arg.Any<SentryLevel>()).Returns(true);
+        var cancellationTokenSource = new CancellationTokenSource();
+        await cancellationTokenSource.CancelAsync();
+        var sut = _fixture.GetSut();
+
+        // Act
+        await sut.FlushAsync(true, cancellationTokenSource.Token);
+        // await Task.Delay(1000);
+
+        // Assert
+        _fixture.Logger.Received(1).Log(SentryLevel.Info, MetricAggregator.FlushShutdownMessage, null);
+    }
 }