Fix race condition in RequestTelemetryLogger

dotnet · Feb 20, 2022 · c43e7ed · c43e7ed
1 parent 1e091e6
commit c43e7ed
Showing 1 changed file with 13 additions and 12 deletions.
diff --git a/src/Features/LanguageServer/Protocol/Handler/RequestExecutionQueue.RequestTelemetryLogger.cs b/src/Features/LanguageServer/Protocol/Handler/RequestExecutionQueue.RequestTelemetryLogger.cs
@@ -4,8 +4,6 @@
 
 using System;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
-using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using Microsoft.CodeAnalysis.Internal.Log;
 
@@ -17,7 +15,7 @@ internal partial class RequestExecutionQueue
         /// Logs metadata on LSP requests (duration, success / failure metrics)
         /// for this particular LSP server instance.
         /// </summary>
-        internal class RequestTelemetryLogger : IDisposable
+        internal sealed class RequestTelemetryLogger : IDisposable
         {
             private const string QueuedDurationKey = "QueuedDuration";
 
@@ -26,7 +24,7 @@ internal class RequestTelemetryLogger : IDisposable
             /// <summary>
             /// Histogram to aggregate the time in queue metrics.
             /// </summary>
-            private HistogramLogAggregator? _queuedDurationLogAggregator;
+            private readonly HistogramLogAggregator _queuedDurationLogAggregator;
 
             /// <summary>
             /// Histogram to aggregate total request duration metrics.
@@ -36,7 +34,7 @@ internal class RequestTelemetryLogger : IDisposable
             /// This provides highly detailed buckets when duration is in MS, but less detailed
             /// when the duration is in terms of seconds or minutes.
             /// </summary>
-            private HistogramLogAggregator? _requestDurationLogAggregator;
+            private readonly HistogramLogAggregator _requestDurationLogAggregator;
 
             /// <summary>
             /// Store request counters in a concurrent dictionary as non-mutating LSP requests can
@@ -46,6 +44,8 @@ internal class RequestTelemetryLogger : IDisposable
 
             private readonly LogAggregator _findDocumentResults;
 
+            private int _disposed;
+
             public RequestTelemetryLogger(string serverTypeName)
             {
                 _serverTypeName = serverTypeName;
@@ -79,10 +79,10 @@ public void UpdateTelemetryData(
             {
                 // Find the bucket corresponding to the queued duration and update the count of durations in that bucket.
                 // This is not broken down per method as time in queue is not specific to an LSP method.
-                _queuedDurationLogAggregator?.IncreaseCount(QueuedDurationKey, Convert.ToDecimal(queuedDuration.TotalMilliseconds));
+                _queuedDurationLogAggregator.IncreaseCount(QueuedDurationKey, Convert.ToDecimal(queuedDuration.TotalMilliseconds));
 
                 // Store the request time metrics per LSP method.
-                _requestDurationLogAggregator?.IncreaseCount(methodName, Convert.ToDecimal(ComputeLogValue(requestDuration.TotalMilliseconds)));
+                _requestDurationLogAggregator.IncreaseCount(methodName, Convert.ToDecimal(ComputeLogValue(requestDuration.TotalMilliseconds)));
                 _requestCounters.GetOrAdd(methodName, (_) => new Counter()).IncrementCount(result);
             }
 
@@ -103,8 +103,12 @@ private static double ComputeLogValue(double durationInMS)
             /// </summary>
             public void Dispose()
             {
-                if (_queuedDurationLogAggregator is null || _queuedDurationLogAggregator.IsEmpty
-                    || _requestDurationLogAggregator is null || _requestDurationLogAggregator.IsEmpty)
+                if (Interlocked.Exchange(ref _disposed, 1) != 0)
+                {
+                    return;
+                }
+
+                if (_queuedDurationLogAggregator.IsEmpty || _requestDurationLogAggregator.IsEmpty)
                 {
                     return;
                 }
@@ -150,10 +154,7 @@ public void Dispose()
                     }
                 }));
 
-                // Clear telemetry we've published in case dispose is called multiple times.
                 _requestCounters.Clear();
-                _queuedDurationLogAggregator = null;
-                _requestDurationLogAggregator = null;
             }
 
             private class Counter