Stabilize distributed trace pipeline tests

src/Microsoft.Diagnostics.Monitoring.EventPipe/DistributedTraces/DistributedTracesPipeline.cs

@@ -29,8 +29,6 @@ protected override MonitoringSourceConfiguration CreateConfiguration()
 
         protected override async Task OnEventSourceAvailable(EventPipeEventSource eventSource, Func<Task> stopSessionAsync, CancellationToken token)
         {
-            await ExecuteActivityLoggerActionAsync((logger) => logger.PipelineStarted(token)).ConfigureAwait(false);
-
             eventSource.Dynamic.All += traceEvent => {
                 try
                 {
@@ -55,6 +53,8 @@ protected override async Task OnEventSourceAvailable(EventPipeEventSource eventS
                 }
             };
 
+            await ExecuteActivityLoggerActionAsync((logger) => logger.PipelineStarted(token)).ConfigureAwait(false);
+
             using EventTaskSource<Action> sourceCompletedTaskSource = new(
                 taskComplete => taskComplete,
                 handler => eventSource.Completed += handler,

src/tests/CommonTestRunner/Debuggees/EventPipeTracee/Program.cs

@@ -34,8 +34,9 @@ public static async Task<int> Main(string[] args)
             bool diagMetrics = args.Any("DiagMetrics".Equals);
             bool duplicateNameMetrics = args.Any("DuplicateNameMetrics".Equals);
             bool useActivitySource = args.Any("UseActivitySource".Equals);
+            bool waitForActivitySourceListener = args.Any("WaitForActivitySourceListener".Equals);
 
-            Console.WriteLine($"{pid} EventPipeTracee: DiagMetrics {diagMetrics} UseActivitySource {useActivitySource}");
+            Console.WriteLine($"{pid} EventPipeTracee: DiagMetrics {diagMetrics} UseActivitySource {useActivitySource} WaitForActivitySourceListener {waitForActivitySourceListener}");
             Console.WriteLine($"{pid} EventPipeTracee: DuplicateNameMetrics {duplicateNameMetrics}");
 
             Console.WriteLine($"{pid} EventPipeTracee: start process");
@@ -76,6 +77,11 @@ public static async Task<int> Main(string[] args)
             // Wait for server to send something
             int input = pipeStream.ReadByte();
 
+            if (waitForActivitySourceListener)
+            {
+                await WaitForActivitySourceListener(activitySource, pid).ConfigureAwait(false);
+            }
+
             Console.WriteLine($"{pid} {DateTime.UtcNow} Starting test body '{input}'");
             Console.Out.Flush();
 
@@ -141,6 +147,33 @@ public static async Task<int> Main(string[] args)
             return 0;
         }
 
+        private static async Task WaitForActivitySourceListener(ActivitySource activitySource, int pid)
+        {
+            if (activitySource == null)
+            {
+                throw new InvalidOperationException("WaitForActivitySourceListener requires UseActivitySource.");
+            }
+
+            TimeSpan timeout = TimeSpan.FromSeconds(30);
+            Stopwatch stopwatch = Stopwatch.StartNew();
+
+            Console.WriteLine($"{pid} EventPipeTracee: waiting for ActivitySource listener");
+            Console.Out.Flush();
+
+            while (!activitySource.HasListeners() && stopwatch.Elapsed < timeout)
+            {
+                await Task.Delay(10).ConfigureAwait(false);
+            }
+
+            if (!activitySource.HasListeners())
+            {
+                throw new TimeoutException($"ActivitySource listener was not installed within {timeout}.");
+            }
+
+            Console.WriteLine($"{pid} EventPipeTracee: ActivitySource listener ready after {stopwatch.Elapsed}");
+            Console.Out.Flush();
+        }
+
         // TODO At some point we may want parameters to choose different test bodies.
         private static async Task TestBodyCore(ILogger customCategoryLogger, ILogger appCategoryLogger, ActivitySource activitySource)
         {

src/tests/Microsoft.Diagnostics.Monitoring.EventPipe/DistributedTracesPipelineUnitTests.cs

@@ -32,9 +32,11 @@ public DistributedTracesPipelineUnitTests(ITestOutputHelper output)
         [SkippableTheory, MemberData(nameof(Configurations))]
         public async Task TestTracesPipeline(TestConfiguration config)
         {
-            TestActivityLogger logger = new();
+            TaskCompletionSource<object?> pipelineStartedSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            TaskCompletionSource<object?> activityLoggedSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            TestActivityLogger logger = new(pipelineStartedSource, activityLoggedSource);
 
-            await using (TestRunner testRunner = await PipelineTestUtilities.StartProcess(config, "TracesRemoteTest UseActivitySource", _output))
+            await using (TestRunner testRunner = await PipelineTestUtilities.StartProcess(config, "TracesRemoteTest UseActivitySource WaitForActivitySourceListener", _output))
             {
                 DiagnosticsClient client = new(testRunner.Pid);
 
@@ -47,7 +49,9 @@ public async Task TestTracesPipeline(TestConfiguration config)
 
                 await PipelineTestUtilities.ExecutePipelineWithTracee(
                     pipeline,
-                    testRunner);
+                    testRunner,
+                    activityLoggedSource,
+                    pipelineStartedSource.Task);
             }
 
             Assert.Single(logger.LoggedActivities);
@@ -83,9 +87,11 @@ await PipelineTestUtilities.ExecutePipelineWithTracee(
         [SkippableTheory, MemberData(nameof(Configurations))]
         public async Task TestTracesPipelineWithSamplingRatio(TestConfiguration config)
         {
-            TestActivityLogger logger = new();
+            TaskCompletionSource<object?> pipelineStartedSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            TaskCompletionSource<object?> activityLoggedSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            TestActivityLogger logger = new(pipelineStartedSource, activityLoggedSource);
 
-            await using (TestRunner testRunner = await PipelineTestUtilities.StartProcess(config, "TracesRemoteTest UseActivitySource", _output))
+            await using (TestRunner testRunner = await PipelineTestUtilities.StartProcess(config, "TracesRemoteTest UseActivitySource WaitForActivitySourceListener", _output))
             {
                 DiagnosticsClient client = new(testRunner.Pid);
 
@@ -98,14 +104,11 @@ public async Task TestTracesPipelineWithSamplingRatio(TestConfiguration config)
 
                 await PipelineTestUtilities.ExecutePipelineWithTracee(
                     pipeline,
-                    testRunner);
+                    testRunner,
+                    activityLoggedSource,
+                    pipelineStartedSource.Task);
             }
 
-            // ParentRatioSampler used to be unavailable on DiagnosticSource < 9,
-            // which caused the sampler portion of the FilterAndPayloadSpec to be
-            // ignored and no activities to flow. Newer net 8 servicing runtimes
-            // honor it, so all supported TFMs now produce a single (un-recorded)
-            // activity here.
             Assert.Single(logger.LoggedActivities);
 
             var activityData = logger.LoggedActivities[0];
@@ -134,16 +137,32 @@ await PipelineTestUtilities.ExecutePipelineWithTracee(
 
         private sealed class TestActivityLogger : IActivityLogger
         {
+            private readonly TaskCompletionSource<object?> _pipelineStartedSource;
+            private readonly TaskCompletionSource<object?> _activityLoggedSource;
+
+            public TestActivityLogger(
+                TaskCompletionSource<object?> pipelineStartedSource,
+                TaskCompletionSource<object?> activityLoggedSource)
+            {
+                _pipelineStartedSource = pipelineStartedSource;
+                _activityLoggedSource = activityLoggedSource;
+            }
+
             public List<(ActivityData, KeyValuePair<string, object?>[])> LoggedActivities { get; } = new();
 
             public void Log(
                 in ActivityData activity,
                 ReadOnlySpan<KeyValuePair<string, object?>> tags)
             {
                 LoggedActivities.Add((activity, tags.ToArray()));
+                _activityLoggedSource.TrySetResult(null);
             }
 
-            public Task PipelineStarted(CancellationToken token) => Task.CompletedTask;
+            public Task PipelineStarted(CancellationToken token)
+            {
+                _pipelineStartedSource.TrySetResult(null);
+                return Task.CompletedTask;
+            }
 
             public Task PipelineStopped(CancellationToken token) => Task.CompletedTask;
         }

src/tests/Microsoft.Diagnostics.Monitoring.EventPipe/PipelineTestUtilities.cs

@@ -37,7 +37,8 @@ await ExecutePipelineWithTracee(
         public static async Task ExecutePipelineWithTracee<T>(
             EventSourcePipeline<T> pipeline,
             TestRunner testRunner,
-            TaskCompletionSource<object> waitTaskSource = null)
+            TaskCompletionSource<object> waitTaskSource = null,
+            Task pipelineStartedTask = null)
             where T : EventSourcePipelineSettings
         {
             using CancellationTokenSource cancellation = new(DefaultPipelineRunTimeout);
@@ -47,7 +48,8 @@ await ExecutePipelineWithTracee(
                 (p, t) => p.StartAsync(t),
                 testRunner,
                 cancellation.Token,
-                waitTaskSource);
+                waitTaskSource,
+                pipelineStartedTask);
         }
 
         public static Task ExecutePipelineWithTracee(
@@ -69,7 +71,8 @@ private static async Task ExecutePipelineWithTracee<TPipeline>(
             Func<TPipeline, CancellationToken, Task<Task>> startPipelineAsync,
             TestRunner testRunner,
             CancellationToken token,
-            TaskCompletionSource<object> waitTaskSource = null)
+            TaskCompletionSource<object> waitTaskSource = null,
+            Task pipelineStartedTask = null)
             where TPipeline : Pipeline
         {
             Task runTask;
@@ -83,6 +86,20 @@ private static async Task ExecutePipelineWithTracee<TPipeline>(
                 throw;
             }
 
+            if (pipelineStartedTask != null)
+            {
+                // The startup callback only completes on successful initialization, so also
+                // observe pipeline failure and the test timeout while waiting for readiness.
+                Task cancellationTask = Task.Delay(Timeout.InfiniteTimeSpan, token);
+                Task completedTask = await Task.WhenAny(pipelineStartedTask, runTask, cancellationTask).ConfigureAwait(false);
+                await completedTask.ConfigureAwait(false);
+
+                if (completedTask == runTask)
+                {
+                    throw new InvalidOperationException("Pipeline completed before signaling startup.");
+                }
+            }
+
             Func<CancellationToken, Task> waitForPipeline = async (cancellationToken) => {
                 // Optionally wait on caller before allowing the pipeline to stop.
                 if (null != waitTaskSource)