Adds direct access to buffer by using a payload protocol

Author: mustii

Sources/Examples/Echo/Model/echo.grpc.swift

@@ -149,3 +149,12 @@ extension Echo_EchoProvider {
   }
 }
 
+extension Echo_EchoRequest: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}
+
+extension Echo_EchoResponse: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}

Sources/Examples/HelloWorld/Model/helloworld.grpc.swift

@@ -84,3 +84,31 @@ extension Helloworld_GreeterProvider {
   }
 }
 
+extension Helloworld_HelloRequest: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {
+        try self.init(serializedData: serializedByteBuffer.readData(length: serializedByteBuffer.readableBytes)!,
+                  extensions: nil,
+                  partial: false,
+                  options: BinaryDecodingOptions())
+    }
+    
+    public func serialize(into buffer: inout ByteBuffer) throws {
+        let data = try serializedData()
+        buffer.writeBytes(data)
+    }
+}
+
+extension Helloworld_HelloReply: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {
+        try self.init(serializedData: serializedByteBuffer.readData(length: serializedByteBuffer.readableBytes)!,
+                  extensions: nil,
+                  partial: false,
+                  options: BinaryDecodingOptions())
+    }
+    
+    public func serialize(into buffer: inout ByteBuffer) throws {
+        let data = try serializedData()
+        buffer.writeBytes(data)
+    }
+}
+

Sources/Examples/RouteGuide/Model/route_guide.grpc.swift

@@ -149,3 +149,28 @@ extension Routeguide_RouteGuideProvider {
   }
 }
 
+extension Routeguide_RouteNote: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}
+
+extension Routeguide_Feature: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}
+
+extension Routeguide_RouteSummary: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}
+
+extension Routeguide_Point: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}
+
+extension Routeguide_Rectangle: Payload {
+    public init(serializedByteBuffer: inout ByteBuffer) throws {}
+    public func serialize(into buffer: inout ByteBuffer) throws {}
+}
+

Sources/GRPC/CallHandlers/BidirectionalStreamingCallHandler.swift

@@ -26,8 +26,8 @@ import Logging
 ///   they can fail the observer block future.
 /// - To close the call and send the status, complete `context.statusPromise`.
 public class BidirectionalStreamingCallHandler<
-  RequestMessage: Message,
-  ResponseMessage: Message
+  RequestMessage: Payload,
+  ResponseMessage: Payload
 >: _BaseCallHandler<RequestMessage, ResponseMessage> {
   public typealias Context = StreamingResponseCallContext<ResponseMessage>
   public typealias EventObserver = (StreamEvent<RequestMessage>) -> Void

Sources/GRPC/CallHandlers/ClientStreamingCallHandler.swift

@@ -34,8 +34,8 @@ enum ClientStreamingHandlerObserverState<Factory, Observer> {
 ///   they can fail the observer block future.
 /// - To close the call and send the response, complete `context.responsePromise`.
 public final class ClientStreamingCallHandler<
-  RequestMessage: Message,
-  ResponseMessage: Message
+  RequestMessage: Payload,
+  ResponseMessage: Payload
 >: _BaseCallHandler<RequestMessage, ResponseMessage> {
   public typealias Context = UnaryResponseCallContext<ResponseMessage>
   public typealias EventObserver = (StreamEvent<RequestMessage>) -> Void

Sources/GRPC/CallHandlers/ServerStreamingCallHandler.swift

@@ -24,8 +24,8 @@ import Logging
 /// - The observer block is implemented by the framework user and calls `context.sendResponse` as needed.
 /// - To close the call and send the status, complete the status future returned by the observer block.
 public final class ServerStreamingCallHandler<
-  RequestMessage: Message,
-  ResponseMessage: Message
+  RequestMessage: Payload,
+  ResponseMessage: Payload
 >: _BaseCallHandler<RequestMessage, ResponseMessage> {
   public typealias EventObserver = (RequestMessage) -> EventLoopFuture<GRPCStatus>
 

Sources/GRPC/CallHandlers/UnaryCallHandler.swift

@@ -25,8 +25,8 @@ import Logging
 /// - To return a response to the client, the framework user should complete that future
 ///   (similar to e.g. serving regular HTTP requests in frameworks such as Vapor).
 public final class UnaryCallHandler<
-  RequestMessage: Message,
-  ResponseMessage: Message
+  RequestMessage: Payload,
+  ResponseMessage: Payload
 >: _BaseCallHandler<RequestMessage, ResponseMessage> {
   public typealias EventObserver = (RequestMessage) -> EventLoopFuture<ResponseMessage>
   private var eventObserver: EventObserver?

Sources/GRPC/CallHandlers/_BaseCallHandler.swift

@@ -23,7 +23,7 @@ import Logging
 ///
 /// Calls through to `processMessage` for individual messages it receives, which needs to be implemented by subclasses.
 /// - Important: This is **NOT** part of the public API.
-public class _BaseCallHandler<RequestMessage: Message, ResponseMessage: Message>: GRPCCallHandler {
+public class _BaseCallHandler<RequestMessage: Payload, ResponseMessage: Payload>: GRPCCallHandler {
   public func makeGRPCServerCodec() -> ChannelHandler {
     return GRPCServerCodec<RequestMessage, ResponseMessage>()
   }

Sources/GRPC/ClientCalls/BaseClientCall.swift

@@ -48,7 +48,7 @@ import Logging
 ///
 /// This class also provides much of the framework user facing functionality via conformance to
 /// `ClientCall`.
-public class BaseClientCall<Request: Message, Response: Message>: ClientCall {
+public class BaseClientCall<Request: Payload, Response: Payload>: ClientCall {
   public typealias RequestMessage = Request
   public typealias ResponseMessage = Response
 

Sources/GRPC/ClientCalls/BidirectionalStreamingCall.swift

@@ -27,7 +27,7 @@ import Logging
 /// - `initialMetadata`: the initial metadata returned from the server,
 /// - `status`: the status of the gRPC call after it has ended,
 /// - `trailingMetadata`: any metadata returned from the server alongside the `status`.
-public final class BidirectionalStreamingCall<RequestMessage: Message, ResponseMessage: Message>
+public final class BidirectionalStreamingCall<RequestMessage: Payload, ResponseMessage: Payload>
   : BaseClientCall<RequestMessage, ResponseMessage>,
     StreamingRequestClientCall {
   private var messageQueue: EventLoopFuture<Void>

Sources/GRPC/ClientCalls/ClientCall.swift

@@ -23,9 +23,9 @@ import SwiftProtobuf
 /// Base protocol for a client call to a gRPC service.
 public protocol ClientCall {
   /// The type of the request message for the call.
-  associatedtype RequestMessage: Message
+  associatedtype RequestMessage: Payload
   /// The type of the response message for the call.
-  associatedtype ResponseMessage: Message
+  associatedtype ResponseMessage: Payload
 
   /// HTTP/2 stream that requests and responses are sent and received on.
   var subchannel: EventLoopFuture<Channel> { get }

Sources/GRPC/ClientCalls/ClientStreamingCall.swift

@@ -28,7 +28,7 @@ import Logging
 /// - `response`: the response from the call,
 /// - `status`: the status of the gRPC call after it has ended,
 /// - `trailingMetadata`: any metadata returned from the server alongside the `status`.
-public final class ClientStreamingCall<RequestMessage: Message, ResponseMessage: Message>
+public final class ClientStreamingCall<RequestMessage: Payload, ResponseMessage: Payload>
   : BaseClientCall<RequestMessage, ResponseMessage>,
     StreamingRequestClientCall,
     UnaryResponseClientCall {

Sources/GRPC/ClientCalls/ServerStreamingCall.swift

@@ -24,7 +24,7 @@ import Logging
 /// - `initialMetadata`: the initial metadata returned from the server,
 /// - `status`: the status of the gRPC call after it has ended,
 /// - `trailingMetadata`: any metadata returned from the server alongside the `status`.
-public final class ServerStreamingCall<RequestMessage: Message, ResponseMessage: Message>: BaseClientCall<RequestMessage, ResponseMessage> {
+public final class ServerStreamingCall<RequestMessage: Payload, ResponseMessage: Payload>: BaseClientCall<RequestMessage, ResponseMessage> {
   public init(
     connection: ClientConnection,
     path: String,

Sources/GRPC/ClientCalls/UnaryCall.swift

@@ -27,7 +27,7 @@ import Logging
 /// - `response`: the response from the unary call,
 /// - `status`: the status of the gRPC call after it has ended,
 /// - `trailingMetadata`: any metadata returned from the server alongside the `status`.
-public final class UnaryCall<RequestMessage: Message, ResponseMessage: Message>
+public final class UnaryCall<RequestMessage: Payload, ResponseMessage: Payload>
   : BaseClientCall<RequestMessage, ResponseMessage>,
     UnaryResponseClientCall {
   public let response: EventLoopFuture<ResponseMessage>

Sources/GRPC/GRPCClientResponseChannelHandler.swift

@@ -24,7 +24,7 @@ import Logging
 /// This includes holding promises for the initial metadata and status of the gRPC call. This handler
 /// is also responsible for error handling, via an error delegate and by appropriately failing the
 /// aforementioned promises.
-internal class GRPCClientResponseChannelHandler<ResponseMessage: Message>: ChannelInboundHandler {
+internal class GRPCClientResponseChannelHandler<ResponseMessage: Payload>: ChannelInboundHandler {
   public typealias InboundIn = _GRPCClientResponsePart<ResponseMessage>
   internal let logger: Logger
   internal var stopwatch: Stopwatch?
@@ -186,7 +186,7 @@ internal class GRPCClientResponseChannelHandler<ResponseMessage: Message>: Chann
 }
 
 /// A channel handler for client calls which receive a single response.
-final class GRPCClientUnaryResponseChannelHandler<ResponseMessage: Message>: GRPCClientResponseChannelHandler<ResponseMessage> {
+final class GRPCClientUnaryResponseChannelHandler<ResponseMessage: Payload>: GRPCClientResponseChannelHandler<ResponseMessage> {
   let responsePromise: EventLoopPromise<ResponseMessage>
 
   internal init(
@@ -236,7 +236,7 @@ final class GRPCClientUnaryResponseChannelHandler<ResponseMessage: Message>: GRP
 }
 
 /// A channel handler for client calls which receive a stream of responses.
-final class GRPCClientStreamingResponseChannelHandler<ResponseMessage: Message>: GRPCClientResponseChannelHandler<ResponseMessage> {
+final class GRPCClientStreamingResponseChannelHandler<ResponseMessage: Payload>: GRPCClientResponseChannelHandler<ResponseMessage> {
   typealias ResponseHandler = (ResponseMessage) -> Void
 
   let responseHandler: ResponseHandler

Sources/GRPC/GRPCClientStateMachine.swift

@@ -66,7 +66,7 @@ enum SendEndOfRequestStreamError: Error {
 /// A state machine for a single gRPC call from the perspective of a client.
 ///
 /// See: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md
-struct GRPCClientStateMachine<Request: Message, Response: Message> {
+struct GRPCClientStateMachine<Request: Payload, Response: Payload> {
   /// The combined state of the request (client) and response (server) streams for an RPC call.
   ///
   /// The following states are not possible:

Sources/GRPC/GRPCServerCodec.swift

@@ -22,7 +22,7 @@ import NIOHTTP1
 /// Incoming gRPC package with a fixed message type.
 ///
 /// - Important: This is **NOT** part of the public API.
-public enum _GRPCServerRequestPart<RequestMessage: Message> {
+public enum _GRPCServerRequestPart<RequestMessage: Payload> {
   case head(HTTPRequestHead)
   case message(RequestMessage)
   case end
@@ -31,14 +31,14 @@ public enum _GRPCServerRequestPart<RequestMessage: Message> {
 /// Outgoing gRPC package with a fixed message type.
 ///
 /// - Important: This is **NOT** part of the public API.
-public enum _GRPCServerResponsePart<ResponseMessage: Message> {
+public enum _GRPCServerResponsePart<ResponseMessage: Payload> {
   case headers(HTTPHeaders)
   case message(ResponseMessage)
   case statusAndTrailers(GRPCStatus, HTTPHeaders)
 }
 
 /// A simple channel handler that translates raw gRPC packets into decoded protobuf messages, and vice versa.
-internal final class GRPCServerCodec<RequestMessage: Message, ResponseMessage: Message> {}
+internal final class GRPCServerCodec<RequestMessage: Payload, ResponseMessage: Payload> {}
 
 extension GRPCServerCodec: ChannelInboundHandler {
   typealias InboundIn = _RawGRPCServerRequestPart
@@ -50,9 +50,8 @@ extension GRPCServerCodec: ChannelInboundHandler {
       context.fireChannelRead(self.wrapInboundOut(.head(requestHead)))
 
     case .message(var message):
-      let messageAsData = message.readData(length: message.readableBytes)!
       do {
-        context.fireChannelRead(self.wrapInboundOut(.message(try RequestMessage(serializedData: messageAsData))))
+        context.fireChannelRead(self.wrapInboundOut(.message(try RequestMessage(serializedByteBuffer: &message))))
       } catch {
         context.fireErrorCaught(GRPCError.DeserializationFailure().captureContext())
       }
@@ -75,8 +74,9 @@ extension GRPCServerCodec: ChannelOutboundHandler {
 
     case .message(let message):
       do {
-        let messageData = try message.serializedData()
-        context.write(self.wrapOutboundOut(.message(messageData)), promise: promise)
+        var buffer = ByteBufferAllocator().buffer(capacity: 0)
+        try message.serialize(into: &buffer)
+        context.write(self.wrapOutboundOut(.message(buffer.readData(length: buffer.readableBytes)!)), promise: promise)
       } catch {
         let error = GRPCError.SerializationFailure().captureContext()
         promise?.fail(error)

Sources/GRPC/Payload.swift

@@ -0,0 +1,7 @@
+import NIO
+import SwiftProtobuf
+
+public protocol Payload {
+    init(serializedByteBuffer: inout NIO.ByteBuffer) throws
+    func serialize(into buffer: inout NIO.ByteBuffer) throws
+}

Sources/GRPC/ReadWriteStates.swift

@@ -57,23 +57,25 @@ enum WriteState {
   /// - Parameter allocator: An allocator to provide a `ByteBuffer` into which the message will be
   ///     written.
   mutating func write(
-    _ message: Message,
+    _ message: Payload,
     allocator: ByteBufferAllocator
   ) -> Result<ByteBuffer, MessageWriteError> {
     switch self {
     case .notWriting:
       return .failure(.cardinalityViolation)
 
     case let .writing(writeArity, contentType, writer):
-      guard let data = try? message.serializedData() else {
+      // Zero is fine: the writer will allocate the correct amount of space.
+      var buffer = allocator.buffer(capacity: 0)
+      
+      do {
+        try message.serialize(into: &buffer)
+      } catch {
         self = .notWriting
         return .failure(.serializationFailed)
       }
-
-      // Zero is fine: the writer will allocate the correct amount of space.
-      var buffer = allocator.buffer(capacity: 0)
-      writer.write(data, into: &buffer)
-
+      var buf = allocator.buffer(capacity: 0)
+      writer.write(buffer.readData(length: buffer.readableBytes)!, into: &buf)
       // If we only expect to write one message then we're no longer writable.
       if case .one = writeArity {
         self = .notWriting
@@ -113,7 +115,7 @@ enum ReadState {
   /// a message has been produced then subsequent calls will result in an error.
   ///
   /// - Parameter buffer: The buffer to read from.
-  mutating func readMessages<MessageType: Message>(
+  mutating func readMessages<MessageType: Payload>(
     _ buffer: inout ByteBuffer,
     as: MessageType.Type = MessageType.self
   ) -> Result<[MessageType], MessageReadError> {
@@ -128,8 +130,7 @@ enum ReadState {
       do {
         while var serializedBytes = try? reader.nextMessage() {
           // Force unwrapping is okay here: we will always be able to read `readableBytes`.
-          let serializedData = serializedBytes.readData(length: serializedBytes.readableBytes)!
-          messages.append(try MessageType(serializedData: serializedData))
+          messages.append(try MessageType(serializedByteBuffer: &serializedBytes))
         }
       } catch {
         self = .notReading

Sources/GRPC/ServerCallContexts/StreamingResponseCallContext.swift

@@ -25,7 +25,7 @@ import Logging
 /// - If `statusPromise` is failed and the error is of type `GRPCStatusTransformable`,
 ///   the result of `error.asGRPCStatus()` will be returned to the client.
 /// - If `error.asGRPCStatus()` is not available, `GRPCStatus.processingError` is returned to the client.
-open class StreamingResponseCallContext<ResponseMessage: Message>: ServerCallContextBase {
+open class StreamingResponseCallContext<ResponseMessage: Payload>: ServerCallContextBase {
   typealias WrappedResponse = _GRPCServerResponsePart<ResponseMessage>
 
   public let statusPromise: EventLoopPromise<GRPCStatus>
@@ -41,7 +41,7 @@ open class StreamingResponseCallContext<ResponseMessage: Message>: ServerCallCon
 }
 
 /// Concrete implementation of `StreamingResponseCallContext` used by our generated code.
-open class StreamingResponseCallContextImpl<ResponseMessage: Message>: StreamingResponseCallContext<ResponseMessage> {
+open class StreamingResponseCallContextImpl<ResponseMessage: Payload>: StreamingResponseCallContext<ResponseMessage> {
   public let channel: Channel
 
   /// - Parameters:
@@ -80,7 +80,7 @@ open class StreamingResponseCallContextImpl<ResponseMessage: Message>: Streaming
 /// Concrete implementation of `StreamingResponseCallContext` used for testing.
 ///
 /// Simply records all sent messages.
-open class StreamingResponseCallContextTestStub<ResponseMessage: Message>: StreamingResponseCallContext<ResponseMessage> {
+open class StreamingResponseCallContextTestStub<ResponseMessage: Payload>: StreamingResponseCallContext<ResponseMessage> {
   open var recordedResponses: [ResponseMessage] = []
 
   open override func sendResponse(_ message: ResponseMessage) -> EventLoopFuture<Void> {

Sources/GRPC/ServerCallContexts/UnaryResponseCallContext.swift

@@ -28,7 +28,7 @@ import Logging
 ///
 /// For unary calls, the response is not actually provided by fulfilling `responsePromise`, but instead by completing
 /// the future returned by `UnaryCallHandler.EventObserver`.
-open class UnaryResponseCallContext<ResponseMessage: Message>: ServerCallContextBase, StatusOnlyCallContext {
+open class UnaryResponseCallContext<ResponseMessage: Payload>: ServerCallContextBase, StatusOnlyCallContext {
   typealias WrappedResponse = _GRPCServerResponsePart<ResponseMessage>
 
   public let responsePromise: EventLoopPromise<ResponseMessage>
@@ -55,7 +55,7 @@ public protocol StatusOnlyCallContext: ServerCallContext {
 }
 
 /// Concrete implementation of `UnaryResponseCallContext` used by our generated code.
-open class UnaryResponseCallContextImpl<ResponseMessage: Message>: UnaryResponseCallContext<ResponseMessage> {
+open class UnaryResponseCallContextImpl<ResponseMessage: Payload>: UnaryResponseCallContext<ResponseMessage> {
   public let channel: Channel
 
   /// - Parameters:
@@ -93,4 +93,4 @@ open class UnaryResponseCallContextImpl<ResponseMessage: Message>: UnaryResponse
 /// Concrete implementation of `UnaryResponseCallContext` used for testing.
 ///
 /// Only provided to make it clear in tests that no "real" implementation is used.
-open class UnaryResponseCallContextTestStub<ResponseMessage: Message>: UnaryResponseCallContext<ResponseMessage> { }
+open class UnaryResponseCallContextTestStub<ResponseMessage: Payload>: UnaryResponseCallContext<ResponseMessage> { }

Sources/GRPC/StreamEvent.swift

@@ -16,7 +16,7 @@
 import SwiftProtobuf
 
 /// An event that can occur on a client-streaming RPC. Provided to the event observer registered for that call.
-public enum StreamEvent<Message: SwiftProtobuf.Message> {
+public enum StreamEvent<Message: Payload> {
   case message(Message)
   case end
   //! FIXME: Also support errors in this type, to propagate them to the event handler.