From e5a90534e95d8d4c4e62698ba7ec8428aa4c62bb Mon Sep 17 00:00:00 2001
From: Noah Falk <noahfalk@microsoft.com>
Date: Tue, 8 Oct 2024 00:50:31 -0700
Subject: [PATCH] Fixes deadlock for IncrementingPollingCounter callbacks

This is a modified backport of https://github.com/dotnet/runtime/pull/105548. It mostly preserves the logic of the original fix in .NET 9 with a few adjustments:
- Added a config switch System.Diagnostics.Tracing.CounterCallbackOnTimerThread that must be set to true to opt-in to the fix behavior. The .NET 9 fix
  was documented as a breaking change because it slighly modifies the timing and thread used for first call to an IncrementingPollingCounter callback. I did not want anyone in 8.0 to opted into this by default.
- The opt-in switch sets the property CounterCallbackOnTimerThread and I added this condition to several of the if checks in the code. Its more than would be strictly necessary just to make it obvious when code reviewing individual methods that the new code paths are unreachable unless the app opts in.
- The original 9.0 change had a bit more refactoring that wasn't essential (renaming a method, removing an unneeded lock() scope) and I removed that here to reduce the code delta.
---
 .../Diagnostics/Tracing/CounterGroup.cs       | 48 ++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/CounterGroup.cs b/src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/CounterGroup.cs
index 71af1ba5b20973..9af21bac674984 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/CounterGroup.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/CounterGroup.cs
@@ -14,6 +14,10 @@ namespace System.Diagnostics.Tracing
 #endif
     internal sealed class CounterGroup
     {
+        // This is the opt-in switch for a .NET 8 backport of https://github.com/dotnet/runtime/pull/105548
+        private static bool CounterCallbackOnTimerThread { get; } =
+            AppContext.TryGetSwitch("System.Diagnostics.Tracing.CounterCallbackOnTimerThread", out bool value) ? value : false;
+
         private readonly EventSource _eventSource;
         private readonly List<DiagnosticCounter> _counters;
         private static readonly object s_counterGroupLock = new object();
@@ -189,6 +193,16 @@ private void DisableTimer()
             Debug.Assert(Monitor.IsEntered(s_counterGroupLock));
             _pollingIntervalInMilliseconds = 0;
             s_counterGroupEnabledList?.Remove(this);
+
+            if (CounterCallbackOnTimerThread && s_needsResetIncrementingPollingCounters.Count > 0)
+            {
+                // if there are any pending callbacks for the timer thread, remove them now that the timer is being disabled
+                foreach (DiagnosticCounter diagnosticCounter in _counters)
+                {
+                    if (diagnosticCounter is IncrementingPollingCounter pollingCounter)
+                        s_needsResetIncrementingPollingCounters.Remove(pollingCounter);
+                }
+            }
         }
 
         private void ResetCounters()
@@ -203,7 +217,17 @@ private void ResetCounters()
                     }
                     else if (counter is IncrementingPollingCounter ipCounter)
                     {
-                        ipCounter.UpdateMetric();
+                        if (CounterCallbackOnTimerThread)
+                        {
+                            // IncrementingPollingCounters will be reset on timer thread
+                            // We need this to avoid deadlocks caused by running IncrementingPollingCounter._totalValueProvider under EventListener.EventListenersLock
+                            // This is an opt-in fix for servicing because there is a small chance this could break apps (most likely tests).
+                            s_needsResetIncrementingPollingCounters.Add(ipCounter);
+                        }
+                        else
+                        {
+                            ipCounter.UpdateMetric();
+                        }
                     }
                     else if (counter is EventCounter eCounter)
                     {
@@ -265,6 +289,10 @@ private void OnTimer()
 
         private static List<CounterGroup>? s_counterGroupEnabledList;
 
+        // Stores a list of IncrementingPollingCounters that need to be reset on the timer thread. This field
+        // is only accessed when CallbackOnTimerThread is true.
+        private static List<IncrementingPollingCounter> s_needsResetIncrementingPollingCounters = [];
+
         private static void PollForValues()
         {
             AutoResetEvent? sleepEvent = null;
@@ -274,6 +302,7 @@ private static void PollForValues()
             // calling into the callbacks can cause a re-entrancy into CounterGroup.Enable()
             // and result in a deadlock. (See https://github.com/dotnet/runtime/issues/40190 for details)
             var onTimers = new List<CounterGroup>();
+            List<IncrementingPollingCounter>? countersToReset = null;
             while (true)
             {
                 int sleepDurationInMilliseconds = int.MaxValue;
@@ -292,7 +321,24 @@ private static void PollForValues()
                         millisecondsTillNextPoll = Math.Max(1, millisecondsTillNextPoll);
                         sleepDurationInMilliseconds = Math.Min(sleepDurationInMilliseconds, millisecondsTillNextPoll);
                     }
+
+                    if (CounterCallbackOnTimerThread && s_needsResetIncrementingPollingCounters.Count > 0)
+                    {
+                        countersToReset = s_needsResetIncrementingPollingCounters;
+                        s_needsResetIncrementingPollingCounters = [];
+                    }
+                }
+
+                if (countersToReset != null)
+                {
+                    foreach (IncrementingPollingCounter counter in countersToReset)
+                    {
+                        counter.UpdateMetric();
+                    }
+
+                    countersToReset = null;
                 }
+
                 foreach (CounterGroup onTimer in onTimers)
                 {
                     onTimer.OnTimer();
