[browser][MT] fix CaptureContextFromParameter and SealJSImportCapturi…

…ng (#96504)

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

@@ -29,7 +29,7 @@ public static void CancelPromise(Task promise)
             }
             _CancelPromise(holder.GCHandle);
 #else
-            // this need to e manually dispatched via holder.ProxyContext, because we don't pass JSObject with affinity
+            // this need to be manually dispatched via holder.ProxyContext, because we don't pass JSObject with affinity
             holder.ProxyContext.SynchronizationContext.Post(static (object? h) =>
             {
                 var holder = (JSHostImplementation.PromiseHolder)h!;
@@ -63,6 +63,7 @@ public static void CancelPromise<T>(Task promise, Action<T> callback, T state)
             _CancelPromise(holder.GCHandle);
             callback.Invoke(state);
 #else
+            // this need to be manually dispatched via holder.ProxyContext, because we don't pass JSObject with affinity
             holder.ProxyContext.SynchronizationContext.Post(_ =>
             {
                 lock (holder.ProxyContext)

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

@@ -197,7 +197,9 @@ public static JSFunctionBinding BindManagedFunction(string fullyQualifiedName, i
             return BindManagedFunctionImpl(fullyQualifiedName, signatureHash, signatures);
         }
 
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void InvokeJSFunction(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
         {
             jsFunction.AssertNotDisposed();
@@ -218,7 +220,9 @@ internal static unsafe void InvokeJSFunction(JSObject jsFunction, Span<JSMarshal
 #endif
         }
 
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void InvokeJSFunctionCurrent(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
         {
             var functionHandle = (int)jsFunction.JSHandle;
@@ -235,7 +239,9 @@ internal static unsafe void InvokeJSFunctionCurrent(JSObject jsFunction, Span<JS
 
 
 #if FEATURE_WASM_THREADS
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void DispatchJSFunctionSync(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
         {
             var args = (nint)Unsafe.AsPointer(ref arguments[0]);
@@ -255,7 +261,9 @@ internal static unsafe void DispatchJSFunctionSync(JSObject jsFunction, Span<JSM
         }
 #endif
 
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void InvokeJSImportImpl(JSFunctionBinding signature, Span<JSMarshalerArgument> arguments)
         {
 #if FEATURE_WASM_THREADS
@@ -311,7 +319,9 @@ internal static unsafe void InvokeJSImportImpl(JSFunctionBinding signature, Span
 #endif
         }
 
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void InvokeJSImportCurrent(JSFunctionBinding signature, Span<JSMarshalerArgument> arguments)
         {
             fixed (JSMarshalerArgument* args = arguments)
@@ -332,7 +342,9 @@ internal static unsafe void InvokeJSImportCurrent(JSFunctionBinding signature, S
 
 #if FEATURE_WASM_THREADS
 
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void DispatchJSImportSync(JSFunctionBinding signature, JSProxyContext targetContext, Span<JSMarshalerArgument> arguments)
         {
             var args = (nint)Unsafe.AsPointer(ref arguments[0]);
@@ -351,7 +363,9 @@ internal static unsafe void DispatchJSImportSync(JSFunctionBinding signature, JS
             }
         }
 
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void DispatchJSImportAsync(JSFunctionBinding signature, JSProxyContext targetContext, Span<JSMarshalerArgument> arguments)
         {
             // this copy is freed in mono_wasm_invoke_import_async
@@ -418,7 +432,9 @@ internal static unsafe void ResolveOrRejectPromise(Span<JSMarshalerArgument> arg
             }
         }
 #else
+#if !DEBUG
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
         internal static unsafe void ResolveOrRejectPromise(JSProxyContext targetContext, Span<JSMarshalerArgument> arguments)
         {
             // this copy is freed in mono_wasm_invoke_import_async

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

@@ -237,6 +237,8 @@ public static void UninstallWebWorkerInterop()
             {
                 SynchronizationContext.SetSynchronizationContext(syncContext.previousSynchronizationContext);
             }
+            JSProxyContext.CurrentThreadContext = null;
+            JSProxyContext.ExecutionContext = null;
             ctx.Dispose();
         }
 

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

@@ -16,13 +16,7 @@ public partial class JSObject : IDisposable
         /// <summary>
         /// Returns true if the proxy was already disposed.
         /// </summary>
-        public bool IsDisposed
-        {
-            get
-            {
-                return _isDisposed;
-            }
-        }
+        public bool IsDisposed => _isDisposed;
 
         /// <summary>
         /// Checks whether the target object or one of its prototypes has a property with the specified name.

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

@@ -134,9 +134,17 @@ public static JSProxyContext SealJSImportCapturing()
             var executionContext = ExecutionContext;
             if (executionContext != null)
             {
-                // we could will call JS on the current thread (or child task), if it has the JS interop installed
+                // we could will call JS on the task's AsyncLocal context, if it has the JS interop installed
                 return executionContext;
             }
+
+            var currentThreadContext = CurrentThreadContext;
+            if (currentThreadContext != null)
+            {
+                // we could will call JS on the current thread (or child task), if it has the JS interop installed
+                return currentThreadContext;
+            }
+
             // otherwise we will call JS on the main thread, which always has JS interop
             return MainThreadContext;
         }
@@ -150,13 +158,13 @@ public static void CaptureContextFromParameter(JSProxyContext parameterContext)
                 Environment.FailFast($"Method only allowed during JSImport capturing phase, ManagedThreadId: {Environment.CurrentManagedThreadId}. {Environment.NewLine} {Environment.StackTrace}");
             }
 
-            var capturedContext = _CapturedOperationContext;
+            var alreadyCapturedContext = _CapturedOperationContext;
 
-            if (capturedContext == null)
+            if (alreadyCapturedContext == null)
             {
-                _CapturedOperationContext = capturedContext;
+                _CapturedOperationContext = parameterContext;
             }
-            else if (parameterContext != capturedContext)
+            else if (parameterContext != alreadyCapturedContext)
             {
                 _CapturedOperationContext = null;
                 _CapturingState = JSImportOperationState.None;

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

@@ -163,7 +163,7 @@ private void Pump()
             }
             catch (Exception e)
             {
-                Environment.FailFast("JSSynchronizationContext.BackgroundJobHandler failed", e);
+                Environment.FailFast($"JSSynchronizationContext.BackgroundJobHandler failed, ManagedThreadId: {Environment.CurrentManagedThreadId}. {Environment.NewLine} {e.StackTrace}");
             }
             finally
             {

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

@@ -51,7 +51,8 @@ public static async Task RunAsync(Func<Task> body, CancellationToken cancellatio
         private static Task<T> RunAsyncImpl<T>(Func<Task<T>> body, CancellationToken cancellationToken)
         {
             var parentContext = SynchronizationContext.Current ?? new SynchronizationContext();
-            var tcs = new TaskCompletionSource<T>();
+            // continuation should not be running synchronously in the JSWebWorker thread because we are about to kill it after we resolve/reject the Task.
+            var tcs = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);
             var capturedContext = SynchronizationContext.Current;
             var t = new Thread(() =>
             {
@@ -88,7 +89,8 @@ private static Task<T> RunAsyncImpl<T>(Func<Task<T>> body, CancellationToken can
         private static Task RunAsyncImpl(Func<Task> body, CancellationToken cancellationToken)
         {
             var parentContext = SynchronizationContext.Current ?? new SynchronizationContext();
-            var tcs = new TaskCompletionSource();
+            // continuation should not be running synchronously in the JSWebWorker thread because we are about to kill it after we resolve/reject the Task.
+            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
             var capturedContext = SynchronizationContext.Current;
             var t = new Thread(() =>
             {

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

@@ -67,11 +67,9 @@ public unsafe void ToJS(Exception? value)
                 if (jse != null && jse.jsException != null)
                 {
                     var jsException = jse.jsException;
-#if !FEATURE_WASM_THREADS
                     jsException.AssertNotDisposed();
-#else
+#if FEATURE_WASM_THREADS
                     var ctx = jsException.ProxyContext;
-                    jsException.AssertNotDisposed();
 
                     if (JSProxyContext.CapturingState == JSProxyContext.JSImportOperationState.JSImportParams)
                     {

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

@@ -40,10 +40,8 @@ public void ToJS(JSObject? value)
             }
             else
             {
-#if !FEATURE_WASM_THREADS
-                value.AssertNotDisposed();
-#else
                 value.AssertNotDisposed();
+#if FEATURE_WASM_THREADS
                 var ctx = value.ProxyContext;
 
                 if (JSProxyContext.CapturingState == JSProxyContext.JSImportOperationState.JSImportParams)

src/libraries/System.Runtime.InteropServices.JavaScript/tests/System.Runtime.InteropServices.JavaScript.UnitTests/System/Runtime/InteropServices/JavaScript/WebWorkerTest.cs

@@ -24,14 +24,15 @@ namespace System.Runtime.InteropServices.JavaScript.Tests
     // thread allocation, many threads
     // TLS
     // ProxyContext flow, child thread, child task
-    // use JSObject after JSWebWorker finished
+    // use JSObject after JSWebWorker finished, especially HTTP
     // JSWebWorker JS setTimeout till after close
     // WS on JSWebWorker
     // Yield will hit event loop 3x
     // HTTP continue on TP
     // WS continue on TP
     // event pipe
     // FS
+    // unit test for problem **7)**
 
     public class WebWorkerTest
     {
@@ -42,7 +43,13 @@ public static IEnumerable<object[]> GetTargetThreads()
             return Enum.GetValues<ExecutorType>().Select(type => new object[] { new Executor(type) });
         }
 
-        public static IEnumerable<object[]> GetTargetThreads2x2()
+        public static IEnumerable<object[]> GetSpecificTargetThreads()
+        {
+            yield return new object[] { new Executor(ExecutorType.JSWebWorker), new Executor(ExecutorType.Main) };
+            yield break;
+        }
+
+        public static IEnumerable<object[]> GetTargetThreads2x()
         {
             return Enum.GetValues<ExecutorType>().SelectMany(
                 type1 => Enum.GetValues<ExecutorType>().Select(
@@ -143,12 +150,13 @@ await executor.Execute(async () =>
         }
 
         [Theory, MemberData(nameof(GetTargetThreads))]
+        [ActiveIssue("https://github.com/dotnet/runtime/issues/96493")]
         public async Task ManagedDelay_ContinueWith(Executor executor)
         {
             await executor.Execute(async () =>
             {
                 executor.AssertTargetThread();
-                await Task.Delay(1).ContinueWith(_ =>
+                await Task.Delay(10).ContinueWith(_ =>
                 {
                     // continue on the context of the Timer's thread pool thread
                     Assert.True(Thread.CurrentThread.IsThreadPoolThread);
@@ -185,5 +193,68 @@ await executor.Execute(async () =>
 
         #endregion
 
+        #region affinity
+
+        private async Task ActionsInDifferentThreads<T>(Executor executor1, Executor executor2, Func<Task, TaskCompletionSource<T>, Task> e1Job, Func<T, Task> e2Job)
+        {
+            TaskCompletionSource<T> readyTCS = new TaskCompletionSource<T>();
+            TaskCompletionSource doneTCS = new TaskCompletionSource();
+
+            var e1 = executor1.Execute(async () =>
+            {
+                await e1Job(doneTCS.Task, readyTCS);
+                if (!readyTCS.Task.IsCompleted)
+                {
+                    readyTCS.SetResult(default);
+                }
+                await doneTCS.Task;
+            });
+
+            var r1 = await readyTCS.Task.ConfigureAwait(true);
+
+            var e2 = executor2.Execute(async () =>
+            {
+
+                executor2.AssertTargetThread();
+
+                await e2Job(r1);
+
+                doneTCS.SetResult();
+            });
+
+            await e2;
+            await e1;
+        }
+
+        [Theory, MemberData(nameof(GetTargetThreads2x))]
+        public async Task JSObject_CapturesAffinity(Executor executor1, Executor executor2)
+        {
+            var e1Job = async (Task e2done, TaskCompletionSource<JSObject> e1State) =>
+            {
+                await WebWorkerTestHelper.InitializeAsync();
+
+                executor1.AssertAwaitCapturedContext();
+
+                var jsState = await WebWorkerTestHelper.PromiseState();
+
+                // share the state with the E2 continuation
+                e1State.SetResult(jsState);
+
+                await e2done;
+
+                // cleanup
+                await WebWorkerTestHelper.DisposeAsync();
+            };
+
+            var e2Job = async (JSObject e1State) =>
+            {
+                bool valid = await WebWorkerTestHelper.PromiseValidateState(e1State);
+                Assert.True(valid);
+            };
+
+            await ActionsInDifferentThreads<JSObject>(executor1, executor2, e1Job, e2Job);
+        }
+
+        #endregion
     }
 }