Add strong partitioning to listeners, connections, and streams (#5376)

* WIP - listener affinity

* WIP - listener events are onto partition thread

* make connections explicitly partitioned, too

* fix winuser build

* update the sidecar, again

* try to fix thread id abstraction

* add array bounds assert

* fix winkernel build (finally?)

* update rust bindings

* fix dotnet, too

* fix more rust

* fix a few bugs

* WIP - test

* add more tests, which all fail

* sync submodules

* WIP

* add docs

* WIP - tests passing w/o partition option

* WIP - listener use-after-free on close notification

* WIP - partition test (WIP) passes

* WIP - more precise tests

* test variable numbers of ecs

* finalize basic test coverage

* fix non-linux test build

* more rust

* fix release build

* fix dotnet

* updating sidecar

* codecheck

* fix winkernel build

* try again to fix winkernel

* STILL fix winkernel

* fix winkernel yet again

* fix win official

* RE-fix

* fix one listener bug

* more listener fixes

* more rust gen

* more fixes

* properly disable the feature on iouring/linux_xdp

* ensure connections aren't implicitly partitioned

* remove flaky/fragile/wrong stateless retry test

* WIP

* more fixes

* improve comment

* fix sidecar

* initialize event for compiler

* fix unref'd param

* make connection async close wait opt-in

* rename listener refcounts for clarity

* add connection close async param

* fixes

* official build fix

* dotnet

* update spinquic for conn close async param

* address PR feedback

* do todo

Author: mtfriesen

CMakeLists.txt

@@ -317,7 +317,7 @@ else()
             set(QUIC_LOGGING_TYPE "lttng")
             message(STATUS "Choosing lttng as default logging type for platform")
         endif()
-        
+
         if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm|aarch64)$")
             set(QUIC_LINUX_XDP_ENABLED OFF CACHE BOOL "XDP not supported on ARM architectures" FORCE)
         endif()
@@ -775,6 +775,10 @@ if (QUIC_USE_SYSTEM_LIBCRYPTO)
     list(APPEND QUIC_COMMON_DEFINES CXPLAT_SYSTEM_CRYPTO)
 endif()
 
+if (QUIC_LINUX_XDP_ENABLED)
+    list(APPEND QUIC_COMMON_DEFINES CXPLAT_LINUX_XDP_ENABLED)
+endif()
+
 if (QUIC_LINUX_IOURING_ENABLED)
     list(APPEND QUIC_COMMON_DEFINES CXPLAT_USE_IO_URING)
 endif()

docs/API.md

@@ -90,7 +90,7 @@ Please note, there is no explicit start/stop API for this library. Each API func
 
 Generally, each app only needs a single registration. The registration represents the execution context where all logic for the app's connections run. The library will create a number of worker threads for each registration, shared for all the connections. This execution context is not shared between different registrations.
 
-A registration is created by calling [RegistrationOpen](api/RegistrationOpen.md) and deleted by calling [RegistrationClose](api/RegistrationClose.md).
+A registration is created by calling [RegistrationOpen](api/RegistrationOpen.md) and deleted by calling [RegistrationClose](api/RegistrationClose.md) or (Preview) [RegistrationClose2](api/RegistrationClose2.md).
 
 ## Configuration
 

docs/Diagnostics.md

@@ -291,6 +291,7 @@ QUIC_PERF_COUNTER_PATH_FAILURE | Total path challenges that fail ever
 QUIC_PERF_COUNTER_SEND_STATELESS_RESET | Total stateless reset packets sent ever
 QUIC_PERF_COUNTER_SEND_STATELESS_RETRY | Total stateless retry packets sent ever
 QUIC_PERF_COUNTER_CONN_LOAD_REJECT | Total connections rejected due to worker load.
+QUIC_PERF_COUNTER_LISTEN_QUEUE_DEPTH | Current listeners queued for processing.
 
 ## Windows Performance Monitor
 

docs/Settings.md

@@ -171,6 +171,7 @@ These parameters are accessed by calling [GetParam](./api/GetParam.md) or [SetPa
 | `QUIC_PARAM_LISTENER_STATS`<br> 1         | QUIC_LISTENER_STATISTICS  | Get-only  | Get statistics specific to this Listener instance.        |
 | `QUIC_PARAM_LISTENER_CIBIR_ID`<br> 2      | uint8_t[]                 | Both      | The CIBIR well-known idenfitier.                          |
 | `QUIC_PARAM_DOS_MODE_EVENTS`<br> 2        | BOOLEAN                   | Both      | The Listener opted in for DoS Mode event.                 |
+| `QUIC_PARAM_LISTENER_PARTITION_INDEX`<br> (preview) | uint16_t           | Both      | The partition to use for listener callback events and incoming connections. |
 
 ## Connection Parameters
 
@@ -204,7 +205,8 @@ These parameters are accessed by calling [GetParam](./api/GetParam.md) or [SetPa
 | `QUIC_PARAM_CONN_STATISTICS_V2_PLAT`<br> 23       | QUIC_STATISTICS_V2            | Get-only  | Connection-level statistics with platform-specific time format, version 2.                |
 | `QUIC_PARAM_CONN_ORIG_DEST_CID` <br> 24           | uint8_t[]                     | Get-only  | The original destination connection ID used by the client to connect to the server.       |
 | `QUIC_PARAM_CONN_SEND_DSCP` <br> 25               | uint8_t                       | Both      | The DiffServ Code Point put in the DiffServ field (formerly TypeOfService/TrafficClass) on packets sent from this connection. |
-| `QUIC_PARAM_CONN_NETWORK_STATISTICS` <br> 20      | QUIC_NETWORK_STATISTICS       | Get-only  | Returns Connection level network statistics |
+| `QUIC_PARAM_CONN_NETWORK_STATISTICS` <br> 32      | QUIC_NETWORK_STATISTICS       | Get-only  | Returns Connection level network statistics |
+| `QUIC_PARAM_CONN_CLOSE_ASYNC` <br> 26      | uint8_t (BOOLEAN)      | Both  | The desired connection close behavior. Defaults to false (synchronous). |
 
 ### QUIC_PARAM_CONN_STATISTICS_V2
 

docs/api/QUIC_API_TABLE.md

@@ -80,6 +80,10 @@ See [RegistrationOpen](RegistrationOpen.md)
 
 See [RegistrationClose](RegistrationClose.md)
 
+`RegistrationClose2`
+
+See (Preview) [RegistrationClose2](RegistrationClose2.md)
+
 `RegistrationShutdown`
 
 See [RegistrationShutdown](RegistrationShutdown.md)

docs/api/RegistrationClose.md

@@ -23,8 +23,9 @@ A registration handle from a previous call to [RegistrationOpen](RegistrationOpe
 
 # Remarks
 
-The application **must** close/delete all child configurations and connection objects before closing the registration. This call **will block** on those outstanding objects being cleaned up. Do no call it on any MsQuic event callback, or it will deadlock.
+The application **must** close/delete all child configurations and connection objects before closing the registration. This call **will block** on those outstanding objects being cleaned up. Do not call it on any MsQuic event callback or a thread that would otherwise be running an external execution context, or it will deadlock.
 
 # See Also
 
 [RegistrationOpen](RegistrationOpen.md)<br>
+(Preview) [RegistrationClose2](RegistrationClose2.md)<br>

docs/api/RegistrationClose2.md

@@ -0,0 +1,54 @@
+RegistrationClose2 function
+======
+
+> **Preview**
+> This routine is in preview and is subject to breaking changes.
+
+Closes an existing registration asynchronously.
+
+# Syntax
+
+```C
+typedef
+_Function_class_(QUIC_REGISTRATION_CLOSE_CALLBACK)
+void
+(QUIC_API QUIC_REGISTRATION_CLOSE_CALLBACK)(
+    _In_opt_ void* Context
+    );
+typedef QUIC_REGISTRATION_CLOSE_CALLBACK *QUIC_REGISTRATION_CLOSE_CALLBACK_HANDLER;
+
+typedef
+_IRQL_requires_max_(PASSIVE_LEVEL)
+void
+(QUIC_API * QUIC_REGISTRATION_CLOSE2_FN)(
+    _In_ _Pre_defensive_ __drv_freesMem(Mem)
+        HQUIC Registration,
+    _In_ _Pre_defensive_ QUIC_REGISTRATION_CLOSE_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context
+    );
+```
+
+# Parameters
+
+`Registration`
+
+A registration handle from a previous call to [RegistrationOpen](RegistrationOpen.md).
+
+`Handler`
+
+A registration close completion handler. It will be invoked exactly once upon completion of the registration close request.
+
+`Context`
+
+The context to provide to the close completion handler.
+
+# Remarks
+
+> **Preview**
+> This routine is in preview and is subject to breaking changes.
+
+The application should close/delete all child configurations and connection objects before closing the registration. This request **will not complete** until those outstanding objects are cleaned up.
+
+# See Also
+
+[RegistrationOpen](RegistrationOpen.md)<br>

docs/api/RegistrationOpen.md

@@ -39,3 +39,4 @@ A caveat to this independence is that until a packet or connection can be determ
 # See Also
 
 [RegistrationClose](RegistrationClose.md)<br>
+(Preview) [RegistrationClose2](RegistrationClose2.md)<br>

docs/api/RegistrationShutdown.md

@@ -44,3 +44,4 @@ The 62-bit error code to indicate to the peer as the reason for the shutdown.
 [RegistrationOpen](RegistrationOpen.md)<br>
 [RegistrationClose](RegistrationClose.md)<br>
 [ConnectionShutdown](ConnectionShutdown.md)<br>
+(Preview) [RegistrationClose2](RegistrationClose2.md)<br>

src/core/api.c

@@ -33,31 +33,12 @@
 _IRQL_requires_max_(DISPATCH_LEVEL)
 QUIC_STATUS
 QUIC_API
-MsQuicConnectionOpen(
-    _In_ _Pre_defensive_ HQUIC RegistrationHandle,
-    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
-    _In_opt_ void* Context,
-    _Outptr_ _At_(*NewConnection, __drv_allocatesMem(Mem)) _Pre_defensive_
-        HQUIC *NewConnection
-    )
-{
-    return
-        MsQuicConnectionOpenInPartition(
-            RegistrationHandle,
-            QuicLibraryGetCurrentPartition()->Index,
-            Handler,
-            Context,
-            NewConnection);
-}
-
-_IRQL_requires_max_(DISPATCH_LEVEL)
-QUIC_STATUS
-QUIC_API
-MsQuicConnectionOpenInPartition(
+QuicConnectionOpenInPartition(
     _In_ _Pre_defensive_ HQUIC RegistrationHandle,
     _In_ uint16_t PartitionIndex,
     _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
     _In_opt_ void* Context,
+    _In_ BOOLEAN Partitioned,
     _Outptr_ _At_(*NewConnection, __drv_allocatesMem(Mem)) _Pre_defensive_
         HQUIC *NewConnection
     )
@@ -98,6 +79,14 @@ MsQuicConnectionOpenInPartition(
         goto Error;
     }
 
+    //
+    // Hard partitioning is only supported on a subset of platforms.
+    //
+#if defined(__linux__) && !defined(CXPLAT_USE_IO_URING) && !defined(CXPLAT_LINUX_XDP_ENABLED)
+    Connection->State.Partitioned = Partitioned;
+#else
+    UNREFERENCED_PARAMETER(Partitioned);
+#endif
     Connection->ClientCallbackHandler = Handler;
     Connection->ClientContext = Context;
 
@@ -114,6 +103,49 @@ MsQuicConnectionOpenInPartition(
     return Status;
 }
 
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicConnectionOpen(
+    _In_ _Pre_defensive_ HQUIC RegistrationHandle,
+    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context,
+    _Outptr_ _At_(*NewConnection, __drv_allocatesMem(Mem)) _Pre_defensive_
+        HQUIC *NewConnection
+    )
+{
+    return
+        QuicConnectionOpenInPartition(
+            RegistrationHandle,
+            QuicLibraryGetCurrentPartition()->Index,
+            Handler,
+            Context,
+            FALSE,
+            NewConnection);
+}
+
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicConnectionOpenInPartition(
+    _In_ _Pre_defensive_ HQUIC RegistrationHandle,
+    _In_ uint16_t PartitionIndex,
+    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context,
+    _Outptr_ _At_(*NewConnection, __drv_allocatesMem(Mem)) _Pre_defensive_
+        HQUIC *NewConnection
+    )
+{
+    return
+        QuicConnectionOpenInPartition(
+            RegistrationHandle,
+            PartitionIndex,
+            Handler,
+            Context,
+            TRUE,
+            NewConnection);
+}
+
 #pragma warning(push)
 #pragma warning(disable:6014) // SAL doesn't understand the free happens on the worker
 _IRQL_requires_max_(PASSIVE_LEVEL)
@@ -125,6 +157,7 @@ MsQuicConnectionClose(
     )
 {
     QUIC_CONNECTION* Connection;
+    BOOLEAN WaitForCompletion = TRUE;
 
     CXPLAT_PASSIVE_CODE();
 
@@ -157,7 +190,7 @@ MsQuicConnectionClose(
 
     CXPLAT_TEL_ASSERT(!Connection->State.HandleClosed);
 
-    if (MsQuicLib.CustomExecutions || IsWorkerThread) {
+    if (IsWorkerThread) {
         //
         // Execute this blocking API call inline if called on the worker thread.
         //
@@ -172,40 +205,49 @@ MsQuicConnectionClose(
 
     } else {
 
-        CXPLAT_EVENT CompletionEvent;
-        QUIC_OPERATION Oper = { 0 };
-        QUIC_API_CONTEXT ApiCtx;
+        CXPLAT_EVENT CompletionEvent = {0};
 
-        Oper.Type = QUIC_OPER_TYPE_API_CALL;
-        Oper.FreeAfterProcess = FALSE;
-        Oper.API_CALL.Context = &ApiCtx;
+        Connection->CloseOper.Type = QUIC_OPER_TYPE_API_CALL;
+        Connection->CloseOper.FreeAfterProcess = FALSE;
+        Connection->CloseOper.API_CALL.Context = &Connection->CloseApiContext;
 
-        ApiCtx.Type = QUIC_API_TYPE_CONN_CLOSE;
-        CxPlatEventInitialize(&CompletionEvent, TRUE, FALSE);
-        ApiCtx.Completed = &CompletionEvent;
-        ApiCtx.Status = NULL;
+        Connection->CloseApiContext.Type = QUIC_API_TYPE_CONN_CLOSE;
+        Connection->CloseApiContext.Status = NULL;
+
+        if (Connection->State.CloseAsync) {
+            Connection->CloseApiContext.Completed = NULL;
+            WaitForCompletion = FALSE;
+        } else {
+            CxPlatEventInitialize(&CompletionEvent, TRUE, FALSE);
+            Connection->CloseApiContext.Completed = &CompletionEvent;
+        }
 
         //
         // Queue the operation and wait for it to be processed.
         //
-        QuicConnQueueOper(Connection, &Oper);
-        QuicTraceEvent(
-            ApiWaitOperation,
-            "[ api] Waiting on operation");
-        CxPlatEventWaitForever(CompletionEvent);
-        CxPlatEventUninitialize(CompletionEvent);
+        QuicConnQueueOper(Connection, &Connection->CloseOper);
+
+        if (WaitForCompletion) {
+            QuicTraceEvent(
+                ApiWaitOperation,
+                "[ api] Waiting on operation");
+            CxPlatEventWaitForever(CompletionEvent);
+            CxPlatEventUninitialize(CompletionEvent);
+        }
     }
 
     //
     // Connection can only be released by the application after the released
     // flag was set, in response to the CONN_CLOSE operation was processed.
     //
-    CXPLAT_TEL_ASSERT(Connection->State.HandleClosed);
+    if (WaitForCompletion) {
+        CXPLAT_TEL_ASSERT(Connection->State.HandleClosed);
 
-    //
-    // Release the reference to the Connection.
-    //
-    QuicConnRelease(Connection, QUIC_CONN_REF_HANDLE_OWNER);
+        //
+        // Release the reference to the Connection.
+        //
+        QuicConnRelease(Connection, QUIC_CONN_REF_HANDLE_OWNER);
+    }
 
 Error: