Enable per-message GSO on Linux (#3436)

Lukasa · web-flow · commit a7da98c3ac60 · 2025-11-07T08:34:33.000Z
Motivation:

NIO has existing GSO support. This GSO support is applied at Channel
scope, which makes it somewhat less than ideal, as it forces the
datagram sizes to be known statically ahead of time. In real
applications this is less common than being able to do this on a
per-superbuffer
context.

Linux supports this already by using sendmsg/sendmmsg control data. We
just need to wire this up in NIO.

An extension to this is that we can also do per-message GRO in Linux.
I'll tackle this in a later patch to keep the diffs small.

Modifications:

- Adding a new field on AddressedEnvelope Metadata.
- Reject setting this field on non-Linux devices.
- On Linux, use this to set the control data.
- Added new tests that validate the behaviour is correct.

Result:

Enable per-message GSO.
diff --git a/Sources/NIOCore/AddressedEnvelope.swift b/Sources/NIOCore/AddressedEnvelope.swift
@@ -39,15 +39,37 @@ public struct AddressedEnvelope<DataType> {
         /// Details of any congestion state.
         public var ecnState: NIOExplicitCongestionNotificationState
         public var packetInfo: NIOPacketInfo?
+        /// The UDP segment size for Generic Segmentation Offload (GSO).
+        ///
+        /// When set, this enables per-message GSO, allowing the kernel to split this datagram
+        /// into multiple segments of the specified size. The maximum segment size is platform-dependent
+        /// and can be queried via `System.udpMaxSegments`.
+        ///
+        /// On non-Linux platforms, writes with a non-nil `segmentSize` will fail with
+        /// `ChannelError.operationUnsupported`. The error will be propagated to the write promise
+        /// (if attached).
+        public var segmentSize: Int?
 
         public init(ecnState: NIOExplicitCongestionNotificationState) {
             self.ecnState = ecnState
             self.packetInfo = nil
+            self.segmentSize = nil
         }
 
         public init(ecnState: NIOExplicitCongestionNotificationState, packetInfo: NIOPacketInfo?) {
             self.ecnState = ecnState
             self.packetInfo = packetInfo
+            self.segmentSize = nil
+        }
+
+        public init(
+            ecnState: NIOExplicitCongestionNotificationState,
+            packetInfo: NIOPacketInfo?,
+            segmentSize: Int?
+        ) {
+            self.ecnState = ecnState
+            self.packetInfo = packetInfo
+            self.segmentSize = segmentSize
         }
     }
 }
diff --git a/Sources/NIOPosix/ControlMessage.swift b/Sources/NIOPosix/ControlMessage.swift
@@ -388,6 +388,25 @@ extension UnsafeOutboundControlBytes {
             break
         }
     }
+
+    /// Appends a UDP segment size control message for Generic Segmentation Offload (GSO).
+    ///
+    /// - Parameter metadata: Metadata from the addressed envelope which may contain a segment size.
+    /// - Warning: This will `fatalError` if called with a segmentSize on non-Linux platforms.
+    ///   Callers should validate platform support before enqueueing writes with segmentSize.
+    internal mutating func appendUDPSegmentSize(metadata: AddressedEnvelope<ByteBuffer>.Metadata?) {
+        guard let metadata = metadata, let segmentSize = metadata.segmentSize else { return }
+
+        #if os(Linux)
+        self.appendGenericControlMessage(
+            level: Posix.SOL_UDP,
+            type: .init(NIOBSDSocket.Option.udp_segment.rawValue),
+            payload: UInt16(segmentSize)
+        )
+        #else
+        fatalError("UDP segment size is only supported on Linux")
+        #endif
+    }
 }
 
 extension AddressedEnvelope.Metadata {
diff --git a/Sources/NIOPosix/PendingDatagramWritesManager.swift b/Sources/NIOPosix/PendingDatagramWritesManager.swift
@@ -148,6 +148,7 @@ private func doPendingDatagramWriteVectorOperation(
 
                 var controlBytes = UnsafeOutboundControlBytes(controlBytes: controlMessageStorage[c])
                 controlBytes.appendExplicitCongestionState(metadata: p.metadata, protocolFamily: protocolFamily)
+                controlBytes.appendUDPSegmentSize(metadata: p.metadata)
                 let controlMessageBytePointer = controlBytes.validControlBytes
 
                 var msg = msghdr()
diff --git a/Sources/NIOPosix/SocketChannel.swift b/Sources/NIOPosix/SocketChannel.swift
@@ -929,6 +929,14 @@ final class DatagramChannel: BaseSocketChannel<Socket>, @unchecked Sendable {
     /// Buffer a write in preparation for a flush.
     private func bufferPendingAddressedWrite(envelope: AddressedEnvelope<ByteBuffer>, promise: EventLoopPromise<Void>?)
     {
+        // Check if segment size is set on a non-Linux platform
+        #if !os(Linux)
+        if envelope.metadata?.segmentSize != nil {
+            promise?.fail(ChannelError.operationUnsupported)
+            return
+        }
+        #endif
+
         // If the socket is connected, check the remote provided matches the connected address.
         if let connectedRemoteAddress = self.remoteAddress {
             guard envelope.remoteAddress == connectedRemoteAddress else {
@@ -977,6 +985,7 @@ final class DatagramChannel: BaseSocketChannel<Socket>, @unchecked Sendable {
                         metadata: metadata,
                         protocolFamily: self.localAddress?.protocol
                     )
+                    controlBytes.appendUDPSegmentSize(metadata: metadata)
                     return try self.socket.sendmsg(
                         pointer: ptr,
                         destinationPtr: destinationPtr,
diff --git a/Sources/NIOPosix/System.swift b/Sources/NIOPosix/System.swift
@@ -542,6 +542,14 @@ internal enum Posix: Sendable {
     static let IPV6_PKTINFO: CInt = CInt(WinSDK.IPV6_PKTINFO)
     #endif
 
+    #if canImport(Darwin)
+    static let SOL_UDP: CInt = CInt(IPPROTO_UDP)
+    #elseif os(Linux) || os(FreeBSD) || os(Android)
+    static let SOL_UDP: CInt = CInt(IPPROTO_UDP)
+    #elseif os(Windows)
+    static let SOL_UDP: CInt = CInt(IPPROTO_UDP)
+    #endif
+
     #if !os(Windows)
     @inline(never)
     public static func shutdown(descriptor: CInt, how: Shutdown) throws {
diff --git a/Tests/NIOPosixTests/DatagramChannelTests.swift b/Tests/NIOPosixTests/DatagramChannelTests.swift
@@ -1860,6 +1860,187 @@ class DatagramChannelTests: XCTestCase {
 
         return hasGoodGROSupport
     }
+
+    // MARK: - Per-Message GSO Tests
+
+    func testLargeScalarWriteWithPerMessageGSO() throws {
+        try XCTSkipUnless(System.supportsUDPSegmentationOffload, "UDP_SEGMENT (GSO) is not supported on this platform")
+
+        // We're going to send one large buffer with per-message GSO metadata.
+        // The kernel should split it into multiple segments.
+        let segmentSize = 1000
+        let segments = 10
+
+        // Form a handful of segments
+        let buffers = (0..<segments).map { i in
+            ByteBuffer(repeating: UInt8(i), count: segmentSize)
+        }
+
+        // Coalesce the segments into a single buffer.
+        var buffer = self.firstChannel.allocator.buffer(capacity: segments * segmentSize)
+        for segment in buffers {
+            buffer.writeImmutableBuffer(segment)
+        }
+
+        // Write the single large buffer with per-message GSO metadata.
+        let writeData = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: buffer,
+            metadata: .init(ecnState: .transportNotCapable, packetInfo: nil, segmentSize: segmentSize)
+        )
+        XCTAssertNoThrow(try self.firstChannel.writeAndFlush(writeData).wait())
+
+        // The receiver will receive separate segments.
+        let receivedBuffers = try self.secondChannel.waitForDatagrams(count: segments)
+        let receivedBytes = receivedBuffers.map { $0.data.readableBytes }.reduce(0, +)
+        XCTAssertEqual(segmentSize * segments, receivedBytes)
+
+        var unusedIndexes = Set(buffers.indices)
+        for envelope in receivedBuffers {
+            if let index = buffers.firstIndex(of: envelope.data) {
+                XCTAssertNotNil(unusedIndexes.remove(index))
+            } else {
+                XCTFail("No matching buffer")
+            }
+        }
+    }
+
+    func testLargeVectorWriteWithPerMessageGSO() throws {
+        try XCTSkipUnless(System.supportsUDPSegmentationOffload, "UDP_SEGMENT (GSO) is not supported on this platform")
+
+        // Similar to the test above, but with multiple writes using different segment sizes.
+        let segmentSize1 = 1000
+        let segments1 = 10
+        let segmentSize2 = 500
+        let segments2 = 5
+
+        // Form segments for first write
+        let buffers1 = (0..<segments1).map { i in
+            ByteBuffer(repeating: UInt8(i), count: segmentSize1)
+        }
+        var buffer1 = self.firstChannel.allocator.buffer(capacity: segments1 * segmentSize1)
+        for segment in buffers1 {
+            buffer1.writeImmutableBuffer(segment)
+        }
+
+        // Form segments for second write
+        let buffers2 = (0..<segments2).map { i in
+            ByteBuffer(repeating: UInt8(100 + i), count: segmentSize2)
+        }
+        var buffer2 = self.firstChannel.allocator.buffer(capacity: segments2 * segmentSize2)
+        for segment in buffers2 {
+            buffer2.writeImmutableBuffer(segment)
+        }
+
+        // Write both buffers with different segment sizes.
+        let writeData1 = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: buffer1,
+            metadata: .init(ecnState: .transportNotCapable, packetInfo: nil, segmentSize: segmentSize1)
+        )
+        let writeData2 = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: buffer2,
+            metadata: .init(ecnState: .transportNotCapable, packetInfo: nil, segmentSize: segmentSize2)
+        )
+        let write1 = self.firstChannel.write(writeData1)
+        let write2 = self.firstChannel.write(writeData2)
+        self.firstChannel.flush()
+        XCTAssertNoThrow(try write1.wait())
+        XCTAssertNoThrow(try write2.wait())
+
+        // The receiver will receive separate segments from both writes.
+        let receivedBuffers = try self.secondChannel.waitForDatagrams(count: segments1 + segments2)
+        XCTAssertEqual(receivedBuffers.count, segments1 + segments2)
+    }
+
+    func testMixedGSOAndNonGSO() throws {
+        try XCTSkipUnless(System.supportsUDPSegmentationOffload, "UDP_SEGMENT (GSO) is not supported on this platform")
+
+        // Send some messages with GSO and some without.
+        let segmentSize = 1000
+        let segments = 5
+
+        // GSO message
+        var gsoBuffer = self.firstChannel.allocator.buffer(capacity: segments * segmentSize)
+        for _ in 0..<segments {
+            gsoBuffer.writeRepeatingByte(1, count: segmentSize)
+        }
+        let gsoEnvelope = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: gsoBuffer,
+            metadata: .init(ecnState: .transportNotCapable, packetInfo: nil, segmentSize: segmentSize)
+        )
+
+        // Non-GSO message
+        var normalBuffer = self.firstChannel.allocator.buffer(capacity: 100)
+        normalBuffer.writeRepeatingByte(2, count: 100)
+        let normalEnvelope = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: normalBuffer
+        )
+
+        // Send both
+        let write1 = self.firstChannel.write(gsoEnvelope)
+        let write2 = self.firstChannel.write(normalEnvelope)
+        self.firstChannel.flush()
+        XCTAssertNoThrow(try write1.wait())
+        XCTAssertNoThrow(try write2.wait())
+
+        // Should receive 5 segments + 1 normal message
+        let received = try self.secondChannel.waitForDatagrams(count: segments + 1)
+        XCTAssertEqual(received.count, segments + 1)
+    }
+
+    func testPerMessageGSOOverridesChannelLevel() throws {
+        try XCTSkipUnless(System.supportsUDPSegmentationOffload, "UDP_SEGMENT (GSO) is not supported on this platform")
+
+        // Set channel-level GSO to one value
+        XCTAssertNoThrow(try self.firstChannel.setOption(.datagramSegmentSize, value: CInt(500)).wait())
+
+        // But use per-message GSO with a different value
+        let segmentSize = 1000
+        let segments = 5
+
+        var buffer = self.firstChannel.allocator.buffer(capacity: segments * segmentSize)
+        for i in 0..<segments {
+            buffer.writeRepeatingByte(UInt8(i), count: segmentSize)
+        }
+
+        let envelope = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: buffer,
+            metadata: .init(ecnState: .transportNotCapable, packetInfo: nil, segmentSize: segmentSize)
+        )
+
+        XCTAssertNoThrow(try self.firstChannel.writeAndFlush(envelope).wait())
+
+        // Should receive segments of size 1000, not 500
+        let received = try self.secondChannel.waitForDatagrams(count: segments)
+        for datagram in received {
+            XCTAssertEqual(datagram.data.readableBytes, segmentSize)
+        }
+    }
+
+    func testPerMessageGSOThrowsOnNonLinux() throws {
+        #if os(Linux)
+        try XCTSkipIf(true, "This test only runs on non-Linux platforms")
+        #else
+        // On non-Linux platforms, setting segmentSize in metadata should throw an error
+        var buffer = self.firstChannel.allocator.buffer(capacity: 1000)
+        buffer.writeRepeatingByte(1, count: 1000)
+
+        let envelope = AddressedEnvelope(
+            remoteAddress: self.secondChannel.localAddress!,
+            data: buffer,
+            metadata: .init(ecnState: .transportNotCapable, packetInfo: nil, segmentSize: 500)
+        )
+
+        XCTAssertThrowsError(try self.firstChannel.writeAndFlush(envelope).wait()) { error in
+            XCTAssertEqual(error as? ChannelError, .operationUnsupported)
+        }
+        #endif
+    }
 }
 
 extension System {
