[ASDisplayNode] Add locking to view and layer in ASDisplayNode (#3179)

* Add locking to view and layer in ASDisplayNode

* Another approach

* Some more

* Address comments

Source/ASDisplayNode.mm

@@ -575,7 +575,7 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
 
 - (BOOL)_locked_shouldLoadViewOrLayer
 {
-  return !(_hierarchyState & ASHierarchyStateRasterized);
+  return !_flags.isDeallocating && !(_hierarchyState & ASHierarchyStateRasterized);
 }
 
 - (UIView *)_locked_viewToLoad
@@ -633,61 +633,39 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   return layer;
 }
 
-- (void)_loadViewOrLayerIsLayerBacked:(BOOL)isLayerBacked
+- (void)_locked_loadViewOrLayerIsLayerBacked:(BOOL)isLayerBacked
 {
-  {
-    ASDN::MutexLocker l(__instanceLock__);
-
-    if (_flags.isDeallocating) {
-      return;
-    }
-
-    if (![self _locked_shouldLoadViewOrLayer]) {
-      return;
-    }
-
-    if (isLayerBacked) {
-      TIME_SCOPED(_debugTimeToCreateView);
-      _layer = [self _locked_layerToLoad];
-      static int ASLayerDelegateAssociationKey;
-
-      /**
-       * CALayer's .delegate property is documented to be weak, but the implementation is actually assign.
-       * Because our layer may survive longer than the node (e.g. if someone else retains it, or if the node
-       * begins deallocation on a background thread and it waiting for the -dealloc call to reach main), the only
-       * way to avoid a dangling pointer is to use a weak proxy.
-       */
-      ASWeakProxy *instance = [ASWeakProxy weakProxyWithTarget:self];
-      _layer.delegate = (id<CALayerDelegate>)instance;
-      objc_setAssociatedObject(_layer, &ASLayerDelegateAssociationKey, instance, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
-    } else {
-      TIME_SCOPED(_debugTimeToCreateView);
-      _view = [self _locked_viewToLoad];
-      _view.asyncdisplaykit_node = self;
-      _layer = _view.layer;
-    }
-    _layer.asyncdisplaykit_node = self;
-
-    self._locked_asyncLayer.asyncDelegate = self;
-  }
-
-  {
-    TIME_SCOPED(_debugTimeToApplyPendingState);
-    [self _applyPendingStateToViewOrLayer];
-  }
-  {
-    TIME_SCOPED(_debugTimeToAddSubnodeViews);
-    [self _addSubnodeViewsAndLayers];
-  }
-  {
-    TIME_SCOPED(_debugTimeForDidLoad);
-    [self _didLoad];
+  if (isLayerBacked) {
+    TIME_SCOPED(_debugTimeToCreateView);
+    _layer = [self _locked_layerToLoad];
+    static int ASLayerDelegateAssociationKey;
+    
+    /**
+     * CALayer's .delegate property is documented to be weak, but the implementation is actually assign.
+     * Because our layer may survive longer than the node (e.g. if someone else retains it, or if the node
+     * begins deallocation on a background thread and it waiting for the -dealloc call to reach main), the only
+     * way to avoid a dangling pointer is to use a weak proxy.
+     */
+    ASWeakProxy *instance = [ASWeakProxy weakProxyWithTarget:self];
+    _layer.delegate = (id<CALayerDelegate>)instance;
+    objc_setAssociatedObject(_layer, &ASLayerDelegateAssociationKey, instance, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+  } else {
+    TIME_SCOPED(_debugTimeToCreateView);
+    _view = [self _locked_viewToLoad];
+    _view.asyncdisplaykit_node = self;
+    _layer = _view.layer;
   }
+  _layer.asyncdisplaykit_node = self;
+  
+  self._locked_asyncLayer.asyncDelegate = self;
 }
 
 - (void)_didLoad
 {
+  ASDisplayNodeAssertMainThread();
+  ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
   ASDisplayNodeLogEvent(self, @"didLoad");
+  TIME_SCOPED(_debugTimeForDidLoad);
   
   [self didLoad];
   
@@ -729,14 +707,47 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
 
 - (UIView *)view
 {
+  ASDN::MutexLocker l(__instanceLock__);
+
   ASDisplayNodeAssert(!_flags.layerBacked, @"Call to -view undefined on layer-backed nodes");
-  if (_flags.layerBacked) {
+  BOOL isLayerBacked = _flags.layerBacked;
+  if (isLayerBacked) {
     return nil;
   }
 
-  if (_view == nil) {
-    ASDisplayNodeAssertMainThread();
-    [self _loadViewOrLayerIsLayerBacked:NO];
+  if (_view != nil) {
+    return _view;
+  }
+
+  if (![self _locked_shouldLoadViewOrLayer]) {
+    return nil;
+  }
+  
+  // Loading a view needs to happen on the main thread
+  ASDisplayNodeAssertMainThread();
+  [self _locked_loadViewOrLayerIsLayerBacked:isLayerBacked];
+  
+  // FIXME: Ideally we'd call this as soon as the node receives -setNeedsLayout
+  // but automatic subnode management would require us to modify the node tree
+  // in the background on a loaded node, which isn't currently supported.
+  if (_pendingViewState.hasSetNeedsLayout) {
+    // Need to unlock before calling setNeedsLayout to avoid deadlocks.
+    // MutexUnlocker will re-lock at the end of scope.
+    ASDN::MutexUnlocker u(__instanceLock__);
+    [self __setNeedsLayout];
+  }
+  
+  [self _locked_applyPendingStateToViewOrLayer];
+  
+  {
+    // The following methods should not be called with a lock
+    ASDN::MutexUnlocker u(__instanceLock__);
+
+    // No need for the lock as accessing the subviews or layers are always happening on main
+    [self _addSubnodeViewsAndLayers];
+    
+    // A subclass hook should never be called with a lock
+    [self _didLoad];
   }
 
   return _view;
@@ -744,14 +755,47 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
 
 - (CALayer *)layer
 {
-  if (_layer == nil) {
-    ASDisplayNodeAssertMainThread();
+  ASDN::MutexLocker l(__instanceLock__);
+  if (_layer != nil) {
+    return _layer;
+  }
+  
+  BOOL isLayerBacked = _flags.layerBacked;
+  if (!isLayerBacked) {
+    // No need for the lock and call the view explicitly in case it needs to be loaded first
+    ASDN::MutexUnlocker u(__instanceLock__);
+    return self.view.layer;
+  }
+  
+  if (![self _locked_shouldLoadViewOrLayer]) {
+    return nil;
+  }
+  
+  // Loading a layer needs to happen on the main thread
+  ASDisplayNodeAssertMainThread();
+  [self _locked_loadViewOrLayerIsLayerBacked:isLayerBacked];
+  
+  // FIXME: Ideally we'd call this as soon as the node receives -setNeedsLayout
+  // but automatic subnode management would require us to modify the node tree
+  // in the background on a loaded node, which isn't currently supported.
+  if (_pendingViewState.hasSetNeedsLayout) {
+    // Need to unlock before calling setNeedsLayout to avoid deadlocks.
+    // MutexUnlocker will re-lock at the end of scope.
+    ASDN::MutexUnlocker u(__instanceLock__);
+    [self __setNeedsLayout];
+  }
+  
+  [self _locked_applyPendingStateToViewOrLayer];
+  
+  {
+    // The following methods should not be called with a lock
+    ASDN::MutexUnlocker u(__instanceLock__);
+
+    // No need for the lock as accessing the subviews or layers are always happening on main
+    [self _addSubnodeViewsAndLayers];
     
-    if (!_flags.layerBacked) {
-      return self.view.layer;
-    }
-    
-    [self _loadViewOrLayerIsLayerBacked:YES];
+    // A subclass hook should never be called with a lock
+    [self _didLoad];
   }
 
   return _layer;
@@ -2793,6 +2837,8 @@ ASDISPLAYNODE_INLINE BOOL nodeIsInRasterizedTree(ASDisplayNode *node) {
 {
   ASDisplayNodeAssertMainThread();
   
+  TIME_SCOPED(_debugTimeToAddSubnodeViews);
+  
   for (ASDisplayNode *node in self.subnodes) {
     [self _addSubnodeSubviewOrSublayer:node];
   }
@@ -3761,17 +3807,16 @@ ASDISPLAYNODE_INLINE BOOL nodeIsInRasterizedTree(ASDisplayNode *node) {
 
 #pragma mark - Pending View State
 
-- (void)_applyPendingStateToViewOrLayer
+- (void)_locked_applyPendingStateToViewOrLayer
 {
   ASDisplayNodeAssertMainThread();
   ASDisplayNodeAssert(self.nodeLoaded, @"must have a view or layer");
 
+  TIME_SCOPED(_debugTimeToApplyPendingState);
+  
   // If no view/layer properties were set before the view/layer were created, _pendingViewState will be nil and the default values
   // for the view/layer are still valid.
-
-  [self applyPendingViewState];
-  
-  __instanceLock__.lock();
+  [self _locked_applyPendingViewState];
   
   if (_flags.displaySuspended) {
     self._locked_asyncLayer.displaySuspended = YES;
@@ -3779,26 +3824,32 @@ ASDISPLAYNODE_INLINE BOOL nodeIsInRasterizedTree(ASDisplayNode *node) {
   if (!_flags.displaysAsynchronously) {
     self._locked_asyncLayer.displaysAsynchronously = NO;
   }
-  
-  __instanceLock__.unlock();
 }
 
 - (void)applyPendingViewState
 {
   ASDisplayNodeAssertMainThread();
+  ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
+  
   ASDN::MutexLocker l(__instanceLock__);
-
   // FIXME: Ideally we'd call this as soon as the node receives -setNeedsLayout
   // but automatic subnode management would require us to modify the node tree
   // in the background on a loaded node, which isn't currently supported.
   if (_pendingViewState.hasSetNeedsLayout) {
-    //Need to unlock before calling setNeedsLayout to avoid deadlocks.
-    //MutexUnlocker will re-lock at the end of scope.
+    // Need to unlock before calling setNeedsLayout to avoid deadlocks.
+    // MutexUnlocker will re-lock at the end of scope.
     ASDN::MutexUnlocker u(__instanceLock__);
     [self __setNeedsLayout];
   }
+  
+  [self _locked_applyPendingViewState];
+}
 
-  if (self.layerBacked) {
+- (void)_locked_applyPendingViewState
+{
+  ASDisplayNodeAssertMainThread();
+
+  if (_flags.layerBacked) {
     [_pendingViewState applyToLayer:self.layer];
   } else {
     BOOL specialPropertiesHandling = ASDisplayNodeNeedsSpecialPropertiesHandlingForFlags(_flags);