Fixed quadratic performance of ListMerger when each executed job creates 2+ new jobs of the same priority

See https://forums.swift.org/t/quadratic-performance-of-the-listmerger-in-specific-use-case/69393

Author: nickolas-pohilets

include/swift/Basic/ListMerger.h

@@ -19,6 +19,7 @@
 #define SWIFT_BASIC_LISTMERGER_H
 
 #include <assert.h>
+#include <tuple>
 
 namespace swift {
 
@@ -62,13 +63,31 @@ namespace swift {
 /// to the merger and before being released except by the merger.
 template <class Node, class NodeTraits>
 class ListMerger {
+public:
+  class LastInsertionPoint {
+    friend class ListMerger;
+    Node node = Node();
+    bool isKnownLastOfEquals = false;
+
+  public:
+    LastInsertionPoint() {}
+
+    void nodeWasRemoved(Node removedNode) {
+      if (node == removedNode) {
+        *this = LastInsertionPoint();
+      }
+    }
+  };
+
+private:
   Node root;
-  Node lastInsertionPoint = Node();
-  bool lastInsertionPointIsKnownLastOfEquals = false;
+  LastInsertionPoint lastInsertionPoint;
+
 public:
   /// Construct a merger with the given sorted list as its current list.
-  ListMerger(Node initialList = Node())
-    : root(initialList) {}
+  ListMerger(Node initialList = Node(),
+             LastInsertionPoint insertionPoint = LastInsertionPoint())
+      : root(initialList), lastInsertionPoint(insertionPoint) {}
 
   /// Add a single node to this merger's current list.
   ///
@@ -86,7 +105,7 @@ class ListMerger {
     Node stopper = Node();
 
     // If we have a previous insertion point, compare against it.
-    if (Node lastIP = lastInsertionPoint) {
+    if (Node lastIP = lastInsertionPoint.node) {
       int comparison = NodeTraits::compare(lastIP, newNode);
 
       // If it compares equal, put the new node immediately after the
@@ -169,7 +188,7 @@ class ListMerger {
 
     // If we have a previous insertion point, check for the presumed-common
     // case that we're inserting something that should immediately follow it.
-    if (auto lastIP = lastInsertionPoint) {
+    if (auto lastIP = lastInsertionPoint.node) {
       lastIP = findLastOfEqualsFromLastIP(lastIP);
 
       // Compare against the next node after lastIP, if it exists.
@@ -246,7 +265,7 @@ class ListMerger {
 
     // If we have a previous insertion point, compare the new root
     // against it.
-    if (Node lastIP = lastInsertionPoint) {
+    if (Node lastIP = lastInsertionPoint.node) {
       int comparison = NodeTraits::compare(lastIP, rootOfNewList);
 
       // If it compares equal, we've got an insertion point where
@@ -341,27 +360,27 @@ class ListMerger {
 
   /// Get the current list that's been built up, and clear the internal
   /// state of this merger.
-  Node release() {
-    Node result = root;
+  std::tuple<Node, LastInsertionPoint> release() {
+    auto result = std::make_tuple(root, lastInsertionPoint);
     root = Node();
-    lastInsertionPoint = Node();
+    lastInsertionPoint = LastInsertionPoint();
     return result;
   }
 
 private:
   /// Set the last point at which we inserted a node, and specify
   /// whether we know it was the last in its sequence of equals.
   void setLastInsertionPoint(Node lastIP, bool knownEndOfEquals) {
-    lastInsertionPoint = lastIP;
-    lastInsertionPointIsKnownLastOfEquals = knownEndOfEquals;
+    lastInsertionPoint.node = lastIP;
+    lastInsertionPoint.isKnownLastOfEquals = knownEndOfEquals;
   }
 
   /// Given the value of lastInsertionPoint (passed in to avoid
   /// reloading it), find the last node in the sequence of equals that
   /// contains it.
   Node findLastOfEqualsFromLastIP(Node lastIP) const {
-    assert(lastIP == lastInsertionPoint);
-    if (!lastInsertionPointIsKnownLastOfEquals)
+    assert(lastIP == lastInsertionPoint.node);
+    if (!lastInsertionPoint.isKnownLastOfEquals)
       return findLastOfEquals(lastIP);
     return lastIP;
   }

stdlib/public/Concurrency/Actor.cpp

@@ -981,6 +981,35 @@ class alignas(sizeof(void *) * 2) ActiveActorStatus {
   }
 };
 
+#if !SWIFT_CONCURRENCY_ACTORS_AS_LOCKS
+
+/// Given that a job is enqueued normally on a default actor, get/set
+/// the next job in the actor's queue.
+static JobRef getNextJobInQueue(Job *job) {
+  return *reinterpret_cast<JobRef *>(job->SchedulerPrivate);
+}
+static void setNextJobInQueue(Job *job, JobRef next) {
+  *reinterpret_cast<JobRef *>(job->SchedulerPrivate) = next;
+}
+
+namespace {
+
+struct JobQueueTraits {
+  static Job *getNext(Job *job) {
+    return getNextJobInQueue(job).getAsPreprocessedJob();
+  }
+  static void setNext(Job *job, Job *next) {
+    setNextJobInQueue(job, JobRef::getPreprocessed(next));
+  }
+  static int compare(Job *lhs, Job *rhs) {
+    return descendingPriorityOrder(lhs->getPriority(), rhs->getPriority());
+  }
+};
+
+} // end anonymous namespace
+
+#endif
+
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION && SWIFT_POINTER_IS_4_BYTES
 #define ACTIVE_ACTOR_STATUS_SIZE (4 * (sizeof(uintptr_t)))
 #else
@@ -1052,10 +1081,13 @@ class DefaultActorImpl : public HeapObject {
   // enforce alignment. This is space that is available for us to use in
   // the future
   alignas(sizeof(ActiveActorStatus)) char StatusStorage[sizeof(ActiveActorStatus)];
+
+  using ListMerger = swift::ListMerger<Job *, JobQueueTraits>;
+  ListMerger::LastInsertionPoint lastInsertionPoint =
+      ListMerger::LastInsertionPoint();
 #endif
   // TODO (rokhinip): Make this a flagset
   bool isDistributedRemoteActor;
-
 public:
   /// Properly construct an actor, except for the heap header.
   void initialize(bool isDistributedRemote = false) {
@@ -1128,6 +1160,10 @@ class DefaultActorImpl : public HeapObject {
   /// It can be done when actor transitions from Idle to Scheduled or
   /// when actor gets a priority override and we schedule a stealer.
   void scheduleActorProcessJob(JobPriority priority);
+
+  Job *preprocessQueue(JobRef start);
+  Job *preprocessQueue(JobRef unprocessedStart, JobRef unprocessedEnd,
+                       Job *existingProcessedJobsToMergeInto);
 #endif /* !SWIFT_CONCURRENCY_ACTORS_AS_LOCKS */
 
   void deallocateUnconditional();
@@ -1203,31 +1239,6 @@ static NonDefaultDistributedActorImpl *asImpl(NonDefaultDistributedActor *actor)
 /*****************************************************************************/
 
 #if !SWIFT_CONCURRENCY_ACTORS_AS_LOCKS
-/// Given that a job is enqueued normally on a default actor, get/set
-/// the next job in the actor's queue.
-static JobRef getNextJobInQueue(Job *job) {
-  return *reinterpret_cast<JobRef*>(job->SchedulerPrivate);
-}
-static void setNextJobInQueue(Job *job, JobRef next) {
-  *reinterpret_cast<JobRef*>(job->SchedulerPrivate) = next;
-}
-
-namespace {
-
-struct JobQueueTraits {
-  static Job *getNext(Job *job) {
-    return getNextJobInQueue(job).getAsPreprocessedJob();
-  }
-  static void setNext(Job *job, Job *next) {
-    setNextJobInQueue(job, JobRef::getPreprocessed(next));
-  }
-  static int compare(Job *lhs, Job *rhs) {
-    return descendingPriorityOrder(lhs->getPriority(), rhs->getPriority());
-  }
-};
-
-} // end anonymous namespace
-
 
 // Called with the actor drain lock held
 //
@@ -1240,15 +1251,14 @@ struct JobQueueTraits {
 // and the previous start. We can then process these jobs and merge them into
 // the already processed list of jobs from the previous iteration of
 // preprocessQueue
-static Job *
-preprocessQueue(JobRef unprocessedStart, JobRef unprocessedEnd, Job *existingProcessedJobsToMergeInto)
-{
+Job *DefaultActorImpl::preprocessQueue(JobRef unprocessedStart,
+                                       JobRef unprocessedEnd,
+                                       Job *existingProcessedJobsToMergeInto) {
   assert(existingProcessedJobsToMergeInto != NULL);
   assert(unprocessedStart.needsPreprocessing());
   assert(unprocessedStart.getAsJob() != unprocessedEnd.getAsJob());
 
   // Build up a list of jobs we need to preprocess
-  using ListMerger = swift::ListMerger<Job*, JobQueueTraits>;
   ListMerger jobsToProcess;
 
   // Get just the prefix list of unprocessed jobs
@@ -1263,19 +1273,20 @@ preprocessQueue(JobRef unprocessedStart, JobRef unprocessedEnd, Job *existingPro
   }
 
   // Finish processing the unprocessed jobs
-  Job *newProcessedJobs = jobsToProcess.release();
+  Job *newProcessedJobs = std::get<0>(jobsToProcess.release());
   assert(newProcessedJobs);
 
-  ListMerger mergedList(existingProcessedJobsToMergeInto);
+  ListMerger mergedList(existingProcessedJobsToMergeInto, lastInsertionPoint);
   mergedList.merge(newProcessedJobs);
-  return mergedList.release();
+  Job *result;
+  std::tie(result, lastInsertionPoint) = mergedList.release();
+  return result;
 }
 
 // Called with the actor drain lock held.
 //
 // Preprocess the queue starting from the top
-static Job *
-preprocessQueue(JobRef start) {
+Job *DefaultActorImpl::preprocessQueue(JobRef start) {
   if (!start) {
     return NULL;
   }
@@ -1288,7 +1299,6 @@ preprocessQueue(JobRef start) {
   // There exist some jobs which haven't been preprocessed
 
   // Build up a list of jobs we need to preprocess
-  using ListMerger = swift::ListMerger<Job*, JobQueueTraits>;
   ListMerger jobsToProcess;
 
   Job *wellFormedListStart = NULL;
@@ -1311,18 +1321,19 @@ preprocessQueue(JobRef start) {
   }
 
   // Finish processing the unprocessed jobs
-  auto processedJobHead = jobsToProcess.release();
+  auto processedJobHead = std::get<0>(jobsToProcess.release());
   assert(processedJobHead);
 
   Job *firstJob = NULL;
   if (wellFormedListStart) {
     // Merge it with already known well formed list if we have one.
-    ListMerger mergedList(wellFormedListStart);
+    ListMerger mergedList(wellFormedListStart, lastInsertionPoint);
     mergedList.merge(processedJobHead);
-    firstJob = mergedList.release();
+    std::tie(firstJob, lastInsertionPoint) = mergedList.release();
   } else {
     // Nothing to merge with, just return the head we already have
     firstJob = processedJobHead;
+    lastInsertionPoint = ListMerger::LastInsertionPoint();
   }
 
   return firstJob;
@@ -1528,6 +1539,7 @@ Job * DefaultActorImpl::drainOne() {
     if (_status().compare_exchange_weak(oldState, newState,
                             /* success */ std::memory_order_relaxed,
                             /* failure */ std::memory_order_relaxed)) {
+      lastInsertionPoint.nodeWasRemoved(firstJob);
       SWIFT_TASK_DEBUG_LOG("Drained first job %p from actor %p", firstJob, this);
       traceActorStateTransition(this, oldState, newState, distributedActorIsRemote);
       concurrency::trace::actor_dequeue(this, firstJob);

stdlib/public/Concurrency/CooperativeGlobalExecutor.inc

@@ -108,7 +108,7 @@ static void swift_task_enqueueGlobalImpl(Job *job) {
 
   JobQueueMerger merger(JobQueue);
   merger.insert(job);
-  JobQueue = merger.release();
+  JobQueue = std::get<0>(merger.release());
 }
 
 /// Enqueues a task on the main executor.
@@ -202,7 +202,7 @@ static void recognizeReadyDelayedJobs() {
     nextDelayedJob = next;
   }
 
-  JobQueue = readyJobs.release();
+  JobQueue = std::get<0>(readyJobs.release());
   DelayedJobQueue = nextDelayedJob;
 }
 

utils/test-list-merger/Makefile

@@ -1,7 +1,7 @@
 SWIFT_SRCROOT=${CURDIR}/../..
 SRCROOT=${SWIFT_SRCROOT}/..
 LLVM_SRCROOT=${SRCROOT}/llvm/
-LLVM_OBJROOT=${SRCROOT}/build/Ninja-DebugAssert/llvm-macosx-x86_64
+LLVM_OBJROOT=${SRCROOT}/build/Ninja-DebugAssert/llvm-macosx-$(shell uname -m)
 
 HEADERS=${SWIFT_SRCROOT}/include/swift/Basic/ListMerger.h
 

utils/test-list-merger/TestListMerger.cpp

@@ -161,7 +161,7 @@ static void runInsertAndMergeTest(llvm::ArrayRef<Instruction> values) {
   assert(!lastMergeEntry && "ended while still building a merge list");
 
   entries.sort(creationOrder);
-  entries.checkSameAs(merger.release());
+  entries.checkSameAs(std::get<0>(merger.release()));
 }
 
 static void runInsertAtFrontTest(llvm::ArrayRef<unsigned> values) {
@@ -171,7 +171,7 @@ static void runInsertAtFrontTest(llvm::ArrayRef<unsigned> values) {
     merger.insertAtFront(entries.create(value));
   }
   entries.sort(reverseCreationOrder);
-  entries.checkSameAs(merger.release());
+  entries.checkSameAs(std::get<0>(merger.release()));
 }
 
 static void runConcreteTests() {