Fix native wait lifetime during asynchronous fence completion

src/ComputeSharp.Win32.D3D12/Windows/um/WinBase/Windows.cs

@@ -9,12 +9,20 @@ namespace ComputeSharp.Win32;
 
 internal static unsafe partial class Windows
 {
-    [DllImport("kernel32", ExactSpelling = true)]
-    public static extern BOOL RegisterWaitForSingleObject([NativeTypeName("PHANDLE")] HANDLE* phNewWaitObject, HANDLE hObject, [NativeTypeName("WAITORTIMERCALLBACK")] delegate* unmanaged<void*, byte, void> Callback, [NativeTypeName("PVOID")] void* Context, [NativeTypeName("ULONG")] uint dwMilliseconds, [NativeTypeName("ULONG")] uint dwFlags);
+    // Runtime marshalling is disabled for this assembly. LibraryImport generates
+    // the last-error capture needed by Marshal.GetLastPInvokeError; adding
+    // SetLastError to DllImport would rely on disabled runtime support.
+    [LibraryImport("kernel32.dll", SetLastError = true)]
+    public static partial BOOL RegisterWaitForSingleObject([NativeTypeName("PHANDLE")] HANDLE* phNewWaitObject, HANDLE hObject, [NativeTypeName("WAITORTIMERCALLBACK")] delegate* unmanaged<void*, byte, void> Callback, [NativeTypeName("PVOID")] void* Context, [NativeTypeName("ULONG")] uint dwMilliseconds, [NativeTypeName("ULONG")] uint dwFlags);
+
+    // The completion argument selects blocking or nonblocking cancellation.
+    // Capture last error to distinguish ERROR_IO_PENDING from cancellation failure.
+    [LibraryImport("kernel32.dll", SetLastError = true)]
+    public static partial BOOL UnregisterWaitEx(HANDLE WaitHandle, HANDLE CompletionEvent);
 
     [DllImport("kernel32", ExactSpelling = true)]
     public static extern BOOL UnregisterWait(HANDLE WaitHandle);
 
     [NativeTypeName("#define INFINITE 0xFFFFFFFF")]
     public const uint INFINITE = 0xFFFFFFFF;
-}
+}

src/ComputeSharp.Win32.D3D12/Windows/um/synchapi/Windows.cs

@@ -9,10 +9,12 @@ namespace ComputeSharp.Win32;
 
 internal static unsafe partial class Windows
 {
-    [DllImport("kernel32", ExactSpelling = true)]
-    public static extern HANDLE CreateEventW([NativeTypeName("LPSECURITY_ATTRIBUTES")] SECURITY_ATTRIBUTES* lpEventAttributes, BOOL bManualReset, BOOL bInitialState, [NativeTypeName("LPCWSTR")] ushort* lpName);
+    // Source-generated last-error capture works with DisableRuntimeMarshalling
+    // and lets fence-wait setup report the native error if event creation fails.
+    [LibraryImport("kernel32.dll", SetLastError = true)]
+    public static partial HANDLE CreateEventW([NativeTypeName("LPSECURITY_ATTRIBUTES")] SECURITY_ATTRIBUTES* lpEventAttributes, BOOL bManualReset, BOOL bInitialState, [NativeTypeName("LPCWSTR")] ushort* lpName);
 
     [DllImport("kernel32", ExactSpelling = true)]
     [return: NativeTypeName("DWORD")]
     public static extern uint WaitForSingleObjectEx(HANDLE hHandle, [NativeTypeName("DWORD")] uint dwMilliseconds, BOOL bAlertable);
-}
+}

src/ComputeSharp/Graphics/Commands/FenceWaitPublication.cs

@@ -0,0 +1,51 @@
+using System;
+using System.Threading;
+
+namespace ComputeSharp.Graphics.Commands;
+
+// The callback may run before RegisterWaitForSingleObject returns its handle.
+// Exactly one of the publisher/callback owns cleanup, after both have arrived.
+internal struct FenceWaitPublication
+{
+    // 0: neither participant arrived; 1: handle published; 2: callback arrived;
+    // 3: cleanup claimed. The first arrival must not access the context again
+    // after returning false, because the other participant may then free it.
+    private int state;
+
+    public bool OnRegistered()
+    {
+        return Arrive(1, 2);
+    }
+
+    public bool OnSignaled()
+    {
+        return Arrive(2, 1);
+    }
+
+    private bool Arrive(int own, int other)
+    {
+        // Interlocked publishes the preceding context writes as well as the state.
+        // Only the second arrival can claim cleanup, once the handle is available
+        // and the callback has observed fence completion.
+        int previous = Interlocked.CompareExchange(ref this.state, own, 0);
+        if (previous == 0)
+        {
+            return false;
+        }
+
+        if (previous == other && Interlocked.CompareExchange(ref this.state, 3, other) == other)
+        {
+            return true;
+        }
+
+        throw new InvalidOperationException("Fence wait completion/publication occurred more than once.");
+    }
+
+    public static bool UnregisterAccepted(bool success, int error, bool insideCallback)
+    {
+        // With a one-shot registration, this callback is the only outstanding
+        // callback. Nonblocking cancellation may report it as still running;
+        // the blocking publisher path must instead confirm completion.
+        return success || (insideCallback && error == 997); // ERROR_IO_PENDING: this callback is still executing.
+    }
+}

src/ComputeSharp/Graphics/GraphicsDevice.Execute.cs

@@ -134,6 +134,17 @@ private struct CallbackContext
         /// The event handle set when the target fence value is reached.
         /// </summary>
         public HANDLE EventHandle;
+
+        /// <summary>
+        /// The registration to cancel before releasing the event and callback context.
+        /// A one-shot wait still requires explicit unregistration.
+        /// </summary>
+        public HANDLE WaitHandle;
+
+        /// <summary>
+        /// Coordinates cleanup when the callback races publication of the registration handle.
+        /// </summary>
+        public FenceWaitPublication Publication;
     }
 
     /// <summary>
@@ -151,39 +162,101 @@ private static WaitForFenceValueTaskSource WaitForFenceAsync(
         ID3D12GraphicsCommandList* d3D12GraphicsCommandList,
         ID3D12CommandAllocator* d3D12CommandAllocator)
     {
-        HANDLE eventHandle = Windows.CreateEventW(null, Windows.FALSE, Windows.FALSE, null);
-
-        device.d3D12ComputeFence.Get()->SetEventOnCompletion(d3D12FenceValue, eventHandle).Assert();
-
-        WaitForFenceValueTaskSource waitForFenceValueTaskSource = WaitForFenceValueTaskSource.Rent();
-        CallbackContext* callbackContext = (CallbackContext*)NativeMemory.Alloc((nuint)sizeof(CallbackContext));
-
-        callbackContext->WaitForFenceValueTaskSourceHandle = GCHandle.Alloc(waitForFenceValueTaskSource);
-        callbackContext->GraphicsDeviceHandle = GCHandle.Alloc(device);
-        callbackContext->D3D12GraphicsCommandList = d3D12GraphicsCommandList;
-        callbackContext->D3D12CommandAllocator = d3D12CommandAllocator;
-        callbackContext->EventHandle = eventHandle;
-
-        HANDLE waitHandle;
+        WaitForFenceValueTaskSource? waitForFenceValueTaskSource = null;
+        CallbackContext* callbackContext = null;
+        bool referenceHeld = false;
+        try
+        {
+            waitForFenceValueTaskSource = WaitForFenceValueTaskSource.Rent();
+            // Zero initialization supplies the handshake's initial state and lets
+            // failure cleanup distinguish acquired resources from uninitialized fields.
+            callbackContext = (CallbackContext*)NativeMemory.AllocZeroed((nuint)sizeof(CallbackContext));
+            if (callbackContext is null)
+            {
+                throw new OutOfMemoryException();
+            }
 
-        device.GetReferenceTracker().DangerousAddRef();
+            callbackContext->WaitForFenceValueTaskSourceHandle = GCHandle.Alloc(waitForFenceValueTaskSource);
+            callbackContext->GraphicsDeviceHandle = GCHandle.Alloc(device);
+            callbackContext->D3D12GraphicsCommandList = d3D12GraphicsCommandList;
+            callbackContext->D3D12CommandAllocator = d3D12CommandAllocator;
+            device.GetReferenceTracker().DangerousAddRef();
+            referenceHeld = true;
+
+            callbackContext->EventHandle = Windows.CreateEventW(null, Windows.FALSE, Windows.FALSE, null);
+            // Do not pass a failed event allocation to SetEventOnCompletion:
+            // a null event has the distinct meaning of a synchronous fence wait.
+            if (callbackContext->EventHandle == default)
+            {
+                throw new Win32Exception(Marshal.GetLastPInvokeError());
+            }
 
-        int result = Windows.RegisterWaitForSingleObject(
-            phNewWaitObject: &waitHandle,
-            hObject: eventHandle,
-            Callback: &WaitForSingleObjectCallbackForWaitForFenceAsync,
-            Context: callbackContext,
-            dwMilliseconds: Windows.INFINITE,
-            dwFlags: 0);
+            device.d3D12ComputeFence.Get()->SetEventOnCompletion(d3D12FenceValue, callbackContext->EventHandle).Assert();
+
+            // Queue at most one callback so that cleanup cannot race a second
+            // invocation. WT_EXECUTEONLYONCE does not release the registration.
+            HANDLE waitHandle;
+            int result = Windows.RegisterWaitForSingleObject(
+                phNewWaitObject: &waitHandle,
+                hObject: callbackContext->EventHandle,
+                Callback: &WaitForSingleObjectCallbackForWaitForFenceAsync,
+                Context: callbackContext,
+                dwMilliseconds: Windows.INFINITE,
+                dwFlags: 0x00000008); // WT_EXECUTEONLYONCE
+            if (result == 0)
+            {
+                throw new Win32Exception(Marshal.GetLastPInvokeError());
+            }
 
-        // The register is successful if the return value is nonzero
-        if (result == 0)
+            // Publish the handle before allowing cleanup. A callback arriving
+            // early records completion and returns without touching it again.
+            callbackContext->WaitHandle = waitHandle;
+        }
+        catch
         {
-            device.GetReferenceTracker().DangerousRelease();
+            // The command list was already submitted/detached. Drain it before
+            // closing an event registered with the fence or recycling its allocator.
+            try
+            {
+                device.d3D12ComputeFence.Get()->SetEventOnCompletion(d3D12FenceValue, default).Assert();
+                device.computeCommandListPool.Return(d3D12GraphicsCommandList, d3D12CommandAllocator);
+            }
+            finally
+            {
+                if (callbackContext is not null)
+                {
+                    if (callbackContext->WaitForFenceValueTaskSourceHandle.IsAllocated)
+                    {
+                        callbackContext->WaitForFenceValueTaskSourceHandle.Free();
+                    }
+
+                    if (callbackContext->GraphicsDeviceHandle.IsAllocated)
+                    {
+                        callbackContext->GraphicsDeviceHandle.Free();
+                    }
+
+                    if (callbackContext->EventHandle != default)
+                    {
+                        _ = Windows.CloseHandle(callbackContext->EventHandle);
+                    }
+
+                    NativeMemory.Free(callbackContext);
+                }
+
+                if (referenceHeld)
+                {
+                    device.GetReferenceTracker().DangerousRelease();
+                }
+            }
 
-            NativeMemory.Free(callbackContext);
+            throw;
+        }
 
-            default(Win32Exception).Throw(E.E_FAIL);
+        // Keep completion outside the setup-failure handler: after publication,
+        // the callback may own cleanup and may already have freed the context.
+        if (callbackContext->Publication.OnRegistered())
+        {
+            CompleteFenceWait(callbackContext, insideCallback: false);
         }
 
         return waitForFenceValueTaskSource;
@@ -199,7 +272,33 @@ private static void WaitForSingleObjectCallbackForWaitForFenceAsync(void* pConte
     {
         CallbackContext* callbackContext = (CallbackContext*)pContext;
 
+        if (callbackContext->Publication.OnSignaled())
+        {
+            CompleteFenceWait(callbackContext, insideCallback: true);
+        }
+    }
+
+    private static void CompleteFenceWait(CallbackContext* callbackContext, bool insideCallback)
+    {
+        // Never block on our own callback: NULL requests nonblocking cancellation.
+        // The early-callback/publisher path uses INVALID_HANDLE_VALUE to join the
+        // callback before freeing its context. Neither path closes the event first.
+        int unregistered = Windows.UnregisterWaitEx(callbackContext->WaitHandle,
+            insideCallback ? default : (HANDLE)(-1));
+        int error = unregistered == 0 ? Marshal.GetLastPInvokeError() : 0;
+        Exception? failure = FenceWaitPublication.UnregisterAccepted(unregistered != 0, error, insideCallback)
+            ? null : new Win32Exception(error);
+
         WaitForFenceValueTaskSource waitForFenceValueTaskSource = Unsafe.As<WaitForFenceValueTaskSource>(callbackContext->WaitForFenceValueTaskSourceHandle.Target)!;
+        if (failure is not null)
+        {
+            // Cancellation did not establish that the OS released the wait.
+            // Quarantine its resources and fault the caller rather than free
+            // a context/event whose native ownership is now uncertain.
+            waitForFenceValueTaskSource.Fail(failure);
+            return;
+        }
+
         GraphicsDevice device = Unsafe.As<GraphicsDevice>(callbackContext->GraphicsDeviceHandle.Target)!;
         ID3D12GraphicsCommandList* d3D12GraphicsCommandList = callbackContext->D3D12GraphicsCommandList;
         ID3D12CommandAllocator* d3D12CommandAllocator = callbackContext->D3D12CommandAllocator;
@@ -277,6 +376,11 @@ public void Complete()
             this.manualResetValueTaskSource.SetResult(null);
         }
 
+        public void Fail(Exception error)
+        {
+            this.manualResetValueTaskSource.SetException(error);
+        }
+
         /// <inheritdoc/>
         public void GetResult(short token)
         {
@@ -306,4 +410,4 @@ public void OnCompleted(Action<object?> continuation, object? state, short token
             this.manualResetValueTaskSource.OnCompleted(continuation, state, token, flags);
         }
     }
-}
+}