Implemented Dispatcher.Yield / Dispatcher.Resume (#19370)

* Refactored DispatcherPriorityAwaitable, implemented Dispatcher.Yield

* Picked changes from https://github.com/AvaloniaUI/Avalonia/pull/14212

* Format/compile

* Update API suppressions

---------

Co-authored-by: Yoh Deadfall <[email protected]>
Co-authored-by: Julien Lebosquain <[email protected]>

api/Avalonia.nupkg.xml

@@ -109,6 +109,42 @@
     <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
     <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
   </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:Avalonia.Threading.DispatcherPriorityAwaitable.get_IsCompleted</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:Avalonia.Threading.DispatcherPriorityAwaitable.GetAwaiter</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:Avalonia.Threading.DispatcherPriorityAwaitable.GetResult</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:Avalonia.Threading.DispatcherPriorityAwaitable.OnCompleted(System.Action)</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:Avalonia.Threading.DispatcherPriorityAwaitable`1.GetAwaiter</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:Avalonia.Threading.DispatcherPriorityAwaitable`1.GetResult</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
   <Suppression>
     <DiagnosticId>CP0002</DiagnosticId>
     <Target>M:Avalonia.Controls.Primitives.IPopupHost.ConfigurePosition(Avalonia.Visual,Avalonia.Controls.PlacementMode,Avalonia.Point,Avalonia.Controls.Primitives.PopupPositioning.PopupAnchor,Avalonia.Controls.Primitives.PopupPositioning.PopupGravity,Avalonia.Controls.Primitives.PopupPositioning.PopupPositionerConstraintAdjustment,System.Nullable{Avalonia.Rect})</Target>
@@ -187,6 +223,30 @@
     <Left>baseline/netstandard2.0/Avalonia.Controls.dll</Left>
     <Right>target/netstandard2.0/Avalonia.Controls.dll</Right>
   </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0007</DiagnosticId>
+    <Target>T:Avalonia.Threading.DispatcherPriorityAwaitable</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0007</DiagnosticId>
+    <Target>T:Avalonia.Threading.DispatcherPriorityAwaitable`1</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0008</DiagnosticId>
+    <Target>T:Avalonia.Threading.DispatcherPriorityAwaitable</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0008</DiagnosticId>
+    <Target>T:Avalonia.Threading.DispatcherPriorityAwaitable`1</Target>
+    <Left>baseline/netstandard2.0/Avalonia.Base.dll</Left>
+    <Right>target/netstandard2.0/Avalonia.Base.dll</Right>
+  </Suppression>
   <Suppression>
     <DiagnosticId>CP0009</DiagnosticId>
     <Target>T:Avalonia.Diagnostics.StyleDiagnostics</Target>

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

@@ -672,4 +672,70 @@ public partial class Dispatcher
     /// </summary>
     public DispatcherPriorityAwaitable<T> AwaitWithPriority<T>(Task<T> task, DispatcherPriority priority) =>
         new(this, task, priority);
+    
+    /// <summary>
+    /// Creates an awaitable object that asynchronously resumes execution on the dispatcher.
+    /// </summary>
+    /// <returns>
+    /// An awaitable object that asynchronously resumes execution on the dispatcher.
+    /// </returns>
+    /// <remarks>
+    /// This method is equivalent to calling the <see cref="Resume(DispatcherPriority)"/> method
+    /// and passing in <see cref="DispatcherPriority.Background"/>.
+    /// </remarks>
+    public DispatcherPriorityAwaitable Resume() =>
+        Resume(DispatcherPriority.Background);
+
+    /// <summary>``
+    /// Creates an awaitable object that asynchronously resumes execution on the dispatcher. The work that occurs
+    /// when control returns to the code awaiting the result of this method is scheduled with the specified priority.
+    /// </summary>
+    /// <param name="priority">The priority at which to schedule the continuation.</param>
+    /// <returns>
+    /// An awaitable object that asynchronously resumes execution on the dispatcher.
+    /// </returns>
+    public DispatcherPriorityAwaitable Resume(DispatcherPriority priority)
+    {
+        DispatcherPriority.Validate(priority, nameof(priority));
+        return new(this, null, priority);
+    }
+
+    /// <summary>
+    /// Creates an awaitable object that asynchronously yields control back to the current dispatcher
+    /// and provides an opportunity for the dispatcher to process other events.
+    /// </summary>
+    /// <returns>
+    /// An awaitable object that asynchronously yields control back to the current dispatcher
+    /// and provides an opportunity for the dispatcher to process other events.
+    /// </returns>
+    /// <remarks>
+    /// This method is equivalent to calling the <see cref="Yield(DispatcherPriority)"/> method
+    /// and passing in <see cref="DispatcherPriority.Background"/>.
+    /// </remarks>
+    /// <exception cref="InvalidOperationException">
+    /// The current thread is not the UI thread.
+    /// </exception>
+    public static DispatcherPriorityAwaitable Yield() =>
+        Yield(DispatcherPriority.Background);
+
+    /// <summary>
+    /// Creates an cawaitable object that asynchronously yields control back to the current dispatcher
+    /// and provides an opportunity for the dispatcher to process other events. The work that occurs when
+    /// control returns to the code awaiting the result of this method is scheduled with the specified priority.
+    /// </summary>
+    /// <param name="priority">The priority at which to schedule the continuation.</param>
+    /// <returns>
+    /// An awaitable object that asynchronously yields control back to the current dispatcher
+    /// and provides an opportunity for the dispatcher to process other events.
+    /// </returns>
+    /// <exception cref="InvalidOperationException">
+    /// The current thread is not the UI thread.
+    /// </exception>
+    public static DispatcherPriorityAwaitable Yield(DispatcherPriority priority)
+    {
+        // TODO12: Update to use Dispatcher.CurrentDispatcher once multi-dispatcher support is merged
+        var current = UIThread;
+        current.VerifyAccess();
+        return UIThread.Resume(priority);
+    }
 }

src/Avalonia.Base/Threading/DispatcherPriorityAwaitable.cs

@@ -1,40 +1,130 @@
 using System;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Avalonia.Threading;
 
-public class DispatcherPriorityAwaitable : INotifyCompletion
+/// <summary>
+///     A simple awaitable type that will return a DispatcherPriorityAwaiter.
+/// </summary>
+public struct DispatcherPriorityAwaitable
 {
     private readonly Dispatcher _dispatcher;
-    private protected readonly Task Task;
+    private readonly Task? _task;
     private readonly DispatcherPriority _priority;
 
-    internal DispatcherPriorityAwaitable(Dispatcher dispatcher, Task task, DispatcherPriority priority)
+    internal DispatcherPriorityAwaitable(Dispatcher dispatcher, Task? task, DispatcherPriority priority)
     {
         _dispatcher = dispatcher;
-        Task = task;
+        _task = task;
         _priority = priority;
     }
-    
-    public void OnCompleted(Action continuation) =>
-        Task.ContinueWith(_ => _dispatcher.Post(continuation, _priority));
 
-    public bool IsCompleted => Task.IsCompleted;
+    public DispatcherPriorityAwaiter GetAwaiter() => new(_dispatcher, _task, _priority);
+}
+
+/// <summary>
+///     A simple awaiter type that will queue the continuation to a dispatcher at a specific priority.
+/// </summary>
+/// <remarks>
+///     This is returned from DispatcherPriorityAwaitable.GetAwaiter()
+/// </remarks>
+public struct DispatcherPriorityAwaiter : INotifyCompletion
+{
+    private readonly Dispatcher _dispatcher;
+    private readonly Task? _task;
+    private readonly DispatcherPriority _priority;
+
+    internal DispatcherPriorityAwaiter(Dispatcher dispatcher, Task? task, DispatcherPriority priority)
+    {
+        _dispatcher = dispatcher;
+        _task = task;
+        _priority = priority;
+    }
+
+    public void OnCompleted(Action continuation)
+    {
+        if(_task == null || _task.IsCompleted)
+            _dispatcher.Post(continuation, _priority);
+        else
+        {
+            var self = this;
+            _task.ConfigureAwait(false).GetAwaiter().OnCompleted(() =>
+            {
+                self._dispatcher.Post(continuation, self._priority);
+            });
+        }
+    }
+
+    /// <summary>
+    /// This always returns false since continuation is requested to be queued to a dispatcher queue
+    /// </summary>
+    public bool IsCompleted => false;
+
+    public void GetResult()
+    {
+        if (_task != null)
+            _task.GetAwaiter().GetResult();
+    }
+}
+
+/// <summary>
+///     A simple awaitable type that will return a DispatcherPriorityAwaiter&lt;T&gt;.
+/// </summary>
+public struct DispatcherPriorityAwaitable<T>
+{
+    private readonly Dispatcher _dispatcher;
+    private readonly Task<T> _task;
+    private readonly DispatcherPriority _priority;
 
-    public void GetResult() => Task.GetAwaiter().GetResult();
+    internal DispatcherPriorityAwaitable(Dispatcher dispatcher, Task<T> task, DispatcherPriority priority)
+    {
+        _dispatcher = dispatcher;
+        _task = task;
+        _priority = priority;
+    }
 
-    public DispatcherPriorityAwaitable GetAwaiter() => this;
+    public DispatcherPriorityAwaiter<T> GetAwaiter() => new(_dispatcher, _task, _priority);
 }
 
-public sealed class DispatcherPriorityAwaitable<T> : DispatcherPriorityAwaitable
+/// <summary>
+///     A simple awaiter type that will queue the continuation to a dispatcher at a specific priority.
+/// </summary>
+/// <remarks>
+///     This is returned from DispatcherPriorityAwaitable&lt;T&gt;.GetAwaiter()
+/// </remarks>
+public struct DispatcherPriorityAwaiter<T> : INotifyCompletion
 {
-    internal DispatcherPriorityAwaitable(Dispatcher dispatcher, Task<T> task, DispatcherPriority priority) : base(
-        dispatcher, task, priority)
+    private readonly Dispatcher _dispatcher;
+    private readonly Task<T> _task;
+    private readonly DispatcherPriority _priority;
+
+    internal DispatcherPriorityAwaiter(Dispatcher dispatcher, Task<T> task, DispatcherPriority priority)
     {
+        _dispatcher = dispatcher;
+        _task = task;
+        _priority = priority;
     }
 
-    public new T GetResult() => ((Task<T>)Task).GetAwaiter().GetResult();
+    public void OnCompleted(Action continuation)
+    {
+        if(_task.IsCompleted)
+            _dispatcher.Post(continuation, _priority);
+        else
+        {
+            var self = this;
+            _task.ConfigureAwait(false).GetAwaiter().OnCompleted(() =>
+            {
+                self._dispatcher.Post(continuation, self._priority);
+            });
+        }
+    }
+    
+    /// <summary>
+    /// This always returns false since continuation is requested to be queued to a dispatcher queue
+    /// </summary>
+    public bool IsCompleted => false;
 
-    public new DispatcherPriorityAwaitable<T> GetAwaiter() => this;
-}
+    public void GetResult() => _task.GetAwaiter().GetResult();
+}

tests/Avalonia.Base.UnitTests/DispatcherTests.cs

@@ -505,4 +505,104 @@ public partial class DispatcherTests
             t.GetAwaiter().GetResult();
         }
     }
+    
+    
+    [Fact]
+    public async Task DispatcherResumeContinuesOnUIThread()
+    {
+        using var services = new DispatcherServices(new SimpleControlledDispatcherImpl());
+
+        var tokenSource = new CancellationTokenSource();
+        var workload = Dispatcher.UIThread.InvokeAsync(
+            async () =>
+            {
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+
+                await Task.Delay(1).ConfigureAwait(false);
+                Assert.False(Dispatcher.UIThread.CheckAccess());
+
+                await Dispatcher.UIThread.Resume();
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+
+                tokenSource.Cancel();
+            });
+
+        Dispatcher.UIThread.MainLoop(tokenSource.Token);
+    }
+
+    [Fact]
+    public async Task DispatcherYieldContinuesOnUIThread()
+    {
+        using var services = new DispatcherServices(new SimpleControlledDispatcherImpl());
+
+        var tokenSource = new CancellationTokenSource();
+        var workload = Dispatcher.UIThread.InvokeAsync(
+            async () =>
+            {
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+
+                await Dispatcher.Yield();
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+
+                tokenSource.Cancel();
+            });
+
+        Dispatcher.UIThread.MainLoop(tokenSource.Token);
+    }
+
+    [Fact]
+    public async Task DispatcherYieldThrowsOnNonUIThread()
+    {
+        using var services = new DispatcherServices(new SimpleControlledDispatcherImpl());
+
+        var tokenSource = new CancellationTokenSource();
+        var workload = Dispatcher.UIThread.InvokeAsync(
+            async () =>
+            {
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+
+                await Task.Delay(1).ConfigureAwait(false);
+                Assert.False(Dispatcher.UIThread.CheckAccess());
+                await Assert.ThrowsAsync<InvalidOperationException>(async () => await Dispatcher.Yield());
+
+                tokenSource.Cancel();
+            });
+
+        Dispatcher.UIThread.MainLoop(tokenSource.Token);
+    }
+
+    [Fact]
+    public async Task AwaitWithPriorityRunsOnUIThread()
+    {
+        static async Task<int> Workload()
+        {
+            await Task.Delay(1).ConfigureAwait(false);
+            Assert.False(Dispatcher.UIThread.CheckAccess());
+
+            return Thread.CurrentThread.ManagedThreadId;
+        }
+
+        using var services = new DispatcherServices(new SimpleControlledDispatcherImpl());
+
+        var tokenSource = new CancellationTokenSource();
+        var workload = Dispatcher.UIThread.InvokeAsync(
+            async () =>
+            {
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+                Task taskWithoutResult = Workload();
+
+                await Dispatcher.UIThread.AwaitWithPriority(taskWithoutResult, DispatcherPriority.Default);
+
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+                Task<int> taskWithResult = Workload();
+
+                await Dispatcher.UIThread.AwaitWithPriority(taskWithResult, DispatcherPriority.Default);
+
+                Assert.True(Dispatcher.UIThread.CheckAccess());
+
+                tokenSource.Cancel();
+            });
+
+        Dispatcher.UIThread.MainLoop(tokenSource.Token);
+    }
 }