Fix CollectionRulePipelineTests failure due to callback registration timing.

src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRulePipelineTests.cs

@@ -62,7 +62,7 @@ public Task CollectionRulePipeline_StartupTriggerTest(TargetFrameworkMoniker app
 
                     // Register first callback before pipeline starts. This callback should be completed before
                     // the pipeline finishes starting.
-                    Task callback1Task = callbackService.WaitWithCancellationAsync(cancellationSource.Token);
+                    Task callback1Task = await callbackService.StartWaitForCallbackAsync(cancellationSource.Token);
 
                     // Startup trigger will cause the the pipeline to complete the start phase
                     // after the action list has been completed.
@@ -72,12 +72,11 @@ public Task CollectionRulePipeline_StartupTriggerTest(TargetFrameworkMoniker app
                     // completed because it was registered after the pipeline had finished starting. Since
                     // the action list is only ever executed once and is executed before the pipeline finishes
                     // starting, thus subsequent invocations of the action list should not occur.
-                    Task callback2Task = callbackService.WaitWithCancellationAsync(cancellationSource.Token);
+                    Task callback2Task = await callbackService.StartWaitForCallbackAsync(cancellationSource.Token);
 
                     // Since the action list was completed before the pipeline finished starting,
                     // the action should have invoked it's callback.
-                    using CancellationTokenSource callbackCancellationSource = new(TimeSpan.FromMilliseconds(50));
-                    await callback1Task.WithCancellation(callbackCancellationSource.Token);
+                    await callback1Task.WithCancellation(cancellationSource.Token);
 
                     // Regardless of the action list constraints, the pipeline should have only
                     // executed the action list once due to the use of a startup trigger.
@@ -132,15 +131,15 @@ public Task CollectionRulePipeline_EventCounterTriggerTest(TargetFrameworkMonike
 
                     // Register first callback before pipeline starts. This callback should be completed after
                     // the pipeline finishes starting.
-                    Task callbackTask = callbackService.WaitWithCancellationAsync(cancellationSource.Token);
+                    Task callbackTask = await callbackService.StartWaitForCallbackAsync(cancellationSource.Token);
 
                     // Start pipeline with EventCounter trigger.
                     await pipeline.StartAsync(cancellationSource.Token);
 
                     await runner.SendCommandAsync(TestAppScenarios.SpinWait.Commands.StartSpin);
 
                     // This should not complete until the trigger conditions are satisfied for the first time.
-                    await callbackTask;
+                    await callbackTask.WithCancellation(cancellationSource.Token);
 
                     VerifyExecutionCount(callbackService, 1);
 

src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRules/Actions/CallbackAction.cs

@@ -70,7 +70,22 @@ public async Task NotifyListeners(CancellationToken token)
             }
         }
 
-        public async Task WaitWithCancellationAsync(CancellationToken token)
+        /// <summary>
+        /// Registers a callback with the Callback action that will complete when
+        /// the Callback action is invoked.
+        /// </summary>
+        /// <returns>
+        /// A <see cref="Task{Task}"/> that completes when the callback has finished
+        /// being registered. The inner <see cref="Task"/> will complete when the callback
+        /// is invoked.
+        /// </returns>
+        /// <remarks>
+        /// Await this method to wait for the callback to be registered; await the inner
+        /// task to wait for the callback to be invoked. The <paramref name="token"/> parameter
+        /// only cancels registration if the registration has not completed; it does not cancel
+        /// the inner task that represents the callback invocation.
+        /// </remarks>
+        public async Task<Task> StartWaitForCallbackAsync(CancellationToken token)
         {
             int id = _nextId++;
             string name = $"Callback{id}";
@@ -89,7 +104,7 @@ public async Task WaitWithCancellationAsync(CancellationToken token)
                 _entriesSemaphore.Release();
             }
 
-            await entry.WithCancellation(token);
+            return entry.Task;
         }
 
         public IReadOnlyCollection<DateTime> ExecutionTimestamps
@@ -126,12 +141,7 @@ public void Complete()
                 _outputHelper.WriteLine("[Callback] End completing {0}.", _name);
             }
 
-            public async Task WithCancellation(CancellationToken token)
-            {
-                _outputHelper.WriteLine("[Test] Begin waiting for {0} completion.", _name);
-                await _completionSource.WithCancellation(token);
-                _outputHelper.WriteLine("[Test] End waiting for {0} completion.", _name);
-            }
+            public Task Task => _completionSource.Task;
         }
     }
 }