Real time client memory leaks fixes

Ably.xcodeproj/project.pbxproj

@@ -69,6 +69,9 @@
 		96E408481A3895E800087F77 /* ARTWebSocketTransport.m in Sources */ = {isa = PBXBuildFile; fileRef = 96E408461A3895E800087F77 /* ARTWebSocketTransport.m */; };
 		D3AD0EBD215E2FB000312105 /* ARTNSString+ARTUtil.h in Headers */ = {isa = PBXBuildFile; fileRef = D3AD0EBB215E2FB000312105 /* ARTNSString+ARTUtil.h */; };
 		D3AD0EBE215E2FB000312105 /* ARTNSString+ARTUtil.m in Sources */ = {isa = PBXBuildFile; fileRef = D3AD0EBC215E2FB000312105 /* ARTNSString+ARTUtil.m */; };
+		D5A22172266F526600C87C42 /* GCDTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = D5A22171266F526600C87C42 /* GCDTest.swift */; };
+		D5A22173266F526600C87C42 /* GCDTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = D5A22171266F526600C87C42 /* GCDTest.swift */; };
+		D5A22174266F526600C87C42 /* GCDTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = D5A22171266F526600C87C42 /* GCDTest.swift */; };
 		D7093C0F219E2DB200723F17 /* Ably.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = D710D45B219495E2008F54AD /* Ably.framework */; };
 		D7093C19219E465300723F17 /* TestUtilities.swift in Sources */ = {isa = PBXBuildFile; fileRef = 856AAC961B6E30C800B07119 /* TestUtilities.swift */; };
 		D7093C1A219E465C00723F17 /* NSObject+TestSuite.m in Sources */ = {isa = PBXBuildFile; fileRef = D780846D1C68B3E50083009D /* NSObject+TestSuite.m */; };
@@ -831,6 +834,7 @@
 		96E408461A3895E800087F77 /* ARTWebSocketTransport.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ARTWebSocketTransport.m; sourceTree = "<group>"; };
 		D3AD0EBB215E2FB000312105 /* ARTNSString+ARTUtil.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; name = "ARTNSString+ARTUtil.h"; path = "Private/ARTNSString+ARTUtil.h"; sourceTree = "<group>"; };
 		D3AD0EBC215E2FB000312105 /* ARTNSString+ARTUtil.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; name = "ARTNSString+ARTUtil.m"; path = "Private/ARTNSString+ARTUtil.m"; sourceTree = "<group>"; };
+		D5A22171266F526600C87C42 /* GCDTest.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = GCDTest.swift; sourceTree = "<group>"; };
 		D7093C0A219E2DB200723F17 /* Ably-macOS-Tests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "Ably-macOS-Tests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
 		D7093C0E219E2DB200723F17 /* Info-macOS.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = "Info-macOS.plist"; sourceTree = "<group>"; };
 		D7093C60219EE1AE00723F17 /* Ably-tvOS-Tests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "Ably-tvOS-Tests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
@@ -1158,6 +1162,7 @@
 				851674EE1B7BA5CD00D35169 /* Stats.swift */,
 				856AAC961B6E30C800B07119 /* TestUtilities.swift */,
 				EB1AE0CD1C5C3A4900D62250 /* Utilities.swift */,
+				D5A22171266F526600C87C42 /* GCDTest.swift */,
 			);
 			path = Spec;
 			sourceTree = "<group>";
@@ -2203,6 +2208,7 @@
 				856AAC991B6E312F00B07119 /* RestClient.swift in Sources */,
 				D746AE2C1BBB625E003ECEF8 /* RestClientChannel.swift in Sources */,
 				D71D30041C5F7B2F002115B0 /* RealtimeClientChannels.swift in Sources */,
+				D5A22172266F526600C87C42 /* GCDTest.swift in Sources */,
 				EB1B53F922F85CE4006A59AC /* ObjectLifetimes.swift in Sources */,
 				D7DF73851EA600240013CD36 /* PushActivationStateMachine.swift in Sources */,
 				D7FC1ECB209CEA2E001E4153 /* Push.swift in Sources */,
@@ -2332,6 +2338,7 @@
 				D7093C1E219E466900723F17 /* RestClient.swift in Sources */,
 				D7093C2B219E466E00723F17 /* Crypto.swift in Sources */,
 				D7093C20219E466E00723F17 /* RestClientChannel.swift in Sources */,
+				D5A22173266F526600C87C42 /* GCDTest.swift in Sources */,
 				EB1B53FA22F85CE4006A59AC /* ObjectLifetimes.swift in Sources */,
 				D7093C2A219E466E00723F17 /* Utilities.swift in Sources */,
 				D798554923EB96C000946BE2 /* DeltaCodec.swift in Sources */,
@@ -2362,6 +2369,7 @@
 				D7093C78219EE26400723F17 /* RestClientChannels.swift in Sources */,
 				D7093C76219EE26400723F17 /* RestClientStats.swift in Sources */,
 				D7093C82219EE26400723F17 /* Crypto.swift in Sources */,
+				D5A22174266F526600C87C42 /* GCDTest.swift in Sources */,
 				EB1B53FB22F85CE4006A59AC /* ObjectLifetimes.swift in Sources */,
 				D7093C7C219EE26400723F17 /* RealtimeClientConnection.swift in Sources */,
 				D798554A23EB96C000946BE2 /* DeltaCodec.swift in Sources */,

Source/ARTEventEmitter.m

@@ -138,15 +138,21 @@ - (void)startTimer {
         NSAssert(false, @"timer is already running");
     }
     _timerIsRunning = true;
+
+    __weak ARTEventListener *weakSelf = self;
     _work = artDispatchScheduled(_timeoutDeadline, [_eventHandler queue], ^{
-        [self timeout];
+        ARTEventListener *strongSelf = weakSelf;
+        if (strongSelf != nil) {
+            [strongSelf timeout];
+        }
     });
 }
 
 - (void)stopTimer {
     artDispatchCancel(_work);
     _timerIsRunning = false;
     _timeoutBlock = nil;
+    _work = nil;
 }
 
 - (void)restartTimer {

Source/Private/ARTGCD.h

@@ -9,7 +9,13 @@
 #import <Foundation/Foundation.h>
 
 @interface ARTScheduledBlockHandle : NSObject
+- (instancetype)initWithDelay:(NSTimeInterval)delay queue:(dispatch_queue_t)queue block:(dispatch_block_t)block;
+- (void)cancel;
 @end
 
 ARTScheduledBlockHandle *artDispatchScheduled(NSTimeInterval seconds, dispatch_queue_t queue, dispatch_block_t block);
-void artDispatchCancel(ARTScheduledBlockHandle *handle);
+static inline void artDispatchCancel(ARTScheduledBlockHandle *handle) {
+    if (handle) {
+        [handle cancel];
+    }
+}

Source/Private/ARTGCD.m

@@ -8,52 +8,81 @@
 
 #import "ARTGCD.h"
 
-@interface ARTScheduledBlockHandle ()
-
-@property (strong, atomic) dispatch_block_t scheduled;
-@property (strong, atomic) dispatch_block_t wrapped;
+@implementation ARTScheduledBlockHandle {
+    dispatch_semaphore_t _semaphore;
+    dispatch_block_t _block;
+    dispatch_block_t _scheduledBlock;
+}
 
-@end
+- (instancetype)initWithDelay:(NSTimeInterval)delay queue:(dispatch_queue_t)queue block:(dispatch_block_t)block {
+    self = [super init];
+    if (self == nil)
+        return nil;
 
-@implementation ARTScheduledBlockHandle
+    // Use a sempaphore to coorindate state. We use a reference here to decouple it when creating a block we'll schedule.
+    dispatch_semaphore_t semaphore = dispatch_semaphore_create(1);
 
-@end
+    __weak ARTScheduledBlockHandle *weakSelf = self;
+    _scheduledBlock = dispatch_block_create(0, ^{
+        dispatch_block_t copiedBlock = nil;
+        dispatch_semaphore_wait(semaphore, DISPATCH_TIME_FOREVER);
+        {
+            // Get a strong reference to self within our semaphore to avoid potential race conditions
+            ARTScheduledBlockHandle *strongSelf = weakSelf;
+            if (strongSelf != nil) {
+                copiedBlock = strongSelf->_block; // copied below
+            }
+        }
+        dispatch_semaphore_signal(semaphore);
 
-ARTScheduledBlockHandle *artDispatchScheduled(NSTimeInterval seconds, dispatch_queue_t queue, dispatch_block_t block) {
-	// We don't pass the block directly; instead, we put it in a property, and
-	// read it back from the property once the timer fires. This gives us the
-	// chance to set the property to nil when cancelling the timer, thus
-	// releasing our retain on the block early. dispatch_block_cancel doesn't do
-	// this, it retains the block even if you cancel the dispatch until the
-	// dispatch time passes. (How this is a good idea escapes me.)
-	//
-	// From Apple's documentation [1]:
-	//
-	// > Release of any resources associated with the block object is delayed
-	// > until execution of the block object is next attempted (or any execution
-	// > already in progress completes).
-	//
-	// https://developer.apple.com/documentation/dispatch/1431058-dispatch_block_cancel
-
-	__block ARTScheduledBlockHandle *handle = [[ARTScheduledBlockHandle alloc] init];
-	handle.wrapped = block;
-
-    handle.scheduled = dispatch_block_create(0, ^{
-        dispatch_block_t wrapped = handle.wrapped;
-        if (wrapped) {
-            wrapped();
+        // If our block is non-nil, our scheduled block was still valid by the time this was invoked
+        if (copiedBlock != nil) {
+            copiedBlock();
         }
     });
-    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(NSEC_PER_SEC * seconds)), queue, handle.scheduled);
-    return handle;
-}
 
-void artDispatchCancel(ARTScheduledBlockHandle *handle) {
-    if (handle) {
-    	dispatch_block_cancel(handle.scheduled);
+    _block = [block copy];
+    _semaphore = semaphore;
+
+    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(NSEC_PER_SEC * delay)), queue, _scheduledBlock);
+
+    return self;
+}
 
-    	// Release the block, since it won't be called and dispatch_block_cancel
-    	// won't release it until its dispatch time passes.
-    	handle.wrapped = nil;
+- (void)cancel {
+    // Cancel within our semaphore for predictable behavior if our block is invoked while we're cancelling
+    dispatch_semaphore_wait(_semaphore, DISPATCH_TIME_FOREVER);
+    {
+        dispatch_block_cancel(_scheduledBlock);
+        _block = nil;
     }
+    dispatch_semaphore_signal(_semaphore);
+}
+
+- (void)dealloc {
+    // Explicitly cancel when we deallocate. This happens implicitly since our scheduled block keeps a weak
+    // reference to self, but we want to make sure that the weak reference can be safely accessed, even if
+    // we're in the midst of deallocating.
+    [self cancel];
+}
+
+@end
+
+ARTScheduledBlockHandle *artDispatchScheduled(NSTimeInterval seconds, dispatch_queue_t queue, dispatch_block_t block) {
+    // We don't pass the block directly; instead, we put it in a property, and
+    // read it back from the property once the timer fires. This gives us the
+    // chance to set the property to nil when cancelling the timer, thus
+    // releasing our retain on the block early. dispatch_block_cancel doesn't do
+    // this, it retains the block even if you cancel the dispatch until the
+    // dispatch time passes. (How this is a good idea escapes me.)
+    //
+    // From Apple's documentation [1]:
+    //
+    // > Release of any resources associated with the block object is delayed
+    // > until execution of the block object is next attempted (or any execution
+    // > already in progress completes).
+    //
+    // https://developer.apple.com/documentation/dispatch/1431058-dispatch_block_cancel
+
+    return [[ARTScheduledBlockHandle alloc] initWithDelay:seconds queue:queue block:block];
 }

Spec/AblySpec-Bridging-Header.h

@@ -4,3 +4,4 @@
 
 #include <asl.h>
 #include "NSObject+TestSuite.h"
+#include "ARTGCD.h"

Spec/GCDTest.swift

@@ -0,0 +1,27 @@
+//
+//  GCDTest.swift
+//  Ably
+//
+//  Created by Mikey on 12/01/2021.
+//  Copyright © 2021 Ably. All rights reserved.
+//
+import XCTest
+
+class GCDTest: XCTestCase {
+    func testScheduledBlockHandleDerefsBlockAfterInvoke() {
+        let invokedExpectation = self.expectation(description: "scheduled block invoked")
+        let obj = NSObject()
+
+        var scheduledBlock = artDispatchScheduled(0, .main) {
+            _ = obj
+            invokedExpectation.fulfill()
+        }
+
+        waitForExpectations(timeout: 2, handler: nil)
+
+        _ = scheduledBlock // silence warning
+        scheduledBlock = nil
+
+        XCTAssertEqual(CFGetRetainCount(obj), 1)
+    }
+}