Merge pull request swiftlang#47 from neonichu/feature/cancellation

<rdar://problem/28676638> Need BuildSystem support for build cancellation

Author: neonichu

include/llbuild/BuildSystem/BuildExecutionQueue.h

@@ -85,6 +85,9 @@ class BuildExecutionQueue {
   /// Add a job to be executed.
   virtual void addJob(QueueJob job) = 0;
 
+  /// Cancel all jobs and subprocesses of this queue.
+  virtual void cancelAllJobs() = 0;
+
   /// @name Execution Interfaces
   ///
   /// These are additional interfaces provided by the execution queue which can

include/llbuild/BuildSystem/BuildSystem.h

@@ -201,6 +201,9 @@ class BuildSystem {
   /// if a cycle was discovered).
   bool build(StringRef target);
 
+  /// Cancel the current build
+  void cancel();
+
   /// @}
 };
 

include/llbuild/BuildSystem/BuildSystemFrontend.h

@@ -92,6 +92,7 @@ class BuildSystemFrontendDelegate : public BuildSystemDelegate {
 
 private:
   void* impl;
+  std::atomic<bool> isCancelled_;
 
   /// Default implementation, cannot be overriden by subclasses.
   virtual void setFileContentsBeingParsed(StringRef buffer) override;
@@ -127,6 +128,12 @@ class BuildSystemFrontendDelegate : public BuildSystemDelegate {
   /// Provides a default cancellation implementation that will cancel when any
   /// command has failed.
   virtual bool isCancelled() override;
+
+  /// Cancels the current build.
+  virtual void cancel();
+
+  /// Reset mutable build state before a new build operation.
+  void resetForBuild();
 
   /// Provides a default handler.
   ///

lib/BuildSystem/BuildSystem.cpp

@@ -281,6 +281,10 @@ static BuildSystemImpl& getBuildSystem(BuildEngine& engine) {
   return static_cast<BuildSystemEngineDelegate*>(
       engine.getDelegate())->getBuildSystem();
 }
+
+static bool isCancelled(BuildEngine& engine) {
+  return getBuildSystem(engine).getCommandInterface().getDelegate().isCancelled();
+}
 
 /// This is the task used to "build" a target, it translates between the request
 /// for building a target key and the requests for all of its nodes.
@@ -328,6 +332,12 @@ class TargetTask : public Task {
   }
 
   virtual void inputsAvailable(BuildEngine& engine) override {
+    // If the build should cancel, do nothing.
+    if (isCancelled(engine)) {
+      engine.taskIsComplete(this, BuildValue::makeSkippedCommand().toData());
+      return;
+    }
+
     if (hasMissingInput) {
       // FIXME: Design the logging and status output APIs.
       auto& system = getBuildSystem(engine);
@@ -2041,3 +2051,10 @@ bool BuildSystem::enableTracing(StringRef path,
 bool BuildSystem::build(StringRef name) {
   return static_cast<BuildSystemImpl*>(impl)->build(name);
 }
+
+void BuildSystem::cancel() {
+  if (impl) {
+    auto buildSystemImpl = static_cast<BuildSystemImpl*>(impl);
+    buildSystemImpl->getCommandInterface().getExecutionQueue().cancelAllJobs();
+  }
+}

lib/BuildSystem/BuildSystemFrontend.cpp

@@ -221,7 +221,7 @@ BuildSystemFrontendDelegate(llvm::SourceMgr& sourceMgr,
                             StringRef name,
                             uint32_t version)
     : BuildSystemDelegate(name, version),
-      impl(new BuildSystemFrontendDelegateImpl(sourceMgr, invocation))
+      impl(new BuildSystemFrontendDelegateImpl(sourceMgr, invocation)), isCancelled_(false)
 {
 
 }
@@ -330,7 +330,19 @@ BuildSystemFrontendDelegate::createExecutionQueue() {
 
 bool BuildSystemFrontendDelegate::isCancelled() {
   // Stop the build after any command failures.
-  return getNumFailedCommands() > 0;
+  return getNumFailedCommands() > 0 || isCancelled_;
+}
+
+void BuildSystemFrontendDelegate::cancel() {
+  // FIXME: We should audit that a build is happening.
+  isCancelled_ = true;
+
+  auto delegateImpl = static_cast<BuildSystemFrontendDelegateImpl*>(impl);
+  delegateImpl->system->cancel();
+}
+
+void BuildSystemFrontendDelegate::resetForBuild() {
+  isCancelled_ = false;
 }
 
 void BuildSystemFrontendDelegate::hadCommandFailure() {

lib/BuildSystem/LaneBasedExecutionQueue.cpp

@@ -16,6 +16,7 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
@@ -28,6 +29,7 @@
 #include <thread>
 #include <vector>
 #include <string>
+#include <unordered_set>
 
 #include <fcntl.h>
 #include <pthread.h>
@@ -64,6 +66,16 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
   std::mutex readyJobsMutex;
   std::condition_variable readyJobsCondition;
 
+  /// The set of spawned processes to terminate if we get cancelled.
+  std::unordered_set<pid_t> spawnedProcesses;
+  std::mutex spawnedProcessesMutex;
+
+  /// Management of cancellation and SIGKILL escalation
+  std::unique_ptr<std::thread> killAfterTimeoutThread = nullptr;
+  std::atomic<bool> cancelled { false };
+  std::condition_variable stopKillingCondition;
+  std::mutex stopKillingMutex;
+
   void executeLane(unsigned laneNumber) {
     // Set the thread name, if available.
 #if defined(__APPLE__)
@@ -78,7 +90,7 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
 #endif
 
     // Execute items from the queue until shutdown.
-    while (true) {
+    while (!cancelled) {
       // Take a job from the ready queue.
       QueueJob job{};
       {
@@ -87,6 +99,10 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
         // While the queue is empty, wait for an item.
         while (readyJobs.empty()) {
           readyJobsCondition.wait(lock);
+
+          if (cancelled) {
+            return;
+          }
         }
 
         // Take an item according to the chosen policy.
@@ -106,6 +122,21 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
     }
   }
 
+  void killAfterTimeout() {
+    std::unique_lock<std::mutex> lock(stopKillingMutex);
+    stopKillingCondition.wait_for(lock, std::chrono::seconds(10));
+    sendSignalToProcesses(SIGKILL);
+  }
+
+  void sendSignalToProcesses(int signal) {
+    std::unique_lock<std::mutex> lock(spawnedProcessesMutex);
+
+    for (pid_t pid: spawnedProcesses) {
+      // We are killing the whole process group here, this depends on us spawning each process in its own group earlier
+      ::kill(-pid, signal);
+    }
+  }
+
 public:
   LaneBasedExecutionQueue(BuildExecutionQueueDelegate& delegate,
                           unsigned numLanes)
@@ -120,20 +151,43 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
 
   virtual ~LaneBasedExecutionQueue() {
     // Shut down the lanes.
-    for (unsigned i = 0; i != numLanes; ++i) {
-      addJob({});
-    }
+    cancelled = true;
+    readyJobsCondition.notify_all();
+
     for (unsigned i = 0; i != numLanes; ++i) {
       lanes[i]->join();
     }
+
+    if (killAfterTimeoutThread) {
+      stopKillingCondition.notify_all();
+      killAfterTimeoutThread->join();
+    }
   }
 
   virtual void addJob(QueueJob job) override {
+    if (cancelled) {
+      // FIXME: We should eventually raise an error here as new work should not be enqueued after cancellation
+      return;
+    }
+
     std::lock_guard<std::mutex> guard(readyJobsMutex);
     readyJobs.push_back(job);
     readyJobsCondition.notify_one();
   }
 
+  virtual void cancelAllJobs() override {
+    auto wasAlreadyCancelled = cancelled.exchange(true);
+    // If we were already cancelled, do nothing.
+    if (wasAlreadyCancelled) {
+      return;
+    }
+
+    readyJobsCondition.notify_all();
+
+    sendSignalToProcesses(SIGINT);
+    killAfterTimeoutThread = llvm::make_unique<std::thread>(&LaneBasedExecutionQueue::killAfterTimeout, this);
+  }
+
   virtual bool
   executeProcess(QueueJobContext* opaqueContext,
                  ArrayRef<StringRef> commandLine,
@@ -274,19 +328,24 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
     }
 
     // Spawn the command.
-    //
-    // FIXME: Need to track spawned processes for the purposes of cancellation.
-    
     pid_t pid;
-    if (posix_spawn(&pid, args[0], /*file_actions=*/&fileActions,
-                    /*attrp=*/&attributes, const_cast<char**>(args.data()),
-                    envp) != 0) {
-      getDelegate().commandProcessHadError(
-          context.job.getForCommand(), handle,
-          Twine("unable to spawn process (") + strerror(errno) + ")");
-      getDelegate().commandProcessFinished(context.job.getForCommand(), handle,
-                                           -1);
-      return false;
+    {
+      // We need to hold the spawn processes lock when we spawn, to ensure that
+      // we don't create a process in between when we are cancelled.
+      std::lock_guard<std::mutex> guard(spawnedProcessesMutex);
+
+      if (posix_spawn(&pid, args[0], /*file_actions=*/&fileActions,
+                      /*attrp=*/&attributes, const_cast<char**>(args.data()),
+                      envp) != 0) {
+        getDelegate().commandProcessHadError(
+            context.job.getForCommand(), handle,
+            Twine("unable to spawn process (") + strerror(errno) + ")");
+        getDelegate().commandProcessFinished(context.job.getForCommand(), handle,
+                                             -1);
+        return false;
+      }
+
+      spawnedProcesses.insert(pid);
     }
 
     posix_spawn_file_actions_destroy(&fileActions);
@@ -323,6 +382,13 @@ class LaneBasedExecutionQueue : public BuildExecutionQueue {
     int status, result = waitpid(pid, &status, 0);
     while (result == -1 && errno == EINTR)
       result = waitpid(pid, &status, 0);
+
+    // Update the set of spawned processes.
+    {
+        std::lock_guard<std::mutex> guard(spawnedProcessesMutex);
+        spawnedProcesses.erase(pid);
+    }
+
     if (result == -1) {
       getDelegate().commandProcessHadError(
           context.job.getForCommand(), handle,

lib/BuildSystem/SwiftTools.cpp

@@ -106,6 +106,12 @@ class SwiftGetVersionCommand : public Command {
 
   virtual void inputsAvailable(BuildSystemCommandInterface& bsci,
                                core::Task* task) override {
+    // If the build should cancel, do nothing.
+    if (bsci.getDelegate().isCancelled()) {
+      bsci.taskIsComplete(task, BuildValue::makeSkippedCommand());
+      return;
+    }
+
     // Dispatch a task to query the compiler version.
     auto fn = [this, &bsci=bsci, task=task](QueueJobContext* context) {
       // Suppress static analyzer false positive on generalized lambda capture