Merge pull request #40 from funcmike/large-payload-support

add first implemention of large payload support without tests

Author: funcmike

Sources/AMQPClient/AMQPChannel.swift

@@ -32,11 +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
 
-    init(channelID: Frame.ChannelID, eventLoop: EventLoop, channel: AMQPChannelHandler) {
+    init(channelID: Frame.ChannelID, eventLoop: EventLoop, channel: AMQPChannelHandler, frameMax: UInt32) {
         ID = channelID
         self.eventLoop = eventLoop
         self.channel = channel
+        self.frameMax = frameMax
     }
 
     /// Close the channel
@@ -88,7 +90,25 @@ public final class AMQPChannel: Sendable {
 
         let header = Frame.Payload.header(.init(classID: classID, weight: 0, bodySize: UInt64(body.readableBytes), properties: properties))
 
-        let result: EventLoopFuture<Void> = channel.send(payloads: [publish, header, .body(body)])
+        let payloads: [Frame.Payload]
+
+        if body.readableBytes <= frameMax {
+            payloads = [publish, header, .body(body)]
+        } else {
+            var parts = [publish, header]
+            var buffer = 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)))
+            }
+
+            payloads = parts
+        }
+
+        let result: EventLoopFuture<Void> = channel.send(payloads: payloads)
         return result.map {
             if self.isConfirmMode.load(ordering: .relaxed) {
                 let count = self.deliveryTag.loadThenWrappingIncrement(ordering: .sequentiallyConsistent)

Sources/AMQPClient/AMQPConnection.swift

@@ -31,15 +31,17 @@ public final class AMQPConnection: Sendable {
 
     public var closeFuture: NIOCore.EventLoopFuture<Void> { connectionHandler.channel.closeFuture }
     public var eventLoop: EventLoop { return connectionHandler.channel.eventLoop }
+    private let frameMax: UInt32
 
     private let connectionHandler: AMQPConnectionHandler
     private var channel: NIOCore.Channel { connectionHandler.channel }
 
     private let state = NIOLockedValueBox(ConnectionState.open)
     private let channels: NIOLockedValueBox<AMQPChannels>
 
-    init(connectionHandler: AMQPConnectionHandler, channelMax: UInt16) {
+    init(connectionHandler: AMQPConnectionHandler, channelMax: UInt16, frameMax: UInt32) {
         self.connectionHandler = connectionHandler
+        self.frameMax = frameMax
         channels = .init(AMQPChannels(channelMax: channelMax))
     }
 
@@ -54,7 +56,8 @@ public final class AMQPConnection: Sendable {
                 connectionHandler.startConnection().map {
                     AMQPConnection(
                         connectionHandler: connectionHandler,
-                        channelMax: $0.channelMax
+                        channelMax: $0.channelMax,
+                        frameMax: $0.frameMax
                     )
                 }
             }
@@ -79,7 +82,7 @@ public final class AMQPConnection: Sendable {
         future.whenFailure { _ in self.channels.withLockedValue { $0.remove(id: channelID) } }
 
         return future.map { channel in
-            let amqpChannel = AMQPChannel(channelID: channelID, eventLoop: self.eventLoop, channel: channel)
+            let amqpChannel = AMQPChannel(channelID: channelID, eventLoop: self.eventLoop, channel: channel, frameMax: self.frameMax)
             self.channels.withLockedValue { $0.add(channel: amqpChannel) }
             return amqpChannel
         }

Sources/AMQPClient/AMQPResponse.swift

@@ -131,6 +131,7 @@ public enum AMQPResponse: Sendable {
 
         public struct Connected: Sendable {
             public let channelMax: UInt16
+            public let frameMax: UInt32
         }
     }
 }

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?)?
+        var nextMessage: (frame: Frame.Method.Basic, properties: Properties?, bodySize: UInt64?, prevBody: ByteBuffer?)?
 
         init(initialResponsePromise: EventLoopPromise<AMQPResponse>) {
             pendingRequests = .init([initialResponsePromise])
@@ -53,6 +53,7 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
     private var context: ChannelHandlerContext!
     private var channels: [Frame.ChannelID: ChannelState] = [:]
     private var channelMax: UInt16 = 0
+    private var frameMax: UInt32 = 0
     private var state: State = .unblocked
 
     private let config: AMQPConnectionConfiguration.Server
@@ -131,14 +132,15 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
                     context.writeAndFlush(wrapOutboundOut(.frame(startOk)), promise: nil)
                 case let .tune(channelMax, frameMax, heartbeat):
                     self.channelMax = channelMax
+                    self.frameMax = frameMax
 
                     let tuneOk = Frame(channelID: frame.channelID,
                                        payload: .method(.connection(.tuneOk(channelMax: channelMax, frameMax: frameMax, heartbeat: heartbeat))))
                     let open = Frame(channelID: frame.channelID, payload: .method(.connection(.open(.init(vhost: config.vhost)))))
 
                     context.writeAndFlush(wrapOutboundOut(.bulk([tuneOk, open])), promise: nil)
                 case .openOk:
-                    channel.fulfilNextPendingRequest(with: .connection(.connected(.init(channelMax: channelMax))))
+                    channel.fulfilNextPendingRequest(with: .connection(.connected(.init(channelMax: channelMax, frameMax: frameMax))))
                 case let .close(close):
                     let closeOk = Frame(channelID: frame.channelID, payload: .method(.connection(.closeOk)))
                     context.writeAndFlush(wrapOutboundOut(.frame(closeOk)), promise: nil)
@@ -205,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)
+                    channel.nextMessage = (frame: basic, nil, nil, nil)
                 case .recoverOk:
                     channel.fulfilNextPendingRequest(with: .channel(.basic(.recovered)))
                 case let .consumeOk(consumerTag):
@@ -265,12 +267,34 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
             }
         case let .header(header):
             channel.nextMessage?.properties = header.properties
+            channel.nextMessage?.bodySize = header.bodySize
         case let .body(body):
-            guard let msg = channel.nextMessage, let properties = msg.properties else {
+            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(
@@ -280,7 +304,7 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
                         deliveryTag: getOk.deliveryTag,
                         properties: properties,
                         redelivered: getOk.redelivered,
-                        body: body
+                        body: completeBody
                     ),
                     messageCount: getOk.messageCount
                 )))))
@@ -292,7 +316,7 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
                         deliveryTag: deliver.deliveryTag,
                         properties: properties,
                         redelivered: deliver.redelivered,
-                        body: body
+                        body: completeBody
                     ),
                     for: deliver.consumerTag
                 )
@@ -303,7 +327,7 @@ internal final class AMQPConnectionMultiplexHandler: ChannelDuplexHandler {
                     exchange: `return`.exchange,
                     routingKey: `return`.routingKey,
                     properties: properties,
-                    body: body
+                    body: completeBody
                 ))
             default:
                 // TODO: take down channel