Add async cleanup APIs for external execution mode

src/core/library.c

@@ -1856,6 +1856,7 @@ MsQuicOpenVersion(
     Api->RegistrationOpen = MsQuicRegistrationOpen;
     Api->RegistrationClose = MsQuicRegistrationClose;
     Api->RegistrationShutdown = MsQuicRegistrationShutdown;
+    Api->RegistrationCloseAsync = MsQuicRegistrationCloseAsync;
 
     Api->ConfigurationOpen = MsQuicConfigurationOpen;
     Api->ConfigurationClose = MsQuicConfigurationClose;
@@ -1932,6 +1933,34 @@ MsQuicClose(
     }
 }
 
+_IRQL_requires_max_(PASSIVE_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicCloseAsync(
+    _In_ _Pre_defensive_ const void* QuicApi,
+    _In_opt_ QUIC_CLOSE_COMPLETE_HANDLER Handler,
+    _In_opt_ void* Context
+    )
+{
+    if (QuicApi != NULL) {
+        QuicTraceLogVerbose(
+            LibraryMsQuicClose,
+            "[ api] MsQuicCloseAsync");
+
+        // There's actually nothing to wait for here, so we can just free the API table
+        // and call the completion handler immediately.
+        CXPLAT_FREE(QuicApi, QUIC_POOL_API);
+        MsQuicRelease();
+        MsQuicLibraryUnload();
+
+        if (Handler != NULL) {
+            Handler(Context);
+        }
+    }
+
+    return QUIC_STATUS_SUCCESS;
+}
+
 _IRQL_requires_max_(DISPATCH_LEVEL)
 QUIC_BINDING*
 QuicLibraryLookupBinding(

src/core/registration.c

@@ -143,6 +143,35 @@ MsQuicRegistrationOpen(
     return Status;
 }
 
+typedef struct QUIC_REGISTRATION_CLOSE_COMPLETE_CONTEXT {
+    QUIC_REGISTRATION_CLOSE_COMPLETE_HANDLER Handler;
+    void* Context;
+    QUIC_REGISTRATION* Registration;
+} QUIC_REGISTRATION_CLOSE_COMPLETE_CONTEXT;
+
+_IRQL_requires_max_(PASSIVE_LEVEL)
+void
+QuicRegistrationRundownComplete(
+    _In_ void* Context
+    )
+{
+    QUIC_REGISTRATION_CLOSE_COMPLETE_CONTEXT* CloseContext =
+        (QUIC_REGISTRATION_CLOSE_COMPLETE_CONTEXT*)Context;
+    QUIC_REGISTRATION* Registration = CloseContext->Registration;
+
+    QuicWorkerPoolUninitialize(Registration->WorkerPool);
+    CxPlatRundownUninitialize(&Registration->Rundown);
+    CxPlatDispatchLockUninitialize(&Registration->ConnectionLock);
+    CxPlatLockUninitialize(&Registration->ConfigLock);
+
+    if (CloseContext->Handler != NULL) {
+        CloseContext->Handler(CloseContext->Context);
+    }
+
+    CXPLAT_FREE(Registration, QUIC_POOL_REGISTRATION);
+    CXPLAT_FREE(CloseContext, QUIC_POOL_GENERIC);
+}
+
 _IRQL_requires_max_(PASSIVE_LEVEL)
 void
 QUIC_API
@@ -187,6 +216,75 @@ MsQuicRegistrationClose(
     }
 }
 
+_IRQL_requires_max_(PASSIVE_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicRegistrationCloseAsync(
+    _In_ _Pre_defensive_ __drv_freesMem(Mem)
+        HQUIC Handle,
+    _In_opt_ QUIC_REGISTRATION_CLOSE_COMPLETE_HANDLER Handler,
+    _In_opt_ void* Context
+    )
+{
+    QUIC_STATUS Status = QUIC_STATUS_INVALID_PARAMETER;
+
+    if (Handle != NULL && Handle->Type == QUIC_HANDLE_TYPE_REGISTRATION) {
+        QuicTraceEvent(
+            ApiEnter,
+            "[ api] Enter %u (%p).",
+            QUIC_TRACE_API_REGISTRATION_CLOSE_ASYNC,
+            Handle);
+
+#pragma prefast(suppress: __WARNING_25024, "Pointer cast already validated.")
+        QUIC_REGISTRATION* Registration = (QUIC_REGISTRATION*)Handle;
+        QUIC_REGISTRATION_CLOSE_COMPLETE_CONTEXT* CloseContext = NULL;
+
+        // Allocate the context for the completion callback
+        CloseContext = CXPLAT_ALLOC_NONPAGED(sizeof(QUIC_REGISTRATION_CLOSE_COMPLETE_CONTEXT), QUIC_POOL_GENERIC);
+        if (CloseContext == NULL) {
+            Status = QUIC_STATUS_OUT_OF_MEMORY;
+            goto Error;
+        }
+
+        CloseContext->Handler = Handler;
+        CloseContext->Context = Context;
+        CloseContext->Registration = Registration;
+
+        QuicTraceEvent(
+            RegistrationCleanup,
+            "[ reg][%p] Cleaning up (async)",
+            Registration);
+
+        if (Registration->ExecProfile != QUIC_EXECUTION_PROFILE_TYPE_INTERNAL) {
+            CxPlatLockAcquire(&MsQuicLib.Lock);
+            CxPlatListEntryRemove(&Registration->Link);
+            CxPlatLockRelease(&MsQuicLib.Lock);
+        }
+
+        Status = QUIC_STATUS_SUCCESS;
+        if (!CxPlatRundownRelease(&Registration->Rundown)) {
+            // No more references, complete synchronously
+            QuicRegistrationRundownComplete(CloseContext);
+        } else {
+            // TODO: Need to implement and use the platform completion callback mechanism
+            // This is a temporary solution that will cause a memory leak if we get here
+            QuicTraceLogWarning(
+                RegistrationCleanupAsyncNotSupported,
+                "[ reg][%p] Async cleanup not fully implemented! Memory will be leaked.",
+                Registration);
+            Status = QUIC_STATUS_NOT_SUPPORTED;
+            goto Error;
+        }
+
+        QuicTraceEvent(
+            ApiExit,
+            "[ api] Exit");
+    }
+
+Error:
+    return Status;
+}
+
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
 QUIC_API

src/inc/msquic.h

@@ -370,6 +370,26 @@ void
 
 typedef QUIC_CREDENTIAL_LOAD_COMPLETE *QUIC_CREDENTIAL_LOAD_COMPLETE_HANDLER;
 
+typedef
+_IRQL_requires_max_(PASSIVE_LEVEL)
+_Function_class_(QUIC_REGISTRATION_CLOSE_COMPLETE)
+void
+(QUIC_API QUIC_REGISTRATION_CLOSE_COMPLETE)(
+    _In_opt_ void* Context
+    );
+
+typedef QUIC_REGISTRATION_CLOSE_COMPLETE *QUIC_REGISTRATION_CLOSE_COMPLETE_HANDLER;
+
+typedef
+_IRQL_requires_max_(PASSIVE_LEVEL)
+_Function_class_(QUIC_CLOSE_COMPLETE)
+void
+(QUIC_API QUIC_CLOSE_COMPLETE)(
+    _In_opt_ void* Context
+    );
+
+typedef QUIC_CLOSE_COMPLETE *QUIC_CLOSE_COMPLETE_HANDLER;
+
 typedef struct QUIC_CERTIFICATE_HASH {
     uint8_t ShaHash[20];
 } QUIC_CERTIFICATE_HASH;
@@ -1101,6 +1121,21 @@ void
         HQUIC Registration
     );
 
+//
+// Asynchronously closes the registration. Instead of synchronizing cleanup,
+// this function registers a callback to be invoked when cleanup is complete.
+// This avoids deadlocks in single-threaded execution environments.
+//
+typedef
+_IRQL_requires_max_(PASSIVE_LEVEL)
+QUIC_STATUS
+(QUIC_API * QUIC_REGISTRATION_CLOSE_ASYNC_FN)(
+    _In_ _Pre_defensive_ __drv_freesMem(Mem)
+        HQUIC Registration,
+    _In_opt_ QUIC_REGISTRATION_CLOSE_COMPLETE_HANDLER Handler,
+    _In_opt_ void* Context
+    );
+
 //
 // Calls shutdown for all connections in this registration. Don't call on a
 // MsQuic callback thread or it might deadlock.
@@ -1809,6 +1844,8 @@ typedef struct QUIC_API_TABLE {
     QUIC_CONNECTION_OPEN_IN_PARTITION_FN
                                         ConnectionOpenInPartition;   // Available from v2.5
 
+    QUIC_REGISTRATION_CLOSE_ASYNC_FN   RegistrationCloseAsync;       // Available from v2.5
+
 #ifdef QUIC_API_ENABLE_PREVIEW_FEATURES
     QUIC_STREAM_PROVIDE_RECEIVE_BUFFERS_FN
                                         StreamProvideReceiveBuffers; // Available from v2.5
@@ -1866,6 +1903,19 @@ MsQuicClose(
     _In_ _Pre_defensive_ const void* QuicApi
     );
 
+//
+// Asynchronously cleans up the function table returned from MsQuicOpenVersion
+// and releases the reference on the API. Calls the provided callback when done.
+//
+_IRQL_requires_max_(PASSIVE_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicCloseAsync(
+    _In_ _Pre_defensive_ const void* QuicApi,
+    _In_opt_ QUIC_CLOSE_COMPLETE_HANDLER Handler,
+    _In_opt_ void* Context
+    );
+
 #endif
 
 _IRQL_requires_max_(PASSIVE_LEVEL)

src/inc/msquic.hpp

@@ -446,6 +446,22 @@ class MsQuicApi : public QUIC_API_TABLE {
             memset(thisTable, 0, sizeof(*thisTable));
         }
     }
+
+    QUIC_STATUS CloseAsync(
+        _In_opt_ QUIC_CLOSE_COMPLETE_HANDLER Handler,
+        _In_opt_ void* Context) noexcept {
+        QUIC_STATUS Status = QUIC_STATUS_INVALID_STATE;
+        if (QUIC_SUCCEEDED(InitStatus)) {
+            Status = MsQuicCloseAsync(ApiTable, Handler, Context);
+            if (QUIC_SUCCEEDED(Status)) {
+                ApiTable = nullptr;
+                QUIC_API_TABLE* thisTable = this;
+                memset(thisTable, 0, sizeof(*thisTable));
+            }
+        }
+        return Status;
+    }
+
     QUIC_STATUS GetInitStatus() const noexcept { return InitStatus; }
     bool IsValid() const noexcept { return QUIC_SUCCEEDED(InitStatus); }
 };
@@ -543,6 +559,27 @@ struct MsQuicRegistration {
         ) noexcept {
         MsQuic->RegistrationShutdown(Handle, Flags, ErrorCode);
     }
+
+    QUIC_STATUS CloseAsync(
+        _In_ QUIC_REGISTRATION_CLOSE_COMPLETE_HANDLER Handler,
+        _In_opt_ void* Context = nullptr
+        ) noexcept {
+        QUIC_STATUS Status = QUIC_STATUS_INVALID_STATE;
+        if (Handle != nullptr) {
+            if (CloseAllConnectionsOnDelete) {
+                MsQuic->RegistrationShutdown(
+                    Handle,
+                    QUIC_CONNECTION_SHUTDOWN_FLAG_SILENT,
+                    1);
+            }
+            Status = MsQuic->RegistrationCloseAsync(Handle, Handler, Context);
+            if (QUIC_SUCCEEDED(Status)) {
+                // Prevent auto-cleanup in destructor, async cleanup will handle it
+                Handle = nullptr;
+            }
+        }
+        return Status;
+    }
 };
 
 class MsQuicAlpn {

src/inc/quic_trace.h

@@ -108,6 +108,7 @@ typedef enum QUIC_TRACE_API_TYPE {
     QUIC_TRACE_API_EXECUTION_CREATE,
     QUIC_TRACE_API_EXECUTION_DELETE,
     QUIC_TRACE_API_EXECUTION_POLL,
+    QUIC_TRACE_API_REGISTRATION_CLOSE_ASYNC,
     QUIC_TRACE_API_COUNT // Must be last
 } QUIC_TRACE_API_TYPE;
 

src/tools/execution/execution_windows.cpp

@@ -61,8 +61,14 @@ void QueueCleanupJob() {
     auto Sqe = new(std::nothrow) QUIC_SQE;
     ZeroMemory(&Sqe->Overlapped, sizeof(Sqe->Overlapped));
     Sqe->Completion = [](QUIC_CQE* Cqe) {
-        printf("Cleaning up...\n");
-        AllDone = true;
+        printf("Cleaning up asynchronously...\n");
+        printf("Note: Using the new MsQuicRegistrationCloseAsync API to avoid deadlocks in the external app-driven execution model.\n");
+        Registration->CloseAsync(
+            [](void* Context) {
+                printf("Registration closed asynchronously\n");
+                AllDone = true;
+            },
+            nullptr);
         delete CONTAINING_RECORD(Cqe->lpOverlapped, QUIC_SQE, Overlapped);
     };
     IOCP->Enqueue(&Sqe->Overlapped);