瀏覽代碼

Reduce allocations in the GRPCClientStateMachine (#653)

Motivation:

There were some unnecessary allocations in the client state machine.

Modifications:

- Use `HPACKHeader`s new `first(name:)` method to avoid allocating an
  array of headers before taking the first.
- Turn LengthPrefixedMessageWriter into a struct.
- Turn LengthPrefixedMessageReader into a struct.

Result:

- Reduce allocations, a small perf gain.
George Barnett 6 年之前
父節點
當前提交
3111de60d6

+ 12 - 13
Sources/GRPC/GRPCClientStateMachine.swift

@@ -544,21 +544,22 @@ extension GRPCClientStateMachine.State {
     // responses as well as a variety of non-GRPC content-types and to omit Status & Status-Message.
     // Implementations must synthesize a Status & Status-Message to propagate to the application
     // layer when this occurs."
-    let statusHeader = headers[":status"].first
-    let responseStatus = statusHeader.flatMap(Int.init).map { code in
-      HTTPResponseStatus(statusCode: code)
-    } ?? .preconditionFailed
+    let responseStatus = headers.first(name: ":status")
+      .flatMap(Int.init)
+      .map { code in
+        HTTPResponseStatus(statusCode: code)
+      } ?? .preconditionFailed
 
     guard responseStatus == .ok else {
       return .failure(.invalidHTTPStatus(responseStatus))
     }
 
-    guard headers["content-type"].first.flatMap(ContentType.init) != nil else {
+    guard headers.first(name: "content-type").flatMap(ContentType.init) != nil else {
       return .failure(.invalidContentType)
     }
 
     // What compression mechanism is the server using, if any?
-    let compression = CompressionMechanism(value: headers[GRPCHeaderName.encoding].first)
+    let compression = CompressionMechanism(value: headers.first(name: GRPCHeaderName.encoding))
 
     // From: https://github.com/grpc/grpc/blob/master/doc/compression.md
     //
@@ -586,13 +587,13 @@ extension GRPCClientStateMachine.State {
   }
 
   private func readStatusCode(from trailers: HPACKHeaders) -> GRPCStatus.Code? {
-    return trailers[GRPCHeaderName.statusCode].first
+    return trailers.first(name: GRPCHeaderName.statusCode)
       .flatMap(Int.init)
       .flatMap(GRPCStatus.Code.init)
   }
 
   private func readStatusMessage(from trailers: HPACKHeaders) -> String? {
-    return trailers[GRPCHeaderName.statusMessage].first
+    return trailers.first(name: GRPCHeaderName.statusMessage)
       .map(GRPCStatusMessageMarshaller.unmarshall)
   }
 
@@ -609,10 +610,8 @@ extension GRPCClientStateMachine.State {
     // one from the ":status".
     //
     // See: https://github.com/grpc/grpc/blob/master/doc/http-grpc-status-mapping.md
-    let statusHeader = trailers[":status"].first
-    guard let status = statusHeader.flatMap(Int.init).map({ HTTPResponseStatus(statusCode: $0) })
-      else {
-        return .failure(.invalidHTTPStatus(nil))
+    guard let status = trailers.first(name: ":status").flatMap(Int.init).map({ HTTPResponseStatus(statusCode: $0) }) else {
+      return .failure(.invalidHTTPStatus(nil))
     }
 
     guard status == .ok else {
@@ -624,7 +623,7 @@ extension GRPCClientStateMachine.State {
       }
     }
 
-    guard trailers["content-type"].first.flatMap(ContentType.init) != nil else {
+    guard trailers.first(name: "content-type").flatMap(ContentType.init) != nil else {
       return .failure(.invalidContentType)
     }
 

+ 5 - 5
Sources/GRPC/LengthPrefixedMessageReader.swift

@@ -30,7 +30,7 @@ import Logging
 ///
 /// - SeeAlso:
 /// [gRPC Protocol](https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md)
-public class LengthPrefixedMessageReader {
+public struct LengthPrefixedMessageReader {
   public typealias Mode = GRPCError.Origin
 
   /// The mechanism that messages will be compressed with.
@@ -79,7 +79,7 @@ public class LengthPrefixedMessageReader {
   }
 
   /// Appends data to the buffer from which messages will be read.
-  public func append(buffer: inout ByteBuffer) {
+  public mutating func append(buffer: inout ByteBuffer) {
     guard buffer.readableBytes > 0 else {
       return
     }
@@ -98,7 +98,7 @@ public class LengthPrefixedMessageReader {
   /// - Returns: A buffer containing a message if one has been read, or `nil` if not enough
   ///   bytes have been consumed to return a message.
   /// - Throws: Throws an error if the compression algorithm is not supported.
-  public func nextMessage() throws -> ByteBuffer? {
+  public mutating func nextMessage() throws -> ByteBuffer? {
     switch try self.processNextState() {
     case .needMoreData:
       self.nilBufferIfPossible()
@@ -116,13 +116,13 @@ public class LengthPrefixedMessageReader {
   /// `nil`s out `buffer` if it exists and has no readable bytes.
   ///
   /// This allows the next call to `append` to avoid writing the contents of the appended buffer.
-  private func nilBufferIfPossible() {
+  private mutating func nilBufferIfPossible() {
     if self.buffer?.readableBytes == 0 {
       self.buffer = nil
     }
   }
 
-  private func processNextState() throws -> ParseResult {
+  private mutating func processNextState() throws -> ParseResult {
     guard self.buffer != nil else {
       return .needMoreData
     }

+ 1 - 1
Sources/GRPC/LengthPrefixedMessageWriter.swift

@@ -16,7 +16,7 @@
 import Foundation
 import NIO
 
-public class LengthPrefixedMessageWriter {
+public struct LengthPrefixedMessageWriter {
   public static let metadataLength = 5
 
   private let compression: CompressionMechanism

+ 1 - 1
Sources/GRPC/ReadWriteStates.swift

@@ -121,7 +121,7 @@ enum ReadState {
     case .notReading:
       return .failure(.cardinalityViolation)
 
-    case let .reading(readArity, reader):
+    case .reading(let readArity, var reader):
       reader.append(buffer: &buffer)
       var messages: [MessageType] = []