Skip to content

Add async cleanup APIs for external execution mode #5127

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 15 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file modified scripts/build.ps1
100644 → 100755
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This file has no changes. Please remove it from this PR.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've removed this file from the PR as it has no changes.

Empty file.
105 changes: 105 additions & 0 deletions src/core/library.c
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,17 @@ QUIC_LIBRARY MsQuicLib = { 0 };

QUIC_TRACE_RUNDOWN_CALLBACK QuicTraceRundown;

// Forward declaration
_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
);

_IRQL_requires_max_(PASSIVE_LEVEL)
void
QuicLibApplyLoadBalancingSetting(
Expand Down Expand Up @@ -1856,6 +1867,7 @@ MsQuicOpenVersion(
Api->RegistrationOpen = MsQuicRegistrationOpen;
Api->RegistrationClose = MsQuicRegistrationClose;
Api->RegistrationShutdown = MsQuicRegistrationShutdown;
Api->RegistrationCloseAsync = MsQuicRegistrationCloseAsync;

Api->ConfigurationOpen = MsQuicConfigurationOpen;
Api->ConfigurationClose = MsQuicConfigurationClose;
Expand Down Expand Up @@ -1932,6 +1944,99 @@ MsQuicClose(
}
}

typedef struct QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT {
QUIC_CLOSE_COMPLETE_HANDLER Handler;
void* Context;
const void* QuicApi; // The API object to be freed
} QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT;

_IRQL_requires_max_(PASSIVE_LEVEL)
void
QuicLibraryCleanupComplete(
_In_ void* Context
)
{
QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT* CloseContext =
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just copy the whole context to a stack variable so you can immediately delete it

QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT CloseContext =
    *((QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT*)Context);
CXPLAT_FREE(CloseContext.QuicApi, QUIC_POOL_API);
CXPLAT_FREE(Context, QUIC_POOL_GENERIC);

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've updated the code to copy the entire context to a stack variable, which allows us to free the context memory before using its contents.

(QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT*)Context;

// Free the API table
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

None of the comments added in this PR follow our style of:

//
// Free the API table
//

Also, many of these comments are unnecessary and merely restate exactly what the code does on the next line. Please remove unnecessary comments and update the style for any ones that you keep.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've updated all the comments to follow the project style and removed unnecessary comments that merely restate what the code does.

CXPLAT_FREE(CloseContext->QuicApi, QUIC_POOL_API);

// Call the user-provided completion handler
if (CloseContext->Handler != NULL) {
CloseContext->Handler(CloseContext->Context);
}

// Free our context
CXPLAT_FREE(CloseContext, QUIC_POOL_GENERIC);
}

_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");

// Allocate context for the cleanup operation
QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT* CloseContext =
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we instead allocate the memory for the QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT with the original API memory allocation? Then we can just use it here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've refactored the code to allocate the context with the original API memory. I created an extended API table structure that includes the close handler and context. This approach eliminates the need for additional memory allocations when calling MsQuicCloseAsync.

CXPLAT_ALLOC_NONPAGED(sizeof(QUIC_LIBRARY_CLOSE_COMPLETE_CONTEXT), QUIC_POOL_GENERIC);
if (CloseContext == NULL) {
return QUIC_STATUS_OUT_OF_MEMORY;
}

CloseContext->Handler = Handler;
CloseContext->Context = Context;
CloseContext->QuicApi = QuicApi;

// Check if this is the last reference which may require blocking operations
BOOLEAN NeedsBlocking = FALSE;

CxPlatLockAcquire(&MsQuicLib.Lock);
NeedsBlocking = (MsQuicLib.OpenRefCount == 1);
CxPlatLockRelease(&MsQuicLib.Lock);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This logic is not correct. If the OpenRefCount is non-zero, it means there are other APIs using the library, so this call can immediately return success. That's not an error condition. If the ref is zero, then it needs to clean everything up (asynchronously).

But several things in the cleanup path (e.g., MsQuicRelease) can block, and they need to be async too. Once only once they complete should the rest of the cleanup code run.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, return QUIC_STATUS_PENDING if the code didn't complete inline and the callback will eventually be called.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've corrected the logic for handling OpenRefCount. Now when it's not the last reference, the function simply completes synchronously without any error condition. The comments have been improved to clarify that the last reference is the one that might require blocking operations.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've updated the code to return QUIC_STATUS_PENDING when the callback will be called asynchronously. The function now returns QUIC_STATUS_SUCCESS when the operation completes synchronously, and QUIC_STATUS_PENDING when the completion is deferred.


if (!NeedsBlocking) {
// No blocking needed, complete synchronously
CXPLAT_FREE(CloseContext, QUIC_POOL_GENERIC);
CXPLAT_FREE(QuicApi, QUIC_POOL_API);
MsQuicRelease();
MsQuicLibraryUnload();

if (Handler != NULL) {
Handler(Context);
}
} 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(
LibraryCleanupAsyncNotSupported,
"[ lib] Async cleanup not fully implemented! Potential deadlock can occur.");

// For now, we have to do the same thing as the existing implementation
// and accept the potential deadlock because the async completion mechanism
// isn't fully implemented yet.
CXPLAT_FREE(CloseContext, QUIC_POOL_GENERIC);
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(
Expand Down
136 changes: 128 additions & 8 deletions src/core/registration.c
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,63 @@ 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);

CXPLAT_FREE(Registration, QUIC_POOL_REGISTRATION);

// Store callback and context locally as we're about to free CloseContext
QUIC_REGISTRATION_CLOSE_COMPLETE_HANDLER Handler = CloseContext->Handler;
void* HandlerContext = CloseContext->Context;

CXPLAT_FREE(CloseContext, QUIC_POOL_GENERIC);

// Call the callback as the very last thing in the function
if (Handler != NULL) {
Handler(HandlerContext);
}
}

// Forward declaration
_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
);

// Helper function for MsQuicRegistrationClose
void
QUIC_API
QuicRegistrationCloseComplete(
void* Context
)
{
CXPLAT_EVENT* Event = (CXPLAT_EVENT*)Context;
CxPlatEventSet(*Event);
}

_IRQL_requires_max_(PASSIVE_LEVEL)
void
QUIC_API
Expand All @@ -158,12 +215,65 @@ MsQuicRegistrationClose(
QUIC_TRACE_API_REGISTRATION_CLOSE,
Handle);

// Create an event to wait on
CXPLAT_EVENT CompletionEvent;
CxPlatEventInitialize(&CompletionEvent, TRUE, FALSE);

// Call the async version with our event as the context
MsQuicRegistrationCloseAsync(
Handle,
QuicRegistrationCloseComplete,
&CompletionEvent);

// Wait for the event to be set by the completion callback
CxPlatEventWaitForever(CompletionEvent);

// Clean up the event
CxPlatEventUninitialize(CompletionEvent);

QuicTraceEvent(
ApiExit,
"[ api] Exit");
}
}

_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);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Put this as a part of the registration so we don't have to allocate in the cleanup path.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've updated the implementation to make the close handler and context part of the registration structure itself, eliminating the need for additional memory allocation during cleanup. This is implemented using a QUIC_REGISTRATION_EX structure that contains the registration and the callback information.

if (CloseContext == NULL) {
Status = QUIC_STATUS_OUT_OF_MEMORY;
goto Exit;
}

CloseContext->Handler = Handler;
CloseContext->Context = Context;
CloseContext->Registration = Registration;

QuicTraceEvent(
RegistrationCleanup,
"[ reg][%p] Cleaning up",
"[ reg][%p] Cleaning up (async)",
Registration);

if (Registration->ExecProfile != QUIC_EXECUTION_PROFILE_TYPE_INTERNAL) {
Expand All @@ -172,19 +282,29 @@ MsQuicRegistrationClose(
CxPlatLockRelease(&MsQuicLib.Lock);
}

CxPlatRundownReleaseAndWait(&Registration->Rundown);

QuicWorkerPoolUninitialize(Registration->WorkerPool);
CxPlatRundownUninitialize(&Registration->Rundown);
CxPlatDispatchLockUninitialize(&Registration->ConnectionLock);
CxPlatLockUninitialize(&Registration->ConfigLock);
Status = QUIC_STATUS_SUCCESS;

QuicTraceEvent(
RegistrationCleanup,
"[ reg][%p] Cleaning up (async)",
Registration);

// Release the rundown - this starts the cleanup process
CxPlatRundownRelease(&Registration->Rundown);

CXPLAT_FREE(Registration, QUIC_POOL_REGISTRATION);
// Call the callback function to continue cleaning up
// For now, this happens synchronously since we don't have a full
// async implementation yet. In the future, the platform would
// queue this for completion when rundown is fully released.
QuicRegistrationRundownComplete(CloseContext);

QuicTraceEvent(
ApiExit,
"[ api] Exit");
}

Exit:
return Status;
}

_IRQL_requires_max_(DISPATCH_LEVEL)
Expand Down
14 changes: 14 additions & 0 deletions src/generated/linux/library.c.clog.h
Original file line number Diff line number Diff line change
Expand Up @@ -626,10 +626,24 @@ tracepoint(CLOG_LIBRARY_C, ApiExit );\



/*----------------------------------------------------------
// Decoder Ring for LibraryCleanupAsyncNotSupported
// [ lib] Async cleanup not fully implemented! Potential deadlock can occur.
// QuicTraceLogWarning(
LibraryCleanupAsyncNotSupported,
"[ lib] Async cleanup not fully implemented! Potential deadlock can occur.");
----------------------------------------------------------*/
#ifndef _clog_2_ARGS_TRACE_LibraryCleanupAsyncNotSupported
#define _clog_2_ARGS_TRACE_LibraryCleanupAsyncNotSupported(uniqueId, encoded_arg_string)\
tracepoint(CLOG_LIBRARY_C, LibraryCleanupAsyncNotSupported );\

#endif


#ifdef __cplusplus
}
#endif
#endif
#ifdef CLOG_INLINE_IMPLEMENTATION
#include "quic.clog_library.c.clog.h.c"
#endif
15 changes: 15 additions & 0 deletions src/generated/linux/library.c.clog.h.lttng.h
Original file line number Diff line number Diff line change
Expand Up @@ -619,3 +619,18 @@ TRACEPOINT_EVENT(CLOG_LIBRARY_C, ApiExit,
TP_FIELDS(
)
)


/*----------------------------------------------------------
// Decoder Ring for LibraryCleanupAsyncNotSupported
// [ lib] Async cleanup not fully implemented! Potential deadlock can occur.
// QuicTraceLogWarning(
LibraryCleanupAsyncNotSupported,
"[ lib] Async cleanup not fully implemented! Potential deadlock can occur.");
----------------------------------------------------------*/
TRACEPOINT_EVENT(CLOG_LIBRARY_C, LibraryCleanupAsyncNotSupported,
TP_ARGS(
),
TP_FIELDS(
)
)
Loading