Fix native wait lifetime during asynchronous fence completion

Author: andrewsil1

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

@@ -9,12 +9,15 @@ 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);
+    [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);
+
+    [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,10 @@ 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);
+    [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,42 @@
+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
+{
+    private int state;
+
+    public bool OnRegistered()
+    {
+        return Arrive(1, 2);
+    }
+
+    public bool OnSignaled()
+    {
+        return Arrive(2, 1);
+    }
+
+    private bool Arrive(int own, int other)
+    {
+        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)
+    {
+        return success || (insideCallback && error == 997); // ERROR_IO_PENDING: this callback is still executing.
+    }
+}

src/ComputeSharp/Graphics/GraphicsDevice.Execute.cs

@@ -134,6 +134,10 @@ private struct CallbackContext
         /// The event handle set when the target fence value is reached.
         /// </summary>
         public HANDLE EventHandle;
+
+        public HANDLE WaitHandle;
+
+        public FenceWaitPublication Publication;
     }
 
     /// <summary>
@@ -151,39 +155,93 @@ 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;
+        WaitForFenceValueTaskSource? waitForFenceValueTaskSource = null;
+        CallbackContext* callbackContext = null;
+        bool referenceHeld = false;
+        try
+        {
+            waitForFenceValueTaskSource = WaitForFenceValueTaskSource.Rent();
+            callbackContext = (CallbackContext*)NativeMemory.AllocZeroed((nuint)sizeof(CallbackContext));
+            if (callbackContext is null)
+            {
+                throw new OutOfMemoryException();
+            }
 
-        HANDLE waitHandle;
+            callbackContext->WaitForFenceValueTaskSourceHandle = GCHandle.Alloc(waitForFenceValueTaskSource);
+            callbackContext->GraphicsDeviceHandle = GCHandle.Alloc(device);
+            callbackContext->D3D12GraphicsCommandList = d3D12GraphicsCommandList;
+            callbackContext->D3D12CommandAllocator = d3D12CommandAllocator;
+            device.GetReferenceTracker().DangerousAddRef();
+            referenceHeld = true;
 
-        device.GetReferenceTracker().DangerousAddRef();
+            callbackContext->EventHandle = Windows.CreateEventW(null, Windows.FALSE, Windows.FALSE, null);
+            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();
+
+            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);
+        if (callbackContext->Publication.OnRegistered())
+        {
+            CompleteFenceWait(callbackContext, insideCallback: false);
         }
 
         return waitForFenceValueTaskSource;
@@ -199,7 +257,32 @@ 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. The early-callback/publisher path
+        // runs outside the callback and joins it before freeing its context.
+        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 +360,11 @@ public void Complete()
             this.manualResetValueTaskSource.SetResult(null);
         }
 
+        public void Fail(Exception error)
+        {
+            this.manualResetValueTaskSource.SetException(error);
+        }
+
         /// <inheritdoc/>
         public void GetResult(short token)
         {