Use NIO scheduleCallback API for connection handler timers (#28)

## Motivation

NIO 2.75.0 added new APIs for scheduling callbacks, which can be used to
implement timers with fewer allocations than the previous APIs[^0].

## Modifications

Replace the timers used in the client and server connection handlers
with a new implementation that uses these new NIO APIs.

## Result

This should have no functional impact. However it reduces the
allocations in the connection handler. How many allocations it reduces
will depend on the exact scenario, but running the `echo` server under
Instruments and sending 1000 requests suggests around 10%.

[^0]: https://github.com/apple/swift-nio/releases/tag/2.75.0

---------

Co-authored-by: George Barnett <[email protected]>

Package.swift

@@ -39,7 +39,7 @@ let dependencies: [Package.Dependency] = [
   ),
   .package(
     url: "https://github.com/apple/swift-nio.git",
-    from: "2.65.0"
+    from: "2.75.0"
   ),
   .package(
     url: "https://github.com/apple/swift-nio-http2.git",

Sources/GRPCNIOTransportCore/Client/Connection/ClientConnectionHandler.swift

@@ -63,17 +63,14 @@ package final class ClientConnectionHandler: ChannelInboundHandler, ChannelOutbo
   /// The `EventLoop` of the `Channel` this handler exists in.
   private let eventLoop: any EventLoop
 
-  /// The maximum amount of time the connection may be idle for. If the connection remains idle
-  /// (i.e. has no open streams) for this period of time then the connection will be gracefully
-  /// closed.
-  private var maxIdleTimer: Timer?
+  /// The timer used to gracefully close idle connections.
+  private var maxIdleTimerHandler: Timer<MaxIdleTimerHandlerView>?
 
-  /// The amount of time to wait before sending a keep alive ping.
-  private var keepaliveTimer: Timer?
+  /// The timer used to send keep-alive pings.
+  private var keepaliveTimerHandler: Timer<KeepaliveTimerHandlerView>?
 
-  /// The amount of time the client has to reply after sending a keep alive ping. Only used if
-  /// `keepaliveTimer` is set.
-  private var keepaliveTimeoutTimer: Timer
+  /// The timer used to detect keep alive timeouts, if keep-alive pings are enabled.
+  private var keepaliveTimeoutHandler: Timer<KeepaliveTimeoutHandlerView>?
 
   /// Opaque data sent in keep alive pings.
   private let keepalivePingData: HTTP2PingData
@@ -110,14 +107,34 @@ package final class ClientConnectionHandler: ChannelInboundHandler, ChannelOutbo
     keepaliveWithoutCalls: Bool
   ) {
     self.eventLoop = eventLoop
-    self.maxIdleTimer = maxIdleTime.map { Timer(delay: $0) }
-    self.keepaliveTimer = keepaliveTime.map { Timer(delay: $0, repeat: true) }
-    self.keepaliveTimeoutTimer = Timer(delay: keepaliveTimeout ?? .seconds(20))
     self.keepalivePingData = HTTP2PingData(withInteger: .random(in: .min ... .max))
     self.state = StateMachine(allowKeepaliveWithoutCalls: keepaliveWithoutCalls)
 
     self.flushPending = false
     self.inReadLoop = false
+    if let maxIdleTime {
+      self.maxIdleTimerHandler = Timer(
+        eventLoop: eventLoop,
+        duration: maxIdleTime,
+        repeating: false,
+        handler: MaxIdleTimerHandlerView(self)
+      )
+    }
+    if let keepaliveTime {
+      let keepaliveTimeout = keepaliveTimeout ?? .seconds(20)
+      self.keepaliveTimerHandler = Timer(
+        eventLoop: eventLoop,
+        duration: keepaliveTime,
+        repeating: true,
+        handler: KeepaliveTimerHandlerView(self)
+      )
+      self.keepaliveTimeoutHandler = Timer(
+        eventLoop: eventLoop,
+        duration: keepaliveTimeout,
+        repeating: false,
+        handler: KeepaliveTimeoutHandlerView(self)
+      )
+    }
   }
 
   package func handlerAdded(context: ChannelHandlerContext) {
@@ -142,8 +159,8 @@ package final class ClientConnectionHandler: ChannelInboundHandler, ChannelOutbo
       promise.succeed()
     }
 
-    self.keepaliveTimer?.cancel()
-    self.keepaliveTimeoutTimer.cancel()
+    self.keepaliveTimerHandler?.cancel()
+    self.keepaliveTimeoutHandler?.cancel()
     context.fireChannelInactive()
   }
 
@@ -222,11 +239,8 @@ package final class ClientConnectionHandler: ChannelInboundHandler, ChannelOutbo
       // Pings are ack'd by the HTTP/2 handler so we only pay attention to acks here, and in
       // particular only those carrying the keep-alive data.
       if ack, data == self.keepalivePingData {
-        let loopBound = LoopBoundView(handler: self, context: context)
-        self.keepaliveTimeoutTimer.cancel()
-        self.keepaliveTimer?.schedule(on: context.eventLoop) {
-          loopBound.keepaliveTimerFired()
-        }
+        self.keepaliveTimeoutHandler?.cancel()
+        self.keepaliveTimerHandler?.start()
       }
 
     case .settings(.settings(_)):
@@ -236,15 +250,8 @@ package final class ClientConnectionHandler: ChannelInboundHandler, ChannelOutbo
       // becoming active is insufficient as, for example, a TLS handshake may fail after
       // establishing the TCP connection, or the server isn't configured for gRPC (or HTTP/2).
       if isInitialSettings {
-        let loopBound = LoopBoundView(handler: self, context: context)
-        self.keepaliveTimer?.schedule(on: context.eventLoop) {
-          loopBound.keepaliveTimerFired()
-        }
-
-        self.maxIdleTimer?.schedule(on: context.eventLoop) {
-          loopBound.maxIdleTimerFired()
-        }
-
+        self.keepaliveTimerHandler?.start()
+        self.maxIdleTimerHandler?.start()
         context.fireChannelRead(self.wrapInboundOut(.ready))
       }
 
@@ -290,29 +297,44 @@ package final class ClientConnectionHandler: ChannelInboundHandler, ChannelOutbo
   }
 }
 
+// Timer handler views.
 extension ClientConnectionHandler {
-  struct LoopBoundView: @unchecked Sendable {
+  struct MaxIdleTimerHandlerView: @unchecked Sendable, NIOScheduledCallbackHandler {
     private let handler: ClientConnectionHandler
-    private let context: ChannelHandlerContext
 
-    init(handler: ClientConnectionHandler, context: ChannelHandlerContext) {
+    init(_ handler: ClientConnectionHandler) {
       self.handler = handler
-      self.context = context
     }
 
-    func keepaliveTimerFired() {
-      self.context.eventLoop.assertInEventLoop()
-      self.handler.keepaliveTimerFired(context: self.context)
+    func handleScheduledCallback(eventLoop: some EventLoop) {
+      self.handler.eventLoop.assertInEventLoop()
+      self.handler.maxIdleTimerFired()
     }
+  }
+
+  struct KeepaliveTimerHandlerView: @unchecked Sendable, NIOScheduledCallbackHandler {
+    private let handler: ClientConnectionHandler
 
-    func keepaliveTimeoutExpired() {
-      self.context.eventLoop.assertInEventLoop()
-      self.handler.keepaliveTimeoutExpired(context: self.context)
+    init(_ handler: ClientConnectionHandler) {
+      self.handler = handler
     }
 
-    func maxIdleTimerFired() {
-      self.context.eventLoop.assertInEventLoop()
-      self.handler.maxIdleTimerFired(context: self.context)
+    func handleScheduledCallback(eventLoop: some EventLoop) {
+      self.handler.eventLoop.assertInEventLoop()
+      self.handler.keepaliveTimerFired()
+    }
+  }
+
+  struct KeepaliveTimeoutHandlerView: @unchecked Sendable, NIOScheduledCallbackHandler {
+    private let handler: ClientConnectionHandler
+
+    init(_ handler: ClientConnectionHandler) {
+      self.handler = handler
+    }
+
+    func handleScheduledCallback(eventLoop: some EventLoop) {
+      self.handler.eventLoop.assertInEventLoop()
+      self.handler.keepaliveTimeoutExpired()
     }
   }
 }
@@ -356,7 +378,7 @@ extension ClientConnectionHandler {
     self.eventLoop.assertInEventLoop()
 
     // Stream created, so the connection isn't idle.
-    self.maxIdleTimer?.cancel()
+    self.maxIdleTimerHandler?.cancel()
     self.state.streamOpened(id)
   }
 
@@ -368,13 +390,10 @@ extension ClientConnectionHandler {
     case .startIdleTimer(let cancelKeepalive):
       // All streams are closed, restart the idle timer, and stop the keep-alive timer (it may
       // not stop if keep-alive is allowed when there are no active calls).
-      let loopBound = LoopBoundView(handler: self, context: context)
-      self.maxIdleTimer?.schedule(on: context.eventLoop) {
-        loopBound.maxIdleTimerFired()
-      }
+      self.maxIdleTimerHandler?.start()
 
       if cancelKeepalive {
-        self.keepaliveTimer?.cancel()
+        self.keepaliveTimerHandler?.cancel()
       }
 
     case .close:
@@ -397,34 +416,31 @@ extension ClientConnectionHandler {
     }
   }
 
-  private func keepaliveTimerFired(context: ChannelHandlerContext) {
-    guard self.state.sendKeepalivePing() else { return }
+  private func keepaliveTimerFired() {
+    guard self.state.sendKeepalivePing(), let context = self.context else { return }
 
     // Cancel the keep alive timer when the client sends a ping. The timer is resumed when the ping
     // is acknowledged.
-    self.keepaliveTimer?.cancel()
+    self.keepaliveTimerHandler?.cancel()
 
     let ping = HTTP2Frame(streamID: .rootStream, payload: .ping(self.keepalivePingData, ack: false))
     context.write(self.wrapOutboundOut(ping), promise: nil)
     self.maybeFlush(context: context)
 
     // Schedule a timeout on waiting for the response.
-    let loopBound = LoopBoundView(handler: self, context: context)
-    self.keepaliveTimeoutTimer.schedule(on: context.eventLoop) {
-      loopBound.keepaliveTimeoutExpired()
-    }
+    self.keepaliveTimeoutHandler?.start()
   }
 
-  private func keepaliveTimeoutExpired(context: ChannelHandlerContext) {
-    guard self.state.beginClosing() else { return }
+  private func keepaliveTimeoutExpired() {
+    guard self.state.beginClosing(), let context = self.context else { return }
 
     context.fireChannelRead(self.wrapInboundOut(.closing(.keepaliveExpired)))
     self.writeAndFlushGoAway(context: context, message: "keepalive_expired")
     context.close(promise: nil)
   }
 
-  private func maxIdleTimerFired(context: ChannelHandlerContext) {
-    guard self.state.beginClosing() else { return }
+  private func maxIdleTimerFired() {
+    guard self.state.beginClosing(), let context = self.context else { return }
 
     context.fireChannelRead(self.wrapInboundOut(.closing(.idle)))
     self.writeAndFlushGoAway(context: context, message: "idle")

Sources/GRPCNIOTransportCore/Internal/Timer.swift

@@ -16,55 +16,54 @@
 
 package import NIOCore
 
-package struct Timer {
-  /// The delay to wait before running the task.
-  private let delay: TimeAmount
-  /// The task to run, if scheduled.
-  private var task: Kind?
-  /// Whether the task to schedule is repeated.
-  private let `repeat`: Bool
+/// A timer backed by `NIOScheduledCallback`.
+package final class Timer<Handler: NIOScheduledCallbackHandler> where Handler: Sendable {
+  /// The event loop on which to run this timer.
+  private let eventLoop: any EventLoop
 
-  private enum Kind {
-    case once(Scheduled<Void>)
-    case repeated(RepeatedTask)
+  /// The duration of the timer.
+  private let duration: TimeAmount
 
-    func cancel() {
-      switch self {
-      case .once(let task):
-        task.cancel()
-      case .repeated(let task):
-        task.cancel()
-      }
-    }
-  }
+  /// Whether this timer should repeat.
+  private let repeating: Bool
+
+  /// The handler to call when the timer fires.
+  private let handler: Handler
+
+  /// The currently scheduled callback if the timer is running.
+  private var scheduledCallback: NIOScheduledCallback?
 
-  package init(delay: TimeAmount, repeat: Bool = false) {
-    self.delay = delay
-    self.task = nil
-    self.repeat = `repeat`
+  package init(eventLoop: any EventLoop, duration: TimeAmount, repeating: Bool, handler: Handler) {
+    self.eventLoop = eventLoop
+    self.duration = duration
+    self.repeating = repeating
+    self.handler = handler
+    self.scheduledCallback = nil
   }
 
-  /// Schedule a task on the given `EventLoop`.
-  package mutating func schedule(
-    on eventLoop: any EventLoop,
-    work: @escaping @Sendable () throws -> Void
-  ) {
-    self.task?.cancel()
+  /// Cancel the timer, if it is running.
+  package func cancel() {
+    self.eventLoop.assertInEventLoop()
+    guard let scheduledCallback = self.scheduledCallback else { return }
+    scheduledCallback.cancel()
+  }
 
-    if self.repeat {
-      let task = eventLoop.scheduleRepeatedTask(initialDelay: self.delay, delay: self.delay) { _ in
-        try work()
-      }
-      self.task = .repeated(task)
-    } else {
-      let task = eventLoop.scheduleTask(in: self.delay, work)
-      self.task = .once(task)
-    }
+  /// Start or restart the timer.
+  package func start() {
+    self.eventLoop.assertInEventLoop()
+    self.scheduledCallback?.cancel()
+    // Only throws if the event loop is shutting down, so we'll just swallow the error here.
+    self.scheduledCallback = try? self.eventLoop.scheduleCallback(in: self.duration, handler: self)
   }
+}
 
-  /// Cancels the task, if one was scheduled.
-  package mutating func cancel() {
-    self.task?.cancel()
-    self.task = nil
+extension Timer: NIOScheduledCallbackHandler, @unchecked Sendable where Handler: Sendable {
+  /// For repeated timer support, the timer itself proxies the callback and restarts the timer.
+  ///
+  /// - NOTE: Users should not call this function directly.
+  package func handleScheduledCallback(eventLoop: some EventLoop) {
+    self.eventLoop.assertInEventLoop()
+    self.handler.handleScheduledCallback(eventLoop: eventLoop)
+    if self.repeating { self.start() }
   }
 }