more fixes

Author: mtfriesen

src/core/api.c

@@ -79,7 +79,12 @@ QuicConnectionOpenInPartition(
         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;
+#endif
     Connection->ClientCallbackHandler = Handler;
     Connection->ClientContext = Context;
 
@@ -150,6 +155,7 @@ MsQuicConnectionClose(
     )
 {
     QUIC_CONNECTION* Connection;
+    BOOLEAN WaitForCompletion = TRUE;
 
     CXPLAT_PASSIVE_CODE();
 
@@ -182,7 +188,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.
         //
@@ -198,39 +204,51 @@ MsQuicConnectionClose(
     } else {
 
         CXPLAT_EVENT CompletionEvent;
-        QUIC_OPERATION Oper = { 0 };
-        QUIC_API_CONTEXT ApiCtx;
 
-        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 (MsQuicLib.CustomExecutions) {
+            //
+            // For custom executions, ConnectionClose completes asynchronously.
+            //
+            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:
 

src/core/connection.c

@@ -5597,8 +5597,8 @@ QuicConnRecvDatagramBatch(
             RecvState->ResetIdleTimeout |= Packet->CompletelyValid;
 
             if (Connection->Registration != NULL && !Connection->Registration->NoPartitioning &&
-                !Connection->State.Partitioned && Path->IsActive && !Path->PartitionUpdated &&
-                Packet->CompletelyValid &&
+                !Path->Binding->Partitioned && !Connection->State.Partitioned && Path->IsActive &&
+                !Path->PartitionUpdated && Packet->CompletelyValid &&
                 (Packets[i]->PartitionIndex % MsQuicLib.PartitionCount) != RecvState->PartitionIndex) {
                 RecvState->PartitionIndex = Packets[i]->PartitionIndex % MsQuicLib.PartitionCount;
                 RecvState->UpdatePartitionId = TRUE;
@@ -7579,6 +7579,14 @@ QuicConnProcessApiOperation(
 
     case QUIC_API_TYPE_CONN_CLOSE:
         QuicConnCloseHandle(Connection);
+        if (MsQuicLib.CustomExecutions && Connection->State.ExternalOwner) {
+            CXPLAT_TEL_ASSERT(Connection->State.HandleClosed);
+            //
+            // Release the external reference to the connection for async
+            // completion.
+            //
+            QuicConnRelease(Connection, QUIC_CONN_REF_HANDLE_OWNER);
+        }
         break;
 
     case QUIC_API_TYPE_CONN_SHUTDOWN:

src/core/connection.h

@@ -559,6 +559,8 @@ typedef struct QUIC_CONNECTION {
     QUIC_OPERATION BackUpOper;
     QUIC_API_CONTEXT BackupApiContext;
     uint16_t BackUpOperUsed;
+    QUIC_OPERATION CloseOper;
+    QUIC_API_CONTEXT CloseApiContext;
 
     //
     // The status code used for indicating transport closed notifications.

src/inc/msquic.hpp

@@ -1149,19 +1149,22 @@ struct MsQuicConnection {
     QUIC_UINT62 AppShutdownErrorCode {0};
     bool HandshakeComplete {false};
     bool HandshakeResumed {false};
+    bool Started {false};
     uint32_t ResumptionTicketLength {0};
     uint8_t* ResumptionTicket {nullptr};
 #ifdef CX_PLATFORM_TYPE
     CxPlatEvent HandshakeCompleteEvent;
     CxPlatEvent ResumptionTicketReceivedEvent;
+    CxPlatEvent ShutdownCompleteEvent {true};
 #endif // CX_PLATFORM_TYPE
 
     MsQuicConnection(
         _In_ const MsQuicRegistration& Registration,
         _In_ MsQuicCleanUpMode CleanUpMode = CleanUpManual,
         _In_ MsQuicConnectionCallback* Callback = NoOpCallback,
         _In_ void* Context = nullptr
-        ) noexcept : CleanUpMode(CleanUpMode), Callback(Callback), Context(Context) {
+        ) noexcept : CleanUpMode(CleanUpMode), Callback(Callback), Context(Context)
+        {
         if (!Registration.IsValid()) {
             InitStatus = Registration.GetInitStatus();
             return;
@@ -1213,6 +1216,11 @@ struct MsQuicConnection {
 
     ~MsQuicConnection() noexcept {
         Close();
+#ifdef CX_PLATFORM_TYPE
+        if (IsValid()) {
+            ShutdownCompleteEvent.WaitForever();
+        }
+#endif // CX_PLATFORM_TYPE
         delete[] ResumptionTicket;
     }
 
@@ -1255,7 +1263,11 @@ struct MsQuicConnection {
             const char* ServerName,
         _In_ uint16_t ServerPort // Host byte order
         ) noexcept {
-        return MsQuic->ConnectionStart(Handle, Config, Family, ServerName, ServerPort);
+        QUIC_STATUS Status = MsQuic->ConnectionStart(Handle, Config, Family, ServerName, ServerPort);
+        if (QUIC_SUCCEEDED(Status)) {
+            Started = true;
+        }
+        return Status;
     }
 
     QUIC_STATUS
@@ -1485,6 +1497,9 @@ struct MsQuicConnection {
         _Inout_ QUIC_CONNECTION_EVENT* Event
         ) noexcept {
         CXPLAT_DBG_ASSERT(pThis);
+        auto DeleteOnExit =
+            Event->Type == QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE &&
+            pThis->CleanUpMode == CleanUpAutoDelete;
         if (Event->Type == QUIC_CONNECTION_EVENT_CONNECTED) {
             pThis->HandshakeComplete = true;
             pThis->HandshakeResumed = Event->CONNECTED.SessionResumed;
@@ -1516,10 +1531,10 @@ struct MsQuicConnection {
 #endif // CX_PLATFORM_TYPE
             }
         }
-        auto DeleteOnExit =
-            Event->Type == QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE &&
-            pThis->CleanUpMode == CleanUpAutoDelete;
         auto Status = pThis->Callback(pThis, pThis->Context, Event);
+        if (Event->Type == QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE) {
+            pThis->ShutdownCompleteEvent.Set();
+        }
         if (DeleteOnExit) {
             delete pThis;
         }