From 5fc8ba1f4f527218d7a824590abd3dc9acab5c71 Mon Sep 17 00:00:00 2001
From: George Barnett <gbarnett@apple.com>
Date: Wed, 24 Jul 2019 16:00:11 +0100
Subject: [PATCH] Track the event observer and factory with state

---
 .../BidirectionalStreamingCallHandler.swift   | 28 +++++++++------
 .../ClientStreamingCallHandler.swift          | 35 +++++++++++++------
 2 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/Sources/GRPC/CallHandlers/BidirectionalStreamingCallHandler.swift b/Sources/GRPC/CallHandlers/BidirectionalStreamingCallHandler.swift
index 80beb8a69..df9820bf0 100644
--- a/Sources/GRPC/CallHandlers/BidirectionalStreamingCallHandler.swift
+++ b/Sources/GRPC/CallHandlers/BidirectionalStreamingCallHandler.swift
@@ -25,11 +25,12 @@ import NIOHTTP1
 ///   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>: BaseCallHandler<RequestMessage, ResponseMessage> {
+  public typealias Context = StreamingResponseCallContext<ResponseMessage>
   public typealias EventObserver = (StreamEvent<RequestMessage>) -> Void
-  private var eventObserver: EventLoopFuture<EventObserver>?
-  private let eventObserverFactory: (StreamingResponseCallContext<ResponseMessage>) -> EventLoopFuture<EventObserver>
+  public typealias EventObserverFactory = (Context) -> EventLoopFuture<EventObserver>
 
-  private var callContext: StreamingResponseCallContext<ResponseMessage>?
+  private var observerState: ClientStreamingHandlerObserverState<EventObserverFactory, EventObserver>
+  private var callContext: Context?
 
   // We ask for a future of type `EventObserver` to allow the framework user to e.g. asynchronously authenticate a call.
   // If authentication fails, they can simply fail the observer future, which causes the call to be terminated.
@@ -37,7 +38,7 @@ public class BidirectionalStreamingCallHandler<RequestMessage: Message, Response
     // Delay the creation of the event observer until `handlerAdded(context:)`, otherwise it is
     // possible for the service to write into the pipeline (by fulfilling the status promise
     // of the call context outside of the observer) before it has been configured.
-    self.eventObserverFactory = eventObserverFactory
+    self.observerState = .pendingCreation(eventObserverFactory)
 
     let context = StreamingResponseCallContextImpl<ResponseMessage>(channel: channel, request: request, errorDelegate: errorDelegate)
     self.callContext = context
@@ -46,15 +47,20 @@ public class BidirectionalStreamingCallHandler<RequestMessage: Message, Response
 
     context.statusPromise.futureResult.whenComplete { _ in
       // When done, reset references to avoid retain cycles.
-      self.eventObserver = nil
       self.callContext = nil
+      self.observerState = .notRequired
     }
   }
 
   public override func handlerAdded(context: ChannelHandlerContext) {
-    guard let callContext = self.callContext else { return }
-    let eventObserver = self.eventObserverFactory(callContext)
-    self.eventObserver = eventObserver
+    guard let callContext = self.callContext,
+      case let .pendingCreation(factory) = self.observerState else {
+      return
+    }
+
+    let eventObserver = factory(callContext)
+    self.observerState = .created(eventObserver)
+
     // Terminate the call if the future providing an observer fails.
     // This is being done _after_ we have been added as a handler to ensure that the `GRPCServerCodec` required to
     // translate our outgoing `GRPCServerResponsePart<ResponseMessage>` message is already present on the channel.
@@ -64,13 +70,15 @@ public class BidirectionalStreamingCallHandler<RequestMessage: Message, Response
 
 
   public override func processMessage(_ message: RequestMessage) {
-    self.eventObserver?.whenSuccess { observer in
+    guard case .created(let eventObserver) = self.observerState else { return }
+    eventObserver.whenSuccess { observer in
       observer(.message(message))
     }
   }
 
   public override func endOfStreamReceived() throws {
-    self.eventObserver?.whenSuccess { observer in
+    guard case .created(let eventObserver) = self.observerState else { return }
+    eventObserver.whenSuccess { observer in
       observer(.end)
     }
   }
diff --git a/Sources/GRPC/CallHandlers/ClientStreamingCallHandler.swift b/Sources/GRPC/CallHandlers/ClientStreamingCallHandler.swift
index 3048f2409..ae687cbb4 100644
--- a/Sources/GRPC/CallHandlers/ClientStreamingCallHandler.swift
+++ b/Sources/GRPC/CallHandlers/ClientStreamingCallHandler.swift
@@ -18,6 +18,14 @@ import SwiftProtobuf
 import NIO
 import NIOHTTP1
 
+/// For calls which support client streaming we need to delay the creation of the event observer
+/// until the handler has been added to the pipeline.
+enum ClientStreamingHandlerObserverState<Factory, Observer> {
+  case pendingCreation(Factory)
+  case created(EventLoopFuture<Observer>)
+  case notRequired
+}
+
 /// Handles client-streaming calls. Forwards incoming messages and end-of-stream events to the observer block.
 ///
 /// - The observer block is implemented by the framework user and fulfills `context.responsePromise` when done.
@@ -25,19 +33,20 @@ import NIOHTTP1
 ///   they can fail the observer block future.
 /// - To close the call and send the response, complete `context.responsePromise`.
 public class ClientStreamingCallHandler<RequestMessage: Message, ResponseMessage: Message>: BaseCallHandler<RequestMessage, ResponseMessage> {
+  public typealias Context = UnaryResponseCallContext<ResponseMessage>
   public typealias EventObserver = (StreamEvent<RequestMessage>) -> Void
-  private var eventObserver: EventLoopFuture<EventObserver>?
-  private let eventObserverFactory: (UnaryResponseCallContext<ResponseMessage>) -> EventLoopFuture<EventObserver>
+  public typealias EventObserverFactory = (Context) -> EventLoopFuture<EventObserver>
 
+  private var observerState: ClientStreamingHandlerObserverState<EventObserverFactory, EventObserver>
   private var callContext: UnaryResponseCallContext<ResponseMessage>?
 
   // We ask for a future of type `EventObserver` to allow the framework user to e.g. asynchronously authenticate a call.
   // If authentication fails, they can simply fail the observer future, which causes the call to be terminated.
-  public init(channel: Channel, request: HTTPRequestHead, errorDelegate: ServerErrorDelegate?, eventObserverFactory: @escaping (UnaryResponseCallContext<ResponseMessage>) -> EventLoopFuture<EventObserver>) {
+  public init(channel: Channel, request: HTTPRequestHead, errorDelegate: ServerErrorDelegate?, eventObserverFactory: @escaping EventObserverFactory) {
     // Delay the creation of the event observer until `handlerAdded(context:)`, otherwise it is
     // possible for the service to write into the pipeline (by fulfilling the response promise
     // of the call context outside of the observer) before it has been configured.
-    self.eventObserverFactory = eventObserverFactory
+    self.observerState = .pendingCreation(eventObserverFactory)
 
     let callContext = UnaryResponseCallContextImpl<ResponseMessage>(channel: channel, request: request, errorDelegate: errorDelegate)
     self.callContext = callContext
@@ -46,15 +55,19 @@ public class ClientStreamingCallHandler<RequestMessage: Message, ResponseMessage
 
     callContext.responsePromise.futureResult.whenComplete { _ in
       // When done, reset references to avoid retain cycles.
-      self.eventObserver = nil
       self.callContext = nil
+      self.observerState = .notRequired
     }
   }
 
   public override func handlerAdded(context: ChannelHandlerContext) {
-    guard let callContext = self.callContext else { return }
-    let eventObserver = self.eventObserverFactory(callContext)
-    self.eventObserver = eventObserver
+    guard let callContext = self.callContext,
+      case let .pendingCreation(factory) = self.observerState else {
+      return
+    }
+
+    let eventObserver = factory(callContext)
+    self.observerState = .created(eventObserver)
 
     // Terminate the call if the future providing an observer fails.
     // This is being done _after_ we have been added as a handler to ensure that the `GRPCServerCodec` required to
@@ -64,13 +77,15 @@ public class ClientStreamingCallHandler<RequestMessage: Message, ResponseMessage
   }
 
   public override func processMessage(_ message: RequestMessage) {
-    self.eventObserver?.whenSuccess { observer in
+    guard case .created(let eventObserver) = self.observerState else { return }
+    eventObserver.whenSuccess { observer in
       observer(.message(message))
     }
   }
 
   public override func endOfStreamReceived() throws {
-    self.eventObserver?.whenSuccess { observer in
+    guard case .created(let eventObserver) = self.observerState else { return }
+    eventObserver.whenSuccess { observer in
       observer(.end)
     }
   }