[SoruceKit] Audit SwiftASTManager to prevent deadlocks

- Ensure that no callbacks are called when either of these mutexes are held:
  - `SwiftASTConsumer::CancellationRequestCallbackAndIsCancelledMtx`
  - `ASTBuildOperation::DependencyStampsMtx`
    - Exception: We allow calls to `getBuferStamp` because that doesn’t claim any locks)
    - No change was necessary here
  - `ASTBuildOperation::ConsumersAndResultMtx`
  - `ASTBuildOperation::CacheMtx`
    - No change was necessary here
  - `ASTBuildOperation::ScheduledConsumersMtx`
    - Make sure we always inform the consumers asynchronously, even if the operation was already cancelled when the consumer got enqueued by introducing `ConsumerNotificationQueue`
- Ensure that all callbacks to `SwiftASTConsumer` (i.e. `requestCancellation` and `handlePrimaryAST`, `failed`, `cancelled`) are performed asyncronously from a queue to avoid deadlocks

Author: ahoppen

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -604,6 +604,13 @@ struct SwiftASTManager::Implementation {
   WorkQueue ASTBuildQueue{ WorkQueue::Dequeuing::Serial,
                            "sourcekit.swift.ASTBuilding" };
 
+  /// Queue on which consumers may be notified about results and cancellation.
+  /// This is essentially just a background queue to which we can jump to inform
+  /// consumers while making sure that no locks are currently claimed.
+  WorkQueue ConsumerNotificationQueue{
+      WorkQueue::Dequeuing::Concurrent,
+      "SwiftASTManager::Implementation::ConsumerNotificationQueue"};
+
   /// Remove all scheduled consumers that don't exist anymore. This is just a
   /// garbage-collection operation to make sure the \c ScheduledConsumers vector
   /// doesn't explode. One should never make assumptions that all consumers in
@@ -769,9 +776,11 @@ void SwiftASTManager::processASTAsync(
       // Cancel any consumers with the same OncePerASTToken.
       for (auto ScheduledConsumer : Impl.ScheduledConsumers) {
         if (ScheduledConsumer.OncePerASTToken == OncePerASTToken) {
-          if (auto Consumer = ScheduledConsumer.Consumer.lock()) {
-            Consumer->requestCancellation();
-          }
+          Impl.ConsumerNotificationQueue.dispatch([ScheduledConsumer]() {
+            if (auto Consumer = ScheduledConsumer.Consumer.lock()) {
+              Consumer->requestCancellation();
+            }
+          });
         }
       }
     }
@@ -781,11 +790,17 @@ void SwiftASTManager::processASTAsync(
   Producer->enqueueConsumer(ASTConsumer, fileSystem, shared_from_this());
 
   auto WeakConsumer = SwiftASTConsumerWeakRef(ASTConsumer);
-  Impl.ReqTracker->setCancellationHandler(CancellationToken, [WeakConsumer] {
-    if (auto Consumer = WeakConsumer.lock()) {
-      Consumer->requestCancellation();
-    }
-  });
+  auto WeakThis = std::weak_ptr<SwiftASTManager>(shared_from_this());
+  Impl.ReqTracker->setCancellationHandler(
+      CancellationToken, [WeakConsumer, WeakThis] {
+        if (auto This = WeakThis.lock()) {
+          This->Impl.ConsumerNotificationQueue.dispatch([WeakConsumer]() {
+            if (auto Consumer = WeakConsumer.lock()) {
+              Consumer->requestCancellation();
+            }
+          });
+        }
+      });
 }
 
 void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
@@ -937,12 +952,14 @@ void ASTBuildOperation::requestConsumerCancellation(
     return;
   }
   Consumers.erase(ConsumerIndex);
-  Consumer->cancelled();
   if (Consumers.empty()) {
     // If there are no more consumers waiting for this result, cancel the AST
     // build.
     CancellationFlag->store(true, std::memory_order_relaxed);
   }
+  ASTManager->Impl.ConsumerNotificationQueue.dispatch([Consumer] {
+    Consumer->cancelled();
+  });
 }
 
 static void collectModuleDependencies(ModuleDecl *TopMod,
@@ -1014,11 +1031,15 @@ void ASTBuildOperation::informConsumer(SwiftASTConsumerRef Consumer) {
                     "more consumers attached to it and should not accept any "
                     "new consumers if the build operation was cancelled. Thus "
                     "this case should never happen.");
-    Consumer->cancelled();
+    ASTManager->Impl.ConsumerNotificationQueue.dispatch([Consumer] {
+      Consumer->cancelled();
+    });
   } else if (Result.AST) {
     Result.AST->Impl.consumeAsync(Consumer, Result.AST);
   } else {
-    Consumer->failed(Result.Error);
+    ASTManager->Impl.ConsumerNotificationQueue.dispatch([Consumer, Error = Result.Error] {
+      Consumer->failed(Error);
+    });
   }
 }
 
@@ -1136,6 +1157,7 @@ void ASTBuildOperation::schedule(WorkQueue Queue) {
                             /*ExpectedOldState=*/State::Running);
         };
 
+        SmallVector<SwiftASTConsumerRef, 4> ConsumersToCancel;
         {
           llvm::sys::ScopedLock L(ConsumersAndResultMtx);
           if (Consumers.empty()) {
@@ -1146,24 +1168,25 @@ void ASTBuildOperation::schedule(WorkQueue Queue) {
           if (CancellationFlag->load(std::memory_order_relaxed)) {
             assert(false && "We should only set the cancellation flag if there "
                             "are no more consumers");
-            for (auto &Consumer : Consumers) {
-              Consumer->cancelled();
-            }
+            ConsumersToCancel = Consumers;
           }
         }
+        for (auto &Consumer : ConsumersToCancel) {
+          Consumer->cancelled();
+        }
 
         std::string Error;
         assert(!Result && "We should only be producing a result once");
         ASTUnitRef AST = buildASTUnit(Error);
-        SmallVector<SwiftASTConsumerRef, 4> LocalConsumers;
+        SmallVector<SwiftASTConsumerRef, 4> ConsumersToInform;
         {
           llvm::sys::ScopedLock L(ConsumersAndResultMtx);
           bool WasCancelled = CancellationFlag->load(std::memory_order_relaxed);
           Result.emplace(AST, Error, WasCancelled);
-          LocalConsumers = Consumers;
+          ConsumersToInform = Consumers;
           Consumers = {};
         }
-        for (auto &Consumer : LocalConsumers) {
+        for (auto &Consumer : ConsumersToInform) {
           informConsumer(Consumer);
         }
         DidFinishCallback();

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -134,8 +134,9 @@ class ASTUnit : public SourceKit::ThreadSafeRefCountedBase<ASTUnit> {
 typedef IntrusiveRefCntPtr<ASTUnit> ASTUnitRef;
 
 class SwiftASTConsumer : public std::enable_shared_from_this<SwiftASTConsumer> {
-  /// Mutex guarding all accesses to CancellationRequestCallback
-  llvm::sys::Mutex CancellationRequestCallbackMtx;
+  /// Mutex guarding all accesses to \c CancellationRequestCallback and \c
+  /// IsCancelled.
+  llvm::sys::Mutex CancellationRequestCallbackAndIsCancelledMtx;
 
   /// A callback that informs the \c ASTBuildOperation, which is producing the
   /// AST for this consumer, that the consumer is no longer of interest. Calling
@@ -159,12 +160,17 @@ class SwiftASTConsumer : public std::enable_shared_from_this<SwiftASTConsumer> {
   /// cause the \c ASTBuildOperation to be cancelled if no other consumer is
   /// depending on it.
   void requestCancellation() {
-    llvm::sys::ScopedLock L(CancellationRequestCallbackMtx);
-    IsCancelled = true;
-    if (CancellationRequestCallback.hasValue()) {
-      (*CancellationRequestCallback)(shared_from_this());
+    Optional<std::function<void(std::shared_ptr<SwiftASTConsumer>)>>
+        CallbackToCall;
+    {
+      llvm::sys::ScopedLock L(CancellationRequestCallbackAndIsCancelledMtx);
+      IsCancelled = true;
+      CallbackToCall = CancellationRequestCallback;
       CancellationRequestCallback = None;
     }
+    if (CallbackToCall.hasValue()) {
+      (*CallbackToCall)(shared_from_this());
+    }
   }
 
   /// Set a cancellation request callback that informs a \c ASTBuildOperation
@@ -176,20 +182,26 @@ class SwiftASTConsumer : public std::enable_shared_from_this<SwiftASTConsumer> {
   /// called, \c NewCallback will be called immediately.
   void setCancellationRequestCallback(
       std::function<void(std::shared_ptr<SwiftASTConsumer>)> NewCallback) {
-    llvm::sys::ScopedLock L(CancellationRequestCallbackMtx);
-    assert(!CancellationRequestCallback.hasValue() &&
-           "Can't set two cancellation callbacks on a SwiftASTConsumer");
-    if (IsCancelled) {
+    bool ShouldCallCallback = false;
+    {
+      llvm::sys::ScopedLock L(CancellationRequestCallbackAndIsCancelledMtx);
+      assert(!CancellationRequestCallback.hasValue() &&
+             "Can't set two cancellation callbacks on a SwiftASTConsumer");
+      if (IsCancelled) {
+        ShouldCallCallback = true;
+      } else {
+        CancellationRequestCallback = NewCallback;
+      }
+    }
+    if (ShouldCallCallback) {
       NewCallback(shared_from_this());
-    } else {
-      CancellationRequestCallback = NewCallback;
     }
   }
 
   /// Removes the cancellation request callback previously set by \c
   /// setCancellationRequestCallback.
   void removeCancellationRequestCallback() {
-    llvm::sys::ScopedLock L(CancellationRequestCallbackMtx);
+    llvm::sys::ScopedLock L(CancellationRequestCallbackAndIsCancelledMtx);
     CancellationRequestCallback = None;
   }