Keep alive connectionstream (#30)

Signed-off-by: ciychodianda <[email protected]>
Signed-off-by: David Nadoba <[email protected]>
Co-authored-by: David Nadoba <[email protected]>

Author: ciychodianda

Package.resolved

@@ -100,7 +100,7 @@ extension ClientSetupConfig {
     public static var defaultMobileToServer: ClientSetupConfig {
         ClientSetupConfig(
             timeBetweenKeepaliveFrames: 30_000,
-            maxLifetime: 30_000,
+            maxLifetime: 60_000,
             metadataEncodingMimeType: "application/octet-stream",
             dataEncodingMimeType: "application/octet-stream"
         )
@@ -109,7 +109,7 @@ extension ClientSetupConfig {
     public static var defaultServerToServer: ClientSetupConfig {
         ClientSetupConfig(
             timeBetweenKeepaliveFrames: 500,
-            maxLifetime: 500,
+            maxLifetime: 2_000,
             metadataEncodingMimeType: "application/octet-stream",
             dataEncodingMimeType: "application/octet-stream"
         )

Sources/RSocketCore/Channel Handler/ClientSetupConfig.swift

@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-import NIO
 import Foundation
+import NIO
 
 /// Information about a client which is about to connect or is connected.
 public struct SetupInfo {
     /// If the connection should honor `LEASE`
     public let honorsLease: Bool
-    
+
     /// version of the client protocol implementation
     public let version: Version
     /**
@@ -72,11 +72,11 @@ public struct SetupInfo {
 
 internal struct SetupValidator {
     internal var maximumClientVersion = Version.v1_0
-    
+
     internal func validate(frame: Frame) throws -> SetupInfo {
         try validateSetup(try getSetupBody(frame))
     }
-    
+
     private func getSetupBody(_ frame: Frame) throws -> SetupFrameBody {
         guard frame.header.streamId == .connection else {
             throw Error.invalidSetup(message: "connection needs to be setup on stream 0")
@@ -89,7 +89,7 @@ internal struct SetupValidator {
         }
         return setup
     }
-    
+
     private func validateSetup(_ setup: SetupFrameBody) throws -> SetupInfo {
         guard setup.version <= maximumClientVersion else {
             throw Error.unsupportedSetup(message: "only version \(maximumClientVersion) and lower are supported")
@@ -184,7 +184,6 @@ internal final class ConnectionEstablishmentHandler: ChannelInboundHandler, Remo
 
     private let setupValidator = SetupValidator()
 
-    
     /// Configure `ConnectionEstablishmentHandler`.  If `shouldAcceptClient` is nil, valid clients are always accepted.
     /// - Parameters:
     ///   - initializeConnection: called after successful handshake and after `shouldAcceptClient`did accept the client.

Sources/RSocketCore/Channel Handler/ConnectionEstablishment.swift

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2015-present the original author or authors.
+ *
+ * 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 Foundation
+import NIO
+
+final class KeepaliveHandler {
+    private var keepAliveHandle: RepeatedTask?
+    
+    /// receive time in **seconds** of the last keepalive frame
+    private var lastReceivedTime: TimeInterval
+    /// time is in **milliseconds**
+    private let timeBetweenKeepaliveFrames: Int32
+    /// time is in **milliseconds**
+    private let maxLifetime: Int32
+    private let connectionSide: ConnectionRole
+    /// returns the current time in **seconds**
+    private let now: () -> TimeInterval
+    
+    init(
+        timeBetweenKeepaliveFrames: Int32,
+        maxLifetime: Int32,
+        connectionSide: ConnectionRole,
+        now: @escaping () -> TimeInterval = { ProcessInfo.processInfo.systemUptime }
+    ) {
+        self.timeBetweenKeepaliveFrames = timeBetweenKeepaliveFrames
+        self.maxLifetime = maxLifetime
+        self.connectionSide = connectionSide
+        self.now = now
+        self.lastReceivedTime = now()
+    }
+}
+
+extension KeepaliveHandler: ChannelInboundHandler {
+    typealias InboundIn = Frame
+    typealias OutboundOut = Frame
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let frame = unwrapInboundIn(data)
+        switch frame.body {
+        case let .keepalive(body):
+            lastReceivedTime = now()
+            if body.respondWithKeepalive {
+                let keepAliveFrame = KeepAliveFrameBody(respondWithKeepalive: false, lastReceivedPosition: 0, data: Data()).asFrame()
+                context.writeAndFlush(self.wrapOutboundOut(keepAliveFrame), promise: nil)
+            }
+        default:
+            break
+        }
+    }
+    
+    func handlerAdded(context: ChannelHandlerContext) {
+        if context.channel.isActive {
+            /// this handler may get added to the pipeline after the channel is already active and `channelActive` is then not called
+            onActive(context: context)
+        }
+    }
+
+    func channelActive(context: ChannelHandlerContext) {
+        onActive(context: context)
+    }
+    
+    func onActive(context: ChannelHandlerContext) {
+        guard timeBetweenKeepaliveFrames > 0 else { return }
+        lastReceivedTime = now()
+        guard connectionSide == .client else { return }
+        
+        keepAliveHandle = context.eventLoop.scheduleRepeatedAsyncTask(
+            initialDelay: .milliseconds(Int64(timeBetweenKeepaliveFrames)),
+            delay: .milliseconds(Int64(timeBetweenKeepaliveFrames))
+        ) { [self] task in
+            let elapsedTimeSinceLastKeepaliveInSeconds =  now() - lastReceivedTime
+            let elapsedTimeSinceLastKeepaliveInMilliseconds = Int32((elapsedTimeSinceLastKeepaliveInSeconds * 1000).rounded(.up))
+            if elapsedTimeSinceLastKeepaliveInMilliseconds >= maxLifetime {
+                let errorFrame = Error.connectionClose(message: "KeepAlive timeout exceeded").asFrame(withStreamId: .connection)
+                context.writeAndFlush(self.wrapOutboundOut(errorFrame), promise: nil)
+                task.cancel()
+                return context.eventLoop.makeFailedFuture(Error.applicationError(message: "KeepAliveHandler Shutdown"))
+            } else {
+                let keepAliveFrame = KeepAliveFrameBody(
+                    respondWithKeepalive: true,
+                    /// we do not support resumability yet, thus do not keep track of `lastReceivedPosition` and just always send 0
+                    lastReceivedPosition: 0,
+                    data: Data()
+                ).asFrame()
+                return context.writeAndFlush(self.wrapOutboundOut(keepAliveFrame))
+            }
+        }
+    }
+
+    func channelInactive(context: ChannelHandlerContext) {
+        keepAliveHandle?.cancel()
+    }
+}

Sources/RSocketCore/Channel Handler/ConnectionStreamHandler.swift

@@ -22,12 +22,12 @@ internal final class SetupWriter: ChannelInboundHandler, RemovableChannelHandler
     typealias OutboundOut = Frame
     private let setup: ClientSetupConfig
     private let connectedPromise: EventLoopPromise<Void>?
-    
+
     internal init(config: ClientSetupConfig, connectedPromise: EventLoopPromise<Void>? = nil) {
         self.setup = config
         self.connectedPromise = connectedPromise
     }
-    
+
     func channelActive(context: ChannelHandlerContext) {
         context.writeAndFlush(self.wrapOutboundOut(SetupFrameBody(
             honorsLease: false,
@@ -39,15 +39,15 @@ internal final class SetupWriter: ChannelInboundHandler, RemovableChannelHandler
             dataEncodingMimeType: setup.dataEncodingMimeType,
             payload: setup.payload
         ).asFrame()), promise: nil)
+        context.fireChannelActive()
         context.channel.pipeline.removeHandler(context: context).eventLoop.assertInEventLoop()
         connectedPromise?.succeed(())
     }
-    
+
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         assertionFailure("should never receive data because we remove this handler right after the channel becomes active")
         /// We need to conform to `ChannelInboundHandler` to get called when the channel becomes active and we remove ourself immediately after the channel becomes active
         /// If, for whatever reason, this method gets called, we just forward the data in release mode
         context.fireChannelRead(data)
     }
 }
-

Sources/RSocketCore/Channel Handler/KeepaliveHandler.swift

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import Foundation
 import NIO
 
 extension ChannelPipeline {
@@ -32,6 +33,7 @@ extension ChannelPipeline {
             responderLateFrameHandler: nil
         )
     }
+
     internal func addRSocketClientHandlers(
         config: ClientSetupConfig,
         responder: RSocket? = nil,
@@ -57,7 +59,7 @@ extension ChannelPipeline {
                 requester: requester,
                 responder: Responder(responderSocket: responder, eventLoop: eventLoop, sendFrame: sendFrame)
             ),
-            ConnectionStreamHandler(),
+            KeepaliveHandler(timeBetweenKeepaliveFrames: config.timeBetweenKeepaliveFrames, maxLifetime: config.maxLifetime, connectionSide: ConnectionRole.client),
         ])
     }
 }
@@ -99,7 +101,7 @@ extension ChannelPipeline {
                         requester: Requester(streamIdGenerator: .server, eventLoop: eventLoop, sendFrame: sendFrame),
                         responder: Responder(responderSocket: responder, eventLoop: eventLoop, sendFrame: sendFrame)
                     ),
-                    ConnectionStreamHandler(),
+                    KeepaliveHandler(timeBetweenKeepaliveFrames: info.timeBetweenKeepaliveFrames, maxLifetime: info.maxLifetime, connectionSide: ConnectionRole.server),
                 ])
             }, shouldAcceptClient: shouldAcceptClient)
         ])

Sources/RSocketCore/Channel Handler/SetupWriter.swift

@@ -33,8 +33,8 @@ extension ServerBootstrap: NIOServerTCPBootstrapProtocol{}
 
 class EndToEndTests: XCTestCase {
     static let defaultClientSetup = ClientSetupConfig(
-        timeBetweenKeepaliveFrames: 500,
-        maxLifetime: 5000,
+        timeBetweenKeepaliveFrames: 100,
+        maxLifetime: 1000,
         metadataEncodingMimeType: "utf8",
         dataEncodingMimeType: "utf8"
     )

Sources/RSocketCore/ChannelPipeline.swift

@@ -60,8 +60,7 @@ final class ConnectionEstablishmentTests: XCTestCase {
 
         self.wait(for: [initializeConnection, shouldAcceptSetup], timeout: 0.1)
     }
-    
-    
+
     func testDeliveryOfExtraMessagesDuringSetup() throws {
         let loop = EmbeddedEventLoop()
         let connectionInitialization = loop.makePromise(of: Void.self)