From 908dbe8fc50b3ce76e896d9ec970a1e04e5f626d Mon Sep 17 00:00:00 2001
From: Ankit Aggarwal <ankit_aggarwal@apple.com>
Date: Sun, 6 Oct 2019 20:08:58 -0700
Subject: [PATCH] Respect numParallelJobs and also install INT handler

---
 Sources/SwiftDriver/Driver/Driver.swift       |  7 ++-
 .../SwiftDriver/Execution/JobExecutor.swift   | 44 +++++++++++++++++--
 Sources/swift-driver/main.swift               | 13 +++++-
 3 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/Sources/SwiftDriver/Driver/Driver.swift b/Sources/SwiftDriver/Driver/Driver.swift
index 252629ec1..1641fbe57 100644
--- a/Sources/SwiftDriver/Driver/Driver.swift
+++ b/Sources/SwiftDriver/Driver/Driver.swift
@@ -358,7 +358,8 @@ extension Driver {
   /// Run the driver.
   public mutating func run(
     resolver: ArgsResolver,
-    executorDelegate: JobExecutorDelegate? = nil
+    executorDelegate: JobExecutorDelegate? = nil,
+    processSet: ProcessSet? = nil
   ) throws {
     // We just need to invoke the corresponding tool if the kind isn't Swift compiler.
     guard driverKind.isSwiftCompiler else {
@@ -389,7 +390,9 @@ extension Driver {
     // Start up an executor and perform the build.
     let jobExecutor = JobExecutor(
         jobs: jobs, resolver: resolver,
-        executorDelegate: executorDelegate
+        executorDelegate: executorDelegate,
+        numParallelJobs: numParallelJobs,
+        processSet: processSet
     )
     try jobExecutor.execute()
   }
diff --git a/Sources/SwiftDriver/Execution/JobExecutor.swift b/Sources/SwiftDriver/Execution/JobExecutor.swift
index e41a7edc8..3fe80e7f6 100644
--- a/Sources/SwiftDriver/Execution/JobExecutor.swift
+++ b/Sources/SwiftDriver/Execution/JobExecutor.swift
@@ -86,14 +86,24 @@ public final class JobExecutor {
     /// Queue for executor delegate.
     let delegateQueue: DispatchQueue = DispatchQueue(label: "org.swift.driver.job-executor-delegate")
 
+    /// Operation queue for executing tasks in parallel.
+    let jobQueue: OperationQueue
+
+    /// The process set to use when launching new processes.
+    let processSet: ProcessSet?
+
     init(
       argsResolver: ArgsResolver,
       producerMap: [VirtualPath: Job],
-      executorDelegate: JobExecutorDelegate
+      executorDelegate: JobExecutorDelegate,
+      jobQueue: OperationQueue,
+      processSet: ProcessSet?
     ) {
       self.producerMap = producerMap
       self.argsResolver = argsResolver
       self.executorDelegate = executorDelegate
+      self.jobQueue = jobQueue
+      self.processSet = processSet
     }
   }
 
@@ -106,14 +116,24 @@ public final class JobExecutor {
   /// The job executor delegate.
   let executorDelegate: JobExecutorDelegate
 
+  /// The number of jobs to run in parallel.
+  let numParallelJobs: Int
+
+  /// The process set to use when launching new processes.
+  let processSet: ProcessSet?
+
   public init(
     jobs: [Job],
     resolver: ArgsResolver,
-    executorDelegate: JobExecutorDelegate
+    executorDelegate: JobExecutorDelegate,
+    numParallelJobs: Int? = nil,
+    processSet: ProcessSet? = nil
   ) {
     self.jobs = jobs
     self.argsResolver = resolver
     self.executorDelegate = executorDelegate
+    self.numParallelJobs = numParallelJobs ?? 1
+    self.processSet = processSet
   }
 
   /// Execute all jobs.
@@ -141,10 +161,16 @@ public final class JobExecutor {
       }
     }
 
+    let jobQueue = OperationQueue()
+    jobQueue.name = "org.swift.driver.job-execution"
+    jobQueue.maxConcurrentOperationCount = numParallelJobs
+
     return Context(
       argsResolver: argsResolver,
       producerMap: producerMap,
-      executorDelegate: executorDelegate
+      executorDelegate: executorDelegate,
+      jobQueue: jobQueue,
+      processSet: processSet
     )
   }
 }
@@ -281,6 +307,13 @@ class ExecuteJobRule: LLBuildRule {
       return engine.taskIsComplete(DriverBuildValue.jobExecution(success: false))
     }
 
+    let context = engine.jobExecutorContext
+    context.jobQueue.addOperation {
+      self.executeJob(engine)
+    }
+  }
+
+  private func executeJob(_ engine: LLTaskBuildEngine) {
     let context = engine.jobExecutorContext
     let resolver = context.argsResolver
     let job = key.job
@@ -295,6 +328,11 @@ class ExecuteJobRule: LLBuildRule {
       let process = try context.executorDelegate.launchProcess(for: job, arguments: arguments)
       pid = Int(process.processID)
 
+      // Add it to the process set if it's a real process.
+      if case let realProcess as TSCBasic.Process = process {
+        try context.processSet?.add(realProcess)
+      }
+
       // Inform the delegate.
       context.delegateQueue.async {
         context.executorDelegate.jobStarted(job: job, arguments: arguments, pid: pid)
diff --git a/Sources/swift-driver/main.swift b/Sources/swift-driver/main.swift
index 3d9100166..b9789d8c5 100644
--- a/Sources/swift-driver/main.swift
+++ b/Sources/swift-driver/main.swift
@@ -1,11 +1,20 @@
 import SwiftDriver
+
 import TSCLibc
-import enum TSCUtility.Diagnostics
+import TSCBasic
+import TSCUtility
+
+var intHandler: InterruptHandler?
 
 do {
+  let processSet = ProcessSet()
+  intHandler = try InterruptHandler {
+    processSet.terminate()
+  }
+
   var driver = try Driver(args: CommandLine.arguments)
   let resolver = try ArgsResolver()
-  try driver.run(resolver: resolver)
+  try driver.run(resolver: resolver, processSet: processSet)
 
   if driver.diagnosticEngine.hasErrors {
     exit(EXIT_FAILURE)