save buffer

Author: glbrntt

Sources/GRPC/GRPCClientChannelHandler.swift

@@ -502,7 +502,10 @@ extension GRPCClientChannelHandler: ChannelOutboundHandler {
     switch self.unwrapOutboundIn(data) {
     case let .head(requestHead):
       // Feed the request into the state machine:
-      switch self.stateMachine.sendRequestHeaders(requestHead: requestHead) {
+      switch self.stateMachine.sendRequestHeaders(
+        requestHead: requestHead,
+        allocator: context.channel.allocator
+      ) {
       case let .success(headers):
         // We're clear to write some headers. Create an appropriate frame and write it.
         let framePayload = HTTP2Frame.FramePayload.headers(.init(headers: headers))
@@ -526,8 +529,7 @@ extension GRPCClientChannelHandler: ChannelOutboundHandler {
       // Feed the request message into the state machine:
       let result = self.stateMachine.sendRequest(
         request.message,
-        compressed: request.compressed,
-        allocator: context.channel.allocator
+        compressed: request.compressed
       )
       switch result {
       case let .success((buffer, maybeBuffer)):

Sources/GRPC/GRPCClientStateMachine.swift

@@ -165,10 +165,11 @@ struct GRPCClientStateMachine {
   ///
   /// - Parameter requestHead: The client request head for the RPC.
   mutating func sendRequestHeaders(
-    requestHead: _GRPCRequestHead
+    requestHead: _GRPCRequestHead,
+    allocator: ByteBufferAllocator
   ) -> Result<HPACKHeaders, SendRequestHeadersError> {
     return self.withStateAvoidingCoWs { state in
-      state.sendRequestHeaders(requestHead: requestHead)
+      state.sendRequestHeaders(requestHead: requestHead, allocator: allocator)
     }
   }
 
@@ -195,11 +196,10 @@ struct GRPCClientStateMachine {
   ///     request will be written.
   mutating func sendRequest(
     _ message: ByteBuffer,
-    compressed: Bool,
-    allocator: ByteBufferAllocator
+    compressed: Bool
   ) -> Result<(ByteBuffer, ByteBuffer?), MessageWriteError> {
     return self.withStateAvoidingCoWs { state in
-      state.sendRequest(message, compressed: compressed, allocator: allocator)
+      state.sendRequest(message, compressed: compressed)
     }
   }
 
@@ -351,7 +351,8 @@ struct GRPCClientStateMachine {
 extension GRPCClientStateMachine.State {
   /// See `GRPCClientStateMachine.sendRequestHeaders(requestHead:)`.
   mutating func sendRequestHeaders(
-    requestHead: _GRPCRequestHead
+    requestHead: _GRPCRequestHead,
+    allocator: ByteBufferAllocator
   ) -> Result<HPACKHeaders, SendRequestHeadersError> {
     let result: Result<HPACKHeaders, SendRequestHeadersError>
 
@@ -369,7 +370,10 @@ extension GRPCClientStateMachine.State {
       result = .success(headers)
 
       self = .clientActiveServerIdle(
-        writeState: pendingWriteState.makeWriteState(messageEncoding: requestHead.encoding),
+        writeState: pendingWriteState.makeWriteState(
+          messageEncoding: requestHead.encoding,
+          allocator: allocator
+        ),
         pendingReadState: .init(arity: responseArity, messageEncoding: requestHead.encoding)
       )
 
@@ -390,18 +394,17 @@ extension GRPCClientStateMachine.State {
   /// See `GRPCClientStateMachine.sendRequest(_:allocator:)`.
   mutating func sendRequest(
     _ message: ByteBuffer,
-    compressed: Bool,
-    allocator: ByteBufferAllocator
+    compressed: Bool
   ) -> Result<(ByteBuffer, ByteBuffer?), MessageWriteError> {
     let result: Result<(ByteBuffer, ByteBuffer?), MessageWriteError>
 
     switch self {
     case .clientActiveServerIdle(var writeState, let pendingReadState):
-      result = writeState.write(message, compressed: compressed, allocator: allocator)
+      result = writeState.write(message, compressed: compressed)
       self = .clientActiveServerIdle(writeState: writeState, pendingReadState: pendingReadState)
 
     case .clientActiveServerActive(var writeState, let readState):
-      result = writeState.write(message, compressed: compressed, allocator: allocator)
+      result = writeState.write(message, compressed: compressed)
       self = .clientActiveServerActive(writeState: writeState, readState: readState)
 
     case .clientClosedServerIdle,

Sources/GRPC/HTTP2ToRawGRPCStateMachine.swift

@@ -303,7 +303,11 @@ extension HTTP2ToRawGRPCStateMachine.State {
     }
 
     // Figure out which encoding we should use for responses.
-    let (writer, responseEncoding) = self.extractResponseEncoding(from: headers, encoding: encoding)
+    let (writer, responseEncoding) = self.extractResponseEncoding(
+      from: headers,
+      encoding: encoding,
+      allocator: allocator
+    )
 
     // Parse the path, and create a call handler.
     guard let path = headers.first(name: ":path") else {
@@ -516,7 +520,8 @@ extension HTTP2ToRawGRPCStateMachine.State {
   /// - Returns: A message writer and the response encoding header to send back to the client.
   private func extractResponseEncoding(
     from headers: HPACKHeaders,
-    encoding: ServerMessageEncoding
+    encoding: ServerMessageEncoding,
+    allocator: ByteBufferAllocator
   ) -> (LengthPrefixedMessageWriter, String?) {
     let writer: LengthPrefixedMessageWriter
     let responseEncoding: String?
@@ -534,12 +539,12 @@ extension HTTP2ToRawGRPCStateMachine.State {
         configuration.enabledAlgorithms.contains($0)
       }
 
-      writer = LengthPrefixedMessageWriter(compression: algorithm)
+      writer = LengthPrefixedMessageWriter(compression: algorithm, allocator: allocator)
       responseEncoding = algorithm?.name
 
     case .disabled:
       // The server doesn't have compression enabled.
-      writer = LengthPrefixedMessageWriter(compression: .none)
+      writer = LengthPrefixedMessageWriter(compression: .none, allocator: allocator)
       responseEncoding = nil
     }
 
@@ -642,11 +647,10 @@ extension HTTP2ToRawGRPCStateMachine {
   static func writeGRPCFramedMessage(
     _ buffer: ByteBuffer,
     compress: Bool,
-    allocator: ByteBufferAllocator,
-    writer: LengthPrefixedMessageWriter
+    writer: inout LengthPrefixedMessageWriter
   ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     do {
-      let buffers = try writer.write(buffer: buffer, allocator: allocator, compressed: compress)
+      let buffers = try writer.write(buffer: buffer, compressed: compress)
       return .success(buffers)
     } catch {
       return .failure(error)
@@ -655,31 +659,27 @@ extension HTTP2ToRawGRPCStateMachine {
 }
 
 extension HTTP2ToRawGRPCStateMachine.RequestOpenResponseOpenState {
-  func send(
+  mutating func send(
     buffer: ByteBuffer,
-    allocator: ByteBufferAllocator,
     compress: Bool
   ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     return HTTP2ToRawGRPCStateMachine.writeGRPCFramedMessage(
       buffer,
       compress: compress,
-      allocator: allocator,
-      writer: self.writer
+      writer: &self.writer
     )
   }
 }
 
 extension HTTP2ToRawGRPCStateMachine.RequestClosedResponseOpenState {
-  func send(
+  mutating func send(
     buffer: ByteBuffer,
-    allocator: ByteBufferAllocator,
     compress: Bool
   ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
     return HTTP2ToRawGRPCStateMachine.writeGRPCFramedMessage(
       buffer,
       compress: compress,
-      allocator: allocator,
-      writer: self.writer
+      writer: &self.writer
     )
   }
 }
@@ -903,12 +903,14 @@ extension HTTP2ToRawGRPCStateMachine {
   }
 
   /// Send a response buffer.
-  func send(
+  mutating func send(
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compress: Bool
   ) -> Result<(ByteBuffer, ByteBuffer?), Error> {
-    return self.state.send(buffer: buffer, allocator: allocator, compress: compress)
+    return self.withStateAvoidingCoWs { state in
+      state.send(buffer: buffer, allocator: allocator, compress: compress)
+    }
   }
 
   /// Send status and trailers.
@@ -1115,7 +1117,7 @@ extension HTTP2ToRawGRPCStateMachine.State {
     }
   }
 
-  func send(
+  mutating func send(
     buffer: ByteBuffer,
     allocator: ByteBufferAllocator,
     compress: Bool
@@ -1129,19 +1131,21 @@ extension HTTP2ToRawGRPCStateMachine.State {
       let error = GRPCError.InvalidState("Response headers must be sent before response message")
       return .failure(error)
 
-    case let .requestOpenResponseOpen(state):
-      return state.send(
+    case var .requestOpenResponseOpen(state):
+      let result = state.send(
         buffer: buffer,
-        allocator: allocator,
         compress: compress
       )
+      self = .requestOpenResponseOpen(state)
+      return result
 
-    case let .requestClosedResponseOpen(state):
-      return state.send(
+    case var .requestClosedResponseOpen(state):
+      let result = state.send(
         buffer: buffer,
-        allocator: allocator,
         compress: compress
       )
+      self = .requestClosedResponseOpen(state)
+      return result
 
     case .requestOpenResponseClosed,
          .requestClosedResponseClosed:

Sources/GRPC/LengthPrefixedMessageWriter.swift

@@ -28,8 +28,13 @@ internal struct LengthPrefixedMessageWriter {
     return self.compression != nil
   }
 
-  init(compression: CompressionAlgorithm? = nil) {
+  /// A scratch buffer that we encode messages into: if the buffer isn't held elsewhere then we
+  /// can avoid having to allocate a new one.
+  private var scratch: ByteBuffer
+
+  init(compression: CompressionAlgorithm? = nil, allocator: ByteBufferAllocator) {
     self.compression = compression
+    self.scratch = allocator.buffer(capacity: 0)
 
     switch self.compression?.algorithm {
     case .none, .some(.identity):
@@ -41,73 +46,62 @@ internal struct LengthPrefixedMessageWriter {
     }
   }
 
-  private func compress(
+  private mutating func compress(
     buffer: ByteBuffer,
-    using compressor: Zlib.Deflate,
-    allocator: ByteBufferAllocator
+    using compressor: Zlib.Deflate
   ) throws -> ByteBuffer {
     // The compressor will allocate the correct size. For now the leading 5 bytes will do.
-    var output = allocator.buffer(capacity: 5)
-
+    self.scratch.clear(minimumCapacity: 5)
     // Set the compression byte.
-    output.writeInteger(UInt8(1))
-
+    self.scratch.writeInteger(UInt8(1))
     // Set the length to zero; we'll write the actual value in a moment.
-    let payloadSizeIndex = output.writerIndex
-    output.writeInteger(UInt32(0))
+    let payloadSizeIndex = self.scratch.writerIndex
+    self.scratch.writeInteger(UInt32(0))
 
     let bytesWritten: Int
 
     do {
       var buffer = buffer
-      bytesWritten = try compressor.deflate(&buffer, into: &output)
+      bytesWritten = try compressor.deflate(&buffer, into: &self.scratch)
     } catch {
       throw error
     }
 
     // Now fill in the message length.
-    output.writePayloadLength(UInt32(bytesWritten), at: payloadSizeIndex)
+    self.scratch.writePayloadLength(UInt32(bytesWritten), at: payloadSizeIndex)
 
     // Finally, the compression context should be reset between messages.
     compressor.reset()
 
-    return output
+    return self.scratch
   }
 
   /// Writes the readable bytes of `buffer` as a gRPC length-prefixed message.
   ///
   /// - Parameters:
   ///   - buffer: The bytes to compress and length-prefix.
-  ///   - allocator: A `ByteBufferAllocator`.
   ///   - compressed: Whether the bytes should be compressed. This is ignored if not compression
   ///     mechanism was configured on this writer.
   /// - Returns: A buffer containing the length prefixed bytes.
-  func write(
+  mutating func write(
     buffer: ByteBuffer,
-    allocator: ByteBufferAllocator,
     compressed: Bool = true
   ) throws -> (ByteBuffer, ByteBuffer?) {
     if compressed, let compressor = self.compressor {
-      let compressedAndFramedPayload = try self.compress(
-        buffer: buffer,
-        using: compressor,
-        allocator: allocator
-      )
+      let compressedAndFramedPayload = try self.compress(buffer: buffer, using: compressor)
       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)
+      self.scratch.clear(minimumCapacity: 5)
+      self.scratch.writeMultipleIntegers(UInt8(0), UInt32(buffer.readableBytes))
+      return (self.scratch, buffer)
     } else {
       // 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.
-      lengthPrefixed.writeImmutableBuffer(buffer)
-      return (lengthPrefixed, nil)
+      self.scratch.clear(minimumCapacity: 5 &+ buffer.readableBytes)
+      self.scratch.writeMultipleIntegers(UInt8(0), UInt32(buffer.readableBytes))
+      self.scratch.writeImmutableBuffer(buffer)
+      return (self.scratch, nil)
     }
   }
 

Sources/GRPC/ReadWriteStates.swift

@@ -30,7 +30,10 @@ struct PendingWriteState {
   /// The 'content-type' being written.
   var contentType: ContentType
 
-  func makeWriteState(messageEncoding: ClientMessageEncoding) -> WriteState {
+  func makeWriteState(
+    messageEncoding: ClientMessageEncoding,
+    allocator: ByteBufferAllocator
+  ) -> WriteState {
     let compression: CompressionAlgorithm?
     switch messageEncoding {
     case let .enabled(configuration):
@@ -39,7 +42,7 @@ struct PendingWriteState {
       compression = nil
     }
 
-    let writer = LengthPrefixedMessageWriter(compression: compression)
+    let writer = LengthPrefixedMessageWriter(compression: compression, allocator: allocator)
     return .writing(self.arity, self.contentType, writer)
   }
 }
@@ -53,25 +56,23 @@ enum WriteState {
   /// more messages to be written.
   case notWriting
 
-  /// Writes a message into a buffer using the `writer` and `allocator`.
+  /// Writes a message into a buffer using the `writer`.
   ///
   /// - Parameter message: The `Message` to write.
-  /// - Parameter allocator: An allocator to provide a `ByteBuffer` into which the message will be
-  ///     written.
   mutating func write(
     _ message: ByteBuffer,
-    compressed: Bool,
-    allocator: ByteBufferAllocator
+    compressed: Bool
   ) -> Result<(ByteBuffer, ByteBuffer?), MessageWriteError> {
     switch self {
     case .notWriting:
       return .failure(.cardinalityViolation)
 
-    case let .writing(writeArity, contentType, writer):
+    case .writing(let writeArity, let contentType, var writer):
+      self = .notWriting
       let buffers: (ByteBuffer, ByteBuffer?)
 
       do {
-        buffers = try writer.write(buffer: message, allocator: allocator, compressed: compressed)
+        buffers = try writer.write(buffer: message, compressed: compressed)
       } catch {
         self = .notWriting
         return .failure(.serializationFailed)