From fda9efafa65556e4ad2e8e5adeb77468703358da Mon Sep 17 00:00:00 2001
From: Huy Nguyen <huy@pinterest.com>
Date: Thu, 4 Feb 2016 13:59:43 -0800
Subject: [PATCH] Add timestamp to notifications of rendering engine and avoid
 race conditions in ASRangeControllerBeta

- Accurately remove notification observer
---
 AsyncDisplayKit/ASDisplayNode.mm              |  5 +++-
 .../Details/ASRangeControllerBeta.mm          | 25 +++++++++++++------
 .../Private/ASDisplayNodeInternal.h           |  1 +
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/AsyncDisplayKit/ASDisplayNode.mm b/AsyncDisplayKit/ASDisplayNode.mm
index cc2f1b3670..a5d33641c8 100644
--- a/AsyncDisplayKit/ASDisplayNode.mm
+++ b/AsyncDisplayKit/ASDisplayNode.mm
@@ -30,6 +30,7 @@
 
 NSInteger const ASDefaultDrawingPriority = ASDefaultTransactionPriority;
 NSString * const ASRenderingEngineDidDisplayScheduledNodesNotification = @"ASRenderingEngineDidDisplayScheduledNodes";
+NSString * const ASRenderingEngineDidDisplayNodesScheduledBeforeTimestamp = @"ASRenderingEngineDidDisplayNodesScheduledBeforeTimestamp";
 
 @interface _ASDisplayNodePosition : NSObject
 
@@ -276,12 +277,14 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
       ASDN::MutexLocker l(displaySchedulerLock);
       displayScheduled = NO;
       NSSet *displayingNodes = [nodesToDisplay copy];
+      CFAbsoluteTime timestamp = CFAbsoluteTimeGetCurrent();
       nodesToDisplay = nil;
       for (ASDisplayNode *node in displayingNodes) {
         [node __recursivelyTriggerDisplayAndBlock:NO];
       }
       [[NSNotificationCenter defaultCenter] postNotificationName:ASRenderingEngineDidDisplayScheduledNodesNotification
-                                                          object:nil];
+                                                          object:displayingNodes
+                                                        userInfo:@{ASRenderingEngineDidDisplayNodesScheduledBeforeTimestamp: [NSNumber numberWithDouble:timestamp]}];
     });
   }
 }
diff --git a/AsyncDisplayKit/Details/ASRangeControllerBeta.mm b/AsyncDisplayKit/Details/ASRangeControllerBeta.mm
index 42d082e6a3..e0b7fef2a1 100644
--- a/AsyncDisplayKit/Details/ASRangeControllerBeta.mm
+++ b/AsyncDisplayKit/Details/ASRangeControllerBeta.mm
@@ -27,6 +27,7 @@
   NSSet<NSIndexPath *> *_allPreviousIndexPaths;
   ASLayoutRangeMode _currentRangeMode;
   BOOL _didRegisterForNotifications;
+  CFAbsoluteTime _pendingDisplayNodesTimestamp;
 }
 
 @end
@@ -48,7 +49,7 @@
 - (void)dealloc
 {
   if (_didRegisterForNotifications) {
-    [[NSNotificationCenter defaultCenter] removeObserver:self];
+    [[NSNotificationCenter defaultCenter] removeObserver:self name:ASRenderingEngineDidDisplayScheduledNodesNotification object:nil];
   }
 }
 
@@ -243,7 +244,11 @@
       }
     }
   }
-  
+
+  if (_didRegisterForNotifications) {
+    _pendingDisplayNodesTimestamp = CFAbsoluteTimeGetCurrent();
+  }
+
   _rangeIsValid = YES;
   _queuedRangeUpdate = NO;
   
@@ -280,7 +285,7 @@
                                                                        currentRangeMode:_currentRangeMode];
     if (_currentRangeMode != nextRangeMode) {
       [[NSNotificationCenter defaultCenter] addObserver:self
-                                               selector:@selector(scheduledNodesDidDisplay)
+                                               selector:@selector(scheduledNodesDidDisplay:)
                                                    name:ASRenderingEngineDidDisplayScheduledNodesNotification
                                                  object:nil];
       _didRegisterForNotifications = YES;
@@ -288,12 +293,16 @@
   }
 }
 
-- (void)scheduledNodesDidDisplay
+- (void)scheduledNodesDidDisplay:(NSNotification *)notification
 {
-  [[NSNotificationCenter defaultCenter] removeObserver:self];
-  _didRegisterForNotifications = NO;
-
-  [self scheduleRangeUpdate];
+  CFAbsoluteTime notificationTimestamp = ((NSNumber *)[notification.userInfo objectForKey:ASRenderingEngineDidDisplayNodesScheduledBeforeTimestamp]).doubleValue;
+  if (_pendingDisplayNodesTimestamp < notificationTimestamp) {
+    // The rendering engine has processed all the nodes this range controller scheduled. Let's schedule a range update
+    [[NSNotificationCenter defaultCenter] removeObserver:self name:ASRenderingEngineDidDisplayScheduledNodesNotification object:nil];
+    _didRegisterForNotifications = NO;
+    
+    [self scheduleRangeUpdate];
+  }
 }
 
 #pragma mark - Cell node view handling
diff --git a/AsyncDisplayKit/Private/ASDisplayNodeInternal.h b/AsyncDisplayKit/Private/ASDisplayNodeInternal.h
index 7bb47cf8a0..6b5b4eb484 100644
--- a/AsyncDisplayKit/Private/ASDisplayNodeInternal.h
+++ b/AsyncDisplayKit/Private/ASDisplayNodeInternal.h
@@ -38,6 +38,7 @@ typedef NS_OPTIONS(NSUInteger, ASDisplayNodeMethodOverrides)
 @class _ASDisplayNodePosition;
 
 FOUNDATION_EXPORT NSString * const ASRenderingEngineDidDisplayScheduledNodesNotification;
+FOUNDATION_EXPORT NSString * const ASRenderingEngineDidDisplayNodesScheduledBeforeTimestamp;
 
 // Allow 2^n increments of begin disabling hierarchy notifications
 #define VISIBILITY_NOTIFICATIONS_DISABLED_BITS 4