[browser][mt] Release all proxies of C# and JS objects on WebWorker

src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/BrowserWebSockets/BrowserWebSocket.cs

@@ -601,7 +601,7 @@ private async Task CancelationHelper(Task jsTask, CancellationToken cancellation
                     FastState = WebSocketState.Aborted;
                     throw new OperationCanceledException(cancellationToken);
                 }
-                if (ex.Message == "OperationCanceledException")
+                if (ex.Message == "Error: OperationCanceledException")
                 {
                     FastState = WebSocketState.Aborted;
                     throw new OperationCanceledException("The operation was cancelled.", ex, cancellationToken);

src/libraries/System.Runtime.InteropServices.JavaScript/src/System/Runtime/InteropServices/JavaScript/Interop/JavaScriptExports.cs

@@ -143,11 +143,8 @@ public static void ReleaseJSOwnedObjectByGCHandle(JSMarshalerArgument* arguments
             {
                 GCHandle handle = (GCHandle)arg_1.slot.GCHandle;
 
-                lock (JSHostImplementation.s_gcHandleFromJSOwnedObject)
-                {
-                    JSHostImplementation.s_gcHandleFromJSOwnedObject.Remove(handle.Target!);
-                    handle.Free();
-                }
+                JSHostImplementation.ThreadJsOwnedObjects.Remove(handle.Target!);
+                handle.Free();
             }
             catch (Exception ex)
             {

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

@@ -172,6 +172,7 @@ public static JSFunctionBinding BindManagedFunction(string fullyQualifiedName, i
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static unsafe void InvokeJSImpl(JSObject jsFunction, Span<JSMarshalerArgument> arguments)
         {
+            ObjectDisposedException.ThrowIf(jsFunction.IsDisposed, jsFunction);
 #if FEATURE_WASM_THREADS
             JSObject.AssertThreadAffinity(jsFunction);
 #endif
@@ -205,10 +206,7 @@ internal static unsafe void InvokeImportImpl(IntPtr fnHandle, Span<JSMarshalerAr
         internal static unsafe JSFunctionBinding BindJSFunctionImpl(string functionName, string moduleName, ReadOnlySpan<JSMarshalerType> signatures)
         {
 #if FEATURE_WASM_THREADS
-            if (JSSynchronizationContext.CurrentJSSynchronizationContext == null)
-            {
-                throw new InvalidOperationException("Please use dedicated worker for working with JavaScript interop. See https://github.com/dotnet/runtime/blob/main/src/mono/wasm/threads.md#JS-interop-on-dedicated-threads ");
-            }
+            JSSynchronizationContext.AssertWebWorkerContext();
 #endif
 
             var signature = JSHostImplementation.GetMethodSignature(signatures);

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

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

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

@@ -32,13 +32,28 @@ public static Dictionary<int, WeakReference<JSObject>> ThreadCsOwnedObjects
         }
 
         // we use this to maintain identity of GCHandle for a managed object
-        public static Dictionary<object, IntPtr> s_gcHandleFromJSOwnedObject = new Dictionary<object, IntPtr>(ReferenceEqualityComparer.Instance);
+#if FEATURE_WASM_THREADS
+        [ThreadStatic]
+#endif
+        private static Dictionary<object, IntPtr>? s_jsOwnedObjects;
+
+        public static Dictionary<object, IntPtr> ThreadJsOwnedObjects
+        {
+            get
+            {
+                s_jsOwnedObjects ??= new Dictionary<object, IntPtr>(ReferenceEqualityComparer.Instance);
+                return s_jsOwnedObjects;
+            }
+        }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static void ReleaseCSOwnedObject(nint jsHandle)
         {
             if (jsHandle != IntPtr.Zero)
             {
+#if FEATURE_WASM_THREADS
+                JSSynchronizationContext.AssertWebWorkerContext();
+#endif
                 ThreadCsOwnedObjects.Remove((int)jsHandle);
                 Interop.Runtime.ReleaseCSOwnedObject(jsHandle);
             }
@@ -64,19 +79,19 @@ public static void ReleaseCSOwnedObject(nint jsHandle)
         public static IntPtr GetJSOwnedObjectGCHandle(object obj, GCHandleType handleType = GCHandleType.Normal)
         {
             if (obj == null)
+            {
                 return IntPtr.Zero;
+            }
 
-            IntPtr result;
-            lock (s_gcHandleFromJSOwnedObject)
+            IntPtr gcHandle;
+            if (ThreadJsOwnedObjects.TryGetValue(obj, out gcHandle))
             {
-                IntPtr gcHandle;
-                if (s_gcHandleFromJSOwnedObject.TryGetValue(obj, out gcHandle))
-                    return gcHandle;
-
-                result = (IntPtr)GCHandle.Alloc(obj, handleType);
-                s_gcHandleFromJSOwnedObject[obj] = result;
-                return result;
+                return gcHandle;
             }
+
+            IntPtr result = (IntPtr)GCHandle.Alloc(obj, handleType);
+            ThreadJsOwnedObjects[obj] = result;
+            return result;
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -188,6 +203,9 @@ public static unsafe void FreeMethodSignatureBuffer(JSFunctionBinding signature)
 
         public static JSObject CreateCSOwnedProxy(nint jsHandle)
         {
+#if FEATURE_WASM_THREADS
+            JSSynchronizationContext.AssertWebWorkerContext();
+#endif
             JSObject? res;
 
             if (!ThreadCsOwnedObjects.TryGetValue((int)jsHandle, out WeakReference<JSObject>? reference) ||
@@ -244,14 +262,55 @@ public static void InstallWebWorkerInterop(bool installJSSynchronizationContext,
 
         public static void UninstallWebWorkerInterop()
         {
-            var ctx = SynchronizationContext.Current as JSSynchronizationContext;
+            var ctx = JSSynchronizationContext.CurrentJSSynchronizationContext;
             var uninstallJSSynchronizationContext = ctx != null;
             if (uninstallJSSynchronizationContext)
             {
-                SynchronizationContext.SetSynchronizationContext(ctx!.previousSynchronizationContext);
-                JSSynchronizationContext.CurrentJSSynchronizationContext = null;
-                ctx.isDisposed = true;
+                try
+                {
+                    foreach (var jsObjectWeak in ThreadCsOwnedObjects.Values)
+                    {
+                        if (jsObjectWeak.TryGetTarget(out var jso))
+                        {
+                            jso.Dispose();
+                        }
+                    }
+                    foreach (var gch in ThreadJsOwnedObjects.Values)
+                    {
+                        GCHandle gcHandle = (GCHandle)gch;
+
+                        // if this is pending promise we reject it
+                        if (gcHandle.Target is TaskCallback holder)
+                        {
+                            unsafe
+                            {
+                                holder.Callback!.Invoke(null);
+                            }
+                        }
+                        gcHandle.Free();
+                    }
+                    SynchronizationContext.SetSynchronizationContext(ctx!.previousSynchronizationContext);
+                    JSSynchronizationContext.CurrentJSSynchronizationContext = null;
+                    ctx.isDisposed = true;
+                }
+                catch(Exception ex)
+                {
+                    Environment.FailFast($"Unexpected error in UninstallWebWorkerInterop, ManagedThreadId: {Thread.CurrentThread.ManagedThreadId}. " + ex);
+                }
+            }
+            else
+            {
+                if (ThreadCsOwnedObjects.Count > 0)
+                {
+                    Environment.FailFast($"There should be no JSObjects proxies on this thread, ManagedThreadId: {Thread.CurrentThread.ManagedThreadId}");
+                }
+                if (ThreadJsOwnedObjects.Count > 0)
+                {
+                    Environment.FailFast($"There should be no JS proxies of managed objects on this thread, ManagedThreadId: {Thread.CurrentThread.ManagedThreadId}");
+                }
             }
+            ThreadCsOwnedObjects.Clear();
+            ThreadJsOwnedObjects.Clear();
             Interop.Runtime.UninstallWebWorkerInterop(uninstallJSSynchronizationContext);
         }
 

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

@@ -27,6 +27,9 @@ public partial class JSObject : IDisposable
         public bool HasProperty(string propertyName)
         {
             ObjectDisposedException.ThrowIf(IsDisposed, this);
+#if FEATURE_WASM_THREADS
+            JSObject.AssertThreadAffinity(this);
+#endif
             return JavaScriptImports.HasProperty(this, propertyName);
         }
 

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

@@ -55,6 +55,16 @@ internal JSSynchronizationContext(Thread targetThread, IntPtr targetThreadId)
         {
         }
 
+        internal static void AssertWebWorkerContext()
+        {
+#if FEATURE_WASM_THREADS
+            if (CurrentJSSynchronizationContext == null)
+            {
+                throw new InvalidOperationException("Please use dedicated worker for working with JavaScript interop. See https://aka.ms/dotnet-JS-interop-threads");
+            }
+#endif
+        }
+
         private JSSynchronizationContext(Thread targetThread, IntPtr targetThreadId, WorkItemQueueType queue)
         {
             TargetThread = targetThread;

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

@@ -49,6 +49,11 @@ public unsafe void ToManaged(out Task? value)
             TaskCompletionSource tcs = new TaskCompletionSource(holder);
             JSHostImplementation.ToManagedCallback callback = (JSMarshalerArgument* arguments_buffer) =>
             {
+                if (arguments_buffer == null)
+                {
+                    tcs.TrySetException(new TaskCanceledException("WebWorker which is origin of the Promise is being terminated."));
+                    return;
+                }
                 ref JSMarshalerArgument arg_2 = ref arguments_buffer[3]; // set by caller when this is SetException call
                 // arg_3 set by caller when this is SetResult call, un-used here
                 if (arg_2.slot.Type != MarshalerType.None)
@@ -88,6 +93,12 @@ public unsafe void ToManaged<T>(out Task<T>? value, ArgumentToManagedCallback<T>
             TaskCompletionSource<T> tcs = new TaskCompletionSource<T>(holder);
             JSHostImplementation.ToManagedCallback callback = (JSMarshalerArgument* arguments_buffer) =>
             {
+                if (arguments_buffer == null)
+                {
+                    tcs.TrySetException(new TaskCanceledException("WebWorker which is origin of the Promise is being terminated."));
+                    return;
+                }
+
                 ref JSMarshalerArgument arg_2 = ref arguments_buffer[3]; // set by caller when this is SetException call
                 ref JSMarshalerArgument arg_3 = ref arguments_buffer[4]; // set by caller when this is SetResult call
                 if (arg_2.slot.Type != MarshalerType.None)

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

@@ -7,7 +7,7 @@
 namespace System.Runtime.InteropServices.JavaScript
 {
     /// <summary>
-    /// This is draft for possible public API of SynchronizationContext
+    /// Extensions of SynchronizationContext which propagate errors and return values
     /// </summary>
     public static class SynchronizationContextExtension
     {