Fix possible retain cycle during update

Instagram · Jan 13, 2017 · f9c2763 · f9c2763
1 parent dd3b9ed
commit f9c2763
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 37 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,8 @@ This release closes the [2.2.0 milestone](https://github.com/Instagram/IGListKit
 
 - Fix bug where emptyView's hidden status is not updated after the number of items is changed with `insertInSectionController:atIndexes:` or related methods. [Peter Edmonston](https://github.com/edmonston) [(#395)](https://github.com/Instagram/IGListKit/pull/395)
 
+- Fixes a possible retain cycle when the data source is deallocated after queueing an update. [Ryan Nystrom](https://github.com/rnystrom) (tbd)
+
 2.1.0
 -----
 

diff --git a/IGListKit.xcodeproj/xcshareddata/xcschemes/IGListKit.xcscheme b/IGListKit.xcodeproj/xcshareddata/xcschemes/IGListKit.xcscheme
@@ -50,6 +50,11 @@
          </BuildableReference>
       </MacroExpansion>
       <AdditionalOptions>
+         <AdditionalOption
+            key = "MallocStackLogging"
+            value = ""
+            isEnabled = "YES">
+         </AdditionalOption>
       </AdditionalOptions>
    </TestAction>
    <LaunchAction

diff --git a/Source/IGListAdapter.m b/Source/IGListAdapter.m
@@ -115,7 +115,7 @@ - (void)setScrollViewDelegate:(id<UIScrollViewDelegate>)scrollViewDelegate {
 - (void)updateAfterPublicSettingsChange {
     id<IGListAdapterDataSource> dataSource = _dataSource;
     if (_collectionView != nil && dataSource != nil) {
-        [self updateObjects:[[dataSource objectsForListAdapter:self] copy] dataSource:dataSource];
+        [self updateObjects:[[dataSource objectsForListAdapter:self] copy]];
     }
 }
 
@@ -260,23 +260,31 @@ - (void)performUpdatesAnimated:(BOOL)animated completion:(IGListUpdaterCompletio
 
     __weak __typeof__(self) weakSelf = self;
     [self.updater performUpdateWithCollectionView:collectionView
-                                               fromObjects:fromObjects
-                                                 toObjects:newItems
-                                                  animated:animated
-                                     objectTransitionBlock:^(NSArray *toObjects) {
-                                         // temporarily capture the item map that we are transitioning from in case
-                                         // there are any item deletes at the same
-                                         weakSelf.previousSectionMap = [weakSelf.sectionMap copy];
+                                      fromObjects:fromObjects
+                                        toObjects:newItems
+                                         animated:animated
+                            objectTransitionBlock:^(NSArray *toObjects) {
+                                // if self or the data source have been deallocated, tell the udpater that object
+                                // changes cannot be applied
+                                if (weakSelf.dataSource == nil) {
+                                    return NO;
+                                }
 
-                                         [weakSelf updateObjects:toObjects dataSource:dataSource];
-                                     } completion:^(BOOL finished) {
-                                         // release the previous items
-                                         weakSelf.previousSectionMap = nil;
+                                // temporarily capture the item map that we are transitioning from in case
+                                // there are any item deletes at the same
+                                weakSelf.previousSectionMap = [weakSelf.sectionMap copy];
 
-                                         if (completion) {
-                                             completion(finished);
-                                         }
-                                     }];
+                                [weakSelf updateObjects:toObjects];
+
+                                return YES;
+                            } completion:^(BOOL finished) {
+                                // release the previous items
+                                weakSelf.previousSectionMap = nil;
+
+                                if (completion) {
+                                    completion(finished);
+                                }
+                            }];
 }
 
 - (void)reloadDataWithCompletion:(nullable IGListUpdaterCompletion)completion {
@@ -297,7 +305,7 @@ - (void)reloadDataWithCompletion:(nullable IGListUpdaterCompletion)completion {
     [self.updater reloadDataWithCollectionView:collectionView reloadUpdateBlock:^{
         // purge all section controllers from the item map so that they are regenerated
         [weakSelf.sectionMap reset];
-        [weakSelf updateObjects:newItems dataSource:dataSource];
+        [weakSelf updateObjects:newItems];
     } completion:completion];
 }
 
@@ -438,9 +446,7 @@ - (CGSize)sizeForSupplementaryViewOfKind:(NSString *)elementKind atIndexPath:(NS
 
 // this method is what updates the "source of truth"
 // this should only be called just before the collection view is updated
-- (void)updateObjects:(NSArray *)objects dataSource:(id<IGListAdapterDataSource>)dataSource {
-    IGParameterAssert(dataSource != nil);
-
+- (void)updateObjects:(NSArray *)objects {
 #if DEBUG
     for (id object in objects) {
         IGAssert([object isEqualToDiffableObject:object], @"Object instance %@ not equal to itself. This will break infra map tables.", object);
@@ -460,6 +466,8 @@ - (void)updateObjects:(NSArray *)objects dataSource:(id<IGListAdapterDataSource>
     id firstObject = objects.firstObject;
     id lastObject = objects.lastObject;
 
+    id<IGListAdapterDataSource> dataSource = self.dataSource;
+
     for (id object in objects) {
         // infra checks to see if a controller exists
         IGListSectionController <IGListSectionType> *sectionController = [map sectionControllerForObject:object];

diff --git a/Source/IGListAdapterUpdater.m b/Source/IGListAdapterUpdater.m
@@ -121,20 +121,22 @@ - (void)performBatchUpdatesWithCollectionView:(UICollectionView *)collectionView
     id<IGListAdapterUpdaterDelegate> delegate = self.delegate;
     NSArray *fromObjects = [self.fromObjects copy];
     NSArray *toObjects = objectsWithDuplicateIdentifiersRemoved(self.toObjects);
-    void (^objectTransitionBlock)(NSArray *) = [self.objectTransitionBlock copy];
+    IGListObjectTransitionBlock objectTransitionBlock = [self.objectTransitionBlock copy];
     NSArray *itemUpdateBlocks = [self.itemUpdateBlocks copy];
     NSArray *completionBlocks = [self.completionBlocks copy];
     const BOOL animated = self.queuedUpdateIsAnimated;
 
     // clean up all state so that new updates can be coalesced while the current update is in flight
     [self cleanupState];
 
-    void (^executeUpdateBlocks)() = ^{
+    // return YES if batch updates should be applied
+    BOOL (^executeUpdateBlocks)() = ^{
         // run the update block so that the adapter can set its items. this makes sure that just before the update is
         // committed that the data source is updated to the /latest/ "toObjects". this makes the data source in sync
         // with the items that the updater is transitioning to
+        BOOL appliedObjectTransition = YES;
         if (objectTransitionBlock != nil) {
-            objectTransitionBlock(toObjects);
+            appliedObjectTransition = objectTransitionBlock(toObjects);
         }
 
         // execute each item update block which should make calls like insert, delete, and reload for index paths
@@ -143,6 +145,8 @@ - (void)performBatchUpdatesWithCollectionView:(UICollectionView *)collectionView
         for (IGListItemUpdateBlock itemUpdateBlock in itemUpdateBlocks) {
             itemUpdateBlock();
         }
+
+        return appliedObjectTransition;
     };
 
     void (^executeCompletionBlocks)(BOOL) = ^(BOOL finished) {
@@ -177,15 +181,17 @@ - (void)performBatchUpdatesWithCollectionView:(UICollectionView *)collectionView
     __block IGListBatchUpdateData *updateData = nil;
 
     void (^updateBlock)() = ^{
-        executeUpdateBlocks();
-
-        updateData = [self flushCollectionView:collectionView
-                                withDiffResult:result
-                                reloadSections:[self.reloadSections copy]
-                              deleteIndexPaths:[self.deleteIndexPaths copy]
-                              insertIndexPaths:[self.insertIndexPaths copy]
-                              reloadIndexPaths:[self.reloadIndexPaths copy]
-                                   fromObjects:fromObjects];
+        // only perform batch updates if the objectTransitionBlock was applied
+        // e.g. toObjects are going to be used to power the collection view after batch updates finishes
+        if (executeUpdateBlocks()) {
+            updateData = [self flushCollectionView:collectionView
+                                    withDiffResult:result
+                                    reloadSections:[self.reloadSections copy]
+                                  deleteIndexPaths:[self.deleteIndexPaths copy]
+                                  insertIndexPaths:[self.insertIndexPaths copy]
+                                  reloadIndexPaths:[self.reloadIndexPaths copy]
+                                       fromObjects:fromObjects];
+        }
 
         [self cleanupUpdateBlockState];
         [self performBatchUpdatesItemBlockApplied];
@@ -332,15 +338,17 @@ - (void)queueUpdateWithCollectionView:(UICollectionView *)collectionView {
     }
 
     __weak __typeof__(self) weakSelf = self;
+    __weak __typeof__(collectionView) weakCollectionView = collectionView;
+
     dispatch_async(dispatch_get_main_queue(), ^{
         if (weakSelf.batchUpdateOrReloadInProgress || ![weakSelf hasChanges]) {
             return;
         }
 
         if (weakSelf.hasQueuedReloadData) {
-            [weakSelf performReloadDataWithCollectionView:collectionView];
+            [weakSelf performReloadDataWithCollectionView:weakCollectionView];
         } else {
-            [weakSelf performBatchUpdatesWithCollectionView:collectionView];
+            [weakSelf performBatchUpdatesWithCollectionView:weakCollectionView];
         }
     });
 }
@@ -371,7 +379,7 @@ - (void)performUpdateWithCollectionView:(UICollectionView *)collectionView
                             fromObjects:(nullable NSArray *)fromObjects
                               toObjects:(nullable NSArray *)toObjects
                                animated:(BOOL)animated
-                  objectTransitionBlock:(void (^)(NSArray *))objectTransitionBlock
+                  objectTransitionBlock:(IGListObjectTransitionBlock)objectTransitionBlock
                              completion:(nullable void (^)(BOOL))completion {
     IGAssertMainThread();
     IGParameterAssert(collectionView != nil);

diff --git a/Source/IGListUpdatingDelegate.h b/Source/IGListUpdatingDelegate.h
@@ -24,8 +24,10 @@ typedef void (^IGListUpdatingCompletion)(BOOL finished);
  A block to be called when the adapter applies changes to the collection view.
 
  @param toObjects The new objects in the collection.
+
+ @return A flag indicating if the transition was applied.
  */
-typedef void (^IGListObjectTransitionBlock)(NSArray *toObjects);
+typedef BOOL (^IGListObjectTransitionBlock)(NSArray *toObjects);
 
 /// A block that contains all of the updates.
 typedef void (^IGListItemUpdateBlock)();

diff --git a/Tests/IGListAdapterE2ETests.m b/Tests/IGListAdapterE2ETests.m
@@ -19,6 +19,8 @@
 #import "IGTestDelegateDataSource.h"
 #import "IGTestObject.h"
 
+#import <objc/runtime.h>
+
 #define genIndexPath(s) [NSIndexPath indexPathForItem:0 inSection:s]
 #define genTestObject(k, v) [[IGTestObject alloc] initWithKey:k value:v]
 
@@ -1150,7 +1152,7 @@ - (void)test_whenDataSourceDeallocatedAfterUpdateQueued_thatUpdateSuccesfullyCom
 
     XCTestExpectation *expectation = genExpectation;
     [self.adapter performUpdatesAnimated:YES completion:^(BOOL finished) {
-        XCTAssertEqual([self.collectionView numberOfSections], 2);
+        XCTAssertEqual([self.collectionView numberOfSections], 1);
         [expectation fulfill];
     }];
 
@@ -1189,4 +1191,38 @@ - (void)test_whenQueuingUpdate_withSectionControllerBatchUpdate_thatSectionContr
     XCTAssertNil(weakSectionController);
 }
 
+- (void)test_whenQueuingUpdate_withSectionControllerBatchUpdate_thatDataSourceNotRetained {
+    __weak id weakDataSource = nil;
+    @autoreleasepool {
+        IGListAdapter *adapter = [[IGListAdapter alloc] initWithUpdater:[IGListAdapterUpdater new] viewController:nil workingRangeSize:0];
+        IGTestDelegateDataSource *dataSource = [IGTestDelegateDataSource new];
+        dataSource.adapter = adapter;
+        IGTestObject *object = genTestObject(@1, @2);
+        dataSource.objects = @[object];
+        IGListCollectionView *collectionView = [[IGListCollectionView alloc] initWithFrame:CGRectMake(0, 0, 100, 100) collectionViewLayout:[UICollectionViewFlowLayout new]];
+        adapter.collectionView = collectionView;
+        adapter.dataSource = dataSource;
+        [collectionView layoutIfNeeded];
+        XCTAssertEqual([collectionView numberOfSections], 1);
+        XCTAssertEqual([collectionView numberOfItemsInSection:0], 2);
+
+        // simulate the data source holding a strong ref to the IGListAdapter
+        // e.g. when a UIViewController owns the adapter and sets `self.adapter.dataSource = self`
+        objc_setAssociatedObject(dataSource, _cmd, adapter, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
+        IGListSectionController<IGListSectionType> *section = [adapter sectionControllerForObject:object];
+        [section.collectionContext performBatchAnimated:YES updates:^{
+            object.value = @3;
+            [section.collectionContext insertInSectionController:section atIndexes:[NSIndexSet indexSetWithIndex:0]];
+        } completion:^(BOOL finished) {}];
+
+        dataSource.objects = @[object, genTestObject(@2, @2)];
+        [adapter performUpdatesAnimated:YES completion:^(BOOL finished) {}];
+
+        weakDataSource = dataSource;
+        XCTAssertNotNil(weakDataSource);
+    }
+    XCTAssertNil(weakDataSource);
+}
+
 @end
diff --git a/Tests/IGListAdapterUpdaterTests.m b/Tests/IGListAdapterUpdaterTests.m
@@ -45,6 +45,7 @@ - (void)setUp {
     __weak __typeof__(self) weakSelf = self;
     self.updateBlock = ^(NSArray *obj) {
         weakSelf.dataSource.sections = obj;
+        return YES;
     };
 }
 
@@ -268,6 +269,8 @@ - (void)test_whenUpdatesAreReentrant_thatUpdatesExecuteSerially {
         preUpdateBlock();
 
         self.dataSource.sections = toObjects;
+
+        return YES;
     } completion:^(BOOL finished) {
         completionCounter++;
         XCTAssertEqual(completionCounter, 1);
@@ -299,6 +302,7 @@ - (void)test_whenQueueingItemUpdates_withBatchUpdate_thatItemUpdateBlockExecutes
                                          animated:YES objectTransitionBlock:^(NSArray * toObjects) {
                                              self.dataSource.sections = toObjects;
                                              sectionUpdateBlockExecuted = YES;
+                                             return YES;
                                          }
                                        completion:nil];
 

diff --git a/Tests/Objects/IGTestDelegateDataSource.h b/Tests/Objects/IGTestDelegateDataSource.h
@@ -20,4 +20,6 @@
 
 @property (nonatomic, copy) void (^cellConfigureBlock)(IGTestDelegateController *);
 
+@property (nonatomic, strong) IGListAdapter *adapter;
+
 @end