[browser][MT] Marshal resolved/unresolved tasks separately (#99347)

Co-authored-by: Ilona Tomkowicz <32700855+ilonatommy@users.noreply.github.com>

Author: pavelsavara

src/libraries/System.Runtime.InteropServices.JavaScript/src/System.Runtime.InteropServices.JavaScript.csproj

@@ -68,6 +68,7 @@
   <ItemGroup Condition="'$(TargetPlatformIdentifier)' == 'browser' and '$(FeatureWasmManagedThreads)' == 'true'">
     <Compile Include="System\Runtime\InteropServices\JavaScript\JSWebWorker.cs" />
     <Compile Include="System\Runtime\InteropServices\JavaScript\JSSynchronizationContext.cs" />
+    <Compile Include="System\Runtime\InteropServices\JavaScript\JSAsyncTaskScheduler.cs" />
   </ItemGroup>
   <ItemGroup Condition="'$(WasmEnableThreads)' == 'true'">
     <ApiCompatSuppressionFile Include="CompatibilitySuppressions.xml;CompatibilitySuppressions.WasmThreads.xml" />

src/libraries/System.Runtime.InteropServices.JavaScript/src/System/Runtime/InteropServices/JavaScript/JSAsyncTaskScheduler.cs

@@ -0,0 +1,43 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+namespace System.Runtime.InteropServices.JavaScript
+{
+    // executes all tasks thru queue, never inline
+    internal sealed class JSAsyncTaskScheduler : TaskScheduler
+    {
+        private readonly JSSynchronizationContext m_synchronizationContext;
+
+        internal JSAsyncTaskScheduler(JSSynchronizationContext synchronizationContext)
+        {
+            m_synchronizationContext = synchronizationContext;
+        }
+
+        protected override void QueueTask(Task task)
+        {
+            m_synchronizationContext.Post((_) =>
+            {
+                if (!TryExecuteTask(task))
+                {
+                    Environment.FailFast("Unexpected failure in JSAsyncTaskScheduler" + Environment.CurrentManagedThreadId);
+                }
+            }, null);
+        }
+
+        // this is the main difference from the SynchronizationContextTaskScheduler
+        protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
+        {
+            return false;
+        }
+
+        protected override IEnumerable<Task>? GetScheduledTasks()
+        {
+            return null;
+        }
+
+        public override int MaximumConcurrencyLevel => 1;
+    }
+}

src/libraries/System.Runtime.InteropServices.JavaScript/src/System/Runtime/InteropServices/JavaScript/JSProxyContext.cs

@@ -40,6 +40,7 @@ private JSProxyContext()
         public int ManagedTID; // current managed thread id
         public bool IsMainThread;
         public JSSynchronizationContext SynchronizationContext;
+        public JSAsyncTaskScheduler? AsyncTaskScheduler;
 
         public static MainThreadingMode MainThreadingMode = MainThreadingMode.DeputyThread;
         public static JSThreadBlockingMode ThreadBlockingMode = JSThreadBlockingMode.NoBlockingWait;
@@ -483,7 +484,7 @@ public static void ReleaseCSOwnedObject(JSObject jso, bool skipJS)
                     {
                         if (IsJSVHandle(jsHandle))
                         {
-                            Environment.FailFast("TODO implement blocking ReleaseCSOwnedObjectSend to make sure the order of FreeJSVHandle is correct.");
+                            Environment.FailFast($"TODO implement blocking ReleaseCSOwnedObjectSend to make sure the order of FreeJSVHandle is correct, ManagedThreadId: {Environment.CurrentManagedThreadId}. {Environment.NewLine} {Environment.StackTrace}");
                         }
 
                         // this is async message, we need to call this as the last thing
@@ -501,7 +502,7 @@ public static void ReleaseCSOwnedObject(JSObject jso, bool skipJS)
             }
         }
 
-#endregion
+        #endregion
 
         #region Dispose
 

src/libraries/System.Runtime.InteropServices.JavaScript/src/System/Runtime/InteropServices/JavaScript/JSSynchronizationContext.cs

@@ -56,6 +56,7 @@ public static unsafe JSSynchronizationContext InstallWebWorkerInterop(bool isMai
             }
 
             var proxyContext = ctx.ProxyContext;
+            proxyContext.AsyncTaskScheduler = new JSAsyncTaskScheduler(ctx);
             JSProxyContext.CurrentThreadContext = proxyContext;
             JSProxyContext.ExecutionContext = proxyContext;
             if (isMainThread)

src/libraries/System.Runtime.InteropServices.JavaScript/src/System/Runtime/InteropServices/JavaScript/Marshaling/JSMarshalerArgument.Task.cs

@@ -7,6 +7,7 @@
 using System.ComponentModel;
 using System.Threading;
 using static System.Runtime.InteropServices.JavaScript.JSHostImplementation;
+using System.Runtime.CompilerServices;
 
 namespace System.Runtime.InteropServices.JavaScript
 {
@@ -140,13 +141,20 @@ internal void ToJSDynamic(Task? value)
         {
             Task? task = value;
 
+            var ctx = ToJSContext;
+            var canMarshalTaskResultOnSameCall = CanMarshalTaskResultOnSameCall(ctx);
+
             if (task == null)
             {
+                if (!canMarshalTaskResultOnSameCall)
+                {
+                    Environment.FailFast("Marshalling null return Task to JS is not supported in MT");
+                }
                 slot.Type = MarshalerType.None;
                 return;
             }
 
-            if (task.IsCompleted)
+            if (canMarshalTaskResultOnSameCall && task.IsCompleted)
             {
                 if (task.Exception != null)
                 {
@@ -172,7 +180,6 @@ internal void ToJSDynamic(Task? value)
                 }
             }
 
-            var ctx = ToJSContext;
 
             if (slot.Type != MarshalerType.TaskPreCreated)
             {
@@ -189,7 +196,9 @@ internal void ToJSDynamic(Task? value)
             var taskHolder = ctx.CreateCSOwnedProxy(slot.JSHandle);
 
 #if FEATURE_WASM_MANAGED_THREADS
-            task.ContinueWith(Complete, taskHolder, CancellationToken.None, TaskContinuationOptions.RunContinuationsAsynchronously, TaskScheduler.FromCurrentSynchronizationContext());
+            // AsyncTaskScheduler will make sure that the resolve message is always sent after this call is completed
+            // that is: synchronous marshaling and eventually message to the target thread, which need to arrive before the resolve message
+            task.ContinueWith(Complete, taskHolder, ctx.AsyncTaskScheduler!);
 #else
             task.ContinueWith(Complete, taskHolder, TaskScheduler.Current);
 #endif
@@ -229,18 +238,18 @@ public void ToJS(Task? value)
         {
             Task? task = value;
             var ctx = ToJSContext;
-            var isCurrentThread = ctx.IsCurrentThread();
+            var canMarshalTaskResultOnSameCall = CanMarshalTaskResultOnSameCall(ctx);
 
             if (task == null)
             {
-                if (!isCurrentThread)
+                if (!canMarshalTaskResultOnSameCall)
                 {
-                    Environment.FailFast("Marshalling null task to JS is not supported in MT");
+                    Environment.FailFast("Marshalling null return Task to JS is not supported in MT");
                 }
                 slot.Type = MarshalerType.None;
                 return;
             }
-            if (isCurrentThread && task.IsCompleted)
+            if (canMarshalTaskResultOnSameCall && task.IsCompleted)
             {
                 if (task.Exception != null)
                 {
@@ -273,7 +282,9 @@ public void ToJS(Task? value)
             var taskHolder = ctx.CreateCSOwnedProxy(slot.JSHandle);
 
 #if FEATURE_WASM_MANAGED_THREADS
-            task.ContinueWith(Complete, taskHolder, CancellationToken.None, TaskContinuationOptions.RunContinuationsAsynchronously, TaskScheduler.FromCurrentSynchronizationContext());
+            // AsyncTaskScheduler will make sure that the resolve message is always sent after this call is completed
+            // that is: synchronous marshaling and eventually message to the target thread, which need to arrive before the resolve message
+            task.ContinueWith(Complete, taskHolder, ctx.AsyncTaskScheduler!);
 #else
             task.ContinueWith(Complete, taskHolder, TaskScheduler.Current);
 #endif
@@ -303,19 +314,19 @@ public void ToJS<T>(Task<T>? value, ArgumentToJSCallback<T> marshaler)
         {
             Task<T>? task = value;
             var ctx = ToJSContext;
-            var isCurrentThread = ctx.IsCurrentThread();
+            var canMarshalTaskResultOnSameCall = CanMarshalTaskResultOnSameCall(ctx);
 
             if (task == null)
             {
-                if (!isCurrentThread)
+                if (!canMarshalTaskResultOnSameCall)
                 {
-                    Environment.FailFast("NULL not supported in MT");
+                    Environment.FailFast("Marshalling null return Task to JS is not supported in MT");
                 }
                 slot.Type = MarshalerType.None;
                 return;
             }
 
-            if (isCurrentThread && task.IsCompleted)
+            if (canMarshalTaskResultOnSameCall && task.IsCompleted)
             {
                 if (task.Exception != null)
                 {
@@ -350,7 +361,9 @@ public void ToJS<T>(Task<T>? value, ArgumentToJSCallback<T> marshaler)
             var taskHolder = ctx.CreateCSOwnedProxy(slot.JSHandle);
 
 #if FEATURE_WASM_MANAGED_THREADS
-            task.ContinueWith(Complete, new HolderAndMarshaler<T>(taskHolder, marshaler), CancellationToken.None, TaskContinuationOptions.RunContinuationsAsynchronously, TaskScheduler.FromCurrentSynchronizationContext());
+            // AsyncTaskScheduler will make sure that the resolve message is always sent after this call is completed
+            // that is: synchronous marshaling and eventually message to the target thread, which need to arrive before the resolve message
+            task.ContinueWith(Complete, new HolderAndMarshaler<T>(taskHolder, marshaler), ctx.AsyncTaskScheduler!);
 #else
             task.ContinueWith(Complete, new HolderAndMarshaler<T>(taskHolder, marshaler), TaskScheduler.Current);
 #endif
@@ -370,6 +383,44 @@ static void Complete(Task<T> task, object? thm)
             }
         }
 
+#if !DEBUG
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
+#if FEATURE_WASM_MANAGED_THREADS
+        // We can't marshal resolved/rejected/null Task.Result directly into current argument when this is marshaling return of JSExport across threads
+        private bool CanMarshalTaskResultOnSameCall(JSProxyContext ctx)
+        {
+            if (slot.Type != MarshalerType.TaskPreCreated)
+            {
+                // this means that we are not in the return value of JSExport
+                // we are marshaling parameter of JSImport
+                return true;
+            }
+
+            if (ctx.IsCurrentThread())
+            {
+                // If the JS and Managed is running on the same thread we can use the args buffer,
+                // because the call is synchronous and the buffer will be processed.
+                // In that case the pre-allocated Promise would be discarded as necessary
+                // and the result will be marshaled by `try_marshal_sync_task_to_js`
+                return true;
+            }
+
+            // Otherwise this is JSExport return value and we can't use the args buffer, because the args buffer arrived in async message and nobody is reading after this.
+            // In such case the JS side already pre-created the Promise and we have to use it, to resolve it in separate call via `mono_wasm_resolve_or_reject_promise_post`
+            // there is JSVHandle in this arg
+            return false;
+        }
+#else
+#pragma warning disable CA1822 // Mark members as static
+        private bool CanMarshalTaskResultOnSameCall(JSProxyContext _)
+        {
+            // in ST build this is always synchronous and we can marshal the result directly
+            return true;
+        }
+#pragma warning restore CA1822 // Mark members as static
+#endif
+
         private sealed record HolderAndMarshaler<T>(JSObject TaskHolder, ArgumentToJSCallback<T> Marshaler);
 
         private static void RejectPromise(JSObject holder, Exception ex)