Rewrite client connection management

Sources/GRPC/ClientConnectivityHandler.swift

@@ -0,0 +1,146 @@
+/*
+ * Copyright 2020, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import NIO
+import NIOHTTP2
+
+internal class ClientConnectivityHandler: ChannelInboundHandler {
+  typealias InboundIn = HTTP2Frame
+
+  private var connectionManager: ConnectionManager
+  private let idleTimeout: TimeAmount
+
+  private var activeStreams = 0
+  private var scheduledIdle: Scheduled<Void>? = nil
+  private var state: State = .notReady
+
+  private enum State {
+    // We haven't marked the connection as "ready" yet.
+    case notReady
+
+    // The connection has been marked as "ready".
+    case ready
+
+    // We called `close` on the channel.
+    case closed
+  }
+
+  init(connectionManager: ConnectionManager, idleTimeout: TimeAmount = .minutes(5)) {
+    self.connectionManager = connectionManager
+    self.idleTimeout = idleTimeout
+  }
+
+  func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+    switch self.state {
+    case .notReady, .ready:
+      if event is NIOHTTP2StreamCreatedEvent {
+        // We have a stream: don't go idle
+        self.scheduledIdle?.cancel()
+        self.scheduledIdle = nil
+
+        self.activeStreams += 1
+      } else if event is StreamClosedEvent {
+        self.activeStreams -= 1
+        // No active streams: go idle soon.
+        if self.activeStreams == 0 {
+          self.scheduleIdleTimeout(context: context)
+        }
+      }
+
+    case .closed:
+      ()
+    }
+
+    context.fireUserInboundEventTriggered(event)
+  }
+
+  func channelActive(context: ChannelHandlerContext) {
+    switch self.state {
+    case .notReady:
+      self.connectionManager.channelActive(channel: context.channel)
+    case .ready, .closed:
+      ()
+    }
+
+    context.fireChannelActive()
+  }
+
+  func channelInactive(context: ChannelHandlerContext) {
+    self.scheduledIdle?.cancel()
+    self.scheduledIdle = nil
+
+    switch self.state {
+    case .notReady, .ready:
+      self.connectionManager.channelInactive()
+    case .closed:
+      ()
+    }
+
+    context.fireChannelInactive()
+  }
+
+  func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+    let frame = self.unwrapInboundIn(data)
+
+    if frame.streamID == .rootStream {
+      switch (self.state, frame.payload) {
+      // We only care about SETTINGS as long as we are in state `.notReady`.
+      case (.notReady, .settings):
+        self.state = .ready
+
+        let remoteAddressDescription = context.channel.remoteAddress.map { "\($0)" } ?? "n/a"
+        self.connectionManager.logger.info("gRPC connection ready", metadata: [
+          "remote_address": "\(remoteAddressDescription)",
+          "event_loop": "\(context.eventLoop)"
+        ])
+
+        // Start the idle timeout.
+        self.scheduleIdleTimeout(context: context)
+
+        // Let the manager know we're ready.
+        self.connectionManager.ready()
+
+      case (.notReady, .goAway),
+           (.ready, .goAway):
+        self.idle(context: context)
+
+      default:
+        ()
+      }
+    }
+
+    context.fireChannelRead(data)
+  }
+
+  private func scheduleIdleTimeout(context: ChannelHandlerContext) {
+    guard self.activeStreams == 0 else {
+      return
+    }
+
+    self.scheduledIdle = context.eventLoop.scheduleTask(in: self.idleTimeout) {
+      self.idle(context: context)
+    }
+  }
+
+  private func idle(context: ChannelHandlerContext) {
+    guard self.activeStreams == 0 else {
+      return
+    }
+
+    self.state = .closed
+    self.connectionManager.idle()
+    context.close(mode: .all, promise: nil)
+  }
+}