[ASLayoutableContext] Properly lock layoutableContextMap

AsyncDisplayKit/ASDisplayNode.mm

@@ -630,9 +630,11 @@ - (BOOL)shouldMeasureWithSizeRange:(ASSizeRange)constrainedSize
     return NO;
   }
 
-  if (ASHierarchyStateIncludesLayoutPending(_hierarchyState)
-      && _pendingTransitionID != ASLayoutableGetCurrentContext().transitionID) {
-    return NO;
+  if (ASHierarchyStateIncludesLayoutPending(_hierarchyState)) {
+    ASLayoutableContext context =  ASLayoutableGetCurrentContext();
+    if (ASLayoutableContextIsNull(context) || _pendingTransitionID != context.transitionID) {
+      return NO;
+    }
   }
 
   // only calculate the size if
@@ -2226,7 +2228,7 @@ - (void)setHierarchyState:(ASHierarchyState)newState
       // Leaving layout pending state, reset related properties
       {
         ASDN::MutexLocker l(_propertyLock);
-        _pendingTransitionID = 0;
+        _pendingTransitionID = ASLayoutableContextInvalidTransitionID;
         _pendingLayoutContext = nil;
       }
     }

AsyncDisplayKit/Layout/ASLayoutable.mm

@@ -9,15 +9,40 @@
 #import "ASLayoutablePrivate.h"
 #import "pthread.h"
 #import <map>
+#import <iterator>
+#import "ASThread.h"
 
-ASLayoutableContext ASLayoutableContextMake(int32_t transitionID, BOOL needsVisualizeNode)
+int32_t const ASLayoutableContextInvalidTransitionID = 0;
+int32_t const ASLayoutableContextDefaultTransitionID = ASLayoutableContextInvalidTransitionID + 1;
+
+static inline ASLayoutableContext _ASLayoutableContextMake(int32_t transitionID, BOOL needsVisualizeNode)
 {
   struct ASLayoutableContext context;
   context.transitionID = transitionID;
   context.needsVisualizeNode = needsVisualizeNode;
   return context;
 }
 
+static inline BOOL _IsValidTransitionID(int32_t transitionID)
+{
+  return transitionID > ASLayoutableContextInvalidTransitionID;
+}
+
+struct ASLayoutableContext const ASLayoutableContextNull = _ASLayoutableContextMake(ASLayoutableContextInvalidTransitionID, NO);
+
+BOOL ASLayoutableContextIsNull(struct ASLayoutableContext context)
+{
+  return !_IsValidTransitionID(context.transitionID);
+}
+
+ASLayoutableContext ASLayoutableContextMake(int32_t transitionID, BOOL needsVisualizeNode)
+{
+  NSCAssert(_IsValidTransitionID(transitionID), @"Invalid transition ID");
+  return _ASLayoutableContextMake(transitionID, needsVisualizeNode);
+}
+
+// Note: This is a non-recursive static lock. If it needs to be recursive, use ASDISPLAYNODE_MUTEX_RECURSIVE_INITIALIZER
+static ASDN::StaticMutex _layoutableContextLock = ASDISPLAYNODE_MUTEX_INITIALIZER;
 static std::map<mach_port_t, ASLayoutableContext> layoutableContextMap;
 
 static inline mach_port_t ASLayoutableGetCurrentContextKey()
@@ -27,14 +52,26 @@ static inline mach_port_t ASLayoutableGetCurrentContextKey()
 
 void ASLayoutableSetCurrentContext(struct ASLayoutableContext context)
 {
-  layoutableContextMap[ASLayoutableGetCurrentContextKey()] = context;
+  const mach_port_t key = ASLayoutableGetCurrentContextKey();
+  ASDN::StaticMutexLocker l(_layoutableContextLock);
+  layoutableContextMap[key] = context;
 }
 
 struct ASLayoutableContext ASLayoutableGetCurrentContext()
 {
-  return layoutableContextMap[ASLayoutableGetCurrentContextKey()];
+  const mach_port_t key = ASLayoutableGetCurrentContextKey();
+  ASDN::StaticMutexLocker l(_layoutableContextLock);
+  const auto it = layoutableContextMap.find(key);
+  if (it != layoutableContextMap.end()) {
+    // Found an interator with above key. "it->first" is the key itself, "it->second" is the context value.
+    return it->second;
+  }
+  return ASLayoutableContextNull;
 }
 
-void ASLayoutableClearCurrentContext() {
-  layoutableContextMap.erase(ASLayoutableGetCurrentContextKey());
+void ASLayoutableClearCurrentContext()
+{
+  const mach_port_t key = ASLayoutableGetCurrentContextKey();
+  ASDN::StaticMutexLocker l(_layoutableContextLock);
+  layoutableContextMap.erase(key);
 }

AsyncDisplayKit/Layout/ASLayoutablePrivate.h

@@ -19,6 +19,14 @@ struct ASLayoutableContext {
   BOOL needsVisualizeNode;
 };
 
+extern int32_t const ASLayoutableContextInvalidTransitionID;
+
+extern int32_t const ASLayoutableContextDefaultTransitionID;
+
+extern struct ASLayoutableContext const ASLayoutableContextNull;
+
+extern BOOL ASLayoutableContextIsNull(struct ASLayoutableContext context);
+
 extern struct ASLayoutableContext ASLayoutableContextMake(int32_t transitionID, BOOL needsVisualizeNode);
 
 extern void ASLayoutableSetCurrentContext(struct ASLayoutableContext context);