Fix DiagnosticListener IsEnabled cyclical dispatch and lost activity hooks

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

@@ -107,9 +107,13 @@ public virtual IDisposable Subscribe(IObserver<KeyValuePair<string, object?>> ob
         /// </param>
         public virtual IDisposable Subscribe(IObserver<KeyValuePair<string, object?>> observer, Func<string, object?, object?, bool>? isEnabled)
         {
-            return isEnabled == null ?
-             SubscribeInternal(observer, null, null, null, null) :
-             SubscribeInternal(observer, name => IsEnabled(name, null, null), isEnabled, null, null);
+            if (isEnabled == null)
+            {
+                return SubscribeInternal(observer, null, null, null, null);
+            }
+
+            Func<string, object?, object?, bool> localIsEnabled = isEnabled;
+            return SubscribeInternal(observer, name => localIsEnabled(name, null, null), isEnabled, null, null);
         }
 
         /// <summary>
@@ -325,7 +329,7 @@ public void Dispose()
                 for (int i = 0; i < 100; i++)
                     GC.KeepAlive("");
 #endif
-                return new DiagnosticSubscription() { Observer = subscriptions.Observer, Owner = subscriptions.Owner, IsEnabled1Arg = subscriptions.IsEnabled1Arg, IsEnabled3Arg = subscriptions.IsEnabled3Arg, Next = Remove(subscriptions.Next, subscription) };
+                return new DiagnosticSubscription() { Observer = subscriptions.Observer, Owner = subscriptions.Owner, IsEnabled1Arg = subscriptions.IsEnabled1Arg, IsEnabled3Arg = subscriptions.IsEnabled3Arg, OnActivityImport = subscriptions.OnActivityImport, OnActivityExport = subscriptions.OnActivityExport, Next = Remove(subscriptions.Next, subscription) };
             }
         }
 

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

@@ -143,9 +143,13 @@ public override void OnActivityExport(Activity activity, object? payload)
         public virtual IDisposable Subscribe(IObserver<KeyValuePair<string, object?>> observer, Func<string, object?, object?, bool>? isEnabled,
             Action<Activity, object?>? onActivityImport = null, Action<Activity, object?>? onActivityExport = null)
         {
-            return isEnabled == null ?
-             SubscribeInternal(observer, null, null, onActivityImport, onActivityExport) :
-             SubscribeInternal(observer, name => IsEnabled(name, null, null), isEnabled, onActivityImport, onActivityExport);
+            if (isEnabled == null)
+            {
+                return SubscribeInternal(observer, null, null, onActivityImport, onActivityExport);
+            }
+
+            Func<string, object?, object?, bool> localIsEnabled = isEnabled;
+            return SubscribeInternal(observer, name => localIsEnabled(name, null, null), isEnabled, onActivityImport, onActivityExport);
         }
     }
 }

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

@@ -717,6 +717,81 @@ public void ActivityImportExport()
             }
         }
 
+        [Fact]
+        public void ActivityImportExportSurvivesOtherSubscriptionRemoval()
+        {
+            using (DiagnosticListener listener = new DiagnosticListener("ActivityImportExportSurvivesRemoval"))
+            {
+                Activity activity = new Activity("MyActivity");
+                object payload = "MyPayload";
+                bool seenActivityImport = false;
+                bool seenActivityExport = false;
+
+                Action<Activity, object> activityImport = (a, p) => seenActivityImport = true;
+                Action<Activity, object> activityExport = (a, p) => seenActivityExport = true;
+
+                // Subscribe the activity-hook observer first. Subscriptions are prepended, so a later
+                // subscription sits ahead of it in the linked list. Removing that deeper (earlier) node
+                // forces this node to be rebuilt, which is where the activity hooks were being dropped.
+                IDisposable plain = listener.Subscribe(
+                    new ObserverToList<TelemData>(new List<KeyValuePair<string, object>>()));
+
+                IDisposable withActivityHooks = listener.Subscribe(
+                    new ObserverToList<TelemData>(new List<KeyValuePair<string, object>>()),
+                    (name, arg1, arg2) => true,
+                    activityImport,
+                    activityExport);
+
+                // Removing the earlier-added (deeper) subscription forces the activity-hook node ahead
+                // of it to be copied.
+                plain.Dispose();
+
+                listener.OnActivityImport(activity, payload);
+                listener.OnActivityExport(activity, payload);
+
+                Assert.True(seenActivityImport, "OnActivityImport hook was lost when another subscription was removed.");
+                Assert.True(seenActivityExport, "OnActivityExport hook was lost when another subscription was removed.");
+
+                withActivityHooks.Dispose();
+            }
+        }
+
+        [Fact]
+        public void IsEnabledOneArgDoesNotRouteThroughThreeArgOverride()
+        {
+            using (var listener = new OverriddenIsEnabledListener("IsEnabledRoutingListener"))
+            {
+                bool seenPredicate = false;
+                Func<string, object, object, bool> predicate = (name, arg1, arg2) =>
+                {
+                    seenPredicate = true;
+                    return true;
+                };
+
+                using (listener.Subscribe(new ObserverToList<TelemData>(new List<KeyValuePair<string, object>>()), predicate))
+                {
+                    Assert.True(listener.IsEnabled("SomeEvent"));
+
+                    // The single-arg IsEnabled must not be dispatched through the three-arg IsEnabled override.
+                    Assert.Equal(0, listener.ThreeArgIsEnabledCount);
+                    Assert.True(seenPredicate);
+                }
+            }
+        }
+
+        private sealed class OverriddenIsEnabledListener : DiagnosticListener
+        {
+            public OverriddenIsEnabledListener(string name) : base(name) { }
+
+            public int ThreeArgIsEnabledCount { get; private set; }
+
+            public override bool IsEnabled(string name, object arg1, object arg2 = null)
+            {
+                ThreeArgIsEnabledCount++;
+                return base.IsEnabled(name, arg1, arg2);
+            }
+        }
+
         #region Helpers
         /// <summary>
         /// Returns the list of active diagnostic listeners.