Avoid exclusivity conflict in DatagramContext

Sources/NIOHTTPServer/Datagrams/HTTP3ConnectionManager.swift

@@ -61,32 +61,16 @@ final class HTTP3ConnectionManager: ChannelInboundHandler {
         }
 
         mutating func finish() {
-            // The connection is going away, so datagrams will never be negotiated if they haven't been already.
-            self.completeNegotiation(datagramsSupported: false)
-
             for datagramStream in self.datagramStreams.values {
                 datagramStream.finish()
             }
             self.datagramStreams.removeAll()
         }
 
-        mutating func receivedSettings(datagramsSupported: Bool) {
-            self.completeNegotiation(datagramsSupported: datagramsSupported)
-        }
-
-        /// Completes the negotiation promise, if it has not been completed already.
-        private mutating func completeNegotiation(datagramsSupported: Bool) {
-            guard let promise = self.datagramsNegotiatedPromise else {
-                // We fulfilled the promise earlier.
-                return
-            }
-            self.datagramsNegotiatedPromise = nil
-
-            if datagramsSupported {
-                promise.succeed()
-            } else {
-                promise.fail(DatagramsNotSupported())
-            }
+        /// Returns the negotiation promise if it has not been taken already.
+        mutating func takeNegotiationPromise() -> EventLoopPromise<Void>? {
+            defer { self.datagramsNegotiatedPromise = nil }
+            return self.datagramsNegotiatedPromise
         }
     }
 
@@ -120,6 +104,20 @@ final class HTTP3ConnectionManager: ChannelInboundHandler {
         self.eventLoop.preconditionInEventLoop()
         self.datagramContext?.deregister(streamID: streamID)
     }
+
+    /// Completes the negotiation promise, if it has not been completed already.
+    private func completeNegotiation(datagramsSupported: Bool) {
+        guard let promise = self.datagramContext?.takeNegotiationPromise() else {
+            // We fulfilled the promise earlier.
+            return
+        }
+
+        if datagramsSupported {
+            promise.succeed()
+        } else {
+            promise.fail(DatagramsNotSupported())
+        }
+    }
     #endif  // UnstableHTTPDatagrams
 
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
@@ -135,7 +133,7 @@ final class HTTP3ConnectionManager: ChannelInboundHandler {
         switch event {
         case let event as ReceivedSettings:
             #if UnstableHTTPDatagrams
-            self.datagramContext?.receivedSettings(datagramsSupported: event.datagramsSupported)
+            self.completeNegotiation(datagramsSupported: event.datagramsSupported)
             #endif
             context.fireUserInboundEventTriggered(event)
 
@@ -151,6 +149,8 @@ final class HTTP3ConnectionManager: ChannelInboundHandler {
 
     func channelInactive(context: ChannelHandlerContext) {
         #if UnstableHTTPDatagrams
+        // The connection is going away, so datagrams will never be negotiated if they haven't been already.
+        self.completeNegotiation(datagramsSupported: false)
         self.datagramContext?.finish()
         self.datagramContext = nil
         #endif

Tests/NIOHTTPServerTests/HTTP3DatagramTests.swift

@@ -113,6 +113,63 @@ struct HTTP3DatagramTests {
         #expect(try await TestHelpers.readDatagram(&secondReader) == nil)
     }
 
+    @Test("A stream opened before SETTINGS can register once datagrams are negotiated")
+    @available(anyAppleOS 27.0, *)
+    func streamOpenedBeforeSettingsRegistersOnNegotiation() async throws {
+        let channel = EmbeddedChannel()
+        let datagramsNegotiatedPromise = channel.eventLoop.makePromise(of: Void.self)
+        let manager = HTTP3ConnectionManager(
+            eventLoop: channel.eventLoop,
+            logger: Self.serverLogger,
+            datagramsNegotiatedPromise: datagramsNegotiatedPromise
+        )
+        try channel.pipeline.syncOperations.addHandler(manager)
+        let loopBoundManager = NIOLoopBound(manager, eventLoop: channel.eventLoop)
+
+        // Register from the negotiation callback.
+        let datagramStreamFuture = datagramsNegotiatedPromise.futureResult.map {
+            let stream = HTTP3UnreliableDatagramStream(
+                streamID: 0,
+                connectionChannel: channel,
+                maxBufferedDatagrams: 16
+            )
+            loopBoundManager.value.register(datagramStream: stream)
+            return stream
+        }
+
+        channel.pipeline.syncOperations.fireUserInboundEventTriggered(ReceivedSettings(datagramsSupported: true))
+
+        let stream = try await datagramStreamFuture.get()
+        var reader = NIOHTTPServer.DatagramReader(iterator: stream.inbound.makeAsyncIterator())
+        try channel.writeInbound(HTTP3Datagram(streamID: 0, payload: ByteBuffer([1])))
+
+        #expect(try await TestHelpers.readDatagram(&reader) == [1])
+    }
+
+    @Test("A stream opened before SETTINGS can deregister when the connection closes")
+    @available(anyAppleOS 27.0, *)
+    func streamOpenedBeforeSettingsDeregistersOnClose() throws {
+        let channel = EmbeddedChannel()
+        let datagramsNegotiatedPromise = channel.eventLoop.makePromise(of: Void.self)
+        let manager = HTTP3ConnectionManager(
+            eventLoop: channel.eventLoop,
+            logger: Self.serverLogger,
+            datagramsNegotiatedPromise: datagramsNegotiatedPromise
+        )
+        try channel.pipeline.syncOperations.addHandler(manager)
+        let loopBoundManager = NIOLoopBound(manager, eventLoop: channel.eventLoop)
+
+        // Deregister from the negotiation callback.
+        datagramsNegotiatedPromise.futureResult.whenComplete { _ in
+            loopBoundManager.value.deregister(streamID: 0)
+        }
+
+        #expect(try channel.finish().isClean)
+        #expect(throws: HTTP3ConnectionManager.DatagramsNotSupported.self) {
+            try datagramsNegotiatedPromise.futureResult.wait()
+        }
+    }
+
     @Test("The oldest datagram is dropped once the buffer is full")
     @available(anyAppleOS 27.0, *)
     func oldestDatagramIsDroppedOnceTheBufferIsFull() async throws {