[cas] Enable the CAS size limiting functionality for Swift caching

rdar://148443907

Author: akyrtzi

Sources/SWBCore/LibSwiftDriver/LibSwiftDriver.swift

@@ -461,6 +461,15 @@ public final class LibSwiftDriver {
         case path(Path)
         case library(libSwiftScanPath: Path)
 
+        public var compilerOrLibraryPath: Path {
+            switch self {
+            case .path(let path):
+                return path
+            case .library(let path):
+                return path
+            }
+        }
+
         public var description: String {
             switch self {
             case .path(let path):
@@ -753,6 +762,14 @@ public final class SwiftCASDatabases {
         self.cas = cas
     }
 
+    public var supportsSizeManagement: Bool { cas.supportsSizeManagement }
+
+    public func getStorageSize() throws -> Int64? { try cas.getStorageSize() }
+
+    public func setSizeLimit(_ size: Int64) throws { try cas.setSizeLimit(size) }
+
+    public func prune() throws { try cas.prune() }
+
     public func queryCacheKey(_ key: String, globally: Bool) async throws -> SwiftCachedCompilation? {
         guard let comp = try await cas.queryCacheKey(key, globally: globally) else { return nil }
         return SwiftCachedCompilation(comp, key: key)

Sources/SWBTaskExecution/DynamicTaskSpecs/CompilationCachingDataPruner.swift

@@ -119,7 +119,7 @@ package final class CompilationCachingDataPruner: Sendable {
                             diagnostic: Diagnostic(
                                 behavior: .remark,
                                 location: .unknown,
-                                data: DiagnosticData("cache size (\(dbSize)) larger than size limit (\(sizeLimit)")
+                                data: DiagnosticData("cache size (\(dbSize)) larger than size limit (\(sizeLimit))")
                             ),
                             for: activityID,
                             signature: signature
@@ -142,6 +142,75 @@ package final class CompilationCachingDataPruner: Sendable {
         }
     }
 
+    package func pruneCAS(
+        _ casDBs: SwiftCASDatabases,
+        key: ClangCachingPruneDataTaskKey,
+        activityReporter: any ActivityReporter,
+        fileSystem fs: any FSProxy
+    ) {
+        let casOpts = key.casOptions
+        guard casOpts.limitingStrategy != .discarded else {
+            return // No need to prune, CAS directory is getting deleted.
+        }
+        let inserted = state.withLock { $0.prunedCASes.insert(key).inserted }
+        guard inserted else {
+            return // already pruned
+        }
+
+        startedAction()
+        let serializer = MsgPackSerializer()
+        key.serialize(to: serializer)
+        let signatureCtx = InsecureHashContext()
+        signatureCtx.add(string: "SwiftCachingPruneData")
+        signatureCtx.add(bytes: serializer.byteString)
+        let signature = signatureCtx.signature
+
+        let casPath = casOpts.casPath.str
+        let swiftscanPath = key.path.str
+
+        // Avoiding the swift concurrency variant because it may lead to starvation when `waitForCompletion()`
+        // blocks on such tasks. Before using a swift concurrency task here make sure there's no deadlock
+        // when setting `LIBDISPATCH_COOPERATIVE_POOL_STRICT`.
+        queue.async {
+            activityReporter.withActivity(
+                ruleInfo: "SwiftCachingPruneData \(casPath) \(swiftscanPath)",
+                executionDescription: "Swift caching pruning \(casPath) using \(swiftscanPath)",
+                signature: signature,
+                target: nil,
+                parentActivity: nil)
+            { activityID in
+                let status: BuildOperationTaskEnded.Status
+                do {
+                    let dbSize = try casDBs.getStorageSize()
+                    let sizeLimit = try computeCASSizeLimit(casOptions: casOpts, dbSize: dbSize.map{Int($0)}, fileSystem: fs)
+                    if casOpts.enableDiagnosticRemarks, let dbSize, let sizeLimit, sizeLimit < dbSize {
+                        activityReporter.emit(
+                            diagnostic: Diagnostic(
+                                behavior: .remark,
+                                location: .unknown,
+                                data: DiagnosticData("cache size (\(dbSize)) larger than size limit (\(sizeLimit))")
+                            ),
+                            for: activityID,
+                            signature: signature
+                        )
+                    }
+                    try casDBs.setSizeLimit(Int64(sizeLimit ?? 0))
+                    try casDBs.prune()
+                    status = .succeeded
+                } catch {
+                    activityReporter.emit(
+                        diagnostic: Diagnostic(behavior: .error, location: .unknown, data: DiagnosticData(error.localizedDescription)),
+                        for: activityID,
+                        signature: signature
+                    )
+                    status = .failed
+                }
+                return status
+            }
+            self.finishedAction()
+        }
+    }
+
     package func pruneCAS(
         _ toolchainCAS: ToolchainCAS,
         key: ClangCachingPruneDataTaskKey,
@@ -188,7 +257,7 @@ package final class CompilationCachingDataPruner: Sendable {
                             diagnostic: Diagnostic(
                                 behavior: .remark,
                                 location: .unknown,
-                                data: DiagnosticData("cache size (\(dbSize)) larger than size limit (\(sizeLimit)")
+                                data: DiagnosticData("cache size (\(dbSize)) larger than size limit (\(sizeLimit))")
                             ),
                             for: activityID,
                             signature: signature

Sources/SWBTaskExecution/TaskActions/SwiftDriverJobTaskAction.swift

@@ -494,6 +494,17 @@ public final class SwiftDriverJobTaskAction: TaskAction, BuildValueValidatingTas
             if let casOpts = payload.casOptions, casOpts.enableIntegratedCacheQueries {
                 let swiftModuleDependencyGraph = dynamicExecutionDelegate.operationContext.swiftModuleDependencyGraph
                 cas = try swiftModuleDependencyGraph.getCASDatabases(casOptions: casOpts, compilerLocation: payload.compilerLocation)
+
+                let casKey = ClangCachingPruneDataTaskKey(
+                    path: payload.compilerLocation.compilerOrLibraryPath,
+                    casOptions: casOpts
+                )
+                dynamicExecutionDelegate.operationContext.compilationCachingDataPruner.pruneCAS(
+                    cas!,
+                    key: casKey,
+                    activityReporter: dynamicExecutionDelegate,
+                    fileSystem: executionDelegate.fs
+                )
             } else {
                 cas = nil
             }

Sources/SWBTestSupport/BuildOperationTester.swift

@@ -631,6 +631,25 @@ package final class BuildOperationTester {
             return nil
         }
 
+        package func getDiagnosticMessageForTask(_ pattern: StringPattern, kind: DiagnosticKind, task: Task) -> String? {
+            for (index, event) in self.events.enumerated() {
+                switch event {
+                case .taskHadEvent(let eventTask, event: .hadDiagnostic(let diagnostic)) where diagnostic.behavior == kind:
+                    guard eventTask == task else {
+                        continue
+                    }
+                    let message = diagnostic.formatLocalizedDescription(.debugWithoutBehavior, task: eventTask)
+                    if pattern ~= message {
+                        _eventList.remove(at: index)
+                        return message
+                    }
+                default:
+                    continue
+                }
+            }
+            return nil
+        }
+
         package func check(_ pattern: StringPattern, kind: BuildOperationTester.DiagnosticKind, failIfNotFound: Bool, sourceLocation: SourceLocation, checkDiagnostic: (Diagnostic) -> Bool) -> Bool {
             let found = (getDiagnosticMessage(pattern, kind: kind, checkDiagnostic: checkDiagnostic) != nil)
             if !found, failIfNotFound {

Tests/SWBBuildSystemTests/SwiftCompilationCachingTests.swift

@@ -126,14 +126,93 @@ fileprivate struct SwiftCompilationCachingTests: CoreBasedTests {
             #expect(try readMetrics("two").contains("\"swiftCacheHits\":\(numCompile),\"swiftCacheMisses\":0"))
         }
     }
+
+    @Test(.requireSDKs(.macOS))
+    func swiftCASLimiting() async throws {
+        try await withTemporaryDirectory { (tmpDirPath: Path) async throws -> Void in
+            let testWorkspace = try await TestWorkspace(
+                "Test",
+                sourceRoot: tmpDirPath.join("Test"),
+                projects: [
+                    TestProject(
+                        "aProject",
+                        groupTree: TestGroup(
+                            "Sources",
+                            children: [
+                                TestFile("main.swift"),
+                            ]),
+                        buildConfigurations: [
+                            TestBuildConfiguration(
+                                "Debug",
+                                buildSettings: [
+                                    "PRODUCT_NAME": "$(TARGET_NAME)",
+                                    "SDKROOT": "macosx",
+                                    "SWIFT_VERSION": swiftVersion,
+                                    "SWIFT_ENABLE_EXPLICIT_MODULES": "YES",
+                                    "SWIFT_ENABLE_COMPILE_CACHE": "YES",
+                                    "COMPILATION_CACHE_ENABLE_DIAGNOSTIC_REMARKS": "YES",
+                                    "COMPILATION_CACHE_LIMIT_SIZE": "1",
+                                    "COMPILATION_CACHE_CAS_PATH": tmpDirPath.join("CompilationCache").str,
+                                    "DSTROOT": tmpDirPath.join("dstroot").str,
+                                ]),
+                        ],
+                        targets: [
+                            TestStandardTarget(
+                                "tool",
+                                type: .framework,
+                                buildPhases: [
+                                    TestSourcesBuildPhase([
+                                        "main.swift",
+                                    ]),
+                                ]
+                            )
+                        ])
+                ])
+            let tester = try await BuildOperationTester(getCore(), testWorkspace, simulated: false)
+
+            try await tester.fs.writeFileContents(tmpDirPath.join("Test/aProject/main.swift")) {
+                $0 <<< "let x = 1\n"
+            }
+
+            try await tester.checkBuild(runDestination: .macOS, persistent: true) { results in
+                results.checkTask(.matchRuleType("SwiftCompile")) { results.checkKeyQueryCacheMiss($0) }
+            }
+            try await tester.checkBuild(runDestination: .macOS, buildCommand: .cleanBuildFolder(style: .regular), body: { _ in })
+
+            // Update the source file and rebuild.
+            try await tester.fs.writeFileContents(tmpDirPath.join("Test/aProject/main.swift")) {
+                $0 <<< "let x = 2\n"
+            }
+            try await tester.checkBuild(runDestination: .macOS, persistent: true) { results in
+                results.checkTask(.matchRuleType("SwiftCompile")) { results.checkKeyQueryCacheMiss($0) }
+            }
+            try await tester.checkBuild(runDestination: .macOS, buildCommand: .cleanBuildFolder(style: .regular), body: { _ in })
+
+            // Revert the source change and rebuild. It should still be a cache miss because of CAS size limiting.
+            try await tester.fs.writeFileContents(tmpDirPath.join("Test/aProject/main.swift")) {
+                $0 <<< "let x = 1\n"
+            }
+            try await tester.checkBuild(runDestination: .macOS, persistent: true) { results in
+                results.checkTask(.matchRuleType("SwiftCompile")) { results.checkKeyQueryCacheMiss($0) }
+            }
+        }
+    }
 }
 
 extension BuildOperationTester.BuildResults {
-    fileprivate func checkKeyQueryCacheMiss(_ task: Task, file: StaticString = #file, line: UInt = #line) {
-        checkRemark(.contains("cache key query miss"))
+    fileprivate func checkKeyQueryCacheMiss(_ task: Task, sourceLocation: SourceLocation = #_sourceLocation) {
+        let found = (getDiagnosticMessageForTask(.contains("cache miss"), kind: .remark, task: task) != nil)
+        guard found else {
+            Issue.record("Unable to find cache miss diagnostic for task \(task)", sourceLocation: sourceLocation)
+            return
+        }
     }
 
-    fileprivate func checkKeyQueryCacheHit(_ task: Task, file: StaticString = #file, line: UInt = #line) {
-        checkRemark(.contains("cache key query hit"))
+    fileprivate func checkKeyQueryCacheHit(_ task: Task, sourceLocation: SourceLocation = #_sourceLocation) {
+        let found = (getDiagnosticMessageForTask(.contains("cache found for key"), kind: .remark, task: task) != nil)
+        guard found else {
+            Issue.record("Unable to find cache hit diagnostic for task \(task)", sourceLocation: sourceLocation)
+            return
+        }
     }
 }