Add timeout configuration if the connection does not receive a response (#101)

Author: madsodgaard

Sources/APNSwift/APNSwiftConfiguration.swift

@@ -168,6 +168,8 @@ public struct APNSwiftConfiguration {
     public var topic: String
     public var environment: Environment
     internal var logger: Logger?
+    /// Optional timeout time if the connection does not receive a response.
+    public var timeout: TimeAmount? = nil
 
     public var url: URL {
         switch environment {
@@ -197,7 +199,8 @@ public struct APNSwiftConfiguration {
         authenticationMethod: AuthenticationMethod,
         topic: String,
         environment: APNSwiftConfiguration.Environment,
-        logger: Logger? = nil
+        logger: Logger? = nil,
+        timeout: TimeAmount? = nil
     ) {
         self.topic = topic
         self.authenticationMethod = authenticationMethod
@@ -206,6 +209,7 @@ public struct APNSwiftConfiguration {
             logger[metadataKey: "origin"] = "APNSwift"
             self.logger = logger
         }
+        self.timeout = timeout
     }
 }
 

Sources/APNSwift/APNSwiftConnection.swift

@@ -200,20 +200,43 @@ public final class APNSwiftConnection: APNSwiftClient {
             request: payload,
             responsePromise: responsePromise
         )
-
         streamPromise.futureResult.cascadeFailure(to: responsePromise)
-        return streamPromise.futureResult.flatMap { stream in
-            logger?.info("Send - sending")
-            return stream.writeAndFlush(context).flatMapErrorThrowing { error in
-                logger?.info("Send - sending - failed - \(error)")
-                responsePromise.fail(error)
-                throw error
+        
+        let timeoutPromise = self.channel.eventLoop.makePromise(of: Void.self)
+        responsePromise.futureResult.cascade(to: timeoutPromise)
+        timeoutPromise.futureResult.cascadeFailure(to: responsePromise)
+        var timeoutTask: Scheduled<Any>? = nil
+        let timeoutTime = configuration.timeout
+        
+        return streamPromise.futureResult
+            .flatMap { stream in
+                logger?.info("Send - sending")
+                if let timeoutTime = timeoutTime {
+                    timeoutTask = stream.eventLoop.scheduleTask(in: timeoutTime) {
+                        logger?.warning("Send - sending - failed - No response was received within the timeout.")
+                        return timeoutPromise.fail(NoResponseWithinTimeoutError())
+                    }
+                } else {
+                    timeoutPromise.succeed(())
+                }
+                
+                return stream.writeAndFlush(context).flatMapErrorThrowing { error in
+                    logger?.info("Send - sending - failed - \(error)")
+                    responsePromise.fail(error)
+                    throw error
+                }
+            }
+            .flatMap {
+                responsePromise
+                    .futureResult
+                    .and(timeoutPromise.futureResult)
+                    .map { _ in () }
+            }
+            .always { _ in
+                timeoutTask?.cancel()
             }
-        }.flatMap {
-            responsePromise.futureResult
-        }
     }
-
+    
     var onClose: EventLoopFuture<Void> {
         logger?.debug("Connection - closed")
         return self.channel.closeFuture

Sources/APNSwift/APNSwiftErrors.swift

@@ -18,6 +18,8 @@ import Foundation
 public struct NoResponseReceivedBeforeConnectionEnded: Error, Equatable {}
 /// An error where a request was made to Apple, but the response body buffer was nil
 public struct NoResponseBodyFromApple: Error, Equatable {}
+/// An error where no the connection received no response within the timeout period
+public struct NoResponseWithinTimeoutError: Error, Equatable {}
 
 public struct APNSwiftError: Equatable {
     public enum ResponseError: Error, Equatable {

Tests/APNSwiftTests/APNSwiftConfigurationTests.swift

@@ -27,7 +27,8 @@ class APNSwiftConfigurationTests: XCTestCase {
                 teamIdentifier: "MY_TEAM_ID"
             ),
             topic: "MY_TOPIC",
-            environment: environment
+            environment: environment,
+            timeout: .seconds(5)
         )
 
         switch environment {
@@ -46,6 +47,7 @@ class APNSwiftConfigurationTests: XCTestCase {
             XCTFail("expected JWT auth method")
         }
         XCTAssertEqual(apnsConfiguration.topic, "MY_TOPIC")
+        XCTAssertEqual(apnsConfiguration.timeout, .seconds(5))
 
     }
 

Tests/APNSwiftTests/APNSwiftRequestTests.swift

@@ -304,7 +304,6 @@ final class APNSwiftRequestTests: XCTestCase {
         // Should have changed
         XCTAssertFalse(newCachedToken == bearerToken.currentBearerToken)
         bearerToken.cancel()
-        
     }
 
     let validAuthKey = """