[Profiler] Add Flaky attribute (#8648)

Author: gleocadie

profiler/test/Datadog.Profiler.IntegrationTests/Exceptions/ExceptionsTest.cs

@@ -9,6 +9,7 @@
 using System.Linq;
 using System.Text;
 using Datadog.Profiler.IntegrationTests.Helpers;
+using Datadog.Profiler.IntegrationTests.Xunit;
 using FluentAssertions;
 using Xunit;
 using Xunit.Abstractions;
@@ -29,6 +30,7 @@ public ExceptionsTest(ITestOutputHelper output)
             _output = output;
         }
 
+        [Flaky("Flaky on ARM64")]
         [TestAppFact("Samples.ExceptionGenerator")]
         public void ThrowExceptionsInParallel(string appName, string framework, string appAssembly)
         {
@@ -145,6 +147,7 @@ public void ThrowExceptionsInParallel(string appName, string framework, string a
             }
         }
 
+        [Flaky("Flaky on ARM64")]
         [TestAppFact("Samples.ExceptionGenerator")]
         public void ThrowExceptionsInParallelWithCustomGetFunctionFromIp(string appName, string framework, string appAssembly)
         {
@@ -262,6 +265,7 @@ public void ThrowExceptionsInParallelWithCustomGetFunctionFromIp(string appName,
             }
         }
 
+        [Flaky("Flaky on ARM64")]
         [Trait("Category", "LinuxOnly")]
         [TestAppFact("Samples.ExceptionGenerator", new[] { "net48", "netcoreapp3.1", "net6.0", "net8.0", })] // FIXME: .NET 9 skipping .NET 9 for now
         public void ThrowExceptionsInParallelWithNewCpuProfiler(string appName, string framework, string appAssembly)
@@ -374,6 +378,7 @@ public void Sampling(string appName, string framework, string appAssembly)
             exceptionCounts.Should().ContainKey("System.InvalidOperationException").WhoseValue.Should().Be(1);
         }
 
+        [Flaky("Flaky on ARM64")]
         [TestAppFact("Samples.ExceptionGenerator", new[] { "net48", "netcoreapp3.1", "net6.0", "net8.0", })] // FIXME: .NET 9 skipping .NET 9 for now
         public void GetExceptionSamplesWithTimestamp(string appName, string framework, string appAssembly)
         {
@@ -385,6 +390,7 @@ public void GetExceptionSamplesWithTimestamp(string appName, string framework, s
             CheckExceptionProfiles(runner, true);
         }
 
+        [Flaky("Flaky on ARM64")]
         [TestAppFact("Samples.ExceptionGenerator", new[] { "net48", "netcoreapp3.1", "net6.0", "net8.0", })] // FIXME: .NET 9 skipping .NET 9 for now
         public void GetExceptionSamplesWithoutTimestamp(string appName, string framework, string appAssembly)
         {

profiler/test/Datadog.Profiler.IntegrationTests/Xunit/DelayedMessageBus.cs

@@ -0,0 +1,41 @@
+// <copyright file="DelayedMessageBus.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2022 Datadog, Inc.
+// </copyright>
+
+using System.Collections.Concurrent;
+using Xunit.Abstractions;
+using Xunit.Sdk;
+
+namespace Datadog.Profiler.IntegrationTests.Xunit;
+
+/// <summary>
+/// Queues messages instead of forwarding them immediately.
+/// On <see cref="Dispose"/>, all queued messages are flushed to the inner bus.
+/// Used by <see cref="ProfilerTestCase"/> to discard intermediate failure messages on retry.
+/// Based on https://github.com/xunit/samples.xunit/blob/main/v2/RetryFactExample/DelayedMessageBus.cs
+/// </summary>
+internal class DelayedMessageBus : IMessageBus
+{
+    private readonly IMessageBus _innerBus;
+    private readonly ConcurrentQueue<IMessageSinkMessage> _messages = new();
+
+    public DelayedMessageBus(IMessageBus innerBus)
+    {
+        _innerBus = innerBus;
+    }
+
+    public bool QueueMessage(IMessageSinkMessage message)
+    {
+        _messages.Enqueue(message);
+        return true;
+    }
+
+    public void Dispose()
+    {
+        while (_messages.TryDequeue(out var message))
+        {
+            _innerBus.QueueMessage(message);
+        }
+    }
+}

profiler/test/Datadog.Profiler.IntegrationTests/Xunit/FlakyAttribute.cs

@@ -0,0 +1,21 @@
+// <copyright file="FlakyAttribute.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2022 Datadog, Inc.
+// </copyright>
+
+using System;
+
+namespace Datadog.Profiler.IntegrationTests.Xunit;
+
+/// <summary>
+/// Marks a test as flaky, so that it is automatically retried by <see cref="ProfilerTestCase"/>.
+/// </summary>
+/// <param name="reason">The reason that this test was marked flaky.</param>
+/// <param name="maxRetries">The maximum number of times a test should be retried (default 5).</param>
+[AttributeUsage(AttributeTargets.Method, Inherited = false)]
+public class FlakyAttribute(string reason, byte maxRetries = 5) : Attribute
+{
+    public string Reason { get; } = reason;
+
+    public byte MaxRetries { get; } = maxRetries;
+}

profiler/test/Datadog.Profiler.IntegrationTests/Xunit/ProfilerTestCase.cs

@@ -4,6 +4,7 @@
 // </copyright>
 
 using System;
+using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit.Abstractions;
@@ -29,12 +30,62 @@ public ProfilerTestCase(
         {
         }
 
-        public override Task<RunSummary> RunAsync(
+        public override async Task<RunSummary> RunAsync(
             IMessageSink diagnosticMessageSink,
             IMessageBus messageBus,
             object[] constructorArguments,
             ExceptionAggregator aggregator,
             CancellationTokenSource cancellationTokenSource)
-           => new ProfilerTestCaseRunner(this, DisplayName, SkipReason, constructorArguments, TestMethodArguments, messageBus, aggregator, cancellationTokenSource).RunAsync();
+        {
+            int attemptsRemaining = 1;
+            var retryReason = string.Empty;
+
+            try
+            {
+                var flakyAttribute = TestMethod.Method.ToRuntimeMethod().GetCustomAttribute<FlakyAttribute>();
+                if (flakyAttribute is not null)
+                {
+                    // First attempt + retries
+                    attemptsRemaining = flakyAttribute.MaxRetries + 1;
+                    retryReason = flakyAttribute.Reason;
+                }
+            }
+            catch (Exception e)
+            {
+                diagnosticMessageSink.OnMessage(new DiagnosticMessage($"ERROR: Looking for FlakyAttribute: {e}"));
+            }
+
+            if (attemptsRemaining <= 1)
+            {
+                return await RunOnceAsync(messageBus);
+            }
+
+            DelayedMessageBus delayedBus = null;
+            try
+            {
+                while (true)
+                {
+                    attemptsRemaining--;
+                    delayedBus = new DelayedMessageBus(messageBus);
+
+                    var summary = await RunOnceAsync(delayedBus);
+
+                    if (summary.Failed == 0 || attemptsRemaining <= 0)
+                    {
+                        return summary;
+                    }
+
+                    diagnosticMessageSink.OnMessage(
+                        new DiagnosticMessage($"RETRYING: {DisplayName} ({attemptsRemaining} attempts remaining, {retryReason})"));
+                }
+            }
+            finally
+            {
+                delayedBus?.Dispose();
+            }
+
+            Task<RunSummary> RunOnceAsync(IMessageBus bus)
+                => new ProfilerTestCaseRunner(this, DisplayName, SkipReason, constructorArguments, TestMethodArguments, bus, aggregator, cancellationTokenSource).RunAsync();
+        }
     }
 }