[browser][MT] JSImport dispatch to target thread via JSSynchronizatio…

…nContext (#96319)

src/libraries/Common/src/Interop/Browser/Interop.Runtime.cs

@@ -12,31 +12,37 @@ internal static partial class Interop
     internal static unsafe partial class Runtime
     {
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
-        internal static extern void ReleaseCSOwnedObject(IntPtr jsHandle);
+        internal static extern void ReleaseCSOwnedObject(nint jsHandle);
         [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern unsafe void BindJSImport(void* signature, out int is_exception, out object result);
-        [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern void InvokeJSFunction(int functionHandle, void* data);
-        [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern void InvokeJSImport(int importHandle, void* data);
-
+        public static extern void InvokeJSFunction(nint functionHandle, nint data);
         [MethodImpl(MethodImplOptions.InternalCall)]
         public static extern unsafe void BindCSFunction(in string fully_qualified_name, int signature_hash, void* signature, out int is_exception, out object result);
         [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern void ResolveOrRejectPromise(void* data);
+        public static extern void ResolveOrRejectPromise(nint data);
 
 #if !ENABLE_JS_INTEROP_BY_VALUE
         [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern IntPtr RegisterGCRoot(IntPtr start, int bytesSize, IntPtr name);
+        public static extern nint RegisterGCRoot(void* start, int bytesSize, IntPtr name);
         [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern void DeregisterGCRoot(IntPtr handle);
+        public static extern void DeregisterGCRoot(nint handle);
 #endif
 
 #if FEATURE_WASM_THREADS
         [MethodImpl(MethodImplOptions.InternalCall)]
-        public static extern void InstallWebWorkerInterop(IntPtr proxyContextGCHandle);
+        public static extern void InstallWebWorkerInterop(nint proxyContextGCHandle);
         [MethodImpl(MethodImplOptions.InternalCall)]
         public static extern void UninstallWebWorkerInterop();
+
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        public static extern void InvokeJSImportSync(nint data, nint signature);
+
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        public static extern void InvokeJSImportAsync(nint data, nint signature);
+#else
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        public static extern unsafe void BindJSImport(void* signature, out int is_exception, out object result);
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        public static extern void InvokeJSImport(int importHandle, nint data);
 #endif
 
         #region Legacy

src/libraries/System.Runtime.InteropServices.JavaScript/gen/JSImportGenerator/JSImportGenerator.cs

@@ -120,9 +120,7 @@ private static MemberDeclarationSyntax PrintGeneratedSource(
 
             FieldDeclarationSyntax sigField = FieldDeclaration(VariableDeclaration(IdentifierName(Constants.JSFunctionSignatureGlobal))
                 .WithVariables(SingletonSeparatedList(VariableDeclarator(Identifier(stub.BindingName)))))
-                .AddModifiers(Token(SyntaxKind.StaticKeyword))
-                .WithAttributeLists(SingletonList(AttributeList(SingletonSeparatedList(
-                    Attribute(IdentifierName(Constants.ThreadStaticGlobal))))));
+                .AddModifiers(Token(SyntaxKind.StaticKeyword));
 
             MemberDeclarationSyntax toPrint = containingSyntaxContext.WrapMembersInContainingSyntaxWithUnsafeModifier(stubMethod, sigField);
             return toPrint;

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

@@ -29,6 +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
             holder.ProxyContext.SynchronizationContext.Post(static (object? h) =>
             {
                 var holder = (JSHostImplementation.PromiseHolder)h!;

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

@@ -30,6 +30,9 @@ internal JSFunctionBinding() { }
         internal static volatile uint nextImportHandle = 1;
         internal int ImportHandle;
         internal bool IsAsync;
+#if DEBUG
+        internal string? FunctionName;
+#endif
 
         [StructLayout(LayoutKind.Sequential, Pack = 4)]
         internal struct JSBindingHeader
@@ -197,15 +200,31 @@ public static JSFunctionBinding BindManagedFunction(string fullyQualifiedName, i
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static unsafe void InvokeJSFunction(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
         {
-            ObjectDisposedException.ThrowIf(jsFunction.IsDisposed, jsFunction);
+            jsFunction.AssertNotDisposed();
+
 #if FEATURE_WASM_THREADS
-            JSObject.AssertThreadAffinity(jsFunction);
+            // if we are on correct thread already, just call it
+            if (jsFunction.ProxyContext.IsCurrentThread())
+            {
+                InvokeJSFunctionCurrent(jsFunction, arguments);
+            }
+            else
+            {
+                DispatchJSFunctionSync(jsFunction, arguments);
+            }
+            // async functions are not implemented
+#else
+            InvokeJSFunctionCurrent(jsFunction, arguments);
 #endif
+        }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe void InvokeJSFunctionCurrent(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
+        {
             var functionHandle = (int)jsFunction.JSHandle;
             fixed (JSMarshalerArgument* ptr = arguments)
             {
-                Interop.Runtime.InvokeJSFunction(functionHandle, ptr);
+                Interop.Runtime.InvokeJSFunction(functionHandle, (nint)ptr);
                 ref JSMarshalerArgument exceptionArg = ref arguments[0];
                 if (exceptionArg.slot.Type != MarshalerType.None)
                 {
@@ -214,12 +233,33 @@ internal static unsafe void InvokeJSFunction(JSObject jsFunction, Span<JSMarshal
             }
         }
 
+
+#if FEATURE_WASM_THREADS
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe void DispatchJSFunctionSync(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
+        {
+            var args = (nint)Unsafe.AsPointer(ref arguments[0]);
+            var functionHandle = jsFunction.JSHandle;
+
+            jsFunction.ProxyContext.SynchronizationContext.Send(static o =>
+            {
+                var state = ((nint functionHandle, nint args))o!;
+                Interop.Runtime.InvokeJSFunction(state.functionHandle, state.args);
+            }, (functionHandle, args));
+
+            ref JSMarshalerArgument exceptionArg = ref arguments[0];
+            if (exceptionArg.slot.Type != MarshalerType.None)
+            {
+                JSHostImplementation.ThrowException(ref exceptionArg);
+            }
+        }
+#endif
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static unsafe void InvokeJSImportImpl(JSFunctionBinding signature, Span<JSMarshalerArgument> arguments)
         {
 #if FEATURE_WASM_THREADS
             var targetContext = JSProxyContext.SealJSImportCapturing();
-            JSProxyContext.AssertIsInteropThread();
             arguments[0].slot.ContextHandle = targetContext.ContextHandle;
             arguments[1].slot.ContextHandle = targetContext.ContextHandle;
 #else
@@ -229,20 +269,36 @@ internal static unsafe void InvokeJSImportImpl(JSFunctionBinding signature, Span
             if (signature.IsAsync)
             {
                 // pre-allocate the result handle and Task
-                var holder = new JSHostImplementation.PromiseHolder(targetContext);
+                var holder = targetContext.CreatePromiseHolder();
                 arguments[1].slot.Type = MarshalerType.TaskPreCreated;
                 arguments[1].slot.GCHandle = holder.GCHandle;
             }
 
-            fixed (JSMarshalerArgument* ptr = arguments)
+#if FEATURE_WASM_THREADS
+            // if we are on correct thread already or this is synchronous call, just call it
+            if (targetContext.IsCurrentThread())
             {
-                Interop.Runtime.InvokeJSImport(signature.ImportHandle, ptr);
-                ref JSMarshalerArgument exceptionArg = ref arguments[0];
-                if (exceptionArg.slot.Type != MarshalerType.None)
+                InvokeJSImportCurrent(signature, arguments);
+
+#if DEBUG
+                if (signature.IsAsync && arguments[1].slot.Type == MarshalerType.None)
                 {
-                    JSHostImplementation.ThrowException(ref exceptionArg);
+                    throw new InvalidOperationException("null Task/Promise return is not supported");
                 }
+#endif
+
             }
+            else if (!signature.IsAsync)
+            {
+                DispatchJSImportSync(signature, targetContext, arguments);
+            }
+            else
+            {
+                DispatchJSImportAsync(signature, targetContext, arguments);
+            }
+#else
+            InvokeJSImportCurrent(signature, arguments);
+
             if (signature.IsAsync)
             {
                 // if js synchronously returned null
@@ -252,22 +308,83 @@ internal static unsafe void InvokeJSImportImpl(JSFunctionBinding signature, Span
                     holderHandle.Free();
                 }
             }
+#endif
         }
 
-        internal static unsafe JSFunctionBinding BindJSImportImpl(string functionName, string moduleName, ReadOnlySpan<JSMarshalerType> signatures)
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe void InvokeJSImportCurrent(JSFunctionBinding signature, Span<JSMarshalerArgument> arguments)
         {
+            fixed (JSMarshalerArgument* args = arguments)
+            {
 #if FEATURE_WASM_THREADS
-            JSProxyContext.AssertIsInteropThread();
+                Interop.Runtime.InvokeJSImportSync((nint)args, (nint)signature.Header);
+#else
+                Interop.Runtime.InvokeJSImport(signature.ImportHandle, (nint)args);
 #endif
+            }
+
+            ref JSMarshalerArgument exceptionArg = ref arguments[0];
+            if (exceptionArg.slot.Type != MarshalerType.None)
+            {
+                JSHostImplementation.ThrowException(ref exceptionArg);
+            }
+        }
+
+#if FEATURE_WASM_THREADS
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe void DispatchJSImportSync(JSFunctionBinding signature, JSProxyContext targetContext, Span<JSMarshalerArgument> arguments)
+        {
+            var args = (nint)Unsafe.AsPointer(ref arguments[0]);
+            var sig = (nint)signature.Header;
+
+            targetContext.SynchronizationContext.Send(static o =>
+            {
+                var state = ((nint args, nint sig))o!;
+                Interop.Runtime.InvokeJSImportSync(state.args, state.sig);
+            }, (args, sig));
+
+            ref JSMarshalerArgument exceptionArg = ref arguments[0];
+            if (exceptionArg.slot.Type != MarshalerType.None)
+            {
+                JSHostImplementation.ThrowException(ref exceptionArg);
+            }
+        }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe void DispatchJSImportAsync(JSFunctionBinding signature, JSProxyContext targetContext, Span<JSMarshalerArgument> arguments)
+        {
+            // this copy is freed in mono_wasm_invoke_import_async
+            var bytes = sizeof(JSMarshalerArgument) * arguments.Length;
+            void* cpy = (void*)Marshal.AllocHGlobal(bytes);
+            void* src = Unsafe.AsPointer(ref arguments[0]);
+            Unsafe.CopyBlock(cpy, src, (uint)bytes);
+            var sig = (nint)signature.Header;
+
+            targetContext.SynchronizationContext.Post(static o =>
+            {
+                var state = ((nint args, nint sig))o!;
+                Interop.Runtime.InvokeJSImportAsync(state.args, state.sig);
+            }, ((nint)cpy, sig));
+
+        }
+
+#endif
+
+        internal static unsafe JSFunctionBinding BindJSImportImpl(string functionName, string moduleName, ReadOnlySpan<JSMarshalerType> signatures)
+        {
             var signature = JSHostImplementation.GetMethodSignature(signatures, functionName, moduleName);
 
+#if !FEATURE_WASM_THREADS
+
             Interop.Runtime.BindJSImport(signature.Header, out int isException, out object exceptionMessage);
             if (isException != 0)
                 throw new JSException((string)exceptionMessage);
 
             JSHostImplementation.FreeMethodSignatureBuffer(signature);
 
+#endif
+
             return signature;
         }
 
@@ -286,18 +403,41 @@ internal static unsafe JSFunctionBinding BindManagedFunctionImpl(string fullyQua
             return signature;
         }
 
+#if !FEATURE_WASM_THREADS
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static unsafe void ResolveOrRejectPromise(Span<JSMarshalerArgument> arguments)
         {
             fixed (JSMarshalerArgument* ptr = arguments)
             {
-                Interop.Runtime.ResolveOrRejectPromise(ptr);
+                Interop.Runtime.ResolveOrRejectPromise((nint)ptr);
                 ref JSMarshalerArgument exceptionArg = ref arguments[0];
                 if (exceptionArg.slot.Type != MarshalerType.None)
                 {
                     JSHostImplementation.ThrowException(ref exceptionArg);
                 }
             }
         }
+#else
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe void ResolveOrRejectPromise(JSProxyContext targetContext, Span<JSMarshalerArgument> arguments)
+        {
+            // this copy is freed in mono_wasm_invoke_import_async
+            var bytes = sizeof(JSMarshalerArgument) * arguments.Length;
+            void* cpy = (void*)Marshal.AllocHGlobal(bytes);
+            void* src = Unsafe.AsPointer(ref arguments[0]);
+            Unsafe.CopyBlock(cpy, src, (uint)bytes);
+
+            // TODO: we could optimize away the work item allocation in JSSynchronizationContext if we synchronously dispatch this when we are already in the right thread.
+
+            // async
+            targetContext.SynchronizationContext.Post(static o =>
+            {
+                var args = (nint)o!;
+                Interop.Runtime.ResolveOrRejectPromise(args);
+            }, (nint)cpy);
+
+            // this never throws directly
+        }
+#endif
     }
 }

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

@@ -21,9 +21,6 @@ public static JSObject GlobalThis
         {
             get
             {
-#if FEATURE_WASM_THREADS
-                JSProxyContext.AssertIsInteropThread();
-#endif
                 return JavaScriptImports.GetGlobalThis();
             }
         }
@@ -35,9 +32,6 @@ public static JSObject DotnetInstance
         {
             get
             {
-#if FEATURE_WASM_THREADS
-                JSProxyContext.AssertIsInteropThread();
-#endif
                 return JavaScriptImports.GetDotnetInstance();
             }
         }
@@ -53,9 +47,6 @@ public static JSObject DotnetInstance
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static Task<JSObject> ImportAsync(string moduleName, string moduleUrl, CancellationToken cancellationToken = default)
         {
-#if FEATURE_WASM_THREADS
-            JSProxyContext.AssertIsInteropThread();
-#endif
             return JSHostImplementation.ImportAsync(moduleName, moduleUrl, cancellationToken);
         }
 

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

@@ -87,13 +87,14 @@ public static void ThrowException(ref JSMarshalerArgument arg)
         public static async Task<JSObject> ImportAsync(string moduleName, string moduleUrl, CancellationToken cancellationToken)
         {
             Task<JSObject> modulePromise = JavaScriptImports.DynamicImport(moduleName, moduleUrl);
-            var wrappedTask = CancelationHelper(modulePromise, cancellationToken);
+            var wrappedTask = CancellationHelper(modulePromise, cancellationToken);
             return await wrappedTask.ConfigureAwait(
                 ConfigureAwaitOptions.ContinueOnCapturedContext |
                 ConfigureAwaitOptions.ForceYielding); // this helps to finish the import before we bind the module in [JSImport]
         }
 
-        public static async Task<JSObject> CancelationHelper(Task<JSObject> jsTask, CancellationToken cancellationToken)
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static async Task<JSObject> CancellationHelper(Task<JSObject> jsTask, CancellationToken cancellationToken)
         {
             if (jsTask.IsCompletedSuccessfully)
             {
@@ -152,6 +153,9 @@ public static unsafe JSFunctionBinding GetMethodSignature(ReadOnlySpan<JSMarshal
             signature.ImportHandle = (int)JSFunctionBinding.nextImportHandle++;
 #endif
 
+#if DEBUG
+            signature.FunctionName = functionName;
+#endif
             for (int i = 0; i < argsCount; i++)
             {
                 var type = signature.Sigs[i] = types[i + 1]._signatureType;