Implemented ExitAllFrames, DisableProcessing and Shutdown, dispatcher callbacks now have proper sync context

src/Avalonia.Base/Threading/AvaloniaSynchronizationContext.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Runtime.ConstrainedExecution;
 using System.Threading;
+using Avalonia.Utilities;
 
 namespace Avalonia.Threading
 {
@@ -9,6 +10,28 @@ namespace Avalonia.Threading
     /// </summary>
     public class AvaloniaSynchronizationContext : SynchronizationContext
     {
+        internal readonly DispatcherPriority Priority;
+        private readonly NonPumpingLockHelper.IHelperImpl? _nonPumpingHelper =
+            AvaloniaLocator.Current.GetService<NonPumpingLockHelper.IHelperImpl>();
+        
+        public AvaloniaSynchronizationContext():  this(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA)
+        {
+            
+        }
+        
+        // This constructor is here to enforce STA behavior for unit tests
+        internal AvaloniaSynchronizationContext(bool isStaThread)
+        {
+            if (_nonPumpingHelper != null 
+                && isStaThread)
+                SetWaitNotificationRequired();
+        }
+
+        public AvaloniaSynchronizationContext(DispatcherPriority priority)
+        {
+            Priority = priority;
+        }
+        
         /// <summary>
         /// Controls if SynchronizationContext should be installed in InstallIfNeeded. Used by Designer.
         /// </summary>
@@ -24,13 +47,13 @@ namespace Avalonia.Threading
                 return;
             }
 
-            SetSynchronizationContext(new AvaloniaSynchronizationContext());
+            SetSynchronizationContext(Dispatcher.UIThread.GetContextWithPriority(DispatcherPriority.Normal));
         }
 
         /// <inheritdoc/>
         public override void Post(SendOrPostCallback d, object? state)
         {
-            Dispatcher.UIThread.Post(d, state, DispatcherPriority.Background);
+            Dispatcher.UIThread.Post(d, state, Priority);
         }
 
         /// <inheritdoc/>
@@ -41,7 +64,50 @@ namespace Avalonia.Threading
             else
                 Dispatcher.UIThread.InvokeAsync(() => d(state), DispatcherPriority.Send).GetAwaiter().GetResult();
         }
+        
+#if !NET6_0_OR_GREATER
+        [PrePrepareMethod]
+#endif
+        public override int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout)
+        {
+            if (
+                _nonPumpingHelper != null
+                && Dispatcher.UIThread.CheckAccess() 
+                && Dispatcher.UIThread.DisabledProcessingCount > 0)
+                return _nonPumpingHelper.Wait(waitHandles, waitAll, millisecondsTimeout);
+            return base.Wait(waitHandles, waitAll, millisecondsTimeout);
+        }
 
+        public record struct RestoreContext : IDisposable
+        {
+            private readonly SynchronizationContext? _oldContext;
+            private bool _needRestore;
 
+            internal RestoreContext(SynchronizationContext? oldContext)
+            {
+                _oldContext = oldContext;
+                _needRestore = true;
+            }
+            
+            public void Dispose()
+            {
+                if (_needRestore)
+                {
+                    SetSynchronizationContext(_oldContext);
+                    _needRestore = false;
+                }
+            }
+        }
+
+        public static RestoreContext Ensure(DispatcherPriority priority)
+        {
+            if (Current is AvaloniaSynchronizationContext avaloniaContext 
+                && avaloniaContext.Priority == priority)
+                return default;
+            var oldContext = Current;
+            Dispatcher.UIThread.VerifyAccess();
+            SetSynchronizationContext(Dispatcher.UIThread.GetContextWithPriority(priority));
+            return new RestoreContext(oldContext);
+        }
     }
 }

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

@@ -107,7 +107,8 @@ public partial class Dispatcher
         // call the callback directly.
         if (!cancellationToken.IsCancellationRequested && priority == DispatcherPriority.Send && CheckAccess())
         {
-            callback();
+            using (AvaloniaSynchronizationContext.Ensure(priority))
+                callback();
             return;
         }
 
@@ -228,7 +229,8 @@ public partial class Dispatcher
         // call the callback directly.
         if (!cancellationToken.IsCancellationRequested && priority == DispatcherPriority.Send && CheckAccess())
         {
-            return callback();
+            using (AvaloniaSynchronizationContext.Ensure(priority))
+                return callback();
         }
 
         // Slow-Path: go through the queue.
@@ -555,10 +557,6 @@ public partial class Dispatcher
     /// <returns>
     ///     An task that completes after the task returned from callback finishes
     /// </returns>
-    /// <remarks>
-    /// DispatcherPriority is only respected for the initial call, any async continuations will be executed through
-    /// SynchronizationContext
-    /// </remarks>
     public Task InvokeTaskAsync(Func<Task> callback, DispatcherPriority priority = default)
     {
         _ = callback ?? throw new ArgumentNullException(nameof(callback));
@@ -580,10 +578,6 @@ public partial class Dispatcher
     /// <returns>
     ///     An task that completes after the task returned from callback finishes
     /// </returns>
-    /// <remarks>
-    /// DispatcherPriority is only respected for the initial call, any async continuations will be executed through
-    /// SynchronizationContext
-    /// </remarks>
     public Task<TResult> InvokeTaskAsync<TResult>(Func<Task<TResult>> action, DispatcherPriority priority = default)
     {
         _ = action ?? throw new ArgumentNullException(nameof(action));

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

@@ -0,0 +1,210 @@
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using Avalonia.Utilities;
+
+namespace Avalonia.Threading;
+
+public partial class Dispatcher
+{
+    internal bool ExitAllFramesRequested { get; private set; }
+    internal bool HasShutdownStarted { get; private set; }
+    internal int DisabledProcessingCount { get; private set; }
+    private bool _hasShutdownFinished;
+    private bool _startingShutdown;
+    
+    private Stack<DispatcherFrame> _frames = new();
+    
+
+    /// <summary>
+    ///     Raised when the dispatcher is shutting down.
+    /// </summary>
+    public event EventHandler? ShutdownStarted;
+
+    /// <summary>
+    ///     Raised when the dispatcher is shut down.
+    /// </summary>
+    public event EventHandler? ShutdownFinished;
+
+    /// <summary>
+    ///     Push an execution frame.
+    /// </summary>
+    /// <param name="frame">
+    ///     The frame for the dispatcher to process.
+    /// </param>
+    public void PushFrame(DispatcherFrame frame)
+    {
+        VerifyAccess();
+        if (_controlledImpl == null)
+            throw new PlatformNotSupportedException();
+        _ = frame ?? throw new ArgumentNullException(nameof(frame));
+
+        if(_hasShutdownFinished) // Dispatcher thread - no lock needed for read
+            throw new InvalidOperationException("Cannot perform requested operation because the Dispatcher shut down");
+
+        if (DisabledProcessingCount > 0)
+            throw new InvalidOperationException(
+                "Cannot perform this operation while dispatcher processing is suspended.");
+
+        try
+        {
+            _frames.Push(frame);
+            using (AvaloniaSynchronizationContext.Ensure(DispatcherPriority.Normal))
+                frame.Run(_controlledImpl);
+        }
+        finally
+        {
+            _frames.Pop();
+            if (_frames.Count == 0)
+            {
+                if (HasShutdownStarted)
+                    ShutdownImpl();
+                else
+                    ExitAllFramesRequested = false;
+            }
+        }
+    }
+    
+    /// <summary>
+    /// Runs the dispatcher's main loop.
+    /// </summary>
+    /// <param name="cancellationToken">
+    /// A cancellation token used to exit the main loop.
+    /// </param>
+    public void MainLoop(CancellationToken cancellationToken)
+    {
+        if (_controlledImpl == null)
+            throw new PlatformNotSupportedException();
+        var frame = new DispatcherFrame();
+        cancellationToken.Register(() => frame.Continue = false);
+        PushFrame(frame);
+    }
+    
+    /// <summary>
+    ///     Requests that all nested frames exit.
+    /// </summary>
+    public void ExitAllFrames()
+    {
+        if (_frames.Count == 0)
+            return;
+        ExitAllFramesRequested = true;
+        foreach (var f in _frames)
+            f.MaybeExitOnDispatcherRequest();
+    }
+    
+    /// <summary>
+    ///     Begins the process of shutting down the dispatcher.
+    /// </summary>
+    public void BeginInvokeShutdown(DispatcherPriority priority) => Post(StartShutdownImpl, priority);
+
+    private void StartShutdownImpl()
+    {
+        if (!_startingShutdown)
+        {
+            // We only need this to prevent reentrancy if the ShutdownStarted event
+            // tries to shut down again.
+            _startingShutdown = true;
+
+            // Call the ShutdownStarted event before we actually mark ourselves
+            // as shutting down.  This is so the handlers can actually do work
+            // when they get this event without throwing exceptions.
+            ShutdownStarted?.Invoke(this, EventArgs.Empty);
+
+            HasShutdownStarted = true;
+            
+            if (_frames.Count > 0)
+                ExitAllFrames();
+            else ShutdownImpl();
+        }
+    }
+
+
+    private void ShutdownImpl()
+    {
+        DispatcherOperation? operation = null;
+        _impl.Timer -= PromoteTimers;
+        _impl.Signaled -= Signaled;
+        do
+        {
+            lock (InstanceLock)
+            {
+                if (_queue.MaxPriority != DispatcherPriority.Invalid)
+                {
+                    operation = _queue.Peek();
+                }
+                else
+                {
+                    operation = null;
+                }
+            }
+
+            if (operation != null)
+            {
+                operation.Abort();
+            }
+        } while (operation != null);
+
+        _impl.UpdateTimer(null);
+        _hasShutdownFinished = true;
+        ShutdownFinished?.Invoke(this, EventArgs.Empty);
+    }
+
+    public record struct DispatcherProcessingDisabled : IDisposable
+    {
+        private readonly SynchronizationContext? _oldContext;
+        
+        private readonly bool _restoreContext;
+        private Dispatcher? _dispatcher;
+
+        internal DispatcherProcessingDisabled(Dispatcher dispatcher)
+        {
+            _dispatcher = dispatcher;
+        }
+
+        internal DispatcherProcessingDisabled(Dispatcher dispatcher, SynchronizationContext? oldContext) : this(
+            dispatcher)
+        {
+            _oldContext = oldContext;
+            _restoreContext = true;
+        }
+        
+        public void Dispose()
+        {
+            if(_dispatcher==null)
+                return;
+            _dispatcher.DisabledProcessingCount--;
+            _dispatcher = null;
+            if (_restoreContext)
+                SynchronizationContext.SetSynchronizationContext(_oldContext);
+        }
+    }
+    
+    /// <summary>
+    ///     Disable the event processing of the dispatcher.
+    /// </summary>
+    /// <remarks>
+    ///     This is an advanced method intended to eliminate the chance of
+    ///     unrelated reentrancy.  The effect of disabling processing is:
+    ///     1) CLR locks will not pump messages internally.
+    ///     2) No one is allowed to push a frame.
+    ///     3) No message processing is permitted.
+    /// </remarks>
+    public DispatcherProcessingDisabled DisableProcessing()
+    {
+        VerifyAccess();
+
+        // Turn off processing.
+        DisabledProcessingCount++;
+        var oldContext = SynchronizationContext.Current;
+        if (oldContext is AvaloniaSynchronizationContext or NonPumpingSyncContext)
+            return new DispatcherProcessingDisabled(this);
+
+        var helper = AvaloniaLocator.Current.GetService<NonPumpingLockHelper.IHelperImpl>();
+        if (helper == null)
+            return new DispatcherProcessingDisabled(this);
+
+        SynchronizationContext.SetSynchronizationContext(new NonPumpingSyncContext(helper, oldContext));
+        return new DispatcherProcessingDisabled(this, oldContext);
+
+    }
+}

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

@@ -49,6 +49,10 @@ public partial class Dispatcher
     
     internal void RunJobs(DispatcherPriority? priority, CancellationToken cancellationToken)
     {
+        if (DisabledProcessingCount > 0)
+            throw new InvalidOperationException(
+                "Cannot perform this operation while dispatcher processing is suspended.");
+        
         priority ??= DispatcherPriority.MinimumActiveValue;
         if (priority < DispatcherPriority.MinimumActiveValue)
             priority = DispatcherPriority.MinimumActiveValue;
@@ -174,7 +178,7 @@ public partial class Dispatcher
         }
     }
 
-    private bool RequestProcessing()
+    internal bool RequestProcessing()
     {
         lock (InstanceLock)
         {

src/Avalonia.Base/Threading/Dispatcher.cs

@@ -18,11 +18,13 @@ public partial class Dispatcher : IDispatcher
 {
     private IDispatcherImpl _impl;
     internal object InstanceLock { get; } = new();
-    private bool _hasShutdownFinished;
     private IControlledDispatcherImpl? _controlledImpl;
     private static Dispatcher? s_uiThread;
     private IDispatcherImplWithPendingInput? _pendingInputImpl;
-    private IDispatcherImplWithExplicitBackgroundProcessing? _backgroundProcessingImpl;
+    private readonly IDispatcherImplWithExplicitBackgroundProcessing? _backgroundProcessingImpl;
+
+    private readonly AvaloniaSynchronizationContext?[] _priorityContexts =
+        new AvaloniaSynchronizationContext?[DispatcherPriority.MaxValue - DispatcherPriority.MinValue + 1];
 
     internal Dispatcher(IDispatcherImpl impl)
     {
@@ -73,50 +75,14 @@ public partial class Dispatcher : IDispatcher
             [MethodImpl(MethodImplOptions.NoInlining)]
             static void ThrowVerifyAccess()
                 => throw new InvalidOperationException("Call from invalid thread");
-
             ThrowVerifyAccess();
         }
     }
 
-    internal void Shutdown()
-    {
-        DispatcherOperation? operation = null;
-        _impl.Timer -= PromoteTimers;
-        _impl.Signaled -= Signaled;
-        do
-        {
-            lock(InstanceLock)
-            {
-                if(_queue.MaxPriority != DispatcherPriority.Invalid)
-                {
-                    operation = _queue.Peek();
-                }
-                else
-                {
-                    operation = null;
-                }
-            }
-
-            if(operation != null)
-            {
-                operation.Abort();
-            }
-        } while(operation != null);
-        _impl.UpdateTimer(null);
-        _hasShutdownFinished = true;
-    }
-
-    /// <summary>
-    /// Runs the dispatcher's main loop.
-    /// </summary>
-    /// <param name="cancellationToken">
-    /// A cancellation token used to exit the main loop.
-    /// </param>
-    public void MainLoop(CancellationToken cancellationToken)
+    internal AvaloniaSynchronizationContext GetContextWithPriority(DispatcherPriority priority)
     {
-        if (_controlledImpl == null)
-            throw new PlatformNotSupportedException();
-        cancellationToken.Register(() => RequestProcessing());
-        _controlledImpl.RunLoop(cancellationToken);
+        DispatcherPriority.Validate(priority, nameof(priority));
+        var index = priority - DispatcherPriority.MinValue;
+        return _priorityContexts[index] ??= new(priority);
     }
 }

src/Avalonia.Base/Threading/DispatcherFrame.cs

@@ -0,0 +1,125 @@
+using System;
+using System.Threading;
+
+namespace Avalonia.Threading;
+
+/// <summary>
+///     Representation of Dispatcher frame.
+/// </summary>
+public class DispatcherFrame
+{
+    private bool _exitWhenRequested;
+    private bool _continue;
+    private bool _isRunning;
+    private CancellationTokenSource? _cancellationTokenSource;
+    
+    /// <summary>
+    ///     Constructs a new instance of the DispatcherFrame class.
+    /// </summary>
+    public DispatcherFrame() : this(true)
+    {
+    }
+
+    public Dispatcher Dispatcher { get; }
+    
+    /// <summary>
+    ///     Constructs a new instance of the DispatcherFrame class.
+    /// </summary>
+    /// <param name="exitWhenRequested">
+    ///     Indicates whether or not this frame will exit when all frames
+    ///     are requested to exit.
+    ///     <p/>
+    ///     Dispatcher frames typically break down into two categories:
+    ///     1) Long running, general purpose frames, that exit only when
+    ///        told to.  These frames should exit when requested.
+    ///     2) Short running, very specific frames that exit themselves
+    ///        when an important criteria is met.  These frames may
+    ///        consider not exiting when requested in favor of waiting
+    ///        for their important criteria to be met.  These frames
+    ///        should have a timeout associated with them.
+    /// </param>
+    public DispatcherFrame(bool exitWhenRequested)
+    {
+        Dispatcher = Dispatcher.UIThread;
+        Dispatcher.VerifyAccess();
+        _exitWhenRequested = exitWhenRequested;
+        _continue = true;
+    }
+
+    /// <summary>
+    ///     Indicates that this dispatcher frame should exit.
+    /// </summary>
+    public bool Continue
+    {
+        get
+        {
+            // This method is free-threaded.
+
+            // First check if this frame wants to continue.
+            bool shouldContinue = _continue;
+            if (shouldContinue)
+            {
+                // This frame wants to continue, so next check if it will
+                // respect the "exit requests" from the dispatcher.
+                if (_exitWhenRequested)
+                {
+                    Dispatcher dispatcher = Dispatcher;
+
+                    // This frame is willing to respect the "exit requests" of
+                    // the dispatcher, so check them.
+                    if (dispatcher.ExitAllFramesRequested || dispatcher.HasShutdownStarted)
+                    {
+                        shouldContinue = false;
+                    }
+                }
+            }
+
+            return shouldContinue;
+        }
+
+        set
+        {
+            // This method is free-threaded.
+            lock (Dispatcher.InstanceLock)
+            {
+                _continue = value;
+                if (!_continue)
+                    _cancellationTokenSource?.Cancel();
+            }
+        }
+    }
+
+    internal void Run(IControlledDispatcherImpl impl)
+    {
+        // Since the actual platform run loop is controlled by a Cancellation token, we are restarting 
+        // it if frame still needs to run
+        while (Continue)
+            RunCore(impl);
+    }
+
+    private void RunCore(IControlledDispatcherImpl impl)
+    {
+        if (_isRunning)
+            throw new InvalidOperationException("This frame is already running");
+        _isRunning = true;
+        try
+        {
+            _cancellationTokenSource = new CancellationTokenSource();
+            // Wake up the dispatcher in case it has pending jobs
+            Dispatcher.RequestProcessing();
+            impl.RunLoop(_cancellationTokenSource.Token);
+        }
+        finally
+        {
+            _isRunning = false;
+            _cancellationTokenSource?.Cancel();
+            _cancellationTokenSource = null;
+        }
+    }
+
+    internal void MaybeExitOnDispatcherRequest()
+    {
+        if (_exitWhenRequested)
+            _cancellationTokenSource?.Cancel();
+    }
+}

src/Avalonia.Base/Threading/DispatcherOperation.cs

@@ -159,7 +159,7 @@ public class DispatcherOperation
                             if (Priority >= DispatcherPriority.MinimumForegroundPriority)
                                 Dispatcher.RunJobs(Priority, cts.Token);
                             else
-                                Dispatcher.MainLoop(cts.Token);
+                                Dispatcher.PushFrame(new DispatcherOperationFrame(this, timeout));
                         }
                         else
                             Dispatcher.RunJobs(DispatcherPriority.MinimumActiveValue, cts.Token);
@@ -175,6 +175,58 @@ public class DispatcherOperation
         GetTask().GetAwaiter().GetResult();
     }
 
+    private class DispatcherOperationFrame : DispatcherFrame
+    {
+        // Note: we pass "exitWhenRequested=false" to the base
+        // DispatcherFrame construsctor because we do not want to exit
+        // this frame if the dispatcher is shutting down. This is
+        // because we may need to invoke operations during the shutdown process.
+        public DispatcherOperationFrame(DispatcherOperation op, TimeSpan timeout) : base(false)
+        {
+            _operation = op;
+
+            // We will exit this frame once the operation is completed or aborted.
+            _operation.Aborted += OnCompletedOrAborted;
+            _operation.Completed += OnCompletedOrAborted;
+
+            // We will exit the frame if the operation is not completed within
+            // the requested timeout.
+            if (timeout.TotalMilliseconds > 0)
+            {
+                _waitTimer = new Timer(_ => Exit(),
+                    null,
+                    timeout,
+                    TimeSpan.FromMilliseconds(-1));
+            }
+
+            // Some other thread could have aborted the operation while we were
+            // setting up the handlers.  We check the state again and mark the
+            // frame as "should not continue" if this happened.
+            if (_operation.Status != DispatcherOperationStatus.Pending)
+            {
+                Exit();
+            }
+        }
+
+        private void Exit()
+        {
+            Continue = false;
+
+            if (_waitTimer != null)
+            {
+                _waitTimer.Dispose();
+            }
+
+            _operation.Aborted -= OnCompletedOrAborted;
+            _operation.Completed -= OnCompletedOrAborted;
+        }
+
+        private void OnCompletedOrAborted(object? sender, EventArgs e) => Exit();
+
+        private DispatcherOperation _operation;
+        private Timer? _waitTimer;
+    }
+
     public Task GetTask() => GetTaskCore();
     
     /// <summary>
@@ -205,7 +257,8 @@ public class DispatcherOperation
 
         try
         {
-            InvokeCore();
+            using (AvaloniaSynchronizationContext.Ensure(Priority))
+                InvokeCore();
         }
         finally
         {

src/Windows/Avalonia.Win32/NonPumpingSyncContext.cs → src/Avalonia.Base/Threading/NonPumpingSyncContext.cs

@@ -2,16 +2,17 @@ using System;
 using System.Runtime.ConstrainedExecution;
 using System.Threading;
 using Avalonia.Utilities;
-using Avalonia.Win32.Interop;
 
-namespace Avalonia.Win32
+namespace Avalonia.Threading
 {
     internal class NonPumpingSyncContext : SynchronizationContext, IDisposable
     {
+        private readonly NonPumpingLockHelper.IHelperImpl _impl;
         private readonly SynchronizationContext? _inner;
 
-        private NonPumpingSyncContext(SynchronizationContext? inner)
+        public NonPumpingSyncContext(NonPumpingLockHelper.IHelperImpl impl, SynchronizationContext? inner)
         {
+            _impl = impl;
             _inner = inner;
             SetWaitNotificationRequired();
             SetSynchronizationContext(this);
@@ -48,15 +49,12 @@ namespace Avalonia.Win32
 #if !NET6_0_OR_GREATER
         [PrePrepareMethod]
 #endif
-        public override int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout)
-        {
-            return UnmanagedMethods.WaitForMultipleObjectsEx(waitHandles.Length, waitHandles, waitAll,
-                millisecondsTimeout, false);
-        }
+        public override int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout) =>
+            _impl.Wait(waitHandles, waitAll, millisecondsTimeout);
 
         public void Dispose() => SetSynchronizationContext(_inner);
 
-        public static IDisposable? Use()
+        internal static IDisposable? Use(NonPumpingLockHelper.IHelperImpl impl)
         {
             var current = Current;
             if (current == null)
@@ -67,12 +65,8 @@ namespace Avalonia.Win32
             if (current is NonPumpingSyncContext)
                 return null;
 
-            return new NonPumpingSyncContext(current);
-        }
-
-        internal class HelperImpl : NonPumpingLockHelper.IHelperImpl
-        {
-            IDisposable? NonPumpingLockHelper.IHelperImpl.Use() => NonPumpingSyncContext.Use();
+            return new NonPumpingSyncContext(impl, current);
         }
+        
     }
 }

src/Avalonia.Base/Utilities/NonPumpingLockHelper.cs

@@ -1,4 +1,5 @@
 using System;
+using Avalonia.Threading;
 
 namespace Avalonia.Utilities
 {
@@ -6,9 +7,15 @@ namespace Avalonia.Utilities
     {
         public interface IHelperImpl
         {
-            IDisposable? Use();
+            int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout);
         }
 
-        public static IDisposable? Use() => AvaloniaLocator.Current.GetService<IHelperImpl>()?.Use();
+        public static IDisposable? Use()
+        {
+            var impl = AvaloniaLocator.Current.GetService<IHelperImpl>();
+            if (impl == null)
+                return null;
+            return NonPumpingSyncContext.Use(impl);
+        }
     }
 }

src/Windows/Avalonia.Win32/Avalonia.Win32.csproj

@@ -28,7 +28,7 @@
     <NoWarn>$(NoWarn);CA1416</NoWarn>
   </PropertyGroup>
   <ItemGroup>
-    <InternalsVisibleTo Include="Avalonia.Win32.Interop, PublicKey=$(AvaloniaPublicKey)"/>
-    <InternalsVisibleTo Include="Avalonia.Direct2D1, PublicKey=$(AvaloniaPublicKey)"/>
+    <InternalsVisibleTo Include="Avalonia.Win32.Interop, PublicKey=$(AvaloniaPublicKey)" />
+    <InternalsVisibleTo Include="Avalonia.Direct2D1, PublicKey=$(AvaloniaPublicKey)" />
   </ItemGroup>
 </Project>

src/Windows/Avalonia.Win32/NonPumpingWaitHelperImpl.cs

@@ -0,0 +1,13 @@
+using System;
+using Avalonia.Utilities;
+using Avalonia.Win32.Interop;
+
+namespace Avalonia.Win32;
+
+internal class NonPumpingWaitHelperImpl : NonPumpingLockHelper.IHelperImpl
+{
+    public static NonPumpingWaitHelperImpl Instance { get; } = new();
+    public int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout) =>
+        UnmanagedMethods.WaitForMultipleObjectsEx(waitHandles.Length, waitHandles, waitAll,
+            millisecondsTimeout, false);
+}

src/Windows/Avalonia.Win32/Win32Platform.cs

@@ -171,7 +171,7 @@ namespace Avalonia.Win32
                     }
                 })
                 .Bind<IPlatformIconLoader>().ToConstant(s_instance)
-                .Bind<NonPumpingLockHelper.IHelperImpl>().ToConstant(new NonPumpingSyncContext.HelperImpl())
+                .Bind<NonPumpingLockHelper.IHelperImpl>().ToConstant(NonPumpingWaitHelperImpl.Instance)
                 .Bind<IMountedVolumeInfoProvider>().ToConstant(new WindowsMountedVolumeInfoProvider())
                 .Bind<IPlatformLifetimeEventsImpl>().ToConstant(s_instance);
             

src/Windows/Avalonia.Win32/WindowImpl.AppWndProc.cs

@@ -579,7 +579,7 @@ namespace Avalonia.Win32
 
                 case WindowsMessage.WM_PAINT:
                     {
-                        using (NonPumpingSyncContext.Use())
+                        using (NonPumpingSyncContext.Use(NonPumpingWaitHelperImpl.Instance))
                         using (_rendererLock.Lock())
                         {
                             if (BeginPaint(_hwnd, out PAINTSTRUCT ps) != IntPtr.Zero)
@@ -602,7 +602,7 @@ namespace Avalonia.Win32
 
                 case WindowsMessage.WM_SIZE:
                     {
-                        using (NonPumpingSyncContext.Use())
+                        using (NonPumpingSyncContext.Use(NonPumpingWaitHelperImpl.Instance))
                         using (_rendererLock.Lock())
                         {
                             // Do nothing here, just block until the pending frame render is completed on the render thread

tests/Avalonia.Base.UnitTests/DispatcherTests.cs

@@ -4,6 +4,7 @@ using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using Avalonia.Threading;
+using Avalonia.Utilities;
 using Xunit;
 namespace Avalonia.Base.UnitTests;
 
@@ -254,4 +255,207 @@ public class DispatcherTests
             Assert.Equal(foreground ? 0 : 1, ((SimpleControlledDispatcherImpl)impl).RunLoopCount);
     }
 
+
+    class DispatcherServices : IDisposable
+    {
+        private readonly IDisposable _scope;
+
+        public DispatcherServices(IDispatcherImpl impl)
+        {
+            _scope = AvaloniaLocator.EnterScope();
+            AvaloniaLocator.CurrentMutable.Bind<IDispatcherImpl>().ToConstant(impl);
+            Dispatcher.ResetForUnitTests();
+            SynchronizationContext.SetSynchronizationContext(null);
+        }
+        
+        public void Dispose()
+        {
+            Dispatcher.ResetForUnitTests();
+            _scope.Dispose();
+            SynchronizationContext.SetSynchronizationContext(null);
+        }
+    }
+    
+    [Fact]
+    public void ExitAllFramesShouldExitAllFramesAndBeAbleToContinue()
+    {
+        using (new DispatcherServices(new SimpleControlledDispatcherImpl()))
+        {
+            var actions = new List<string>();
+            var disp = Dispatcher.UIThread;
+            disp.Post(() =>
+            {
+                actions.Add("Nested frame");
+                Dispatcher.UIThread.MainLoop(CancellationToken.None);
+                actions.Add("Nested frame exited");
+            });
+            disp.Post(() =>
+            {
+                actions.Add("ExitAllFrames");
+                disp.ExitAllFrames();
+            });
+
+
+            disp.MainLoop(CancellationToken.None);
+            
+            Assert.Equal(new[] { "Nested frame", "ExitAllFrames", "Nested frame exited" }, actions);
+            actions.Clear();
+            
+            var secondLoop = new CancellationTokenSource();
+            disp.Post(() =>
+            {
+                actions.Add("Callback after exit");
+                secondLoop.Cancel();
+            });
+            disp.MainLoop(secondLoop.Token);
+            Assert.Equal(new[] { "Callback after exit" }, actions);
+        }
+    }
+    
+        
+    [Fact]
+    public void ShutdownShouldExitAllFramesAndNotAllowNewFrames()
+    {
+        using (new DispatcherServices(new SimpleControlledDispatcherImpl()))
+        {
+            var actions = new List<string>();
+            var disp = Dispatcher.UIThread;
+            disp.Post(() =>
+            {
+                actions.Add("Nested frame");
+                Dispatcher.UIThread.MainLoop(CancellationToken.None);
+                actions.Add("Nested frame exited");
+            });
+            disp.Post(() =>
+            {
+                actions.Add("Shutdown");
+                disp.BeginInvokeShutdown(DispatcherPriority.Normal);
+            });
+            
+            disp.Post(() =>
+            {
+                actions.Add("Nested frame after shutdown");
+                // This should exit immediately and not run any jobs
+                Dispatcher.UIThread.MainLoop(CancellationToken.None);
+                actions.Add("Nested frame after shutdown exited");
+            });
+            
+            var criticalFrameAfterShutdown = new DispatcherFrame(false);
+            disp.Post(() =>
+            {
+                actions.Add("Critical frame after shutdown");
+                
+                Dispatcher.UIThread.PushFrame(criticalFrameAfterShutdown);
+                actions.Add("Critical frame after shutdown exited");
+            });
+            disp.Post(() =>
+            {
+                actions.Add("Stop critical frame");
+                criticalFrameAfterShutdown.Continue = false;
+            });
+
+            disp.MainLoop(CancellationToken.None);
+
+            Assert.Equal(new[]
+            {
+                "Nested frame", 
+                "Shutdown",
+                // Normal nested frames are supposed to exit immediately
+                "Nested frame after shutdown", "Nested frame after shutdown exited",
+                // if frame is configured to not answer dispatcher requests, it should be allowed to run
+                "Critical frame after shutdown", "Stop critical frame", "Critical frame after shutdown exited",
+                // After 3-rd level frames have exited, the normal nested frame exits too
+                "Nested frame exited"
+            }, actions);
+            actions.Clear();
+            
+            disp.Post(()=>actions.Add("Frame after shutdown finished"));
+            Assert.Throws<InvalidOperationException>(() => disp.MainLoop(CancellationToken.None));
+            Assert.Empty(actions);
+        }
+    }
+
+    class WaitHelper : SynchronizationContext, NonPumpingLockHelper.IHelperImpl
+    {
+        public int WaitCount;
+        public int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout)
+        {
+            WaitCount++;
+            return base.Wait(waitHandles, waitAll, millisecondsTimeout);
+        }
+    }
+    
+    [Fact]
+    public void DisableProcessingShouldStopProcessing()
+    {
+        using (new DispatcherServices(new SimpleControlledDispatcherImpl()))
+        {
+            var helper = new WaitHelper();
+            AvaloniaLocator.CurrentMutable.Bind<NonPumpingLockHelper.IHelperImpl>().ToConstant(helper);
+            using (Dispatcher.UIThread.DisableProcessing())
+            {
+                Assert.True(SynchronizationContext.Current is NonPumpingSyncContext);
+                Assert.Throws<InvalidOperationException>(() => Dispatcher.UIThread.MainLoop(CancellationToken.None));
+                Assert.Throws<InvalidOperationException>(() => Dispatcher.UIThread.RunJobs());
+            }
+
+            var avaloniaContext = new AvaloniaSynchronizationContext(true);
+            SynchronizationContext.SetSynchronizationContext(avaloniaContext);
+
+            var waitHandle = new ManualResetEvent(true);
+            
+            helper.WaitCount = 0;
+            waitHandle.WaitOne(100);
+            Assert.Equal(0, helper.WaitCount);
+            using (Dispatcher.UIThread.DisableProcessing())
+            {
+                Assert.Equal(avaloniaContext, SynchronizationContext.Current);
+                waitHandle.WaitOne(100);
+                Assert.Equal(1, helper.WaitCount);
+            }
+        }
+    }
+
+    [Fact]
+    public void DispatcherOperationsHaveContextWithProperPriority()
+    {
+        using (new DispatcherServices(new SimpleControlledDispatcherImpl()))
+        {
+            SynchronizationContext.SetSynchronizationContext(null);
+            var disp = Dispatcher.UIThread;
+            var priorities = new List<DispatcherPriority>();
+
+            void DumpCurrentPriority() =>
+                priorities.Add(((AvaloniaSynchronizationContext)SynchronizationContext.Current!).Priority);
+                
+                
+            disp.Post(DumpCurrentPriority, DispatcherPriority.Normal);
+            disp.Post(DumpCurrentPriority, DispatcherPriority.Loaded);
+            disp.Post(DumpCurrentPriority, DispatcherPriority.Input);
+            disp.Post(() =>
+            {
+                DumpCurrentPriority();
+                disp.ExitAllFrames();
+            }, DispatcherPriority.Background);
+            disp.MainLoop(CancellationToken.None);
+
+            disp.Invoke(DumpCurrentPriority, DispatcherPriority.Send);
+            disp.Invoke(() =>
+            {
+                DumpCurrentPriority();
+                return 1;
+            }, DispatcherPriority.Send);
+
+            Assert.Equal(
+                new[]
+                {
+                    DispatcherPriority.Normal, DispatcherPriority.Loaded, DispatcherPriority.Input, DispatcherPriority.Background,
+                    DispatcherPriority.Send, DispatcherPriority.Send,
+                },
+                priorities);
+
+
+        }
+    }
+
 }