Fix RdTask GetAwaiter behaviour (return to original context after await)

add test
JetBrains · Mar 16, 2024 · 1a04cda · 1a04cda
1 parent 401a19e
commit 1a04cda
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 22 deletions.
diff --git a/rd-net/RdFramework/Tasks/RdTaskEx.cs b/rd-net/RdFramework/Tasks/RdTaskEx.cs
@@ -122,27 +122,8 @@ public static Task<T> AsTask<T>(this IRdTask<T> task)
       });
       return tcs.Task;
     }
-
-    [PublicAPI]
-    public static RdTaskAwaiter<T> GetAwaiter<T>(this IRdTask<T> task) => new(task.Result);
-
-    public readonly struct RdTaskAwaiter<T> : INotifyCompletion
-    {
-      private readonly IReadonlyProperty<RdTaskResult<T>> myResult;
-
-      internal RdTaskAwaiter(IReadonlyProperty<RdTaskResult<T>> result)
-      {
-        myResult = result;
-      }
 
-      public bool IsCompleted => myResult.Maybe.HasValue;
-
-      public T GetResult() => myResult.Value.Unwrap();
-
-      public void OnCompleted(Action continuation)
-      {
-        myResult.Change.AdviseOnce(Lifetime.Eternal, _ => continuation());
-      }
-    }
+    [PublicAPI]
+    public static TaskAwaiter<T> GetAwaiter<T>(this IRdTask<T> task) => task.AsTask().GetAwaiter();
   }
 }
diff --git a/rd-net/Test.RdFramework/AsyncRdTaskTest.cs b/rd-net/Test.RdFramework/AsyncRdTaskTest.cs
@@ -3,8 +3,10 @@
 using System.Threading.Tasks;
 using JetBrains.Collections.Viewable;
 using JetBrains.Core;
+using JetBrains.Rd.Base;
 using JetBrains.Rd.Impl;
 using JetBrains.Rd.Tasks;
+using JetBrains.Threading;
 using JetBrains.Util;
 using NUnit.Framework;
 
@@ -31,7 +33,7 @@ public void BindableRdCallListUseRdTaskTest()
     BindableRdCallListTest(TaskKind.Rd);
   }
 
-  private enum TaskKind
+  public enum TaskKind
   {
     System,
     Rd,
@@ -96,6 +98,41 @@ private void BindableRdCallListTest(TaskKind taskKind)
       Assert.IsTrue(bindClientTask.Wait(Timeout(TimeSpan.FromSeconds(10))));
     }
   }
+
+  [Test]
+  [TestCase(TaskKind.Rd)]
+  [TestCase(TaskKind.System)]
+  public void TestRdTaskAwaiter(TaskKind kind)
+  {
+    var rdTask = new RdTask<Unit>();
+    var scheduler = new TaskSchedulerWrapper(new ConcurrentExclusiveSchedulerPair(TaskScheduler.Default).ExclusiveScheduler, false);
+
+    var task = TestLifetime.StartAsync(scheduler.AsTaskScheduler(), async () =>
+    {
+      scheduler.AssertThread();
+
+      TestLifetime.Start(scheduler.AsTaskScheduler(), () =>
+      {
+        scheduler.AssertThread();
+        TestLifetime.Start(TaskScheduler.Default, () =>
+        {
+          rdTask.ResultInternal.Set(RdTaskResult<Unit>.Success(Unit.Instance));
+        }).NoAwait();
+      }).NoAwait();
+      
+      _ = kind switch
+      {
+        TaskKind.System => await rdTask.AsTask(),
+        TaskKind.Rd     => await rdTask,
+        _ => throw new ArgumentOutOfRangeException(nameof(kind), kind, null)
+      };
+      
+      scheduler.AssertThread();
+    });
+
+    task.Wait(TimeSpan.FromSeconds(10));
+    Assert.IsTrue(task.IsCompleted);
+  } 
 
   private static TimeSpan Timeout(TimeSpan timeout)
   {