Avoid copies of large payloads

glbrntt · glbrntt · commit 564f522fccc1 · 2022-12-07T12:24:31.000Z
Motivation:

Messages are prefixed with a 5-byte header. Currently messages of all
sizes are written into a new buffer with their header. For smaller
payloads this is good: we avoid the extra allocations associated with
creating HTTP/2 frames. For large payloads the cost of the copy
outweighs the cost of extra allocations.

Modifications:

- For messages larger than 8KB emit an extra HTTP/2 DATA frame
  containing just the message header.

Result:

Better performance for large payloads.
diff --git a/Sources/GRPC/GRPCClientChannelHandler.swift b/Sources/GRPC/GRPCClientChannelHandler.swift
@@ -530,15 +530,26 @@ extension GRPCClientChannelHandler: ChannelOutboundHandler {
         allocator: context.channel.allocator
       )
       switch result {
-      case let .success(buffer):
-        // We're clear to send a message; wrap it up in an HTTP/2 frame.
-        let framePayload = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(buffer)))
+      case let .success((buffer, maybeBuffer)):
+        let frame1 = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(buffer)))
         self.logger.trace("writing HTTP2 frame", metadata: [
           MetadataKey.h2Payload: "DATA",
           MetadataKey.h2DataBytes: "\(buffer.readableBytes)",
           MetadataKey.h2EndStream: "false",
         ])
-        context.write(self.wrapOutboundOut(framePayload), promise: promise)
+        // If there's a second buffer, attach the promise to the second write.
+        let promise1 = maybeBuffer == nil ? promise : nil
+        context.write(self.wrapOutboundOut(frame1), promise: promise1)
+
+        if let maybeBuffer = maybeBuffer {
+          let frame2 = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(maybeBuffer)))
+          self.logger.trace("writing HTTP2 frame", metadata: [
+            MetadataKey.h2Payload: "DATA",
+            MetadataKey.h2DataBytes: "\(maybeBuffer.readableBytes)",
+            MetadataKey.h2EndStream: "false",
+          ])
+          context.write(self.wrapOutboundOut(frame2), promise: promise)
+        }
 
       case let .failure(writeError):
         switch writeError {
diff --git a/Sources/GRPC/GRPCClientStateMachine.swift b/Sources/GRPC/GRPCClientStateMachine.swift
@@ -197,7 +197,7 @@ struct GRPCClientStateMachine {
     _ message: ByteBuffer,
     compressed: Bool,
     allocator: ByteBufferAllocator
-  ) -> Result<ByteBuffer, MessageWriteError> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), MessageWriteError> {
     return self.withStateAvoidingCoWs { state in
       state.sendRequest(message, compressed: compressed, allocator: allocator)
     }
@@ -392,8 +392,8 @@ extension GRPCClientStateMachine.State {
     _ message: ByteBuffer,
     compressed: Bool,
     allocator: ByteBufferAllocator
-  ) -> Result<ByteBuffer, MessageWriteError> {
-    let result: Result<ByteBuffer, MessageWriteError>
+  ) -> Result<(ByteBuffer, ByteBuffer?), MessageWriteError> {
+    let result: Result<(ByteBuffer, ByteBuffer?), MessageWriteError>
 
     switch self {
     case .clientActiveServerIdle(var writeState, let pendingReadState):
diff --git a/Sources/GRPC/HTTP2ToRawGRPCServerCodec.swift b/Sources/GRPC/HTTP2ToRawGRPCServerCodec.swift
@@ -295,9 +295,16 @@ internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServe
     )
 
     switch writeBuffer {
-    case let .success(buffer):
-      let payload = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(buffer)))
-      self.context.write(self.wrapOutboundOut(payload), promise: promise)
+    case let .success((buffer, maybeBuffer)):
+      if let maybeBuffer = maybeBuffer {
+        let payload1 = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(buffer)))
+        self.context.write(self.wrapOutboundOut(payload1), promise: nil)
+        let payload2 = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(maybeBuffer)))
+        self.context.write(self.wrapOutboundOut(payload2), promise: promise)
+      } else {
+        let payload = HTTP2Frame.FramePayload.data(.init(data: .byteBuffer(buffer)))
+        self.context.write(self.wrapOutboundOut(payload), promise: promise)
+      }
       if metadata.flush {
         self.markFlushPoint()
       }
diff --git a/Sources/GRPC/HTTP2ToRawGRPCStateMachine.swift b/Sources/GRPC/HTTP2ToRawGRPCStateMachine.swift
@@ -644,10 +644,10 @@ extension HTTP2ToRawGRPCStateMachine {
     compress: Bool,
     allocator: ByteBufferAllocator,
     writer: LengthPrefixedMessageWriter
-  ) -> Result<ByteBuffer, Error> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     do {
-      let prefixed = try writer.write(buffer: buffer, allocator: allocator, compressed: compress)
-      return .success(prefixed)
+      let buffers = try writer.write(buffer: buffer, allocator: allocator, compressed: compress)
+      return .success(buffers)
     } catch {
       return .failure(error)
     }
@@ -659,7 +659,7 @@ extension HTTP2ToRawGRPCStateMachine.RequestOpenResponseOpenState {
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compress: Bool
-  ) -> Result<ByteBuffer, Error> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     return HTTP2ToRawGRPCStateMachine.writeGRPCFramedMessage(
       buffer,
       compress: compress,
@@ -674,7 +674,7 @@ extension HTTP2ToRawGRPCStateMachine.RequestClosedResponseOpenState {
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compress: Bool
-  ) -> Result<ByteBuffer, Error> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     return HTTP2ToRawGRPCStateMachine.writeGRPCFramedMessage(
       buffer,
       compress: compress,
@@ -907,7 +907,7 @@ extension HTTP2ToRawGRPCStateMachine {
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compress: Bool
-  ) -> Result<ByteBuffer, Error> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     return self.state.send(buffer: buffer, allocator: allocator, compress: compress)
   }
 
@@ -1119,7 +1119,7 @@ extension HTTP2ToRawGRPCStateMachine.State {
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compress: Bool
-  ) -> Result<ByteBuffer, Error> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     switch self {
     case .requestIdleResponseIdle:
       preconditionFailure("Invalid state: the request stream is still closed")
diff --git a/Sources/GRPC/LengthPrefixedMessageWriter.swift b/Sources/GRPC/LengthPrefixedMessageWriter.swift
@@ -86,99 +86,42 @@ internal struct LengthPrefixedMessageWriter {
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compressed: Bool = true
-  ) throws -> ByteBuffer {
+  ) throws -> (ByteBuffer, ByteBuffer?) {
     if compressed, let compressor = self.compressor {
-      return try self.compress(buffer: buffer, using: compressor, allocator: allocator)
-    } else if buffer.readerIndex >= 5 {
-      // We're not compressing and we have enough bytes before the reader index that we can write
-      // over with the compression byte and length.
-      var buffer = buffer
-
-      // Get the size of the message.
-      let messageSize = buffer.readableBytes
-
-      // Move the reader index back 5 bytes. This is okay: we validated the `readerIndex` above.
-      buffer.moveReaderIndex(to: buffer.readerIndex - 5)
-
-      // Fill in the compression byte and message length.
-      buffer.setInteger(UInt8(0), at: buffer.readerIndex)
-      buffer.setInteger(UInt32(messageSize), at: buffer.readerIndex + 1)
-
-      // The message bytes are already in place, we're done.
-      return buffer
+      let compressedAndFramedPayload = try self.compress(
+        buffer: buffer,
+        using: compressor,
+        allocator: allocator
+      )
+      return (compressedAndFramedPayload, nil)
+    } else if buffer.readableBytes > Self.singleBufferSizeLimit {
+      // Buffer is larger than the limit for emitting a single buffer: create a second buffer
+      // containing just the message header.
+      var prefixed = allocator.buffer(capacity: 5)
+      prefixed.writeMultipleIntegers(UInt8(0), UInt32(buffer.readableBytes))
+      return (prefixed, buffer)
     } else {
-      // We're not compressing and we don't have enough space before the message bytes passed in.
-      // We need a new buffer.
-      var lengthPrefixed = allocator.buffer(capacity: 5 + buffer.readableBytes)
-
-      // Write the compression byte.
-      lengthPrefixed.writeInteger(UInt8(0))
-
-      // Write the message length.
-      lengthPrefixed.writeInteger(UInt32(buffer.readableBytes))
-
+      // We're not compressing and the message is within our single buffer size limit.
+      var lengthPrefixed = allocator.buffer(capacity: 5 &+ buffer.readableBytes)
+      // Write the compression byte and message length.
+      lengthPrefixed.writeMultipleIntegers(UInt8(0), UInt32(buffer.readableBytes))
       // Write the message.
-      var buffer = buffer
-      lengthPrefixed.writeBuffer(&buffer)
-
-      return lengthPrefixed
+      lengthPrefixed.writeImmutableBuffer(buffer)
+      return (lengthPrefixed, nil)
     }
   }
 
-  /// Writes the data into a `ByteBuffer` as a gRPC length-prefixed message.
+  /// Message size above which we emit two buffers: one containing the header and one with the
+  /// actual message bytes. At or below the limit we copy the message into a new buffer containing
+  /// both the header and the message.
   ///
-  /// - Parameters:
-  ///   - payload: The payload to serialize and write.
-  ///   - buffer: The buffer to write the message into.
-  /// - Returns: A `ByteBuffer` containing a gRPC length-prefixed message.
-  /// - Precondition: `compression.supported` is `true`.
-  /// - Note: See `LengthPrefixedMessageReader` for more details on the format.
-  func write(
-    _ payload: GRPCPayload,
-    into buffer: inout ByteBuffer,
-    compressed: Bool = true
-  ) throws {
-    buffer.reserveCapacity(buffer.writerIndex + LengthPrefixedMessageWriter.metadataLength)
-
-    if compressed, let compressor = self.compressor {
-      // Set the compression byte.
-      buffer.writeInteger(UInt8(1))
-
-      // Leave a gap for the length, we'll set it in a moment.
-      let payloadSizeIndex = buffer.writerIndex
-      buffer.moveWriterIndex(forwardBy: MemoryLayout<UInt32>.size)
-
-      var messageBuf = ByteBufferAllocator().buffer(capacity: 0)
-      try payload.serialize(into: &messageBuf)
-
-      // Compress the message.
-      let bytesWritten = try compressor.deflate(&messageBuf, into: &buffer)
-
-      // Now fill in the message length.
-      buffer.writePayloadLength(UInt32(bytesWritten), at: payloadSizeIndex)
-
-      // Finally, the compression context should be reset between messages.
-      compressor.reset()
-    } else {
-      // We could be using 'identity' compression, but since the result is the same we'll just
-      // say it isn't compressed.
-      buffer.writeInteger(UInt8(0))
-
-      // Leave a gap for the length, we'll set it in a moment.
-      let payloadSizeIndex = buffer.writerIndex
-      buffer.moveWriterIndex(forwardBy: MemoryLayout<UInt32>.size)
-
-      let payloadPrefixedBytes = buffer.readableBytes
-      // Writes the payload into the buffer
-      try payload.serialize(into: &buffer)
-
-      // Calculates the Written bytes with respect to the prefixed ones
-      let bytesWritten = buffer.readableBytes - payloadPrefixedBytes
-
-      // Write the message length.
-      buffer.writePayloadLength(UInt32(bytesWritten), at: payloadSizeIndex)
-    }
-  }
+  /// Using two buffers avoids expensive copies of large messages. For smaller messages the copy
+  /// is cheaper than the additional allocations and overhead required to send an extra HTTP/2 DATA
+  /// frame.
+  ///
+  /// The value of 8192 was chosen empirically. We subtract the length of the message header
+  /// as `ByteBuffer` reserve capacity in powers of two and want to avoid overallocating.
+  private static let singleBufferSizeLimit = 8192 - 5
 }
 
 extension ByteBuffer {
diff --git a/Sources/GRPC/ReadWriteStates.swift b/Sources/GRPC/ReadWriteStates.swift
@@ -62,16 +62,16 @@ enum WriteState {
     _ message: ByteBuffer,
     compressed: Bool,
     allocator: ByteBufferAllocator
-  ) -> Result<ByteBuffer, MessageWriteError> {
+  ) -> Result<(ByteBuffer, ByteBuffer?), MessageWriteError> {
     switch self {
     case .notWriting:
       return .failure(.cardinalityViolation)
 
     case let .writing(writeArity, contentType, writer):
-      let buffer: ByteBuffer
+      let buffers: (ByteBuffer, ByteBuffer?)
 
       do {
-        buffer = try writer.write(buffer: message, allocator: allocator, compressed: compressed)
+        buffers = try writer.write(buffer: message, allocator: allocator, compressed: compressed)
       } catch {
         self = .notWriting
         return .failure(.serializationFailed)
@@ -84,7 +84,7 @@ enum WriteState {
         self = .writing(writeArity, contentType, writer)
       }
 
-      return .success(buffer)
+      return .success(buffers)
     }
   }
 }
diff --git a/Tests/GRPCTests/GRPCClientStateMachineTests.swift b/Tests/GRPCTests/GRPCClientStateMachineTests.swift
@@ -39,7 +39,16 @@ class GRPCClientStateMachineTests: GRPCTestCase {
     let buffer = self.allocator.buffer(string: message)
 
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    return try writer.write(buffer: buffer, allocator: self.allocator, compressed: false)
+    var (buffer1, buffer2) = try writer.write(
+      buffer: buffer,
+      allocator: self.allocator,
+      compressed: false
+    )
+
+    if var buffer2 = buffer2 {
+      buffer1.writeBuffer(&buffer2)
+    }
+    return buffer1
   }
 
   /// Writes a message into the given `buffer`.
@@ -148,8 +157,9 @@ extension GRPCClientStateMachineTests {
       ByteBuffer(string: request),
       compressed: false,
       allocator: self.allocator
-    ).assertSuccess { buffer in
-      var buffer = buffer
+    ).assertSuccess { buffers in
+      var buffer = buffers.0
+      XCTAssertNil(buffers.1)
       // Remove the length and compression flag prefix.
       buffer.moveReaderIndex(forwardBy: 5)
       let data = buffer.readString(length: buffer.readableBytes)!
@@ -1149,13 +1159,14 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer:
     let message = ByteBuffer(string: "Hello!")
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var buffer = try writer.write(buffer: message, allocator: self.allocator)
+    var buffers = try writer.write(buffer: message, allocator: self.allocator)
+    XCTAssertNil(buffers.1)
     // And some extra junk bytes:
     let bytes: [UInt8] = [0x00]
-    buffer.writeBytes(bytes)
+    buffers.0.writeBytes(bytes)
 
     var state: ReadState = .one()
-    state.readMessages(&buffer, maxLength: .max).assertFailure {
+    state.readMessages(&buffers.0, maxLength: .max).assertFailure {
       XCTAssertEqual($0, .leftOverBytes)
     }
     state.assertNotReading()
@@ -1165,12 +1176,14 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer twice:
     let message = ByteBuffer(string: "Hello!")
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var buffer = try writer.write(buffer: message, allocator: self.allocator)
-    var second = try writer.write(buffer: message, allocator: self.allocator)
-    buffer.writeBuffer(&second)
+    var buffers1 = try writer.write(buffer: message, allocator: self.allocator)
+    var buffers2 = try writer.write(buffer: message, allocator: self.allocator)
+    XCTAssertNil(buffers1.1)
+    XCTAssertNil(buffers2.1)
+    buffers1.0.writeBuffer(&buffers2.0)
 
     var state: ReadState = .one()
-    state.readMessages(&buffer, maxLength: .max).assertFailure {
+    state.readMessages(&buffers1.0, maxLength: .max).assertFailure {
       XCTAssertEqual($0, .cardinalityViolation)
     }
     state.assertNotReading()
@@ -1180,7 +1193,8 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer twice:
     let message = ByteBuffer(string: "Hello!")
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var buffer = try writer.write(buffer: message, allocator: self.allocator)
+    var (buffer, other) = try writer.write(buffer: message, allocator: self.allocator)
+    XCTAssertNil(other)
 
     var state: ReadState = .one()
     state.readMessages(&buffer, maxLength: .max).assertSuccess {
@@ -1195,7 +1209,8 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer twice:
     let message = ByteBuffer(string: "Hello!")
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var buffer = try writer.write(buffer: message, allocator: self.allocator)
+    var (buffer, other) = try writer.write(buffer: message, allocator: self.allocator)
+    XCTAssertNil(other)
 
     var state: ReadState = .many()
     state.readMessages(&buffer, maxLength: .max).assertSuccess {
@@ -1211,9 +1226,9 @@ class ReadStateTests: GRPCTestCase {
     let message = ByteBuffer(string: "Hello!")
     let writer = LengthPrefixedMessageWriter(compression: .none)
 
-    var first = try writer.write(buffer: message, allocator: self.allocator)
-    var second = try writer.write(buffer: message, allocator: self.allocator)
-    var third = try writer.write(buffer: message, allocator: self.allocator)
+    var (first, _) = try writer.write(buffer: message, allocator: self.allocator)
+    var (second, _) = try writer.write(buffer: message, allocator: self.allocator)
+    var (third, _) = try writer.write(buffer: message, allocator: self.allocator)
 
     first.writeBuffer(&second)
     first.writeBuffer(&third)
diff --git a/Tests/GRPCTests/LengthPrefixedMessageWriterTests.swift b/Tests/GRPCTests/LengthPrefixedMessageWriterTests.swift
