Socket Disconnect when no remaining subscriptions

AppSyncRealTimeClient/ConnectionProvider/AppsyncRealtimeConnection/RealtimeConnectionProvider.swift

@@ -30,7 +30,8 @@ public class RealtimeConnectionProvider: ConnectionProvider {
 
     let serialWriteQueue = DispatchQueue(label: "com.amazonaws.AppSyncRealTimeConnectionProvider.writeQueue")
 
-    public init(for url: URL, websocket: AppSyncWebsocketProvider) {
+    public init(for url: URL,
+                websocket: AppSyncWebsocketProvider) {
         self.url = url
         self.websocket = websocket
     }
@@ -93,15 +94,26 @@ public class RealtimeConnectionProvider: ConnectionProvider {
     }
 
     public func addListener(identifier: String, callback: @escaping ConnectionProviderCallback) {
-        serialCallbackQueue.async {
-            self.listeners[identifier] = callback
+        serialCallbackQueue.async { [weak self] in
+            self?.listeners[identifier] = callback
         }
-
     }
 
     public func removeListener(identifier: String) {
-        serialCallbackQueue.async {
+        serialCallbackQueue.async { [weak self] in
+            guard let self = self else {
+                return
+            }
+
             self.listeners.removeValue(forKey: identifier)
+
+            self.serialConnectionQueue.async { [weak self] in
+                guard let self = self else {
+                    return
+                }
+                self.status = .notConnected
+                self.websocket.disconnect()
+            }
         }
     }
 

AppSyncRealTimeClient/ConnectionProvider/ConnectionProviderFactory.swift

@@ -26,7 +26,8 @@ public struct ConnectionProviderFactory {
         return provider
     }
 
-    static func createConnectionProvider(for url: URL, connectionType: SubscriptionConnectionType) -> ConnectionProvider {
+    static func createConnectionProvider(for url: URL,
+                                         connectionType: SubscriptionConnectionType) -> ConnectionProvider {
         switch connectionType {
         case .appSyncRealtime:
             let websocketProvider = StarscreamAdapter()

AppSyncRealTimeClient/Websocket/Starscream/StarscreamAdapter.swift

@@ -27,6 +27,7 @@ public class StarscreamAdapter: AppSyncWebsocketProvider {
 
     public func disconnect() {
         socket?.disconnect()
+        socket = nil
     }
 
     public func write(message: String) {

AppSyncRealTimeClientIntegrationTests/AppSyncRealTimeClientIntegrationTests.swift

@@ -12,6 +12,15 @@ class AppSyncRealTimeClientIntegrationTests: XCTestCase {
 
     var url: URL!
     var apiKey: String!
+    let requestString = """
+        subscription onCreate {
+          onCreateTodo{
+            id
+            description
+            name
+          }
+        }
+        """
 
     override func setUp() {
         do {
@@ -49,19 +58,9 @@ class AppSyncRealTimeClientIntegrationTests: XCTestCase {
         let subscribeSuccess = expectation(description: "subscribe successfully")
         let authInterceptor = APIKeyAuthInterceptor(apiKey)
         let connectionProvider = ConnectionProviderFactory.createConnectionProvider(for: url,
-                                                                                authInterceptor: authInterceptor,
-                                                                                connectionType: .appSyncRealtime)
-
+                                                                                    authInterceptor: authInterceptor,
+                                                                                    connectionType: .appSyncRealtime)
         let subscriptionConnection = AppSyncSubscriptionConnection(provider: connectionProvider)
-        let requestString = """
-                            subscription onCreate {
-                              onCreateTodo{
-                                id
-                                description
-                                name
-                              }
-                            }
-                            """
         _ = subscriptionConnection.subscribe(requestString: requestString, variables: nil) { (event, item) in
 
             switch event {
@@ -83,5 +82,88 @@ class AppSyncRealTimeClientIntegrationTests: XCTestCase {
 
         wait(for: [subscribeSuccess], timeout: TestCommonConstants.networkTimeout)
     }
-}
 
+    /// The purpose of this test is to ensure that all websockets can be successfully created, exercised and terminated
+    /// while keeping a single connection provider in memory. Specifically, the following test exercises the following:
+    /// 1. Create a new connection provider
+    /// 2. Create multiple subscriptions
+    /// 3. Unsubscribe the subscriptions
+    /// 4. Sleep to make sure the asynchronous process to disconnect the socket is executed
+    /// 5. Ensure the socket is disconnected
+    /// 6. Repeat Steps 2-5 with the existing connection provider.
+    ///
+    /// - Given: Connected subscriptions
+    /// - When:
+    ///    - All subscription items are unsubscribed
+    /// - Then:
+    ///    - Underlying websocket is disconnected
+    func testAllSubscriptionsCancelledShouldDisconnectTheWebsocket() {
+        let connectedInvoked = expectation(description: "Connection established")
+        connectedInvoked.expectedFulfillmentCount = 3
+
+        let authInterceptor = APIKeyAuthInterceptor(apiKey)
+        let connectionProvider = ConnectionProviderFactory.createConnectionProvider(for: url,
+                                                                                    authInterceptor: authInterceptor,
+                                                                                    connectionType: .appSyncRealtime)
+        let subscriptionConnection1 = AppSyncSubscriptionConnection(provider: connectionProvider)
+        let item1 = subscriptionConnection1.subscribe(requestString: requestString, variables: nil) { (event, item) in
+            if case let .connection(state) = event {
+                if case .connected = state {
+                    connectedInvoked.fulfill()
+                }
+            }
+        }
+        let subscriptionConnection2 = AppSyncSubscriptionConnection(provider: connectionProvider)
+        let item2 = subscriptionConnection2.subscribe(requestString: requestString, variables: nil) { (event, item) in
+            if case let .connection(state) = event {
+                if case .connected = state {
+                    connectedInvoked.fulfill()
+                }
+            }
+        }
+        let subscriptionConnection3 = AppSyncSubscriptionConnection(provider: connectionProvider)
+        let item3 = subscriptionConnection3.subscribe(requestString: requestString, variables: nil) { (event, item) in
+            if case let .connection(state) = event {
+                if case .connected = state {
+                    connectedInvoked.fulfill()
+                }
+            }
+        }
+
+        XCTAssertNotNil(item1)
+        XCTAssertNotNil(item2)
+        XCTAssertNotNil(item3)
+        wait(for: [connectedInvoked], timeout: TestCommonConstants.networkTimeout)
+
+        guard let realTimeConnectionProvider = connectionProvider as? RealtimeConnectionProvider else {
+            XCTFail("Could not retrieve concrete connection provider")
+            return
+        }
+        XCTAssertEqual(realTimeConnectionProvider.status, .connected)
+
+        subscriptionConnection1.unsubscribe(item: item1)
+        subscriptionConnection2.unsubscribe(item: item2)
+        subscriptionConnection3.unsubscribe(item: item3)
+
+        // Sleep is required here as disconnecting the connection provider is done asynchronously on the connection
+        // queue for the very last unsubscribe. This means we need to "pull" for the status to ensure the system is operating correctly by sleeping
+        // and checking that the status is .notConnected
+        sleep(5)
+        XCTAssertEqual(realTimeConnectionProvider.status, .notConnected)
+
+        let newConnectedInvoked = expectation(description: "Connection established")
+        let subscriptionConnection4 = AppSyncSubscriptionConnection(provider: connectionProvider)
+        let newItem = subscriptionConnection4.subscribe(requestString: requestString, variables: nil) { (event, item) in
+            if case let .connection(state) = event {
+                if case .connected = state {
+                    newConnectedInvoked.fulfill()
+                }
+            }
+        }
+        wait(for: [newConnectedInvoked], timeout: TestCommonConstants.networkTimeout)
+        XCTAssertEqual(realTimeConnectionProvider.status, .connected)
+        subscriptionConnection4.unsubscribe(item: newItem)
+        sleep(5)
+        XCTAssertEqual(realTimeConnectionProvider.status, .notConnected)
+    }
+}

CHANGELOG.md

@@ -1,5 +1,11 @@
 # AppSync RealTime Client for iOS
 
+## 1.1.6
+
+### Improvements
+
+- Socket Disconnect when no remaining subscriptions. See [PR #8](https://github.com/aws-amplify/aws-appsync-realtime-client-ios/pull/8)
+
 ## 1.1.5
 
 ### Bug fix