large payload tests and performance improvements

Author: sliemeobn

Sources/AMQPClient/AMQPChannel.swift

@@ -32,13 +32,13 @@ public final class AMQPChannel: Sendable {
     private let isConfirmMode = ManagedAtomic(false)
     private let isTxMode = ManagedAtomic(false)
     private let deliveryTag = ManagedAtomic(UInt64(1))
-    private let frameMax: UInt32
+    private let frameMax: Int
 
     init(channelID: Frame.ChannelID, eventLoop: EventLoop, channel: AMQPChannelHandler, frameMax: UInt32) {
         ID = channelID
         self.eventLoop = eventLoop
         self.channel = channel
-        self.frameMax = frameMax
+        self.frameMax = Int(frameMax)
     }
 
     /// Close the channel
@@ -90,22 +90,19 @@ public final class AMQPChannel: Sendable {
 
         let header = Frame.Payload.header(.init(classID: classID, weight: 0, bodySize: UInt64(body.readableBytes), properties: properties))
 
-        let payloads: [Frame.Payload]
+        var payloads: [Frame.Payload]
 
         if body.readableBytes <= frameMax {
             payloads = [publish, header, .body(body)]
         } else {
-            var parts = [publish, header]
-            var buffer = body
+            payloads = [publish, header]
+            var body = body
 
-            while(buffer.readableBytes > 0) {
-                guard let bytes = buffer.readBytes(length: frameMax < buffer.readableBytes ? Int(frameMax) : buffer.readableBytes) else {
-                    preconditionFailure("invalid bytes read")
-                }
-                parts.append(.body(.init(bytes: bytes)))
+            while body.readableBytes > 0 {
+                // slice is always valid
+                let slice = body.readSlice(length: min(frameMax, body.readableBytes))!
+                payloads.append(Frame.Payload.body(slice))
             }
-
-            payloads = parts
         }
 
         let result: EventLoopFuture<Void> = channel.send(payloads: payloads)

Sources/AMQPClient/ChannelHandlers/AMQPConnectionMultiplexHandler.swift

@@ -42,7 +42,7 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
         // NOTE: this can be extended to keep some state of the open request so a response can be verified against its request
         var pendingRequests: Deque<EventLoopPromise<AMQPResponse>>
         weak var eventHandler: AMQPChannelHandler?
-        var nextMessage: (frame: Frame.Method.Basic, properties: Properties?, bodySize: UInt64?, prevBody: ByteBuffer?)?
+        var nextMessage: PartialDelivery?
 
         init(initialResponsePromise: EventLoopPromise<AMQPResponse>) {
             pendingRequests = .init([initialResponsePromise])
@@ -207,7 +207,7 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
                     channel.fulfilNextPendingRequest(with: .channel(.message(.get())))
                 case .deliver, .getOk, .return:
                     // TODO: wrap this away more nicely, assert message must be nil
-                    channel.nextMessage = (frame: basic, nil, nil, nil)
+                    channel.nextMessage = PartialDelivery(method: basic)
                 case .recoverOk:
                     channel.fulfilNextPendingRequest(with: .channel(.basic(.recovered)))
                 case let .consumeOk(consumerTag):
@@ -266,72 +266,56 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
                 }
             }
         case let .header(header):
-            channel.nextMessage?.properties = header.properties
-            channel.nextMessage?.bodySize = header.bodySize
+            channel.nextMessage?.setHeader(header)
         case let .body(body):
-            guard let msg = channel.nextMessage, let properties = msg.properties, let bodySize = msg.bodySize else {
-                // TODO: take down channel
-                return
-            }
-
-            let prevSize = msg.prevBody?.readableBytes ?? 0
-            if (prevSize + body.readableBytes < bodySize) {
-                if var prevBody = msg.prevBody  {
-                    prevBody.writeImmutableBuffer(body)
-                    channel.nextMessage?.prevBody = prevBody
-                } else {
-                    channel.nextMessage?.prevBody = body
-                }
-
-                return
-            }
-            
-            let completeBody: ByteBuffer
-            
-            if var prevBody = msg.prevBody {
-                prevBody.writeImmutableBuffer(body)
-                completeBody = prevBody
-            } else {
-                completeBody = body
-            }
-
-            switch msg.frame {
-            case let .getOk(getOk):
-                channel.fulfilNextPendingRequest(with: .channel(.message(.get(.init(
-                    message: .init(
-                        exchange: getOk.exchange,
-                        routingKey: getOk.routingKey,
-                        deliveryTag: getOk.deliveryTag,
-                        properties: properties,
-                        redelivered: getOk.redelivered,
-                        body: completeBody
-                    ),
-                    messageCount: getOk.messageCount
-                )))))
-            case let .deliver(deliver):
-                channel.eventHandler?.receiveDelivery(
-                    .init(
-                        exchange: deliver.exchange,
-                        routingKey: deliver.routingKey,
-                        deliveryTag: deliver.deliveryTag,
+            // TODO: take down channel
+            guard channel.nextMessage != nil else { return }
+
+            // written like this to avoid unnecessary copies
+            channel.nextMessage!.addBody(body)
+
+            if channel.nextMessage!.isComplete {
+                let (method, properties, completeBody) = channel.nextMessage!.asCompletedMessage()
+                channel.nextMessage = nil
+
+                switch method {
+                case let .getOk(getOk):
+                    channel.fulfilNextPendingRequest(with: .channel(.message(.get(.init(
+                        message: .init(
+                            exchange: getOk.exchange,
+                            routingKey: getOk.routingKey,
+                            deliveryTag: getOk.deliveryTag,
+                            properties: properties,
+                            redelivered: getOk.redelivered,
+                            body: completeBody
+                        ),
+                        messageCount: getOk.messageCount
+                    )))))
+                case let .deliver(deliver):
+                    channel.eventHandler?.receiveDelivery(
+                        .init(
+                            exchange: deliver.exchange,
+                            routingKey: deliver.routingKey,
+                            deliveryTag: deliver.deliveryTag,
+                            properties: properties,
+                            redelivered: deliver.redelivered,
+                            body: completeBody
+                        ),
+                        for: deliver.consumerTag
+                    )
+                case let .return(`return`):
+                    channel.eventHandler?.receiveReturn(.init(
+                        replyCode: `return`.replyCode,
+                        replyText: `return`.replyText,
+                        exchange: `return`.exchange,
+                        routingKey: `return`.routingKey,
                         properties: properties,
-                        redelivered: deliver.redelivered,
                         body: completeBody
-                    ),
-                    for: deliver.consumerTag
-                )
-            case let .return(`return`):
-                channel.eventHandler?.receiveReturn(.init(
-                    replyCode: `return`.replyCode,
-                    replyText: `return`.replyText,
-                    exchange: `return`.exchange,
-                    routingKey: `return`.routingKey,
-                    properties: properties,
-                    body: completeBody
-                ))
-            default:
-                // TODO: take down channel
-                preconditionUnexpectedFrame(frame)
+                    ))
+                default:
+                    // TODO: take down channel
+                    preconditionUnexpectedFrame(frame)
+                }
             }
         case .heartbeat:
             let heartbeat = Frame(channelID: frame.channelID, payload: .heartbeat)
@@ -400,3 +384,41 @@ extension AMQPConnectionMultiplexHandler.ChannelState {
         eventHandler?.reportAsClosed(error: error)
     }
 }
+
+private struct PartialDelivery {
+    let method: Frame.Method.Basic
+    private var header: Frame.Header?
+    private var payload: ByteBuffer?
+
+    init(method: Frame.Method.Basic) {
+        self.method = method
+    }
+
+    var isComplete: Bool { header != nil && header!.bodySize <= (payload?.readableBytes ?? 0) }
+
+    // NOTE: should be made throwing with validation for a more restrictive protocol implementation
+    mutating func setHeader(_ header: consuming Frame.Header) {
+        // validate that self.header == nil
+        self.header = header
+    }
+
+    // NOTE: should be made throwing with validation for a more restrictive protocol implementation
+    mutating func addBody(_ buffer: consuming ByteBuffer) {
+        guard let header else { return } // probably should take channel down
+
+        if payload == nil {
+            buffer.reserveCapacity(Int(header.bodySize))
+            payload = consume buffer
+        } else {
+            payload!.writeBuffer(&buffer)
+        }
+    }
+
+    func asCompletedMessage() -> (Frame.Method.Basic, Properties, ByteBuffer) {
+        // NOTE: this could be made a consuming func once partial is possible I think
+        assert(isComplete)
+
+        // header and payloads are guaranteed to be non-nil after isComplete
+        return (method, header!.properties, payload!)
+    }
+}

Tests/AMQPClientTests/AMQPChannelLargePayloadsTest.swift

@@ -0,0 +1,96 @@
+
+@testable import AMQPClient
+import AMQPProtocol
+import NIO
+import NIOPosix
+import XCTest
+
+final class AMQPChannelLargePayloadsTest: XCTestCase {
+    let testExchange = "lage_payloads"
+    var connection: AMQPConnection!
+    var channel: AMQPChannel!
+
+    override func setUp() async throws {
+        connection = try await AMQPConnection.connect(use: MultiThreadedEventLoopGroup.singleton.next(), from: .init(connection: .plain, server: .init()))
+        channel = try await connection.openChannel()
+        try await channel.exchangeDeclare(name: testExchange, type: "fanout")
+    }
+
+    override func tearDown() async throws {
+        try await channel?.close()
+        try await connection?.close()
+    }
+
+    func testPublishLargePayloads() async throws {
+        try await channel.basicPublish(from: ByteBuffer(megaBytes: 50), exchange: testExchange, routingKey: "")
+        try await channel.basicPublish(from: ByteBuffer(megaBytes: 10), exchange: testExchange, routingKey: "")
+    }
+
+    func testGetLargePayloads() async throws {
+        let queueName = "lp-test-get"
+        let payload1 = ByteBuffer(megaBytes: 20)
+        let payload2 = ByteBuffer(megaBytes: 30, repeating: UInt8(ascii: "b"))
+
+        try await channel.queueDeclare(name: queueName, exclusive: true)
+        try await channel.queueBind(queue: queueName, exchange: testExchange)
+
+        try await channel.basicPublish(from: payload1, exchange: testExchange, routingKey: "")
+        try await channel.basicPublish(from: payload2, exchange: testExchange, routingKey: "")
+        let getResponse1 = try await channel.basicGet(queue: queueName, noAck: true)
+        let getResponse2 = try await channel.basicGet(queue: queueName, noAck: true)
+
+        XCTAssertEqual(getResponse1?.message.body, payload1)
+        XCTAssertEqual(getResponse1?.messageCount, 1)
+        XCTAssertEqual(getResponse2?.message.body, payload2)
+        XCTAssertEqual(getResponse2?.messageCount, 0)
+    }
+
+    func testConsumeLargePayloads() async throws {
+        let queueName = "lp-test-consume"
+        let payload1 = ByteBuffer(megaBytes: 20, repeating: UInt8(ascii: "a"))
+        let payload2 = ByteBuffer(megaBytes: 10, repeating: UInt8(ascii: "b"))
+
+        try await channel.queueDeclare(name: queueName, exclusive: true)
+        try await channel.queueBind(queue: queueName, exchange: testExchange)
+
+        var messages = try await channel.basicConsume(queue: queueName, noAck: true).prefix(2).makeAsyncIterator()
+
+        try await channel.basicPublish(from: payload1, exchange: testExchange, routingKey: "")
+        try await channel.basicPublish(from: payload2, exchange: testExchange, routingKey: "")
+
+        let received1 = try await messages.next()
+        let received2 = try await messages.next()
+        let received3 = try await messages.next()
+
+        XCTAssertEqual(received1?.body, payload1)
+        XCTAssertEqual(received2?.body, payload2)
+        XCTAssertNil(received3)
+    }
+
+    func testReturnLargePayloads() async throws {
+        let payload1 = ByteBuffer(megaBytes: 20, repeating: UInt8(ascii: "a"))
+        let payload2 = ByteBuffer(megaBytes: 10, repeating: UInt8(ascii: "b"))
+
+        try await channel.exchangeDeclare(name: "no-binding", type: "fanout")
+
+        var returns = try await channel.returnConsume(named: "test").prefix(2).makeAsyncIterator()
+        
+
+        try await channel.basicPublish(from: payload1, exchange: "no-binding", routingKey: "", mandatory: true)
+        try await channel.basicPublish(from: payload2, exchange: "no-binding", routingKey: "", mandatory: true)
+
+        let received1 = try await returns.next()
+        let received2 = try await returns.next()
+        let received3 = try await returns.next()
+
+        XCTAssertEqual(received1?.body, payload1)
+        XCTAssertEqual(received2?.body, payload2)
+        XCTAssertNil(received3)
+    }
+}
+
+extension ByteBuffer {
+    init(megaBytes: Int, repeating: UInt8 = .init(ascii: "a")) {
+        self.init(repeating: repeating, count: megaBytes * 1024 * 1024)
+    }
+}