Fix no-context JSReference leak and TSFN post/release race (follow-ups to #492)

src/NodeApi/Interop/JSSynchronizationContext.cs

@@ -254,6 +254,10 @@ internal sealed unsafe class JSTsfnSynchronizationContext : JSSynchronizationCon
     private readonly JSRuntime _runtime;
     private readonly napi_env _env;
     private readonly JSThreadSafeFunction _tsfn;
+
+    // Coordinates in-flight NonBlockingCall invocations with TSFN release during teardown, so that
+    // a poster can never call into the TSFN after (or while) it is being released.
+    private readonly TsfnCallGate _callGate = new();
     private GCHandle _cleanupHandle;
 
     public JSTsfnSynchronizationContext()
@@ -298,6 +302,11 @@ public override void Dispose()
 
         base.Dispose();
 
+        // Prevent any new NonBlockingCall from starting and wait for in-flight calls to finish
+        // before releasing the TSFN. This closes the post-then-release race in which a poster could
+        // call into the TSFN after the environment cleanup hook has released it.
+        _callGate.Close();
+
         // Destroy TSFN by releasing last thread use count.
         // TSFN is deleted after this point and must not be used.
         _tsfn.Release();
@@ -348,9 +357,16 @@ public override void CloseAsyncScope()
 
     public override void Post(SendOrPostCallback callback, object? state)
     {
-        if (IsDisposed) return;
+        if (!_callGate.TryEnter()) return;
 
-        _tsfn.NonBlockingCall(() => callback(state));
+        try
+        {
+            _tsfn.NonBlockingCall(() => callback(state));
+        }
+        finally
+        {
+            _callGate.Exit();
+        }
     }
 
     public override void Send(SendOrPostCallback callback, object? state)
@@ -361,14 +377,25 @@ public override void Send(SendOrPostCallback callback, object? state)
             return;
         }
 
-        if (IsDisposed) return;
-
         using ManualResetEvent syncEvent = new(false);
-        _tsfn.NonBlockingCall(() =>
+
+        // The gate is held only around the native NonBlockingCall, not around the wait below: the
+        // posted callback runs on the JS thread and could otherwise deadlock TSFN release.
+        if (!_callGate.TryEnter()) return;
+
+        try
         {
-            callback(state);
-            syncEvent.Set();
-        });
+            _tsfn.NonBlockingCall(() =>
+            {
+                callback(state);
+                syncEvent.Set();
+            });
+        }
+        finally
+        {
+            _callGate.Exit();
+        }
+
         syncEvent.WaitOne();
     }
 }

src/NodeApi/Interop/TsfnCallGate.cs

@@ -0,0 +1,97 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Threading;
+
+namespace Microsoft.JavaScript.NodeApi.Interop;
+
+/// <summary>
+/// Coordinates in-flight native thread-safe-function (TSFN) calls with TSFN release, so that a
+/// caller cannot invoke the TSFN after (or while) it is being released.
+/// </summary>
+/// <remarks>
+/// <para>
+/// <see cref="JSTsfnSynchronizationContext"/> posts work to the JS thread by calling into a native
+/// TSFN, and releases that TSFN from an environment cleanup hook during teardown. Without
+/// coordination there is a race: a poster can observe the context as not-yet-disposed, then the
+/// cleanup hook can release the TSFN, and then the poster calls into the released TSFN &mdash; a
+/// native use-after-release that a managed <c>try/catch</c> cannot guard.
+/// </para>
+/// <para>
+/// This gate closes that window. A caller wraps each native TSFN call in
+/// <see cref="TryEnter"/>/<see cref="Exit"/>; the owner calls <see cref="Close"/> before releasing
+/// the TSFN. <see cref="Close"/> atomically prevents any further <see cref="TryEnter"/> from
+/// succeeding and then waits for all in-flight calls to finish, guaranteeing that no native call
+/// is in progress and none can start once <see cref="Close"/> returns.
+/// </para>
+/// </remarks>
+internal sealed class TsfnCallGate
+{
+    // High bit marks the gate as closed; the remaining bits count in-flight calls. Because Enter
+    // never increments once the closed bit is set, and Exit is only called after a successful
+    // Enter, the count part never borrows into the closed bit.
+    private const int ClosedFlag = unchecked((int)0x80000000);
+    private const int CountMask = 0x7FFFFFFF;
+
+    private int _state;
+
+    /// <summary>
+    /// Gets a value indicating whether the gate has been closed.
+    /// </summary>
+    public bool IsClosed => (Volatile.Read(ref _state) & ClosedFlag) != 0;
+
+    /// <summary>
+    /// Attempts to enter the gate for a single native call. When this returns true the caller must
+    /// pair it with exactly one call to <see cref="Exit"/> once the native call has returned.
+    /// </summary>
+    /// <returns>True if the call may proceed; false if the gate is closed and the call must be
+    /// skipped.</returns>
+    public bool TryEnter()
+    {
+        int state = Volatile.Read(ref _state);
+        while ((state & ClosedFlag) == 0)
+        {
+            int updated = Interlocked.CompareExchange(ref _state, state + 1, state);
+            if (updated == state)
+            {
+                return true;
+            }
+
+            state = updated;
+        }
+
+        return false;
+    }
+
+    /// <summary>
+    /// Marks completion of a native call previously admitted by <see cref="TryEnter"/>.
+    /// </summary>
+    public void Exit() => Interlocked.Decrement(ref _state);
+
+    /// <summary>
+    /// Closes the gate so that no further calls can enter, then blocks until all in-flight calls
+    /// have exited. After this returns it is safe to release the underlying TSFN. Calling this more
+    /// than once is safe.
+    /// </summary>
+    public void Close()
+    {
+        int state = Volatile.Read(ref _state);
+        while ((state & ClosedFlag) == 0)
+        {
+            int updated = Interlocked.CompareExchange(
+                ref _state, state | ClosedFlag, state);
+            if (updated == state)
+            {
+                break;
+            }
+
+            state = updated;
+        }
+
+        SpinWait spin = default;
+        while ((Volatile.Read(ref _state) & CountMask) != 0)
+        {
+            spin.SpinOnce();
+        }
+    }
+}

src/NodeApi/JSReference.cs

@@ -2,9 +2,11 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using Microsoft.JavaScript.NodeApi.Interop;
+using Microsoft.JavaScript.NodeApi.Runtime;
 using static Microsoft.JavaScript.NodeApi.Runtime.JSRuntime;
 
 namespace Microsoft.JavaScript.NodeApi;
@@ -32,6 +34,30 @@ public class JSReference : IDisposable
     private readonly napi_env _env;
     private readonly JSRuntimeContext? _context;
 
+    /// <summary>
+    /// Environment-scoped queue of no-context references whose finalizers could not release them
+    /// on the JS thread, keyed by the owning environment handle.
+    /// </summary>
+    /// <remarks>
+    /// A no-context reference has no <see cref="JSSynchronizationContext"/>, so when its finalizer
+    /// runs on the GC finalizer thread (which has no JS scope) it cannot marshal
+    /// <c>DeleteReference</c> back to the owning JS thread. Rather than leaking the
+    /// <c>napi_ref</c> until the environment is destroyed, the finalizer enqueues the handle here.
+    /// The deletion is then performed the next time a scope for the same environment is active on
+    /// its JS thread (see <see cref="JSValueScope"/>), and a final drain runs when that
+    /// environment's root scope is disposed. Deferred handles are always drained while the
+    /// environment is still alive, so the queued <c>napi_ref</c> remains valid.
+    /// </remarks>
+    private static readonly ConcurrentDictionary<napi_env, ConcurrentQueue<napi_ref>>
+        s_pendingFinalizerDeletions = new();
+
+    /// <summary>
+    /// Total number of deferred deletions pending across all environments. Used as a cheap
+    /// fast-path gate so that the common case (nothing pending) avoids a dictionary lookup on
+    /// every scope entry.
+    /// </summary>
+    private static int s_pendingFinalizerDeletionCount;
+
     /// <summary>
     /// Creates a new instance of a <see cref="JSReference"/> that holds a strong or weak
     /// reference to a JS value.
@@ -388,17 +414,21 @@ private void DisposeFromFinalizer()
             {
                 // A no-context reference (for example one created from the native host scope) can
                 // only be deleted on the JS thread. CurrentOrNull is thread-static, so on the real
-                // GC finalizer thread it is null and this delete is skipped; the napi_ref is then
-                // reclaimed when the JS environment is destroyed. The guarded delete still runs if
-                // Dispose(disposing: false) is ever invoked on the owning JS thread. A no-context
-                // scope has no synchronization context, so the finalizer cannot marshal the delete
-                // to the JS thread; doing so would require an env-scoped cleanup queue in the
-                // native host (tracked as a follow-up).
+                // GC finalizer thread it is null and the inline delete below is skipped; the
+                // handle is instead deferred to be deleted the next time a scope for this
+                // environment is active on its JS thread (see JSValueScope). The guarded delete
+                // still runs directly if Dispose(disposing: false) is ever invoked on the owning
+                // JS thread. This releases the napi_ref promptly instead of leaking it until the
+                // JS environment is destroyed.
                 JSValueScope? scope = JSValueScope.CurrentOrNull;
                 if (scope != null && scope.UncheckedEnvironmentHandle == _env)
                 {
                     scope.Runtime.DeleteReference(_env, _handle);
                 }
+                else
+                {
+                    EnqueuePendingDeletion(_env, _handle);
+                }
             }
             else
             {
@@ -425,6 +455,78 @@ private void DisposeFromFinalizer()
         }
     }
 
+    /// <summary>
+    /// Records a no-context reference handle whose deletion was deferred because its finalizer
+    /// could not run on the owning JS thread. The handle is deleted later by
+    /// <see cref="DrainPendingDeletions"/> or <see cref="RemovePendingDeletions"/>.
+    /// </summary>
+    private static void EnqueuePendingDeletion(napi_env env, napi_ref handle)
+    {
+        s_pendingFinalizerDeletions
+            .GetOrAdd(env, _ => new ConcurrentQueue<napi_ref>())
+            .Enqueue(handle);
+        Interlocked.Increment(ref s_pendingFinalizerDeletionCount);
+    }
+
+    /// <summary>
+    /// Deletes any no-context reference handles that were deferred by a finalizer for the given
+    /// environment. Must be called on the JS thread that owns <paramref name="env"/> with an
+    /// active scope, so that <c>DeleteReference</c> is valid.
+    /// </summary>
+    /// <param name="env">The environment whose deferred deletions should be released.</param>
+    /// <param name="runtime">The JS runtime for the environment.</param>
+    internal static void DrainPendingDeletions(napi_env env, JSRuntime runtime)
+    {
+        // Fast path: avoid a dictionary lookup on every scope entry when nothing is pending.
+        if (Volatile.Read(ref s_pendingFinalizerDeletionCount) == 0)
+        {
+            return;
+        }
+
+        if (!s_pendingFinalizerDeletions.TryGetValue(
+            env, out ConcurrentQueue<napi_ref>? queue))
+        {
+            return;
+        }
+
+        DeletePendingHandles(env, runtime, queue);
+    }
+
+    /// <summary>
+    /// Performs a final drain of deferred deletions for an environment that is being torn down and
+    /// removes its queue, so the registry does not retain entries for environments that no longer
+    /// exist. Must be called on the owning JS thread while the environment is still alive.
+    /// </summary>
+    /// <param name="env">The environment being torn down.</param>
+    /// <param name="runtime">The JS runtime for the environment.</param>
+    internal static void RemovePendingDeletions(napi_env env, JSRuntime runtime)
+    {
+        if (!s_pendingFinalizerDeletions.TryRemove(
+            env, out ConcurrentQueue<napi_ref>? queue))
+        {
+            return;
+        }
+
+        DeletePendingHandles(env, runtime, queue);
+    }
+
+    private static void DeletePendingHandles(
+        napi_env env, JSRuntime runtime, ConcurrentQueue<napi_ref> queue)
+    {
+        while (queue.TryDequeue(out napi_ref handle))
+        {
+            Interlocked.Decrement(ref s_pendingFinalizerDeletionCount);
+            try
+            {
+                runtime.DeleteReference(env, handle);
+            }
+            catch
+            {
+                // Best effort: the environment may already be gone.
+            }
+        }
+    }
+
     ~JSReference()
     {
         // An exception escaping a finalizer terminates the process. Dispose(bool) is virtual, so a

src/NodeApi/JSValueScope.cs

@@ -359,13 +359,28 @@ public JSValueScope(
             CurrentOrNull = previousScope;
             throw;
         }
+
+        // Release any no-context reference deletions that a finalizer had to defer because it could
+        // not run on the JS thread. This scope is active on the owning JS thread, so deleting the
+        // deferred handles here is valid. The call is a cheap no-op when nothing is pending.
+        JSReference.DrainPendingDeletions(_env, Runtime);
     }
 
     public void Dispose()
     {
         if (IsDisposed) return;
         IsDisposed = true;
 
+        if (ScopeType == JSValueScopeType.Root || ScopeType == JSValueScopeType.NoContext)
+        {
+            // This environment-scoped scope is going away, so perform a final drain of any deferred
+            // no-context reference deletions and remove the environment's queue, preventing the
+            // registry from retaining entries for environments that no longer exist. The scope is
+            // still current here, so deletion runs on the owning JS thread while the environment is
+            // alive.
+            JSReference.RemovePendingDeletions(_env, Runtime);
+        }
+
         if (ScopeType != JSValueScopeType.NoContext)
         {
             napi_env env = RuntimeContext.EnvironmentHandle;

src/NodeApi/NodeApi.csproj

@@ -28,6 +28,10 @@
     <PackageReference Include="Microsoft.Bcl.AsyncInterfaces" />
   </ItemGroup>
 
+  <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.JavaScript.NodeApi.Test" />
+  </ItemGroup>
+
   <Target Name="SetRuntimeConfigValues" BeforeTargets="GenerateBuildRuntimeConfigurationFiles">
     <ItemGroup>
       <RuntimeHostConfigurationOption Include="System.Runtime.InteropServices.EnableConsumingManagedCodeFromNativeHosting" Value="true" />