Fix HeadlessUnitTestSession creation race condition (#12979)

* Fix HeadlessUnitTestSession creation race condition

* Bind Compositor/MediaContext to a fixed UI thread

* Fix dead lock in AvaloniaTestCase

* Rename Compositor.UIThreadDispatcher to Dispatcher

src/Avalonia.Base/Media/MediaContext.Clock.cs

@@ -3,8 +3,6 @@ using System.Collections.Generic;
 using System.Diagnostics;
 using Avalonia.Animation;
 using Avalonia.Reactive;
-using Avalonia.Threading;
-using Avalonia.Utilities;
 
 namespace Avalonia.Media;
 
@@ -33,12 +31,12 @@ internal partial class MediaContext
         public IDisposable Subscribe(IObserver<TimeSpan> observer)
         {
             _parent.ScheduleRender(false);
-            Dispatcher.UIThread.VerifyAccess();
+            _parent._dispatcher.VerifyAccess();
             _observers.Add(observer);
             _newObservers.Add(observer);
             return Disposable.Create(() =>
             {
-                Dispatcher.UIThread.VerifyAccess();
+                _parent._dispatcher.VerifyAccess();
                 _observers.Remove(observer);
             });
         }
@@ -79,4 +77,4 @@ internal partial class MediaContext
     }
 
     public void RequestAnimationFrame(Action<TimeSpan> action) => _clock.RequestAnimationFrame(action);
-}
+}

src/Avalonia.Base/Media/MediaContext.Compositor.cs

@@ -1,6 +1,4 @@
-using System;
 using System.Linq;
-using System.Threading.Tasks;
 using Avalonia.Platform;
 using Avalonia.Rendering.Composition;
 using Avalonia.Rendering.Composition.Transport;
@@ -22,7 +20,7 @@ partial class MediaContext
         _requestedCommits.Remove(compositor);
         _pendingCompositionBatches[compositor] = commit;
         commit.Processed.ContinueWith(_ =>
-            Dispatcher.UIThread.Post(() => CompositionBatchFinished(compositor, commit), DispatcherPriority.Send));
+            _dispatcher.Post(() => CompositionBatchFinished(compositor, commit), DispatcherPriority.Send));
         return commit;
     }
     
@@ -147,4 +145,4 @@ partial class MediaContext
         // TODO: maybe skip the full render here?
         ScheduleRender(true);
     }
-}
+}

src/Avalonia.Base/Media/MediaContext.cs

@@ -1,8 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.Linq;
-using System.Threading.Tasks;
-using Avalonia.Animation;
 using Avalonia.Layout;
 using Avalonia.Rendering;
 using Avalonia.Rendering.Composition;
@@ -23,6 +20,7 @@ internal partial class MediaContext : ICompositorScheduler
     private readonly Action _inputMarkerHandler;
     private readonly HashSet<Compositor> _requestedCommits = new();
     private readonly Dictionary<Compositor, CompositionBatch> _pendingCompositionBatches = new();
+    private readonly Dispatcher _dispatcher;
     private record  TopLevelInfo(Compositor Compositor, CompositingRenderer Renderer, ILayoutManager LayoutManager);
 
     private List<Action>? _invokeOnRenderCallbacks;
@@ -38,11 +36,12 @@ internal partial class MediaContext : ICompositorScheduler
 
     private Dictionary<object, TopLevelInfo> _topLevels = new();
 
-    private MediaContext()
+    private MediaContext(Dispatcher dispatcher)
     {
         _render = Render;
         _inputMarkerHandler = InputMarkerHandler;
         _clock = new(this);
+        _dispatcher = dispatcher;
         _animationsTimer.Tick += (_, _) =>
         {
             _animationsTimer.Stop();
@@ -58,7 +57,7 @@ internal partial class MediaContext : ICompositorScheduler
             // and need to do a full reset for unit tests
             var context = AvaloniaLocator.Current.GetService<MediaContext>();
             if (context == null)
-                AvaloniaLocator.CurrentMutable.Bind<MediaContext>().ToConstant(context = new());
+                AvaloniaLocator.CurrentMutable.Bind<MediaContext>().ToConstant(context = new(Dispatcher.UIThread));
             return context;
         }
     }
@@ -84,7 +83,7 @@ internal partial class MediaContext : ICompositorScheduler
         
         if (_inputMarkerOp == null)
         {
-            _inputMarkerOp = Dispatcher.UIThread.InvokeAsync(_inputMarkerHandler, DispatcherPriority.Input);
+            _inputMarkerOp = _dispatcher.InvokeAsync(_inputMarkerHandler, DispatcherPriority.Input);
             _inputMarkerAddedAt = _time.Elapsed;
         }
         else if (!now && (_time.Elapsed - _inputMarkerAddedAt).TotalSeconds > MaxSecondsWithoutInput)
@@ -93,7 +92,7 @@ internal partial class MediaContext : ICompositorScheduler
         }
         
 
-        _nextRenderOp = Dispatcher.UIThread.InvokeAsync(_render, priority);
+        _nextRenderOp = _dispatcher.InvokeAsync(_render, priority);
     }
     
     /// <summary>

src/Avalonia.Base/Rendering/Composition/CompositionDrawingSurface.cs

@@ -56,7 +56,7 @@ public class CompositionDrawingSurface : CompositionSurface, IDisposable
 
     ~CompositionDrawingSurface()
     {
-        Dispatcher.UIThread.Post(Dispose);
+        Compositor.Dispatcher.Post(Dispose);
     }
 
     public new void Dispose() => base.Dispose();

src/Avalonia.Base/Rendering/Composition/Compositor.cs

@@ -1,13 +1,10 @@
 using System;
 using System.Collections.Generic;
-using System.IO;
-using System.Numerics;
 using System.Threading.Tasks;
 using Avalonia.Animation.Easings;
 using Avalonia.Media;
 using Avalonia.Metadata;
 using Avalonia.Platform;
-using Avalonia.Rendering.Composition.Animations;
 using Avalonia.Rendering.Composition.Server;
 using Avalonia.Rendering.Composition.Transport;
 using Avalonia.Threading;
@@ -40,6 +37,8 @@ namespace Avalonia.Rendering.Composition
 
         internal IEasing DefaultEasing { get; }
 
+        internal Dispatcher Dispatcher { get; }
+
         private DiagnosticTextRenderer? DiagnosticTextRenderer
         {
             get
@@ -69,16 +68,18 @@ namespace Avalonia.Rendering.Composition
         }
 
         internal Compositor(IRenderLoop loop, IPlatformGraphics? gpu, bool useUiThreadForSynchronousCommits = false)
-            : this(loop, gpu, useUiThreadForSynchronousCommits, MediaContext.Instance, false)
+            : this(loop, gpu, useUiThreadForSynchronousCommits, MediaContext.Instance, false, Dispatcher.UIThread)
         {
         }
 
         internal Compositor(IRenderLoop loop, IPlatformGraphics? gpu,
             bool useUiThreadForSynchronousCommits,
-            ICompositorScheduler scheduler, bool reclaimBuffersImmediately)
+            ICompositorScheduler scheduler, bool reclaimBuffersImmediately,
+            Dispatcher dispatcher)
         {
             Loop = loop;
             UseUiThreadForSynchronousCommits = useUiThreadForSynchronousCommits;
+            Dispatcher = dispatcher;
             _batchMemoryPool = new(reclaimBuffersImmediately);
             _batchObjectPool = new(reclaimBuffersImmediately);
             _server = new ServerCompositor(loop, gpu, _batchObjectPool, _batchMemoryPool);
@@ -99,14 +100,14 @@ namespace Avalonia.Rendering.Composition
         /// <returns>A CompositionBatch object that provides batch lifetime information</returns>
         public CompositionBatch RequestCompositionBatchCommitAsync()
         {
-            Dispatcher.UIThread.VerifyAccess();
+            Dispatcher.VerifyAccess();
             if (_nextCommit == null)
             {
                 _nextCommit = new ();
                 var pending = _pendingBatch;
                 if (pending != null)
                     pending.Processed.ContinueWith(
-                        _ => Dispatcher.UIThread.Post(_triggerCommitRequested, DispatcherPriority.Send));
+                        _ => Dispatcher.Post(_triggerCommitRequested, DispatcherPriority.Send));
                 else
                     _triggerCommitRequested();
             }
@@ -130,7 +131,7 @@ namespace Avalonia.Rendering.Composition
         
         CompositionBatch CommitCore()
         {
-            Dispatcher.UIThread.VerifyAccess();
+            Dispatcher.VerifyAccess();
             using var noPump = NonPumpingLockHelper.Use();
             
             var commit = _nextCommit ??= new();
@@ -197,7 +198,7 @@ namespace Avalonia.Rendering.Composition
 
         internal void RegisterForSerialization(ICompositorSerializable compositionObject)
         {
-            Dispatcher.UIThread.VerifyAccess();
+            Dispatcher.VerifyAccess();
             if(_objectSerializationHashSet.Add(compositionObject))
                 _objectSerializationQueue.Enqueue(compositionObject);
             RequestCommitAsync();
@@ -217,14 +218,14 @@ namespace Avalonia.Rendering.Composition
         /// </summary>
         public void RequestCompositionUpdate(Action action)
         {
-            Dispatcher.UIThread.VerifyAccess();
+            Dispatcher.VerifyAccess();
             _invokeBeforeCommitWrite.Enqueue(action);
             RequestCommitAsync();
         }
 
         internal void PostServerJob(Action job)
         {
-            Dispatcher.UIThread.VerifyAccess();
+            Dispatcher.VerifyAccess();
             _pendingServerCompositorJobs.Add(job);
             RequestCommitAsync();
         }

src/Avalonia.Base/Threading/Dispatcher.Queue.cs

@@ -108,6 +108,7 @@ public partial class Dispatcher
                 job = s_uiThread._queue.Peek();
             if (job == null || job.Priority <= DispatcherPriority.Inactive)
             {
+                s_uiThread.ShutdownImpl();
                 s_uiThread = null;
                 return;
             }

src/Headless/Avalonia.Headless.XUnit/AvaloniaTestCase.cs

@@ -1,9 +1,7 @@
 using System;
 using System.ComponentModel;
-using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
-using Avalonia.Threading;
 using Xunit.Abstractions;
 using Xunit.Sdk;
 
@@ -37,10 +35,11 @@ internal class AvaloniaTestCase : XunitTestCase
 
         // We need to block the XUnit thread to ensure its concurrency throttle is effective.
         // See https://github.com/AArnott/Xunit.StaFact/pull/55#issuecomment-826187354 for details.
-        var runSummary = AvaloniaTestCaseRunner
-            .RunTest(session, this, DisplayName, SkipReason, constructorArguments,
-                TestMethodArguments, messageBus, aggregator, cancellationTokenSource)
-            .GetAwaiter().GetResult();
+        var runSummary =
+            Task.Run(() => AvaloniaTestCaseRunner.RunTest(session, this, DisplayName, SkipReason, constructorArguments,
+                TestMethodArguments, messageBus, aggregator, cancellationTokenSource))
+            .GetAwaiter()
+            .GetResult();
 
         return Task.FromResult(runSummary);
     }

src/Headless/Avalonia.Headless.XUnit/AvaloniaTestCaseRunner.cs

@@ -11,15 +11,17 @@ namespace Avalonia.Headless.XUnit;
 
 internal class AvaloniaTestCaseRunner : XunitTestCaseRunner
 {
+    private readonly HeadlessUnitTestSession _session;
     private readonly Action? _onAfterTestInvoked;
 
     public AvaloniaTestCaseRunner(
-        Action? onAfterTestInvoked,
+        HeadlessUnitTestSession session, Action? onAfterTestInvoked,
         IXunitTestCase testCase, string displayName, string skipReason, object[] constructorArguments,
         object[] testMethodArguments, IMessageBus messageBus, ExceptionAggregator aggregator,
         CancellationTokenSource cancellationTokenSource) : base(testCase, displayName, skipReason, constructorArguments,
         testMethodArguments, messageBus, aggregator, cancellationTokenSource)
     {
+        _session = session;
         _onAfterTestInvoked = onAfterTestInvoked;
     }
 
@@ -29,43 +31,46 @@ internal class AvaloniaTestCaseRunner : XunitTestCaseRunner
         CancellationTokenSource cancellationTokenSource)
     {
         var afterTest = () => Dispatcher.UIThread.RunJobs();
-        return session.Dispatch(async () =>
-        {
-            var runner = new AvaloniaTestCaseRunner(afterTest, testCase, displayName,
-                skipReason, constructorArguments, testMethodArguments, messageBus, aggregator, cancellationTokenSource);
-            return await runner.RunAsync();
-        }, cancellationTokenSource.Token);
+
+        var runner = new AvaloniaTestCaseRunner(session, afterTest, testCase, displayName,
+            skipReason, constructorArguments, testMethodArguments, messageBus, aggregator, cancellationTokenSource);
+        return runner.RunAsync();
     }
-    
+
     protected override XunitTestRunner CreateTestRunner(ITest test, IMessageBus messageBus, Type testClass,
         object[] constructorArguments,
         MethodInfo testMethod, object[] testMethodArguments, string skipReason,
         IReadOnlyList<BeforeAfterTestAttribute> beforeAfterAttributes,
         ExceptionAggregator aggregator, CancellationTokenSource cancellationTokenSource)
     {
-        return new AvaloniaTestRunner(_onAfterTestInvoked, test, messageBus, testClass, constructorArguments,
+        return new AvaloniaTestRunner(_session, _onAfterTestInvoked, test, messageBus, testClass, constructorArguments,
             testMethod, testMethodArguments, skipReason, beforeAfterAttributes, aggregator, cancellationTokenSource);
     }
 
     private class AvaloniaTestRunner : XunitTestRunner
     {
+        private readonly HeadlessUnitTestSession _session;
         private readonly Action? _onAfterTestInvoked;
 
         public AvaloniaTestRunner(
-            Action? onAfterTestInvoked,
+            HeadlessUnitTestSession session, Action? onAfterTestInvoked,
             ITest test, IMessageBus messageBus, Type testClass, object[] constructorArguments, MethodInfo testMethod,
             object[] testMethodArguments, string skipReason,
             IReadOnlyList<BeforeAfterTestAttribute> beforeAfterAttributes, ExceptionAggregator aggregator,
             CancellationTokenSource cancellationTokenSource) : base(test, messageBus, testClass, constructorArguments,
             testMethod, testMethodArguments, skipReason, beforeAfterAttributes, aggregator, cancellationTokenSource)
         {
+            _session = session;
             _onAfterTestInvoked = onAfterTestInvoked;
         }
 
         protected override Task<decimal> InvokeTestMethodAsync(ExceptionAggregator aggregator)
         {
-            return new AvaloniaTestInvoker(_onAfterTestInvoked, Test, MessageBus, TestClass, ConstructorArguments,
-                TestMethod, TestMethodArguments, BeforeAfterAttributes, aggregator, CancellationTokenSource).RunAsync();
+            return _session.Dispatch(
+                () => new AvaloniaTestInvoker(_onAfterTestInvoked, Test, MessageBus, TestClass,
+                    ConstructorArguments, TestMethod, TestMethodArguments, BeforeAfterAttributes, aggregator,
+                    CancellationTokenSource).RunAsync(),
+                CancellationTokenSource.Token);
         }
     }
 

src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs

@@ -3,13 +3,10 @@ using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Reflection;
-using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
-using Avalonia.Controls.Platform;
 using Avalonia.Metadata;
 using Avalonia.Reactive;
-using Avalonia.Rendering;
 using Avalonia.Threading;
 
 namespace Avalonia.Headless;
@@ -22,7 +19,7 @@ namespace Avalonia.Headless;
 [Unstable("This API is experimental and might be unstable. Use on your risk. API might or might not be changed in a minor update.")]
 public sealed class HeadlessUnitTestSession : IDisposable
 {
-    private static readonly ConcurrentDictionary<Assembly, HeadlessUnitTestSession> s_session = new();
+    private static readonly Dictionary<Assembly, HeadlessUnitTestSession> s_session = new();
 
     private readonly AppBuilder _appBuilder;
     private readonly CancellationTokenSource _cancellationTokenSource;
@@ -61,7 +58,7 @@ public sealed class HeadlessUnitTestSession : IDisposable
 
     /// <summary>
     /// Dispatch method queues an async operation on the dispatcher thread, creates a new application instance,
-    /// setting app avalonia services, and runs <see cref="action"/> parameter.
+    /// setting app avalonia services, and runs <paramref name="action"/> parameter.
     /// </summary>
     /// <param name="action">Action to execute on the dispatcher thread with avalonia services.</param>
     /// <param name="cancellationToken">Cancellation token to cancel execution.</param>
@@ -80,20 +77,21 @@ public sealed class HeadlessUnitTestSession : IDisposable
         var tcs = new TaskCompletionSource<TResult>();
         _queue.Add(() =>
         {
-            using var application = EnsureApplication();
-
             var cts = new CancellationTokenSource();
             using var globalCts = token.Register(s => ((CancellationTokenSource)s!).Cancel(), cts, true);
             using var localCts = cancellationToken.Register(s => ((CancellationTokenSource)s!).Cancel(), cts, true);
 
             try
             {
+                using var application = EnsureApplication();
+
                 var task = action();
                 task.ContinueWith((_, s) => ((CancellationTokenSource)s!).Cancel(), cts,
                     TaskScheduler.FromCurrentSynchronizationContext());
 
                 if (cts.IsCancellationRequested)
                 {
+                    tcs.TrySetCanceled(cts.Token);
                     return;
                 }
 
@@ -202,13 +200,23 @@ public sealed class HeadlessUnitTestSession : IDisposable
         Justification = "AvaloniaTestApplicationAttribute attribute should preserve type information.")]
     public static HeadlessUnitTestSession GetOrStartForAssembly(Assembly? assembly)
     {
-        return s_session.GetOrAdd(assembly ?? typeof(HeadlessUnitTestSession).Assembly, a =>
+        assembly ??= typeof(HeadlessUnitTestSession).Assembly;
+
+        lock (s_session)
         {
-            var appBuilderEntryPointType = a.GetCustomAttribute<AvaloniaTestApplicationAttribute>()
-                ?.AppBuilderEntryPointType;
-            return appBuilderEntryPointType is not null ?
-                StartNew(appBuilderEntryPointType) :
-                StartNew(typeof(Application));
-        });
+            if (!s_session.TryGetValue(assembly, out var session))
+            {
+                var appBuilderEntryPointType = assembly.GetCustomAttribute<AvaloniaTestApplicationAttribute>()
+                    ?.AppBuilderEntryPointType;
+
+                session = appBuilderEntryPointType is not null ?
+                    StartNew(appBuilderEntryPointType) :
+                    StartNew(typeof(Application));
+
+                s_session.Add(assembly, session);
+            }
+
+            return session;
+        }
     }
 }

tests/Avalonia.Headless.XUnit.UnitTests/AssemblyInfo.cs

@@ -2,6 +2,5 @@
 global using Avalonia.Headless.XUnit;
 using Avalonia.Headless;
 using Avalonia.Headless.UnitTests;
-using Avalonia.Headless.XUnit;
 
 [assembly: AvaloniaTestApplication(typeof(TestApplication))]

tests/Avalonia.RenderTests/TestBase.cs

@@ -108,7 +108,7 @@ namespace Avalonia.Direct2D1.RenderTests
             var timer = new ManualRenderTimer();
 
             var compositor = new Compositor(new RenderLoop(timer), null, true,
-                new DispatcherCompositorScheduler(), true);
+                new DispatcherCompositorScheduler(), true, Dispatcher.UIThread);
             using (var writableBitmap = factory.CreateWriteableBitmap(pixelSize, dpiVector, factory.DefaultPixelFormat, factory.DefaultAlphaFormat))
             {
                 var root = new TestRenderRoot(dpiVector.X / 96, null!);

tests/Avalonia.UnitTests/CompositorTestServices.cs

@@ -46,7 +46,7 @@ public class CompositorTestServices : IDisposable
             AvaloniaLocator.CurrentMutable.Bind<IRenderTimer>().ToConstant(Timer);
 
             Compositor = new Compositor(new RenderLoop(Timer), null,
-                true, new DispatcherCompositorScheduler(), true);
+                true, new DispatcherCompositorScheduler(), true, Dispatcher.UIThread);
             var impl = new TopLevelImpl(Compositor, size ?? new Size(1000, 1000));
             TopLevel = new EmbeddableControlRoot(impl)
             {

tests/Avalonia.UnitTests/RendererMocks.cs

@@ -18,7 +18,7 @@ namespace Avalonia.UnitTests
 
         public static Compositor CreateDummyCompositor() =>
             new(new RenderLoop(new CompositorTestServices.ManualRenderTimer()), null, false,
-                new CompositionCommitScheduler(), true);
+                new CompositionCommitScheduler(), true, Dispatcher.UIThread);
 
         class CompositionCommitScheduler : ICompositorScheduler
         {