[ASDisplayNode] Use Weak Proxy to Avoid Dangling CALayer.delegate (#2249)

* Add weak proxy between node and layer to avoid dangling layer problem

* Add failing test case for dangling CALayer.delegate pointer issue

* Add docs

* Remove asynctransactioncontainer forwarding preprocessor macro

* Improve comments

* Remove asyncTransactionContainerStateDidChange callback

AsyncDisplayKit/ASDisplayNode.mm

@@ -431,8 +431,6 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
 
   _view = nil;
   _subnodes = nil;
-  if (_flags.layerBacked)
-    _layer.delegate = nil;
   _layer = nil;
 
   // TODO: Remove this? If supernode isn't already nil, this method isn't dealloc-safe anyway.
@@ -448,7 +446,7 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
 
 - (void)__unloadNode
 {
-  ASDisplayNodeAssertThreadAffinity(self);
+  ASDisplayNodeAssertMainThread();
   ASDisplayNodeAssert([self isNodeLoaded], @"Implementation shouldn't call __unloadNode if not loaded: %@", self);
   ASDisplayNodeAssert(_flags.synchronous == NO, @"Node created using -initWithViewBlock:/-initWithLayerBlock: cannot be unloaded. Node: %@", self);
   ASDN::MutexLocker l(__instanceLock__);
@@ -545,8 +543,17 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   if (isLayerBacked) {
     TIME_SCOPED(_debugTimeToCreateView);
     _layer = [self _layerToLoad];
-    // Surpress warning for Base SDK > 10.0
-    _layer.delegate = (id<CALayerDelegate>)self;
+    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 _viewToLoad];

AsyncDisplayKit/Details/ASWeakProxy.m

@@ -11,6 +11,8 @@
 //
 
 #import "ASWeakProxy.h"
+#import "ASObjectDescriptionHelpers.h"
+#import "ASAssert.h"
 
 @implementation ASWeakProxy
 
@@ -32,24 +34,37 @@
   return _target;
 }
 
-- (NSMethodSignature *)methodSignatureForSelector:(SEL)aSelector
+- (BOOL)respondsToSelector:(SEL)aSelector
 {
-  // Check for a compiled definition for the selector
-  NSMethodSignature *methodSignature = [[_target class] instanceMethodSignatureForSelector:aSelector];
-  
+  return [_target respondsToSelector:aSelector];
+}
+
+/// Strangely, this method doesn't get forwarded by ObjC.
+- (BOOL)isKindOfClass:(Class)aClass
+{
+  return [_target isKindOfClass:aClass];
+}
+
+- (NSString *)description
+{
+  return ASObjectDescriptionMake(self, @[@{ @"target": _target ?: (id)kCFNull }]);
+}
+
+- (NSMethodSignature *)methodSignatureForSelector:(SEL)sel
+{
+  ASDisplayNodeAssertNil(_target, @"ASWeakProxy got %@ when its target is still alive, which is unexpected.", NSStringFromSelector(_cmd));
   // Unfortunately, in order to get this object to work properly, the use of a method which creates an NSMethodSignature
   // from a C string. -methodSignatureForSelector is called when a compiled definition for the selector cannot be found.
   // This is the place where we have to create our own dud NSMethodSignature. This is necessary because if this method
   // returns nil, a selector not found exception is raised. The string argument to -signatureWithObjCTypes: outlines
   // the return type and arguments to the message. To return a dud NSMethodSignature, pretty much any signature will
-  // suffice. Since the -forwardInvocation call will do nothing if the delegate does not respond to the selector,
+  // suffice. Since the -forwardInvocation call will do nothing if the target does not respond to the selector,
   // the dud NSMethodSignature simply gets us around the exception.
-  return methodSignature ?: [NSMethodSignature signatureWithObjCTypes:"@^v^c"];
+  return [NSMethodSignature signatureWithObjCTypes:"@^v^c"];
 }
-
 - (void)forwardInvocation:(NSInvocation *)invocation
 {
-  // If we are down here this means _target where nil. Just don't do anything to prevent a crash
+  ASDisplayNodeAssertNil(_target, @"ASWeakProxy got %@ when its target is still alive, which is unexpected.", NSStringFromSelector(_cmd));
 }
 
 @end

AsyncDisplayKit/Details/Transactions/_ASAsyncTransactionContainer.h

@@ -47,12 +47,6 @@ typedef NS_ENUM(NSUInteger, ASAsyncTransactionContainerState) {
  */
 - (void)asyncdisplaykit_cancelAsyncTransactions;
 
-/**
- @summary Invoked when the asyncdisplaykit_asyncTransactionContainerState property changes.
- @desc You may want to override this in a CALayer or UIView subclass to take appropriate action (such as hiding content while it renders).
- */
-- (void)asyncdisplaykit_asyncTransactionContainerStateDidChange;
-
 @end
 
 @interface CALayer (ASDisplayNodeAsyncTransactionContainer) <ASDisplayNodeAsyncTransactionContainer>

AsyncDisplayKit/Details/Transactions/_ASAsyncTransactionContainer.m

@@ -14,10 +14,6 @@
 #import "_ASAsyncTransactionGroup.h"
 #import <objc/runtime.h>
 
-#ifndef ASASYNCTRANSACTIONCONTAINER_FORWARD_STATE_CHANGE
-#define ASASYNCTRANSACTIONCONTAINER_FORWARD_STATE_CHANGE 1
-#endif
-
 static const char *ASDisplayNodeAssociatedTransactionsKey = "ASAssociatedTransactions";
 static const char *ASDisplayNodeAssociatedCurrentTransactionKey = "ASAssociatedCurrentTransaction";
 
@@ -83,16 +79,6 @@ static const char *ASAsyncTransactionIsContainerKey = "ASTransactionIsContainer"
   }
 }
 
-- (void)asyncdisplaykit_asyncTransactionContainerStateDidChange
-{
-#if ASASYNCTRANSACTIONCONTAINER_FORWARD_STATE_CHANGE
-  id delegate = self.delegate;
-  if ([delegate respondsToSelector:@selector(asyncdisplaykit_asyncTransactionContainerStateDidChange)]) {
-    [delegate asyncdisplaykit_asyncTransactionContainerStateDidChange];
-  }
-#endif
-}
-
 - (_ASAsyncTransaction *)asyncdisplaykit_asyncTransaction
 {
   _ASAsyncTransaction *transaction = self.asyncdisplaykit_currentAsyncLayerTransaction;
@@ -102,19 +88,18 @@ static const char *ASAsyncTransactionIsContainerKey = "ASTransactionIsContainer"
       transactions = [NSHashTable hashTableWithOptions:NSPointerFunctionsObjectPointerPersonality];
       self.asyncdisplaykit_asyncLayerTransactions = transactions;
     }
+    __weak CALayer *weakSelf = self;
     transaction = [[_ASAsyncTransaction alloc] initWithCallbackQueue:dispatch_get_main_queue() completionBlock:^(_ASAsyncTransaction *completedTransaction, BOOL cancelled) {
+      __strong CALayer *self = weakSelf;
+      if (self == nil) {
+        return;
+      }
       [transactions removeObject:completedTransaction];
       [self asyncdisplaykit_asyncTransactionContainerDidCompleteTransaction:completedTransaction];
-      if ([transactions count] == 0) {
-        [self asyncdisplaykit_asyncTransactionContainerStateDidChange];
-      }
     }];
     [transactions addObject:transaction];
     self.asyncdisplaykit_currentAsyncLayerTransaction = transaction;
     [self asyncdisplaykit_asyncTransactionContainerWillBeginTransaction:transaction];
-    if ([transactions count] == 1) {
-      [self asyncdisplaykit_asyncTransactionContainerStateDidChange];
-    }
   }
   [[_ASAsyncTransactionGroup mainTransactionGroup] addTransactionContainer:self];
   return transaction;

AsyncDisplayKit/Details/_ASDisplayView.mm

@@ -330,12 +330,6 @@
 }
 #endif
 
-- (void)asyncdisplaykit_asyncTransactionContainerStateDidChange
-{
-  ASDisplayNode *node = _asyncdisplaykit_node; // Create strong reference to weak ivar.
-  [node asyncdisplaykit_asyncTransactionContainerStateDidChange];
-}
-
 - (void)tintColorDidChange
 {
     ASDisplayNode *node = _asyncdisplaykit_node; // Create strong reference to weak ivar.

AsyncDisplayKit/Private/ASDisplayNode+UIViewBridge.mm

@@ -962,8 +962,4 @@ nodeProperty = nodeValueExpr; _setToViewOnly(viewAndPendingViewStateProperty, vi
   [_layer asyncdisplaykit_cancelAsyncTransactions];
 }
 
-- (void)asyncdisplaykit_asyncTransactionContainerStateDidChange
-{
-}
-
 @end

AsyncDisplayKitTests/ASDisplayNodeTests.m

@@ -1964,6 +1964,25 @@ static bool stringContainsPointer(NSString *description, id p) {
   XCTAssertNoThrow([nodeView removeFromSuperview]);
 }
 
+// Running on main thread
+// Cause retain count of node to fall to zero synchronously on a background thread (pausing main thread)
+// ASDealloc2MainObject queues actual call to -dealloc to occur on the main thread
+// Continue execution on main, before the dealloc can run, to dealloc the host view
+// Node is in an invalid state (about to dealloc, not valid to retain) but accesses to sublayer delegates
+// causes attempted retain — unless weak variable works correctly
+- (void)testThatLayerDelegateDoesntDangleAndCauseCrash
+{
+  NS_VALID_UNTIL_END_OF_SCOPE UIView *host = [[UIView alloc] init];
+
+  __block NS_VALID_UNTIL_END_OF_SCOPE ASDisplayNode *node = [[ASDisplayNode alloc] init];
+  node.layerBacked = YES;
+
+  [host addSubnode:node];
+  [self executeOffThread:^{
+    node = nil;
+  }];
+  host = nil; // <- Would crash here, when UIView accesses its sublayers' delegates in -dealloc.
+}
 
 - (void)testThatSubnodeGetsInterfaceStateSetIfRasterized
 {