From c70d7b3692df9f3f1e33d8433ec7d0719296bdd9 Mon Sep 17 00:00:00 2001
From: Huy Nguyen <huy@pinterest.com>
Date: Thu, 17 Sep 2020 16:30:39 -0700
Subject: [PATCH] Add an experiment that makes ASDataController to do
 everything on main thread

- Under this experiment, ASDataController will allocate and layout all nodes on the main thread. This helps to avoid deadlocks that would otherwise occur if some of the node allocations or layouts caused ASDataController's backgroud queue to block on main thread. As a bonus, this experiment also helps to measure how much perf wins we get from doing the work off main.
---
 Source/ASExperimentalFeatures.h    |  1 +
 Source/ASExperimentalFeatures.mm   |  3 ++-
 Source/Details/ASDataController.mm | 33 +++++++++++++++++++++++++-----
 Tests/ASConfigurationTests.mm      |  2 ++
 4 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/Source/ASExperimentalFeatures.h b/Source/ASExperimentalFeatures.h
index 4e70df6b7..6407d4c70 100644
--- a/Source/ASExperimentalFeatures.h
+++ b/Source/ASExperimentalFeatures.h
@@ -29,6 +29,7 @@ typedef NS_OPTIONS(NSUInteger, ASExperimentalFeatures) {
   ASExperimentalDrawingGlobal = 1 << 8,                                     // exp_drawing_global
   ASExperimentalOptimizeDataControllerPipeline = 1 << 9,                    // exp_optimize_data_controller_pipeline
   ASExperimentalDisableGlobalTextkitLock = 1 << 10,                         // exp_disable_global_textkit_lock
+  ASExperimentalMainThreadOnlyDataController = 1 << 11,                     // exp_main_thread_only_data_controller
   ASExperimentalFeatureAll = 0xFFFFFFFF
 };
 
diff --git a/Source/ASExperimentalFeatures.mm b/Source/ASExperimentalFeatures.mm
index 71313955e..fe4ebc9bf 100644
--- a/Source/ASExperimentalFeatures.mm
+++ b/Source/ASExperimentalFeatures.mm
@@ -22,7 +22,8 @@
                                       @"exp_dispatch_apply",
                                       @"exp_drawing_global",
                                       @"exp_optimize_data_controller_pipeline",
-                                    @"exp_disable_global_textkit_lock"]));
+                                      @"exp_disable_global_textkit_lock",
+                                      @"exp_main_thread_only_data_controller"]));
   if (flags == ASExperimentalFeatureAll) {
     return allNames;
   }
diff --git a/Source/Details/ASDataController.mm b/Source/Details/ASDataController.mm
index 76e58d8be..184b5cd61 100644
--- a/Source/Details/ASDataController.mm
+++ b/Source/Details/ASDataController.mm
@@ -622,12 +622,9 @@ - (void)updateWithChangeSet:(_ASHierarchyChangeSet *)changeSet
   os_log_debug(ASCollectionLog(), "New content: %@", newMap.smallDescription);
 
   Class<ASDataControllerLayoutDelegate> layoutDelegateClass = [self.layoutDelegate class];
-  ++_editingTransactionGroupCount;
-  dispatch_group_async(_editingTransactionGroup, _editingTransactionQueue, ^{
-    __block __unused os_activity_scope_state_s preparationScope = {}; // unused if deployment target < iOS10
-    as_activity_scope_enter(as_activity_create("Prepare nodes for collection update", AS_ACTIVITY_CURRENT, OS_ACTIVITY_FLAG_DEFAULT), &preparationScope);
 
-    // Step 3: Call the layout delegate if possible. Otherwise, allocate and layout all elements
+  // Step 3: Call the layout delegate if possible. Otherwise, allocate and layout all elements
+  dispatch_block_t step3 = ^{
     if (canDelegate) {
       [layoutDelegateClass calculateLayoutWithContext:layoutContext];
     } else {
@@ -644,6 +641,32 @@ - (void)updateWithChangeSet:(_ASHierarchyChangeSet *)changeSet
       }
       [self _allocateNodesFromElements:elementsToProcess];
     }
+  };
+
+  // Step 3 can be done on the main thread or on _editingTransactionQueue
+  // depending on an experiment.
+  BOOL mainThreadOnly = ASActivateExperimentalFeature(ASExperimentalMainThreadOnlyDataController);
+  if (mainThreadOnly) {
+    // Doing it on main means _editingTransactionQueue will be empty all the time
+    // and blocking on it should be super quick/no-op.
+    //
+    // 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();
+  }
+
+  ++_editingTransactionGroupCount;
+  dispatch_group_async(_editingTransactionGroup, _editingTransactionQueue, ^{
+    __block __unused os_activity_scope_state_s preparationScope = {}; // unused if deployment target < iOS10
+    as_activity_scope_enter(as_activity_create("Prepare nodes for collection update", AS_ACTIVITY_CURRENT, OS_ACTIVITY_FLAG_DEFAULT), &preparationScope);
+
+    if (!mainThreadOnly) {
+      step3();
+    }
 
     // Step 4: Inform the delegate on main thread
     [self->_mainSerialQueue performBlockOnMainThread:^{
diff --git a/Tests/ASConfigurationTests.mm b/Tests/ASConfigurationTests.mm
index b3287fc31..8f66f0753 100644
--- a/Tests/ASConfigurationTests.mm
+++ b/Tests/ASConfigurationTests.mm
@@ -50,6 +50,8 @@ + (NSArray *)names {
     @"exp_dispatch_apply",
     @"exp_drawing_global",
     @"exp_optimize_data_controller_pipeline",
+    @"exp_disable_global_textkit_lock",
+    @"exp_main_thread_only_data_controller"
   ];
 }