Move the logic for reading the payload to GRPCServerCodec, and fixed the logic

Author: mustii

Sources/GRPC/GRPCServerCodec.swift

@@ -41,6 +41,7 @@ public enum _GRPCServerResponsePart<ResponsePayload: GRPCPayload> {
 internal final class GRPCServerCodec<RequestPayload: GRPCPayload, ResponsePayload: GRPCPayload> {
   // 1-byte for compression flag, 4-bytes for message length.
   private let protobufMetadataSize = 5
+  var messageWriter = LengthPrefixedMessageWriter(compression: .none)
 }
 
 extension GRPCServerCodec: ChannelInboundHandler {
@@ -77,9 +78,9 @@ extension GRPCServerCodec: ChannelOutboundHandler {
 
     case .message(let message):
       do {
-        var buffer = ByteBufferAllocator().buffer(capacity: 0)
-        try message.serialize(into: &buffer)
-        context.write(self.wrapOutboundOut(.message(buffer)), promise: promise)
+        var responseBuffer = context.channel.allocator.buffer(capacity: protobufMetadataSize)
+        try messageWriter.write(message, into: &responseBuffer)
+        context.write(self.wrapOutboundOut(.message(responseBuffer)), promise: promise)
       } catch {
         let error = GRPCError.SerializationFailure().captureContext()
         promise?.fail(error)

Sources/GRPC/HTTP1ToRawGRPCServerCodec.swift

@@ -58,6 +58,9 @@ public final class HTTP1ToRawGRPCServerCodec {
   private let logger: Logger
   private let accessLog: Logger
   private var stopwatch: Stopwatch?
+  
+  // 1-byte for compression flag, 4-bytes for message length.
+  private let protobufMetadataSize = 5
 
   // The following buffers use force unwrapping explicitly. With optionals, developers
   // are encouraged to unwrap them using guard-else statements. These don't work cleanly
@@ -88,7 +91,6 @@ public final class HTTP1ToRawGRPCServerCodec {
     }
   }
 
-  var messageWriter = LengthPrefixedMessageWriter(compression: .none)
   var messageReader: LengthPrefixedMessageReader
 }
 
@@ -257,27 +259,29 @@ extension HTTP1ToRawGRPCServerCodec: ChannelOutboundHandler {
         self.logger.error("invalid state '\(self.outboundState)' while writing message", metadata: ["message": "\(messageBytes)"])
         return
       }
-      var messageBytes = messageBytes
+      
       if contentType == .text {
         precondition(self.responseTextBuffer != nil)
 
         // Store the response into an independent buffer. We can't return the message directly as
         // it needs to be aggregated with all the responses plus the trailers, in order to have
         // the base64 response properly encoded in a single byte stream.
-        //
-        // We can try! here because the server does not support compression at the moment and this
-        // only throws when compression fails.
-        try! messageWriter.write(&messageBytes, into: &self.responseTextBuffer)
-
+        
+        var messageBytes = messageBytes
+        // disregards the first 5 bytes since the rewritting them to the buffer doesnt require them
+        messageBytes.moveReaderIndex(forwardBy: protobufMetadataSize)
+        // Writes the compression bit as 'Zero' since the server cant compress data
+        self.responseTextBuffer.writeInteger(UInt8(0))
+        
+        // copying the old data to the responseTextBuffer
+        self.responseTextBuffer.writeInteger(UInt32(messageBytes.readableBytes))
+        self.responseTextBuffer.writeBytes(messageBytes.readableBytesView)
+        
         // Since we stored the written data, mark the write promise as successful so that the
         // ServerStreaming provider continues sending the data.
         promise?.succeed(())
       } else {
-        var responseBuffer = context.channel.allocator.buffer(capacity: LengthPrefixedMessageWriter.metadataLength)
-        // We can try! here because the server does not support compression at the moment and this
-        // only throws when compression fails.
-        try! messageWriter.write(&messageBytes, into: &responseBuffer)
-        context.write(self.wrapOutboundOut(.body(.byteBuffer(responseBuffer))), promise: promise)
+        context.write(self.wrapOutboundOut(.body(.byteBuffer(messageBytes))), promise: promise)
       }
       self.outboundState = .expectingBodyOrStatus
 

Sources/GRPC/LengthPrefixedMessageWriter.swift

@@ -49,39 +49,53 @@ internal struct LengthPrefixedMessageWriter {
   /// - 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: inout ByteBuffer, into buffer: inout ByteBuffer, disableCompression: Bool = false) throws {
-    buffer.reserveCapacity(LengthPrefixedMessageWriter.metadataLength + payload.readableBytes)
-
+  func write(_ payload: GRPCPayload, into buffer: inout ByteBuffer, disableCompression: Bool = false) throws {
+    
+    buffer.reserveCapacity(LengthPrefixedMessageWriter.metadataLength)
+    
     if !disableCompression, 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)
 
+      // Build the message by serializing it into a Buffer first
+      // Since deflate can only take a ByteBuffer
+      var messageBuf = ByteBufferAllocator().buffer(capacity: 0)
+      try payload.serialize(into: &messageBuf)
+      
       // Compress the message.
-      var payload = payload
-      let bytesWritten = try compressor.deflate(&payload, into: &buffer)
+      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 {
+      let startBufferIndex = buffer.readableBytes
       // 'identity' compression has no compressor but should still set the compression bit set
       // unless we explicitly disable compression.
       if self.compression?.algorithm == .identity && !disableCompression {
         buffer.writeInteger(UInt8(1))
       } else {
         buffer.writeInteger(UInt8(0))
       }
+      
+      // Leave a gap for the length, we'll set it in a moment.
+      let payloadSizeIndex = buffer.writerIndex
+      buffer.writeInteger(UInt32(0))
+
+      // Writes the payload into the buffer
+      try payload.serialize(into: &buffer)
+      
+      // Calculates the Written bytes with respect to the prefixed ones
+      let writtenFrame = buffer.readableBytes - LengthPrefixedMessageWriter.metadataLength - startBufferIndex
 
       // Write the message length.
-      buffer.writeInteger(UInt32(payload.readableBytes))
-      // And the message bytes.
-      buffer.writeBytes(payload.readableBytesView)
+      buffer.writePayloadLength(UInt32(writtenFrame), at: payloadSizeIndex)
     }
   }
 }

Sources/GRPC/ReadWriteStates.swift

@@ -63,19 +63,12 @@ enum WriteState {
       return .failure(.cardinalityViolation)
 
     case let .writing(writeArity, contentType, writer):
-      // Zero is fine: the writer will allocate the correct amount of space.
-      var messageBuffer = allocator.buffer(capacity: 0)
-      do {
-        try message.serialize(into: &messageBuffer)
-      } catch {
-        self = .notWriting
-        return .failure(.serializationFailed)
-      }
       // Zero is fine: the writer will allocate the correct amount of space.
       var buffer = allocator.buffer(capacity: 0)
       do {
-        try writer.write(&messageBuffer, into: &buffer, disableCompression: disableCompression)
+        try writer.write(message, into: &buffer)
       } catch {
+        self = .notWriting
         return .failure(.serializationFailed)
       }
 

Tests/GRPCTests/GRPCClientStateMachineTests.swift

@@ -52,10 +52,8 @@ class GRPCClientStateMachineTests: GRPCTestCase {
 
   /// Writes a message into the given `buffer`.
   func writeMessage<T: GRPCPayload>(_ message: T, into buffer: inout ByteBuffer) throws {
-    var messageBuffer = allocator.buffer(capacity: 0)
-    try message.serialize(into: &messageBuffer)
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    try writer.write(&messageBuffer, into: &buffer)
+    try writer.write(message, into: &buffer)
   }
 
   /// Returns a minimally valid `HPACKHeaders` for a response.
@@ -876,10 +874,8 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer:
     let message = Echo_EchoRequest.with { $0.text = "Hello!" }
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var messageBuffer = self.allocator.buffer(capacity: 0)
-    try message.serialize(into: &messageBuffer)
     var buffer = self.allocator.buffer(capacity: 0)
-    try writer.write(&messageBuffer, into: &buffer)
+    try writer.write(message, into: &buffer)
     // And some extra junk bytes:
     let bytes: [UInt8] = [0x00]
     buffer.writeBytes(bytes)
@@ -896,11 +892,9 @@ class ReadStateTests: GRPCTestCase {
     let message = Echo_EchoRequest.with { $0.text = "Hello!" }
     let writer = LengthPrefixedMessageWriter(compression: .none)
     var buffer = self.allocator.buffer(capacity: 0)
-    var messageBuffer = self.allocator.buffer(capacity: 0)
-    try message.serialize(into: &messageBuffer)
-    try writer.write(&messageBuffer, into: &buffer)
-    try writer.write(&messageBuffer, into: &buffer)
-
+    try writer.write(message, into: &buffer)
+    try writer.write(message, into: &buffer)
+    
     var state: ReadState = .one()
     state.readMessages(&buffer, as: Echo_EchoRequest.self).assertFailure {
       XCTAssertEqual($0, .cardinalityViolation)
@@ -912,10 +906,8 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer twice:
     let message = Echo_EchoRequest.with { $0.text = "Hello!" }
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var messageBuffer = self.allocator.buffer(capacity: 0)
-    try message.serialize(into: &messageBuffer)
     var buffer = self.allocator.buffer(capacity: 0)
-    try writer.write(&messageBuffer, into: &buffer)
+    try writer.write(message, into: &buffer)
 
     var state: ReadState = .one()
     state.readMessages(&buffer, as: Echo_EchoRequest.self).assertSuccess {
@@ -930,10 +922,8 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer twice:
     let message = Echo_EchoRequest.with { $0.text = "Hello!" }
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var messageBuffer = self.allocator.buffer(capacity: 0)
-    try message.serialize(into: &messageBuffer)
     var buffer = self.allocator.buffer(capacity: 0)
-    try writer.write(&messageBuffer, into: &buffer)
+    try writer.write(message, into: &buffer)
 
     var state: ReadState = .many()
     state.readMessages(&buffer, as: Echo_EchoRequest.self).assertSuccess {
@@ -948,12 +938,10 @@ class ReadStateTests: GRPCTestCase {
     // Write a message into the buffer twice:
     let message = Echo_EchoRequest.with { $0.text = "Hello!" }
     let writer = LengthPrefixedMessageWriter(compression: .none)
-    var messageBuffer = self.allocator.buffer(capacity: 0)
-    try message.serialize(into: &messageBuffer)
     var buffer = self.allocator.buffer(capacity: 0)
-    try writer.write(&messageBuffer, into: &buffer)
-    try writer.write(&messageBuffer, into: &buffer)
-    try writer.write(&messageBuffer, into: &buffer)
+    try writer.write(message, into: &buffer)
+    try writer.write(message, into: &buffer)
+    try writer.write(message, into: &buffer)
 
     var state: ReadState = .many()
     state.readMessages(&buffer, as: Echo_EchoRequest.self).assertSuccess {