[SwiftCaching] Improve swift caching batch build performance (#301)

Improve swift caching build performance, especially batch build
performance by:
* Use more accurate dependency tracking by requesting cache key querying
  jobs in parallel when requesting all the explicit module dependencies.
  This allows build system to plan cache key querying jobs in the gaps
  of explicit module build, thus the compilation job can start sooner.
* Instead of issuing multiple key queries per batched job, issue one key
  querying job for multiple keys instead. This can dramatically reduce
  the number of jobs planned thus improves the efficiency. Also in the
  cache miss cases, if the first key hits the cache miss, there is no
  need to continue querying other keys since the job needs to be run
  anyway.

rdar://146429825

Author: cachemeifyoucan

Sources/SWBTaskExecution/DynamicTaskSpecs/SwiftCachingDynamicTaskSpec.swift

@@ -23,12 +23,12 @@ final class SwiftCachingKeyQueryDynamicTaskSpec: DynamicTaskSpec {
             type: self,
             payload: nil,
             forTarget: dynamicTask.target,
-            ruleInfo: ["SwiftCachingKeyQuery", swiftCachingKeyQueryTaskKey.cacheKey],
-            commandLine: ["builtin-swiftCachingKeyQuery", .literal(ByteString(encodingAsUTF8: swiftCachingKeyQueryTaskKey.cacheKey))],
+            ruleInfo: ["SwiftCachingKeyQuery", swiftCachingKeyQueryTaskKey.cacheKeys.description],
+            commandLine: ["builtin-swiftCachingKeyQuery", .literal(ByteString(encodingAsUTF8: swiftCachingKeyQueryTaskKey.cacheKeys.description))],
             environment: dynamicTask.environment,
             workingDirectory: dynamicTask.workingDirectory,
             showEnvironment: dynamicTask.showEnvironment,
-            execDescription: "Swift caching query key \(swiftCachingKeyQueryTaskKey.cacheKey)",
+            execDescription: "Swift caching query key \(swiftCachingKeyQueryTaskKey.cacheKeys)",
             priority: .network,
             isDynamic: true
         )
@@ -57,12 +57,12 @@ final class SwiftCachingMaterializeKeyDynamicTaskSpec: DynamicTaskSpec {
             type: self,
             payload: nil,
             forTarget: dynamicTask.target,
-            ruleInfo: ["SwiftCachingKeyMaterializer", swiftCachingTaskKey.cacheKey],
-            commandLine: ["builtin-swiftCachingKeyMaterializer", .literal(ByteString(encodingAsUTF8: swiftCachingTaskKey.cacheKey))],
+            ruleInfo: ["SwiftCachingKeyMaterializer", swiftCachingTaskKey.cacheKeys.description],
+            commandLine: ["builtin-swiftCachingKeyMaterializer", .literal(ByteString(encodingAsUTF8: swiftCachingTaskKey.cacheKeys.description))],
             environment: dynamicTask.environment,
             workingDirectory: dynamicTask.workingDirectory,
             showEnvironment: dynamicTask.showEnvironment,
-            execDescription: "Swift caching materialize key \(swiftCachingTaskKey.cacheKey)",
+            execDescription: "Swift caching materialize key \(swiftCachingTaskKey.cacheKeys)",
             priority: .network,
             isDynamic: true
         )

Sources/SWBTaskExecution/DynamicTaskSpecs/SwiftCachingTaskKeys.swift

@@ -15,32 +15,32 @@ public import SWBUtil
 
 public struct SwiftCachingKeyQueryTaskKey: Serializable, CustomDebugStringConvertible {
     let casOptions: CASOptions
-    let cacheKey: String
+    let cacheKeys: [String]
     let compilerLocation: LibSwiftDriver.CompilerLocation
 
-    init(casOptions: CASOptions, cacheKey: String, compilerLocation: LibSwiftDriver.CompilerLocation) {
+    init(casOptions: CASOptions, cacheKeys: [String], compilerLocation: LibSwiftDriver.CompilerLocation) {
         self.casOptions = casOptions
-        self.cacheKey = cacheKey
+        self.cacheKeys = cacheKeys
         self.compilerLocation = compilerLocation
     }
 
     public func serialize<T>(to serializer: T) where T : Serializer {
         serializer.serializeAggregate(3) {
             serializer.serialize(casOptions)
-            serializer.serialize(cacheKey)
+            serializer.serialize(cacheKeys)
             serializer.serialize(compilerLocation)
         }
     }
 
     public init(from deserializer: any Deserializer) throws {
         try deserializer.beginAggregate(3)
         casOptions = try deserializer.deserialize()
-        cacheKey = try deserializer.deserialize()
+        cacheKeys = try deserializer.deserialize()
         compilerLocation = try deserializer.deserialize()
     }
 
     public var debugDescription: String {
-        "<SwiftCachingKeyQuery casOptions=\(casOptions) cacheKey=\(cacheKey) compilerLocation=\(compilerLocation)>"
+        "<SwiftCachingKeyQuery casOptions=\(casOptions) cacheKey=\(cacheKeys) compilerLocation=\(compilerLocation)>"
     }
 }
 

Sources/SWBTaskExecution/TaskActions/SwiftCachingKeyQueryTaskAction.swift

@@ -45,9 +45,15 @@ public final class SwiftCachingKeyQueryTaskAction: TaskAction {
 
             do {
                 // Request global query to get maximum search space that the database supports
-                let cacheHit = try await cas.queryCacheKey(key.cacheKey, globally: true) != nil
-                if key.casOptions.enableDiagnosticRemarks {
-                    outputDelegate.remark("cache key query \(cacheHit ? "hit" : "miss")")
+                for cacheKey in key.cacheKeys {
+                    let cacheHit = try await cas.queryCacheKey(cacheKey, globally: true) != nil
+                    if key.casOptions.enableDiagnosticRemarks {
+                        outputDelegate.remark("cache key query \(cacheHit ? "hit" : "miss")")
+                    }
+                    guard cacheHit else {
+                        // return on first failure.
+                        return .succeeded
+                    }
                 }
             } catch {
                 guard !key.casOptions.enableStrictCASErrors else { throw error }

Sources/SWBTaskExecution/TaskActions/SwiftCachingMaterializeKeyTaskAction.swift

@@ -41,7 +41,7 @@ public final class SwiftCachingMaterializeKeyTaskAction: TaskAction {
         /// The action is in its initial state, and has not yet performed any work.
         case initial
         /// Waiting for the caching key query to finish.
-        case waitingForKeyQuery(jobTaskIDBase: UInt, cas: SwiftCASDatabases, key: String)
+        case waitingForKeyQuery(jobTaskIDBase: UInt, cas: SwiftCASDatabases, keys: [String])
         /// Waiting for the outputs to finish downloading.
         case waitingForOutputDownloads
         /// Not waiting for any other dependency
@@ -69,22 +69,22 @@ public final class SwiftCachingMaterializeKeyTaskAction: TaskAction {
             guard let cas = try swiftModuleDependencyGraph.getCASDatabases(casOptions: taskKey.casOptions, compilerLocation: taskKey.compilerLocation) else {
                 throw StubError.error("unable to use CAS databases")
             }
-            let jobTaskIDBase: UInt = 0
-            if let cachedComp = try cas.queryLocalCacheKey(taskKey.cacheKey) {
-                if try requestCompilationOutputs(cachedComp,
-                                                 dynamicExecutionDelegate: dynamicExecutionDelegate,
-                                                 cacheKey: taskKey.cacheKey,
-                                                 jobTaskIDBase: jobTaskIDBase) != 0 {
-                    state = .waitingForOutputDownloads
+            var missingKeys: [String] = []
+            var compOutputs: [SwiftCachedCompilation] = []
+            try taskKey.cacheKeys.forEach { key in
+                if let output = try cas.queryLocalCacheKey(key) {
+                    compOutputs.append(output)
                 } else {
-                    state = .done
+                    missingKeys.append(key)
                 }
-            } else {
-                let key = SwiftCachingKeyQueryTaskKey(casOptions: taskKey.casOptions, cacheKey: taskKey.cacheKey, compilerLocation: taskKey.compilerLocation)
+            }
+
+            if !missingKeys.isEmpty {
+                let key = SwiftCachingKeyQueryTaskKey(casOptions: taskKey.casOptions, cacheKeys: missingKeys, compilerLocation: taskKey.compilerLocation)
                 dynamicExecutionDelegate.requestDynamicTask(
                     toolIdentifier: SwiftCachingKeyQueryTaskAction.toolIdentifier,
                     taskKey: .swiftCachingKeyQuery(key),
-                    taskID: jobTaskIDBase,
+                    taskID: 0,
                     singleUse: true,
                     workingDirectory: Path(""),
                     environment: .init(),
@@ -93,7 +93,15 @@ public final class SwiftCachingMaterializeKeyTaskAction: TaskAction {
                     showEnvironment: false,
                     reason: .wasCompilationCachingQuery
                 )
-                state = .waitingForKeyQuery(jobTaskIDBase: jobTaskIDBase + 1, cas: cas, key: taskKey.cacheKey)
+                state = .waitingForKeyQuery(jobTaskIDBase: 1, cas: cas, keys: taskKey.cacheKeys)
+            } else {
+                if try requestCompilationOutputs(compOutputs,
+                                                 dynamicExecutionDelegate: dynamicExecutionDelegate,
+                                                 jobTaskIDBase: 1) != 0 {
+                    state = .waitingForOutputDownloads
+                } else {
+                    state = .done
+                }
             }
         } catch {
             state = .executionError(error)
@@ -111,15 +119,17 @@ public final class SwiftCachingMaterializeKeyTaskAction: TaskAction {
         switch state {
         case .initial:
             state = .executionError(StubError.error("taskDependencyReady unexpectedly called in initial state"))
-        case .waitingForKeyQuery(jobTaskIDBase: let jobTaskIDBase, cas: let cas, key: let key):
+        case .waitingForKeyQuery(jobTaskIDBase: let jobTaskIDBase, cas: let cas, keys: let keys):
             do {
-                guard let cachedComp = try cas.queryLocalCacheKey(key) else {
+                let cachedComps = try keys.compactMap {
+                    try cas.queryLocalCacheKey($0)
+                }
+                guard cachedComps.count == keys.count else {
                     state = .done
                     return // compilation key not found.
                 }
-                if try requestCompilationOutputs(cachedComp,
+                if try requestCompilationOutputs(cachedComps,
                                                  dynamicExecutionDelegate: dynamicExecutionDelegate,
-                                                 cacheKey: key,
                                                  jobTaskIDBase: jobTaskIDBase) != 0 {
                     state = .waitingForOutputDownloads
                 } else {
@@ -139,32 +149,33 @@ public final class SwiftCachingMaterializeKeyTaskAction: TaskAction {
     }
 
     private func requestCompilationOutputs(
-        _ cachedComp: SwiftCachedCompilation,
+        _ cachedComps: [SwiftCachedCompilation],
         dynamicExecutionDelegate: any DynamicTaskExecutionDelegate,
-        cacheKey: String,
         jobTaskIDBase: UInt
     ) throws -> UInt {
-        let requested = try cachedComp.getOutputs().reduce(into: UInt(0)) { (numRequested, output) in
-            if output.isMaterialized { return }
-            let outputMaterializeKey = SwiftCachingOutputMaterializerTaskKey(casOptions: taskKey.casOptions,
-                                                                             casID: output.casID,
-                                                                             outputKind: output.kindName,
-                                                                             compilerLocation: taskKey.compilerLocation)
-            dynamicExecutionDelegate.requestDynamicTask(
-                toolIdentifier: SwiftCachingOutputMaterializerTaskAction.toolIdentifier,
-                taskKey: .swiftCachingOutputMaterializer(outputMaterializeKey),
-                taskID: jobTaskIDBase + numRequested,
-                singleUse: true,
-                workingDirectory: Path(""),
-                environment: .init(),
-                forTarget: nil,
-                priority: .network,
-                showEnvironment: false,
-                reason: .wasCompilationCachingQuery
-            )
-            numRequested += 1
+        let numRequested = try cachedComps.reduce(into: UInt(0)) { (numRequested, cachedComp) in
+            try cachedComp.getOutputs().forEach { output in
+                if output.isMaterialized { return }
+                let outputMaterializeKey = SwiftCachingOutputMaterializerTaskKey(casOptions: taskKey.casOptions,
+                                                                                 casID: output.casID,
+                                                                                 outputKind: output.kindName,
+                                                                                 compilerLocation: taskKey.compilerLocation)
+                dynamicExecutionDelegate.requestDynamicTask(
+                    toolIdentifier: SwiftCachingOutputMaterializerTaskAction.toolIdentifier,
+                    taskKey: .swiftCachingOutputMaterializer(outputMaterializeKey),
+                    taskID: jobTaskIDBase + numRequested,
+                    singleUse: true,
+                    workingDirectory: Path(""),
+                    environment: .init(),
+                    forTarget: nil,
+                    priority: .network,
+                    showEnvironment: false,
+                    reason: .wasCompilationCachingQuery
+                )
+                numRequested += 1
+            }
         }
-        return requested
+        return numRequested
     }
 
     public override func performTaskAction(