Lock up to yogaRoot during layout to avoid deadlocks. (#1356)

* Lock up to yogaRoot during layout to avoid dead lock.

1) lock to root for tree
2) lock self to change parent (& consequently root)
3) Implement ASLocking (tryLock) on ASNodeController
4) add lockPair to try-lock node & controller together
5) lock controllers if they exist in lockToRoot...

Disable some asserts due to lock to root. :(

LL# No commands remaining.

* Add macro so non-Yoga still builds :)

* wut

Source/ASDisplayNode+Layout.mm

@@ -16,6 +16,8 @@
 #import <AsyncDisplayKit/ASLayout.h>
 #import <AsyncDisplayKit/ASLayoutElementStylePrivate.h>
 #import <AsyncDisplayKit/ASLog.h>
+#import <AsyncDisplayKit/ASNodeController+Beta.h>
+#import <AsyncDisplayKit/ASDisplayNode+Yoga.h>
 
 #pragma mark - ASDisplayNode (ASLayoutElement)
 
@@ -65,7 +67,7 @@
 
 - (ASLayout *)layoutThatFits:(ASSizeRange)constrainedSize parentSize:(CGSize)parentSize
 {
-  ASDN::MutexLocker l(__instanceLock__);
+  ASScopedLockSelfOrToRoot();
 
   // If one or multiple layout transitions are in flight it still can happen that layout information is requested
   // on other threads. As the pending and calculated layout to be updated in the layout transition in here just a
@@ -240,11 +242,11 @@ ASLayoutElementStyleExtensibilityForwarding
   }
 }
 
+// TODO It would be easier to work with if we could `ASAssertUnlocked` here, but we
+// cannot due to locking to root in `_u_measureNodeWithBoundsIfNecessary`.
 - (void)_rootNodeDidInvalidateSize
 {
   ASDisplayNodeAssertThreadAffinity(self);
-  ASAssertUnlocked(__instanceLock__);
-  
   __instanceLock__.lock();
   
   // We are the root node and need to re-flow the layout; at least one child needs a new size.
@@ -270,11 +272,21 @@ ASLayoutElementStyleExtensibilityForwarding
   }
 }
 
+// TODO
+// We should remove this logic, which is relatively new, and instead
+// rely on the parent / host of the root node to do this size change. That's always been the
+// expectation with other node containers like ASTableView, ASCollectionView, ASViewController, etc.
+// E.g. in ASCellNode the _interactionDelegate is a Table or Collection that will resize in this
+// case. By resizing without participating with the parent, we could get cases where our parent size
+// does not match, especially if there is a size constraint that is applied at that level.
+//
+// In general a node should never need to set its own size, instead allowing its parent to do so -
+// even in the root case. Anyhow this is a separate / pre-existing issue, but I think it could be
+// causing real issues in cases of resizing nodes.
 - (void)displayNodeDidInvalidateSizeNewSize:(CGSize)size
 {
   ASDisplayNodeAssertThreadAffinity(self);
-  ASAssertUnlocked(__instanceLock__);
+
-  
   // The default implementation of display node changes the size of itself to the new size
   CGRect oldBounds = self.bounds;
   CGSize oldSize = oldBounds.size;
@@ -295,9 +307,9 @@ ASLayoutElementStyleExtensibilityForwarding
 
 - (void)_u_measureNodeWithBoundsIfNecessary:(CGRect)bounds
 {
-  ASAssertUnlocked(__instanceLock__);
+  // ASAssertUnlocked(__instanceLock__);
-  
+  ASScopedLockSelfOrToRoot();
-  ASDN::MutexLocker l(__instanceLock__);
+
   // Check if we are a subnode in a layout transition.
   // In this case no measurement is needed as it's part of the layout transition
   if ([self _locked_isLayoutTransitionInvalid]) {
@@ -455,7 +467,7 @@ ASLayoutElementStyleExtensibilityForwarding
 - (void)_layoutSublayouts
 {
   ASDisplayNodeAssertThreadAffinity(self);
-  ASAssertUnlocked(__instanceLock__);
+  // ASAssertUnlocked(__instanceLock__);
   
   ASLayout *layout;
   {
@@ -620,7 +632,7 @@ ASLayoutElementStyleExtensibilityForwarding
     NSUInteger newLayoutVersion = _layoutVersion;
     ASLayout *newLayout;
     {
-      ASDN::MutexLocker l(__instanceLock__);
+      ASScopedLockSelfOrToRoot();
 
       ASLayoutElementContext *ctx = [[ASLayoutElementContext alloc] init];
       ctx.transitionID = transitionID;
@@ -1009,3 +1021,61 @@ ASLayoutElementStyleExtensibilityForwarding
 }
 
 @end
+
+#pragma mark -
+#pragma mark - ASDisplayNode (YogaLayout)
+
+@implementation ASDisplayNode (YogaInternal)
+
+#pragma mark -
+#pragma mark - ASDisplayNode (Yoga)
+
+- (BOOL)locked_shouldLayoutFromYogaRoot {
+#if YOGA
+  YGNodeRef yogaNode = _style.yogaNode;
+  BOOL hasYogaParent = (_yogaParent != nil);
+  BOOL hasYogaChildren = (_yogaChildren.count > 0);
+  BOOL usesYoga = (yogaNode != NULL && (hasYogaParent || hasYogaChildren));
+  if (usesYoga) {
+    if ([self shouldHaveYogaMeasureFunc] == NO) {
+      return YES;
+    } else {
+      return NO;
+    }
+  } else {
+    return NO;
+  }
+#else
+  return NO;
+#endif
+}
+
+- (ASLockSet)lockToRootIfNeededForLayout {
+  ASLockSet lockSet = ASLockSequence(^BOOL(ASAddLockBlock addLock) {
+    if (!addLock(self)) {
+      return NO;
+    }
+    if (self.nodeController && !addLock(self.nodeController)) {
+      return NO;
+    }
+#if YOGA
+    if (![self locked_shouldLayoutFromYogaRoot]) {
+      return YES;
+    }
+    ASDisplayNode *parent = _supernode;
+    while (parent) {
+      if (!addLock(parent)) {
+        return NO;
+      }
+      if (parent.nodeController && !addLock(parent.nodeController)) {
+        return NO;
+      }
+      parent = parent->_supernode;
+    }
+#endif
+    return true;
+  });
+  return lockSet;
+}
+
+@end

Source/ASDisplayNode+Yoga.h

@@ -33,6 +33,12 @@ AS_EXTERN void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullab
 // Will walk up the Yoga tree and returns the root node
 - (ASDisplayNode *)yogaRoot;
 
+/**
+ * @discussion Attempts(spinning) to lock all node up to root node when yoga is enabled.
+ * This will lock self when yoga is not enabled;
+ */
+- (ASLockSet)lockToRootIfNeededForLayout;
+
 @end
 
 
@@ -47,6 +53,11 @@ AS_EXTERN void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullab
 - (void)calculateLayoutFromYogaRoot:(ASSizeRange)rootConstrainedSize;
 /// For internal usage only
 - (void)invalidateCalculatedYogaLayout;
+/**
+ * @discussion return true only when yoga enabled and the node is in yoga tree and the node is
+ * not leaf that implemented measure function.
+ */
+- (BOOL)locked_shouldLayoutFromYogaRoot;
 
 @end
 
@@ -79,4 +90,9 @@ AS_EXTERN void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullab
 
 NS_ASSUME_NONNULL_END
 
+// When Yoga is enabled, there are several points where we want to lock the tree to the root but otherwise (without Yoga)
+// will want to simply lock self.
+#define ASScopedLockSelfOrToRoot() ASScopedLockSet lockSet = [self lockToRootIfNeededForLayout]
+#else
+#define ASScopedLockSelfOrToRoot() ASLockScopeSelf()
 #endif

Source/ASDisplayNode+Yoga.mm

@@ -46,7 +46,7 @@
 
 - (void)setYogaChildren:(NSArray *)yogaChildren
 {
-  ASLockScope(self.yogaRoot);
+  ASScopedLockSelfOrToRoot();
   for (ASDisplayNode *child in [_yogaChildren copy]) {
     // Make sure to un-associate the YGNodeRef tree before replacing _yogaChildren
     // If this becomes a performance bottleneck, it can be optimized by not doing the NSArray removals here.
@@ -66,7 +66,7 @@
 
 - (void)addYogaChild:(ASDisplayNode *)child
 {
-  ASLockScope(self.yogaRoot);
+  ASScopedLockSelfOrToRoot();
   [self _locked_addYogaChild:child];
 }
 
@@ -77,7 +77,7 @@
 
 - (void)removeYogaChild:(ASDisplayNode *)child
 {
-  ASLockScope(self.yogaRoot);
+  ASScopedLockSelfOrToRoot();
   [self _locked_removeYogaChild:child];
 }
 
@@ -95,7 +95,7 @@
 
 - (void)insertYogaChild:(ASDisplayNode *)child atIndex:(NSUInteger)index
 {
-  ASLockScope(self.yogaRoot);
+  ASScopedLockSelfOrToRoot();
   [self _locked_insertYogaChild:child atIndex:index];
 }
 
@@ -129,6 +129,7 @@
 
 - (void)setYogaParent:(ASDisplayNode *)yogaParent
 {
+  ASLockScopeSelf();
   if (_yogaParent == yogaParent) {
     return;
   }
@@ -184,7 +185,7 @@
 
 - (void)setupYogaCalculatedLayout
 {
-  ASLockScopeSelf();
+  ASScopedLockSelfOrToRoot();
 
   YGNodeRef yogaNode = self.style.yogaNode;
   uint32_t childCount = YGNodeGetChildCount(yogaNode);
@@ -194,7 +195,7 @@
 
   ASLayout *rawSublayouts[childCount];
   int i = 0;
-  for (ASDisplayNode *subnode in self.yogaChildren) {
+  for (ASDisplayNode *subnode in _yogaChildren) {
     rawSublayouts[i++] = [subnode layoutForYogaNode];
   }
   const auto sublayouts = [NSArray<ASLayout *> arrayByTransferring:rawSublayouts count:childCount];
@@ -251,10 +252,11 @@
 
 - (BOOL)shouldHaveYogaMeasureFunc
 {
+  ASLockScopeSelf();
   // Size calculation via calculateSizeThatFits: or layoutSpecThatFits:
   // For these nodes, we assume they may need custom Baseline calculation too.
   // This will be used for ASTextNode, as well as any other node that has no Yoga children
-  BOOL isLeafNode = (self.yogaChildren.count == 0);
+  BOOL isLeafNode = (_yogaChildren.count == 0);
   BOOL definesCustomLayout = [self implementsLayoutMethod];
   return (isLeafNode && definesCustomLayout);
 }
@@ -296,31 +298,24 @@
   // - This node is a Yoga tree root: it has no yogaParent, but has yogaChildren.
   // - This node is a Yoga tree node: it has both a yogaParent and yogaChildren.
   // - This node is a Yoga tree leaf: it has a yogaParent, but no yogaChidlren.
-  YGNodeRef yogaNode = _style.yogaNode;
+  if ([self locked_shouldLayoutFromYogaRoot]) {
-  BOOL hasYogaParent = (_yogaParent != nil);
+    // If we're a yoga root, tree node, or leaf with no measure func (e.g. spacer), then
-  BOOL hasYogaChildren = (_yogaChildren.count > 0);
+    // initiate a new Yoga calculation pass from root.
-  BOOL usesYoga = (yogaNode != NULL && (hasYogaParent || hasYogaChildren));
+    as_activity_create_for_scope("Yoga layout calculation");
-  if (usesYoga) {
+    if (self.yogaLayoutInProgress == NO) {
-    // This node has some connection to a Yoga tree.
+      ASYogaLog("Calculating yoga layout from root %@, %@", self,
-    if ([self shouldHaveYogaMeasureFunc] == NO) {
+              NSStringFromASSizeRange(constrainedSize));
-      // If we're a yoga root, tree node, or leaf with no measure func (e.g. spacer), then
+      [self calculateLayoutFromYogaRoot:constrainedSize];
-      // initiate a new Yoga calculation pass from root.
-
-      as_activity_create_for_scope("Yoga layout calculation");
-      if (self.yogaLayoutInProgress == NO) {
-        ASYogaLog("Calculating yoga layout from root %@, %@", self, NSStringFromASSizeRange(constrainedSize));
-        l.unlock();
-        [self calculateLayoutFromYogaRoot:constrainedSize];
-        l.lock();
-      } else {
-        ASYogaLog("Reusing existing yoga layout %@", _yogaCalculatedLayout);
-      }
-      ASDisplayNodeAssert(_yogaCalculatedLayout, @"Yoga node should have a non-nil layout at this stage: %@", self);
-      return _yogaCalculatedLayout;
     } else {
-      // If we're a yoga leaf node with custom measurement function, proceed with normal layout so layoutSpecs can run (e.g. ASButtonNode).
+      ASYogaLog("Reusing existing yoga layout %@", _yogaCalculatedLayout);
-      ASYogaLog("PROCEEDING past Yoga check to calculate ASLayout for: %@", self);
     }
+    ASDisplayNodeAssert(_yogaCalculatedLayout,
+            @"Yoga node should have a non-nil layout at this stage: %@", self);
+    return _yogaCalculatedLayout;
+  } else {
+    // If we're a yoga leaf node with custom measurement function, proceed with normal layout so
+    // layoutSpecs can run (e.g. ASButtonNode).
+    ASYogaLog("PROCEEDING past Yoga check to calculate ASLayout for: %@", self);
   }
 
   // Delegate to layout spec layout for nodes that do not support Yoga
@@ -345,7 +340,6 @@
     }
   }];
 
-  ASLockScopeSelf();
 
   // Prepare all children for the layout pass with the current Yoga tree configuration.
   ASDisplayNodePerformBlockOnEveryYogaChild(self, ^(ASDisplayNode *_Nonnull node) {

Source/ASDisplayNode.mm

@@ -1046,7 +1046,7 @@ ASSynthesizeLockingMethodsWithMutex(__instanceLock__);
 - (void)__layout
 {
   ASDisplayNodeAssertThreadAffinity(self);
-  ASAssertUnlocked(__instanceLock__);
+  // ASAssertUnlocked(__instanceLock__);
   
   BOOL loaded = NO;
   {
@@ -1096,7 +1096,7 @@ ASSynthesizeLockingMethodsWithMutex(__instanceLock__);
 - (void)_layoutDidFinish
 {
   ASDisplayNodeAssertMainThread();
-  ASAssertUnlocked(__instanceLock__);
+  // ASAssertUnlocked(__instanceLock__);
   ASDisplayNodeAssertTrue(self.isNodeLoaded);
   [self layoutDidFinish];
 }
@@ -1169,7 +1169,7 @@ ASSynthesizeLockingMethodsWithMutex(__instanceLock__);
 {
   // Hook for subclasses
   ASDisplayNodeAssertMainThread();
-  ASAssertUnlocked(__instanceLock__);
+  // ASAssertUnlocked(__instanceLock__);
   ASDisplayNodeAssertTrue(self.isNodeLoaded);
   [self enumerateInterfaceStateDelegates:^(id<ASInterfaceStateDelegate> del) {
     [del nodeDidLayout];
@@ -2684,7 +2684,6 @@ ASDISPLAYNODE_INLINE BOOL subtreeIsRasterized(ASDisplayNode *node) {
 {
   ASDisplayNodeAssertMainThread();
   ASDisplayNodeAssert(!_flags.isEnteringHierarchy, @"Should not cause recursive __enterHierarchy");
-  ASAssertUnlocked(__instanceLock__);
   ASDisplayNodeLogEvent(self, @"enterHierarchy");
   
   // Profiling has shown that locking this method is beneficial, so each of the property accesses don't have to lock and unlock.
@@ -2733,7 +2732,6 @@ ASDISPLAYNODE_INLINE BOOL subtreeIsRasterized(ASDisplayNode *node) {
 {
   ASDisplayNodeAssertMainThread();
   ASDisplayNodeAssert(!_flags.isExitingHierarchy, @"Should not cause recursive __exitHierarchy");
-  ASAssertUnlocked(__instanceLock__);
   ASDisplayNodeLogEvent(self, @"exitHierarchy");
   
   // Profiling has shown that locking this method is beneficial, so each of the property accesses don't have to lock and unlock.

Source/ASNodeController+Beta.h

@@ -12,7 +12,7 @@
 
 /* ASNodeController is currently beta and open to change in the future */
 @interface ASNodeController<__covariant DisplayNodeType : ASDisplayNode *>
-    : NSObject <ASInterfaceStateDelegate, NSLocking>
+    : NSObject <ASInterfaceStateDelegate, ASLocking>
 
 @property (nonatomic, strong /* may be weak! */) DisplayNodeType node;
 
@@ -44,6 +44,12 @@
 
 - (void)hierarchyDisplayDidFinish ASDISPLAYNODE_REQUIRES_SUPER;
 
+/**
+ * @discussion Attempts (via ASLockSequence, a backing-off spinlock similar to
+ * std::lock()) to lock both the node and its ASNodeController, if one exists.
+ */
+- (ASLockSet)lockPair;
+
 @end
 
 @interface ASDisplayNode (ASNodeController)

Source/ASNodeController+Beta.mm

@@ -89,6 +89,20 @@
 
 - (void)hierarchyDisplayDidFinish {}
 
+- (ASLockSet)lockPair {
+  ASLockSet lockSet = ASLockSequence(^BOOL(ASAddLockBlock addLock) {
+    if (!addLock(_node)) {
+      return NO;
+    }
+    if (!addLock(self)) {
+      return NO;
+    }
+    return YES;
+  });
+
+  return lockSet;
+}
+
 #pragma mark NSLocking
 
 - (void)lock
@@ -101,6 +115,11 @@
   __instanceLock__.unlock();
 }
 
+- (BOOL)tryLock
+{
+  return __instanceLock__.try_lock();
+}
+
 @end
 
 @implementation ASDisplayNode (ASNodeController)

Source/ASScrollNode.mm

@@ -10,10 +10,12 @@
 #import <AsyncDisplayKit/ASScrollNode.h>
 #import <AsyncDisplayKit/ASDisplayNodeExtras.h>
 #import <AsyncDisplayKit/ASDisplayNode+FrameworkPrivate.h>
+#import <AsyncDisplayKit/ASDisplayNode+Beta.h>
 #import <AsyncDisplayKit/ASDisplayNode+Subclasses.h>
 #import <AsyncDisplayKit/ASLayout.h>
 #import <AsyncDisplayKit/_ASDisplayLayer.h>
 #import <AsyncDisplayKit/ASThread.h>
+#import <AsyncDisplayKit/ASDisplayNode+Yoga.h>
 
 @interface ASScrollView : UIScrollView
 @end
@@ -79,7 +81,7 @@
                      restrictedToSize:(ASLayoutElementSize)size
                  relativeToParentSize:(CGSize)parentSize
 {
-  ASLockScopeSelf();  // Lock for using our instance variables.
+  ASScopedLockSelfOrToRoot();
 
   ASSizeRange contentConstrainedSize = constrainedSize;
   if (ASScrollDirectionContainsVerticalDirection(_scrollableDirections)) {