Allow multiple EventCounters sessions and make sure EventListeners always send a Disable command

src/libraries/System.Diagnostics.Tracing/tests/BasicEventSourceTest/TestCallbacks.cs

@@ -0,0 +1,90 @@
+// 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.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Xunit;
+#if USE_MDT_EVENTSOURCE
+using Microsoft.Diagnostics.Tracing;
+#else
+using System.Diagnostics.Tracing;
+#endif
+using System.Reflection;
+
+namespace BasicEventSourceTests
+{
+    public class TestsEventSourceCallbacks
+    {
+        /// <summary>
+        /// Validates that the EventProvider AppDomain.ProcessExit handler does not keep the EventProvider instance
+        /// alive.
+        /// </summary>
+        [Fact]
+        public void Test_EventSource_Lifetime()
+        {
+            using (var source = new CallbacksTestEventSource())
+            {
+                bool isDisabledInDelegate = false;
+                source.EventCommandExecuted += (sender, args) =>
+                {
+                    if (args.Command == EventCommand.Disable)
+                    {
+                        EventSource eventSource = (EventSource)sender;
+                        isDisabledInDelegate = !eventSource.IsEnabled();
+                    }
+                };
+
+                using (var listener = new CallbacksEventListener())
+                {
+                    source.Event();
+                }
+
+                if (!source._isDisabledInCallback)
+                {
+                    Assert.Fail("EventSource was still enabled in OnEventCommand callback");
+                }
+
+                if (!isDisabledInDelegate)
+                {
+                    Assert.Fail("EventSource was still enabled in EventCommandExecuted delegate");
+                }
+            }
+        }
+
+        private class CallbacksEventListener : EventListener
+        {
+            protected override void OnEventSourceCreated(EventSource eventSource)
+            {
+                base.OnEventSourceCreated(eventSource);
+
+                if (eventSource.Name.Equals("TestsEventSourceCallbacks.CallbacksTestEventSource"))
+                {
+                    EnableEvents(eventSource, EventLevel.Verbose);
+                }
+            }
+        }
+
+        [EventSource(Name = "TestsEventSourceCallbacks.CallbacksTestEventSource")]
+        private class CallbacksTestEventSource : EventSource
+        {
+            internal bool _isDisabledInCallback;
+
+            [Event(1)]
+            public void Event()
+            {
+                WriteEvent(1);
+            }
+
+            [NonEvent]
+            protected override void OnEventCommand(EventCommandEventArgs command)
+            {
+                base.OnEventCommand(command);
+
+                _isDisabledInCallback = !IsEnabled();
+            }
+        }
+    }
+}

src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/CounterGroup.cs

@@ -46,24 +46,56 @@ private void RegisterCommandCallback()
 
         private void OnEventSourceCommand(object? sender, EventCommandEventArgs e)
         {
-            if (e.Command == EventCommand.Enable || e.Command == EventCommand.Update)
-            {
-                Debug.Assert(e.Arguments != null);
+            // Should only be enable or disable
+            Debug.Assert(e.Command == EventCommand.Enable || e.Command == EventCommand.Disable);
 
-                if (e.Arguments.TryGetValue("EventCounterIntervalSec", out string? valueStr) && float.TryParse(valueStr, out float value))
+            lock (s_counterGroupLock)      // Lock the CounterGroup
+            {
+                if (e.Command == EventCommand.Enable)
                 {
-                    lock (s_counterGroupLock)      // Lock the CounterGroup
+                    Debug.Assert(e.Arguments != null);
+
+                    if (!e.Arguments.TryGetValue("EventCounterIntervalSec", out string? valueStr)
+                        || !float.TryParse(valueStr, out float intervalValue))
+                    {
+                        // Command is Enable but no EventCounterIntervalSec arg so ignore
+                        return;
+                    }
+
+                    // Sending an Enabled with EventCounterIntervalSec <=0 is a signal that we should immediately turn
+                    // off counters
+                    if (intervalValue <= 0)
+                    {
+                        DisableTimer();
+                    }
+                    else
                     {
-                        EnableTimer(value);
+                        EnableTimer(intervalValue);
                     }
                 }
-            }
-            else if (e.Command == EventCommand.Disable)
-            {
-                lock (s_counterGroupLock)
+                else
                 {
-                    DisableTimer();
+                    Debug.Assert(e.Command == EventCommand.Disable);
+                    // Since we allow sessions to send multiple Enable commands to update the interval, we cannot
+                    // rely on ref counting to determine when to enable and disable counters. You will get an arbitrary
+                    // number of Enables and one Disable per session.
+                    //
+                    // Previously we would turn off counters when we received any Disable command, but that meant that any
+                    // session could turn off counters for all other sessions. To get to a good place we now will only
+                    // turn off counters once the EventSource that provides the counters is disabled. We can then end up
+                    // keeping counters on too long in certain circumstances - if one session enables counters, then a second
+                    // session enables the EventSource but not counters we will stay on until both sessions terminate, even
+                    // if the first session terminates first.
+                    if (!_eventSource.IsEnabled())
+                    {
+                        DisableTimer();
+                    }
                 }
+
+                Debug.Assert((s_counterGroupEnabledList == null && !_eventSource.IsEnabled())
+                                || (_eventSource.IsEnabled() && s_counterGroupEnabledList!.Contains(this))
+                                || (_pollingIntervalInMilliseconds == 0 && !s_counterGroupEnabledList!.Contains(this))
+                                || (!_eventSource.IsEnabled() && !s_counterGroupEnabledList!.Contains(this)));
             }
         }
 
@@ -118,12 +150,9 @@ internal static CounterGroup GetCounterGroup(EventSource eventSource)
 
         private void EnableTimer(float pollingIntervalInSeconds)
         {
+            Debug.Assert(pollingIntervalInSeconds > 0);
             Debug.Assert(Monitor.IsEntered(s_counterGroupLock));
-            if (pollingIntervalInSeconds <= 0)
-            {
-                DisableTimer();
-            }
-            else if (_pollingIntervalInMilliseconds == 0 || pollingIntervalInSeconds * 1000 < _pollingIntervalInMilliseconds)
+            if (_pollingIntervalInMilliseconds == 0 || pollingIntervalInSeconds * 1000 < _pollingIntervalInMilliseconds)
             {
                 _pollingIntervalInMilliseconds = (int)(pollingIntervalInSeconds * 1000);
                 ResetCounters(); // Reset statistics for counters before we start the thread.

src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/EventSource.cs

@@ -2629,9 +2629,6 @@ internal void DoCommand(EventCommandEventArgs commandArgs)
                         m_eventSourceEnabled = true;
                     }
 
-                    this.OnEventCommand(commandArgs);
-                    this.m_eventCommandExecuted?.Invoke(this, commandArgs);
-
                     if (!commandArgs.enable)
                     {
                         // If we are disabling, maybe we can turn on 'quick checks' to filter
@@ -2663,6 +2660,9 @@ internal void DoCommand(EventCommandEventArgs commandArgs)
                             m_eventSourceEnabled = false;
                         }
                     }
+
+                    this.OnEventCommand(commandArgs);
+                    this.m_eventCommandExecuted?.Invoke(this, commandArgs);
                 }
                 else
                 {
@@ -4219,6 +4219,38 @@ internal static void DisposeOnShutdown()
             }
         }
 
+        // If an EventListener calls Dispose without calling DisableEvents first we want to issue the Disable command now
+        private static void CallDisableEventsIfNecessary(EventDispatcher eventDispatcher, EventSource eventSource)
+        {
+#if DEBUG
+            // Disable validation of EventSource/EventListener connections in case a call to EventSource.AddListener
+            // causes a recursive call into this method.
+            bool previousValue = s_ConnectingEventSourcesAndListener;
+            s_ConnectingEventSourcesAndListener = true;
+            try
+            {
+#endif
+                if (eventDispatcher.m_EventEnabled == null)
+                {
+                    return;
+                }
+
+                for (int i = 0; i < eventDispatcher.m_EventEnabled.Length; ++i)
+                {
+                    if (eventDispatcher.m_EventEnabled[i])
+                    {
+                        eventDispatcher.m_Listener.DisableEvents(eventSource);
+                    }
+                }
+#if DEBUG
+            }
+            finally
+            {
+                s_ConnectingEventSourcesAndListener = previousValue;
+            }
+#endif
+        }
+
         /// <summary>
         /// Helper used in code:Dispose that removes any references to 'listenerToRemove' in any of the
         /// eventSources in the appdomain.
@@ -4230,14 +4262,38 @@ private static void RemoveReferencesToListenerInEventSources(EventListener liste
             Debug.Assert(Monitor.IsEntered(EventListener.EventListenersLock));
             // Foreach existing EventSource in the appdomain
             Debug.Assert(s_EventSources != null);
-            foreach (WeakReference<EventSource> eventSourceRef in s_EventSources)
+
+            // First pass to call DisableEvents
+            WeakReference<EventSource>[] eventSourcesSnapshot = s_EventSources.ToArray();
+            foreach (WeakReference<EventSource> eventSourceRef in eventSourcesSnapshot)
             {
                 if (eventSourceRef.TryGetTarget(out EventSource? eventSource))
                 {
-                    Debug.Assert(eventSource.m_Dispatchers != null);
+                    EventDispatcher? cur = eventSource.m_Dispatchers;
+                    while (cur != null)
+                    {
+                        if (cur.m_Listener == listenerToRemove)
+                        {
+                            CallDisableEventsIfNecessary(cur!, eventSource);
+                        }
+
+                        cur = cur.m_Next;
+                    }
+                }
+            }
+
+            // DisableEvents can call back to user code and we have to start over since s_EventSources and
+            // eventSource.m_Dispatchers could have mutated
+            foreach (WeakReference<EventSource> eventSourceRef in s_EventSources)
+            {
+                if (eventSourceRef.TryGetTarget(out EventSource? eventSource)
+                    && eventSource.m_Dispatchers != null)
+                {
                     // Is the first output dispatcher the dispatcher we are removing?
                     if (eventSource.m_Dispatchers.m_Listener == listenerToRemove)
+                    {
                         eventSource.m_Dispatchers = eventSource.m_Dispatchers.m_Next;
+                    }
                     else
                     {
                         // Remove 'listenerToRemove' from the eventSource.m_Dispatchers linked list.
@@ -4267,6 +4323,7 @@ private static void RemoveReferencesToListenerInEventSources(EventListener liste
 #endif // FEATURE_PERFTRACING
         }
 
+
         /// <summary>
         /// Checks internal consistency of EventSources/Listeners.
         /// </summary>

src/tests/tracing/eventlistener/eventlistenerenabledisable.cs

@@ -0,0 +1,106 @@
+// 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.IO;
+using System.Diagnostics.Tracing;
+using System.Threading;
+using Tracing.Tests.Common;
+
+namespace Tracing.Tests
+{
+    [EventSource(Name = "Tracing.Tests.EnableDisableEventSource")]
+    class EnableDisableEventSource : EventSource
+    {
+        internal int _enables;
+        internal int _disables;
+
+        public EnableDisableEventSource() : base(true) { }
+
+        [Event(1)]
+        internal void TestEvent() { this.WriteEvent(1); }
+
+        [NonEvent]
+        protected override void OnEventCommand(EventCommandEventArgs command)
+        {
+            base.OnEventCommand(command);
+
+            if (command.Command == EventCommand.Enable)
+            {
+                Interlocked.Increment(ref _enables);
+            }
+            else if (command.Command == EventCommand.Disable)
+            {
+                Interlocked.Increment(ref _disables);
+            }
+            else
+            {
+                throw new Exception($"Saw unexpected command {command.Command} in OnEventCommand");
+            }
+        }
+    }
+
+    internal sealed class EnableDisableListener : EventListener
+    {
+        private EventSource? _target;
+
+        protected override void OnEventSourceCreated(EventSource source)
+        {
+            if (source.Name.Equals("Tracing.Tests.EnableDisableEventSource"))
+            {
+                _target = source;
+                EnableEvents(_target, EventLevel.Verbose);
+            }
+        }
+
+        public void Disable()
+        {
+            DisableEvents(_target);
+        }
+
+        public void Dispose(bool disableEvents)
+        {
+            base.Dispose();
+
+            if (disableEvents)
+            {
+                // Dispose should call DisableEvents for us, this should be ignored after Dispose()
+                Disable();
+            }
+        }
+    }
+
+    class EventListenerEnableDisableTest
+    {
+        static int Main()
+        {
+            bool pass = false;
+            using(var source = new EnableDisableEventSource())
+            {
+                // Testing three scenarios:
+                //      listener1 calls EnableEvents but never calls DisableEvents
+                //      listener2 calls EnableEvents and calls DisableEvents outside of Dispose
+                //      listener3 calls EnableEvents and calls DisableEvents inside of Dispose
+                // 
+                // We should get an Enable and Disable for all of them
+                using (var listener1 = new EnableDisableListener())
+                using (var listener2 = new EnableDisableListener())
+                {
+                    var listener3 = new EnableDisableListener();
+                    source.TestEvent();
+                    listener3.Dispose(true);
+
+                    listener2.Disable();
+                }
+
+                if (source._enables == 3 && source._disables == 3)
+                {
+                    return 100;
+                }
+
+                Console.WriteLine($"Unexpected enable/disable count _enables={source._enables} _disables={source._disables}");
+                return -1;
+            }
+        }
+    }
+}

src/tests/tracing/eventlistener/eventlistenerenabledisable.csproj

@@ -0,0 +1,10 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="eventlistenerenabledisable.cs" />
+    <ProjectReference Include="../common/common.csproj" />
+  </ItemGroup>
+</Project>