More on ASDataController's main-thread-only mode (TextureGroup#1915)

Follow up on TextureGroup#1911: it's not enough to execute step 3 on the main thread because -_allocateNodesFromElements: uses ASDispatchApply to offload the work to other threads. So this diff adds a flag to tell that method to do everything serially on the calling thread.

Author: nguyenhuy

Source/Details/ASDataController.mm

@@ -129,7 +129,16 @@ - (void)setLayoutDelegate:(id<ASDataControllerLayoutDelegate>)layoutDelegate
 
 #pragma mark - Cell Layout
 
+/**
+ * Allocates and layouts nodes from the given collection elements, and blocks the current thread while doing so.
+ *
+ * @param elements The elements from which nodes can be allocated and laid out.
+ * @param strictlyOnCurrentThread Whether or not all the work must be done strictly on the current thread.
+ * YES means all nodes will be allocated and laid out serially on the current thread.
+ * NO means the work can be offloaded to other thread(s), potentially reduce the blocking time on the calling thread.
+ */
 - (void)_allocateNodesFromElements:(NSArray<ASCollectionElement *> *)elements
+           strictlyOnCurrentThread:(BOOL)strictlyOnCurrentThread
 {
   NSUInteger nodeCount = elements.count;
   __weak id<ASDataControllerSource> weakDataSource = _dataSource;
@@ -142,12 +151,7 @@ - (void)_allocateNodesFromElements:(NSArray<ASCollectionElement *> *)elements
   {
     as_activity_create_for_scope("Data controller batch");
 
-    dispatch_queue_t queue = dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0);
-    NSUInteger threadCount = 0;
-    if ([_dataSource dataControllerShouldSerializeNodeCreation:self]) {
-      threadCount = 1;
-    }
-    ASDispatchApply(nodeCount, queue, threadCount, ^(size_t i) {
+    void(^work)(size_t) = ^(size_t i) {
       __strong id<ASDataControllerSource> strongDataSource = weakDataSource;
       if (strongDataSource == nil) {
         return;
@@ -166,7 +170,20 @@ - (void)_allocateNodesFromElements:(NSArray<ASCollectionElement *> *)elements
       if (ASSizeRangeHasSignificantArea(sizeRange)) {
         [self _layoutNode:node withConstrainedSize:sizeRange];
       }
-    });
+    };
+    
+    if (strictlyOnCurrentThread) {
+      for (NSUInteger i = 0; i < nodeCount; i++) {
+        work(i);
+      }
+    } else {
+      dispatch_queue_t queue = dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0);
+      NSUInteger threadCount = 0;
+      if ([_dataSource dataControllerShouldSerializeNodeCreation:self]) {
+        threadCount = 1;
+      }
+      ASDispatchApply(nodeCount, queue, threadCount, work);
+    }
   }
 
   ASSignpostEnd(DataControllerBatch, self, "count: %lu", (unsigned long)nodeCount);
@@ -620,8 +637,11 @@ - (void)updateWithChangeSet:(_ASHierarchyChangeSet *)changeSet
   Class<ASDataControllerLayoutDelegate> layoutDelegateClass = [self.layoutDelegate class];
 
   // Step 3: Call the layout delegate if possible. Otherwise, allocate and layout all elements
-  dispatch_block_t step3 = ^{
+  void (^step3)(BOOL) = ^(BOOL strictlyOnCurrentThread){
     if (canDelegate) {
+      // Don't pass strictlyOnCurrentThread to the layout delegate. Instead give it
+      // total control over its threading behavior, as long as it blocks the
+      // calling thread while preparing the layout (which is part of the API contract).
       [layoutDelegateClass calculateLayoutWithContext:layoutContext];
     } else {
       const auto elementsToProcess = [[NSMutableArray<ASCollectionElement *> alloc] init];
@@ -635,21 +655,24 @@ - (void)updateWithChangeSet:(_ASHierarchyChangeSet *)changeSet
           [elementsToProcess addObject:element];
         }
       }
-      [self _allocateNodesFromElements:elementsToProcess];
+      [self _allocateNodesFromElements:elementsToProcess
+               strictlyOnCurrentThread:strictlyOnCurrentThread];
     }
   };
 
   // Step 3 can be done on the main thread or on _editingTransactionQueue
   // depending on an experiment.
   BOOL mainThreadOnly = ASActivateExperimentalFeature(ASExperimentalMainThreadOnlyDataController);
   if (mainThreadOnly) {
-    // We'll still dispatch to _editingTransactionQueue only to schedule a block
+    // In main-thread-only mode allocate and layout all nodes serially on the main thread.
+    //
+    // After this step, we'll still dispatch to _editingTransactionQueue only to schedule a block
     // to _mainSerialQueue to execute next steps. This is not optimized because
     // in theory we can skip _editingTransactionQueue entirely, but it's much safer
     // because change sets will still flow through the pipeline in pretty the same way
     // (main thread -> _editingTransactionQueue -> _mainSerialQueue) and so
     // any methods that block on _editingTransactionQueue will still work.
-    step3();
+    step3(YES);
   }
 
   ++_editingTransactionGroupCount;
@@ -658,7 +681,7 @@ - (void)updateWithChangeSet:(_ASHierarchyChangeSet *)changeSet
     as_activity_scope_enter(as_activity_create("Prepare nodes for collection update", AS_ACTIVITY_CURRENT, OS_ACTIVITY_FLAG_DEFAULT), &preparationScope);
 
     if (!mainThreadOnly) {
-      step3();
+      step3(NO);
     }
 
     // Step 4: Inform the delegate on main thread