Merge pull request #58978 from ahoppen/pr/no-deadlock-cancellation-bigger-change

[SoruceKit] Audit SwiftASTManager to prevent deadlocks

Author: ahoppen

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -1182,6 +1182,7 @@ void ASTBuildOperation::schedule(WorkQueue Queue) {
                             /*ExpectedOldState=*/State::Running);
         };
 
+        SmallVector<SwiftASTConsumerRef, 4> ConsumersToCancel;
         {
           llvm::sys::ScopedLock L(ConsumersAndResultMtx);
           if (Consumers.empty()) {
@@ -1192,24 +1193,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

@@ -135,8 +135,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
@@ -160,12 +161,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.has_value()) {
-      (*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.has_value()) {
+      (*CallbackToCall)(shared_from_this());
+    }
   }
 
   /// Set a cancellation request callback that informs a \c ASTBuildOperation
@@ -177,20 +183,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.has_value() &&
-           "Can't set two cancellation callbacks on a SwiftASTConsumer");
-    if (IsCancelled) {
+    bool ShouldCallCallback = false;
+    {
+      llvm::sys::ScopedLock L(CancellationRequestCallbackAndIsCancelledMtx);
+      assert(!CancellationRequestCallback.has_value() &&
+             "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;
   }