Uninitialize LossDetection earlier to cleanup Datagram frame (#4316)

* Uninitialize LossDetection earlier to cleanup pending Datagram before removeing client callback

* move shutdown event indication to the end of func

* add QuicDatagramSendShutdown

* check SendEnabled before touching ApiQueueLock

* Uninitialize ApiQueueLock later

* discard datagram if metadata allocation failed

* Renaming Datagram APIs

* move discard indication to QuicSentPacketPoolReturnPacketMetadata

* clear hander after ack

src/core/connection.c

@@ -369,6 +369,7 @@ QuicConnFree(
     QuicOperationQueueUninitialize(&Connection->OperQ);
     QuicStreamSetUninitialize(&Connection->Streams);
     QuicSendBufferUninitialize(&Connection->SendBuffer);
+    QuicDatagramSendShutdown(&Connection->Datagram);
     QuicDatagramUninitialize(&Connection->Datagram);
     if (Connection->Configuration != NULL) {
         QuicConfigurationRelease(Connection->Configuration);
@@ -1328,26 +1329,6 @@ QuicConnOnShutdownComplete(
         Connection,
         Connection->State.ShutdownCompleteTimedOut);
 
-    if (Connection->State.ExternalOwner) {
-
-        QUIC_CONNECTION_EVENT Event;
-        Event.Type = QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE;
-        Event.SHUTDOWN_COMPLETE.HandshakeCompleted =
-            Connection->State.Connected;
-        Event.SHUTDOWN_COMPLETE.PeerAcknowledgedShutdown =
-            !Connection->State.ShutdownCompleteTimedOut;
-        Event.SHUTDOWN_COMPLETE.AppCloseInProgress =
-            Connection->State.HandleClosed;
-
-        QuicTraceLogConnVerbose(
-            IndicateConnectionShutdownComplete,
-            Connection,
-            "Indicating QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE");
-        (void)QuicConnIndicateEvent(Connection, &Event);
-
-        Connection->ClientCallbackHandler = NULL;
-    }
-
     //
     // Clean up any pending state that is irrelevant now.
     //
@@ -1370,8 +1351,29 @@ QuicConnOnShutdownComplete(
     QuicTimerWheelRemoveConnection(&Connection->Worker->TimerWheel, Connection);
     QuicLossDetectionUninitialize(&Connection->LossDetection);
     QuicSendUninitialize(&Connection->Send);
+    QuicDatagramSendShutdown(&Connection->Datagram);
 
-    if (!Connection->State.ExternalOwner) {
+    if (Connection->State.ExternalOwner) {
+
+        QUIC_CONNECTION_EVENT Event;
+        Event.Type = QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE;
+        Event.SHUTDOWN_COMPLETE.HandshakeCompleted =
+            Connection->State.Connected;
+        Event.SHUTDOWN_COMPLETE.PeerAcknowledgedShutdown =
+            !Connection->State.ShutdownCompleteTimedOut;
+        Event.SHUTDOWN_COMPLETE.AppCloseInProgress =
+            Connection->State.HandleClosed;
+
+        QuicTraceLogConnVerbose(
+            IndicateConnectionShutdownComplete,
+            Connection,
+            "Indicating QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE");
+        (void)QuicConnIndicateEvent(Connection, &Event);
+
+        // This need to be later than QuicLossDetectionUninitialize to indicate
+        // status change of Datagram frame for an app to free its buffer
+        Connection->ClientCallbackHandler = NULL;
+    } else {
         //
         // If the connection was never indicated to the application, then the
         // "owner" ref still resides with the stack and needs to be released.

src/core/datagram.c

@@ -144,7 +144,6 @@ QuicDatagramUninitialize(
     _In_ QUIC_DATAGRAM* Datagram
     )
 {
-    QuicDatagramSendShutdown(Datagram);
     CXPLAT_DBG_ASSERT(Datagram->SendQueue == NULL);
     CXPLAT_DBG_ASSERT(Datagram->ApiQueue == NULL);
     CxPlatDispatchLockUninitialize(&Datagram->ApiQueueLock);

src/core/loss_detection.c

@@ -408,7 +408,7 @@ QuicLossDetectionOnPacketSent(
             "Sent packet metadata",
             SIZEOF_QUIC_SENT_PACKET_METADATA(TempSentPacket->FrameCount));
         QuicLossDetectionRetransmitFrames(LossDetection, TempSentPacket, FALSE);
-        QuicSentPacketMetadataReleaseFrames(TempSentPacket);
+        QuicSentPacketMetadataReleaseFrames(TempSentPacket, Connection);
         return;
     }
 
@@ -623,6 +623,7 @@ QuicLossDetectionOnPacketAcknowledged(
                 Packet->Flags.SuspectedLost ?
                     QUIC_DATAGRAM_SEND_ACKNOWLEDGED_SPURIOUS :
                     QUIC_DATAGRAM_SEND_ACKNOWLEDGED);
+            Packet->Frames[i].DATAGRAM.ClientContext = NULL;
             break;
 
         case QUIC_FRAME_HANDSHAKE_DONE:
@@ -882,7 +883,7 @@ QuicLossDetectionRetransmitFrames(
     Packet->Flags.SuspectedLost = TRUE;
 
     if (ReleasePacket) {
-        QuicSentPacketPoolReturnPacketMetadata(&Connection->Worker->SentPacketPool, Packet);
+        QuicSentPacketPoolReturnPacketMetadata(Packet, Connection);
     }
 
     return NewDataQueued;
@@ -898,17 +899,6 @@ QuicLossDetectionOnPacketDiscarded(
 {
     QUIC_CONNECTION* Connection = QuicLossDetectionGetConnection(LossDetection);
 
-    for (uint8_t i = 0; i < Packet->FrameCount; i++) {
-        switch (Packet->Frames[i].Type) {
-
-        case QUIC_FRAME_DATAGRAM:
-        case QUIC_FRAME_DATAGRAM_1:
-            QuicDatagramIndicateSendStateChange(
-                Connection,
-                &Packet->Frames[i].DATAGRAM.ClientContext,
-                QUIC_DATAGRAM_SEND_LOST_DISCARDED);
-            break;
-        }
     if (Packet->Flags.IsMtuProbe && DiscardedForLoss) {
         uint8_t PathIndex;
         QUIC_PATH* Path = QuicConnGetPathByID(Connection, Packet->PathId, &PathIndex);
@@ -921,9 +911,8 @@ QuicLossDetectionOnPacketDiscarded(
             QuicMtuDiscoveryProbePacketDiscarded(&Path->MtuDiscovery, Connection, PacketMtu);
         }
     }
-    }
 
-    QuicSentPacketPoolReturnPacketMetadata(&Connection->Worker->SentPacketPool, Packet);
+    QuicSentPacketPoolReturnPacketMetadata(Packet, Connection);
 }
 
 //
@@ -1145,7 +1134,7 @@ QuicLossDetectionDiscardPackets(
 
             QuicLossDetectionOnPacketAcknowledged(LossDetection, EncryptLevel, Packet, TRUE, TimeNow, 0);
 
-            QuicSentPacketPoolReturnPacketMetadata(&Connection->Worker->SentPacketPool, Packet);
+            QuicSentPacketPoolReturnPacketMetadata(Packet, Connection);
 
             Packet = NextPacket;
 
@@ -1194,7 +1183,7 @@ QuicLossDetectionDiscardPackets(
 
             QuicLossDetectionOnPacketAcknowledged(LossDetection, EncryptLevel, Packet, TRUE, TimeNow, 0);
 
-            QuicSentPacketPoolReturnPacketMetadata(&Connection->Worker->SentPacketPool, Packet);
+            QuicSentPacketPoolReturnPacketMetadata(Packet, Connection);
 
             Packet = NextPacket;
 
@@ -1628,7 +1617,7 @@ QuicLossDetectionProcessAckBlocks(
     while (AckedPacketsIterator != NULL) {
         QUIC_SENT_PACKET_METADATA* PacketMeta = AckedPacketsIterator;
         AckedPacketsIterator = AckedPacketsIterator->Next;
-        QuicSentPacketPoolReturnPacketMetadata(&Connection->Worker->SentPacketPool, PacketMeta);
+        QuicSentPacketPoolReturnPacketMetadata(PacketMeta, Connection);
     }
 
     //

src/core/packet_builder.c

@@ -154,7 +154,7 @@ QuicPacketBuilderCleanup(
         QuicLossDetectionUpdateTimer(&Builder->Connection->LossDetection, FALSE);
     }
 
-    QuicSentPacketMetadataReleaseFrames(Builder->Metadata);
+    QuicSentPacketMetadataReleaseFrames(Builder->Metadata, Builder->Connection);
 
     CxPlatSecureZeroMemory(Builder->HpMask, sizeof(Builder->HpMask));
 }

src/core/sent_packet_metadata.c

@@ -23,7 +23,8 @@ Abstract:
 
 void
 QuicSentPacketMetadataReleaseFrames(
-    _In_ QUIC_SENT_PACKET_METADATA* Metadata
+    _In_ QUIC_SENT_PACKET_METADATA* Metadata,
+    _In_ QUIC_CONNECTION* Connection
     )
 {
     for (uint8_t i = 0; i < Metadata->FrameCount; i++) {
@@ -50,6 +51,15 @@ QuicSentPacketMetadataReleaseFrames(
             QuicStreamSentMetadataDecrement(Metadata->Frames[i].RELIABLE_RESET_STREAM.Stream);
             break;
 #pragma warning(pop)
+        case QUIC_FRAME_DATAGRAM:
+        case QUIC_FRAME_DATAGRAM_1:
+            if (Metadata->Frames[i].DATAGRAM.ClientContext != NULL) {
+                QuicDatagramIndicateSendStateChange(
+                    Connection,
+                    &Metadata->Frames[i].DATAGRAM.ClientContext,
+                    QUIC_DATAGRAM_SEND_LOST_DISCARDED);
+            }
+            break;
         default:
             //
             // Nothing to clean up for other frame types.
@@ -109,8 +119,8 @@ QuicSentPacketPoolGetPacketMetadata(
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
 QuicSentPacketPoolReturnPacketMetadata(
-    _In_ QUIC_SENT_PACKET_POOL* Pool,
+    _In_ QUIC_SENT_PACKET_METADATA* Metadata,
-    _In_ QUIC_SENT_PACKET_METADATA* Metadata
+    _In_ QUIC_CONNECTION* Connection
     )
 {
     _Analysis_assume_(
@@ -121,6 +131,6 @@ QuicSentPacketPoolReturnPacketMetadata(
     Metadata->Flags.Freed = TRUE;
 #endif
 
-    QuicSentPacketMetadataReleaseFrames(Metadata);
+    QuicSentPacketMetadataReleaseFrames(Metadata, Connection);
-    CxPlatPoolFree(Pool->Pools + Metadata->FrameCount - 1, Metadata);
+    CxPlatPoolFree(Connection->Worker->SentPacketPool.Pools + Metadata->FrameCount - 1, Metadata);
 }

src/core/sent_packet_metadata.h

@@ -187,7 +187,8 @@ QuicPacketTraceType(
 
 void
 QuicSentPacketMetadataReleaseFrames(
-    _In_ QUIC_SENT_PACKET_METADATA* Metadata
+    _In_ QUIC_SENT_PACKET_METADATA* Metadata,
+    _In_ QUIC_CONNECTION* Connection
     );
 
 //
@@ -243,6 +244,6 @@ QuicSentPacketPoolGetPacketMetadata(
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
 QuicSentPacketPoolReturnPacketMetadata(
-    _In_ QUIC_SENT_PACKET_POOL* Pool,
+    _In_ QUIC_SENT_PACKET_METADATA* Metadata,
-    _In_ QUIC_SENT_PACKET_METADATA* Metadata
+    _In_ QUIC_CONNECTION* Connection
     );