Merge pull request #797 from owenv/owenv/description-refcounting

Add API to allow clients to manage build description lifecycle

Author: owenv

Sources/SWBBuildService/BuildDescriptionMessages.swift

@@ -54,7 +54,8 @@ fileprivate extension Request {
                 message.buildDescriptionID,
                 request: buildRequest,
                 buildRequestContext: buildRequestContext,
-                workspaceContext: workspaceContext
+                workspaceContext: workspaceContext,
+                retain: false
             ), clientDelegate: clientDelegate, constructionDelegate: operation
         )?.buildDescription
 
@@ -220,3 +221,11 @@ struct IndexBuildSettingsMsg: MessageHandler {
         return IndexBuildSettingsResponse(compilerArguments: compilerArguments)
     }
 }
+
+struct ReleaseBuildDescriptionMsg: MessageHandler {
+    func handle(request: Request, message: ReleaseBuildDescriptionRequest) async throws -> VoidResponse {
+        let session = try request.session(for: message)
+        session.buildDescriptionManager.releaseBuildDescription(id: message.buildDescriptionID)
+        return VoidResponse()
+    }
+}

Sources/SWBBuildService/BuildOperationMessages.swift

@@ -165,6 +165,9 @@ final class ActiveBuild: ActiveBuildOperation {
     /// Whether this operation is intended only for creating and reporting the build description.
     let onlyCreatesBuildDescription: Bool
 
+    /// Whether this operation should retain the build description it uses.
+    let retainBuildDescription: Bool
+
     /// The current state of the build.
     var state: State
 
@@ -194,6 +197,7 @@ final class ActiveBuild: ActiveBuildOperation {
         self.id = request.buildService.nextBuildOperationID()
 
         self.onlyCreatesBuildDescription = message.onlyCreateBuildDescription
+        self.retainBuildDescription = message.retainBuildDescription == true
 
         self.session = try request.session(for: message)
         guard let workspaceContext = session.workspaceContext else {
@@ -420,7 +424,7 @@ final class ActiveBuild: ActiveBuildOperation {
             let preparationDelegate = self.preparationProgressDelegate!
             let clientDelegate = ClientExchangeDelegate(request: self.request, session: self.session)
             // FIXME: We should have a channel for reporting errors here which don't make it look like there was an internal service error. E.g., if we fail to create or write the build description or manifest because of some error outside of our control, we should simply report that and not make it look like we might have a bug.
-            let description = try await MacroNamespace.withExpressionInterningEnabled { try await self.session.buildDescriptionManager.getBuildDescription(planRequest, clientDelegate: clientDelegate, constructionDelegate: preparationDelegate) }
+            let description = try await MacroNamespace.withExpressionInterningEnabled { try await self.session.buildDescriptionManager.getBuildDescription(planRequest, retained: retainBuildDescription, clientDelegate: clientDelegate, constructionDelegate: preparationDelegate) }
             return description
         } catch {
             self.abortBuild(error)
@@ -445,7 +449,7 @@ final class ActiveBuild: ActiveBuildOperation {
 
         do {
             let clientDelegate = ClientExchangeDelegate(request: self.request, session: self.session)
-            let descRequest = BuildDescriptionManager.BuildDescriptionRequest.cachedOnly(buildDescriptionID, request: self.buildRequest, buildRequestContext: self.buildRequestContext, workspaceContext: self.workspaceContext)
+            let descRequest = BuildDescriptionManager.BuildDescriptionRequest.cachedOnly(buildDescriptionID, request: self.buildRequest, buildRequestContext: self.buildRequestContext, workspaceContext: self.workspaceContext, retain: self.retainBuildDescription)
             let retrievedBuildDescription = try await self.session.buildDescriptionManager.getNewOrCachedBuildDescription(descRequest, clientDelegate: clientDelegate, constructionDelegate: self.preparationProgressDelegate!)
             return retrievedBuildDescription?.buildDescription
         } catch {

Sources/SWBBuildService/DocumentationInfo.swift

@@ -51,7 +51,7 @@ extension BuildDescriptionManager {
         // Get the complete build description.
         let buildDescription: BuildDescription
         do {
-            if let retrievedBuildDescription = try await getBuildDescription(planRequest, clientDelegate: delegate.clientDelegate, constructionDelegate: delegate) {
+            if let retrievedBuildDescription = try await getBuildDescription(planRequest, retained: false, clientDelegate: delegate.clientDelegate, constructionDelegate: delegate) {
                 buildDescription = retrievedBuildDescription
             } else {
                 // If we don't receive a build description it means we were cancelled.

Sources/SWBBuildService/LocalizationInfo.swift

@@ -68,7 +68,7 @@ extension BuildDescriptionManager {
 
         let buildDescription: BuildDescription
         do {
-            if let retrievedBuildDescription = try await getNewOrCachedBuildDescription(.cachedOnly(descriptionID, request: buildRequest, buildRequestContext: buildRequestContext, workspaceContext: workspaceContext), clientDelegate: delegate.clientDelegate, constructionDelegate: delegate)?.buildDescription {
+            if let retrievedBuildDescription = try await getNewOrCachedBuildDescription(.cachedOnly(descriptionID, request: buildRequest, buildRequestContext: buildRequestContext, workspaceContext: workspaceContext, retain: false), clientDelegate: delegate.clientDelegate, constructionDelegate: delegate)?.buildDescription {
                 buildDescription = retrievedBuildDescription
             } else {
                 // If we don't receive a build description it means we were cancelled.

Sources/SWBBuildService/Messages.swift

@@ -585,7 +585,7 @@ fileprivate enum ResultOrErrorMessage<T> {
 fileprivate func getIndexBuildDescriptionFromID(buildDescriptionID: BuildDescriptionID, request: Request, session: Session, buildRequest: BuildRequest, buildRequestContext: BuildRequestContext, workspaceContext: WorkspaceContext, constructionDelegate: any BuildDescriptionConstructionDelegate) async -> ResultOrErrorMessage<BuildDescription> {
     let clientDelegate = ClientExchangeDelegate(request: request, session: session)
     do {
-        let descRequest = BuildDescriptionManager.BuildDescriptionRequest.cachedOnly(buildDescriptionID, request: buildRequest, buildRequestContext: buildRequestContext, workspaceContext: workspaceContext)
+        let descRequest = BuildDescriptionManager.BuildDescriptionRequest.cachedOnly(buildDescriptionID, request: buildRequest, buildRequestContext: buildRequestContext, workspaceContext: workspaceContext, retain: false)
         guard let retrievedBuildDescription = try await session.buildDescriptionManager.getNewOrCachedBuildDescription(descRequest, clientDelegate: clientDelegate, constructionDelegate: constructionDelegate) else {
             // If we don't receive a build description it means we were cancelled.
             return .failed(VoidResponse())
@@ -727,7 +727,7 @@ extension MessageHandler {
                     let clientDelegate = ClientExchangeDelegate(request: request, session: session)
                     let buildDescription: BuildDescription
                     do {
-                        if let retrievedBuildDescription = try await session.buildDescriptionManager.getBuildDescription(planRequest, clientDelegate: clientDelegate, constructionDelegate: operation) {
+                        if let retrievedBuildDescription = try await session.buildDescriptionManager.getBuildDescription(planRequest, retained: false, clientDelegate: clientDelegate, constructionDelegate: operation) {
                             buildDescription = retrievedBuildDescription
                         } else {
                             // If we don't receive a build description it means we were cancelled.
@@ -1626,6 +1626,7 @@ package struct ServiceMessageHandlers: ServiceExtension {
         service.registerMessageHandler(BuildDescriptionConfiguredTargetsMsg.self)
         service.registerMessageHandler(BuildDescriptionConfiguredTargetSourcesMsg.self)
         service.registerMessageHandler(IndexBuildSettingsMsg.self)
+        service.registerMessageHandler(ReleaseBuildDescriptionMsg.self)
 
         service.registerMessageHandler(MacroEvaluationMsg.self)
         service.registerMessageHandler(AllExportedMacrosAndValuesMsg.self)

Sources/SWBBuildService/PreviewInfo.swift

@@ -123,6 +123,7 @@ extension BuildDescriptionManager {
         do {
             if let retrievedBuildDescription = try await getBuildDescription(
                 planRequest,
+                retained: false,
                 clientDelegate: delegate.clientDelegate,
                 constructionDelegate: delegate
             ) {

Sources/SWBProtocol/BuildDescriptionMessages.swift

@@ -213,6 +213,24 @@ public struct IndexBuildSettingsResponse: Message, SerializableCodable, Equatabl
     }
 }
 
+public struct ReleaseBuildDescriptionRequest: SessionMessage, RequestMessage, SerializableCodable, Equatable {
+    public typealias ResponseMessage = VoidResponse
+
+    public static let name = "RELEASE_BUILD_DESCRIPTION"
+
+    public var sessionHandle: String
+
+    public let buildDescriptionID: BuildDescriptionID
+
+    public init(
+        sessionHandle: String,
+        buildDescriptionID: BuildDescriptionID
+    ) {
+        self.sessionHandle = sessionHandle
+        self.buildDescriptionID = buildDescriptionID
+    }
+}
+
 // MARK: Registering messages
 
 let buildDescriptionMessages: [any Message.Type] = [
@@ -222,4 +240,5 @@ let buildDescriptionMessages: [any Message.Type] = [
     BuildDescriptionConfiguredTargetSourcesResponse.self,
     IndexBuildSettingsRequest.self,
     IndexBuildSettingsResponse.self,
+    ReleaseBuildDescriptionRequest.self,
 ]

Sources/SWBProtocol/BuildOperationMessages.swift

@@ -263,19 +263,24 @@ public struct CreateBuildRequest: SessionChannelBuildMessage, RequestMessage, Se
     /// If true, the build operation will be completed after the build description is created and reported.
     public let onlyCreateBuildDescription: Bool
 
+    /// If true, the build description for this build wil be kept alive in memory, even if it would otherwise be evicted from caches.
+    public let retainBuildDescription: Bool?
+
+    @available(*, deprecated, renamed: "init(sessionHandle:responseChannel:request:onlyCreateBuildDescription:retainBuildDescription:)")
     public init(sessionHandle: String, responseChannel: UInt64, request: BuildRequestMessagePayload, onlyCreateBuildDescription: Bool) {
         self.sessionHandle = sessionHandle
         self.responseChannel = responseChannel
         self.request = request
         self.onlyCreateBuildDescription = onlyCreateBuildDescription
+        self.retainBuildDescription = nil
     }
 
-    public init(fromLegacy deserializer: any Deserializer) throws {
-        try deserializer.beginAggregate(4)
-        self.sessionHandle = try deserializer.deserialize()
-        self.responseChannel = try deserializer.deserialize()
-        self.request = try deserializer.deserialize()
-        self.onlyCreateBuildDescription = try deserializer.deserialize()
+    public init(sessionHandle: String, responseChannel: UInt64, request: BuildRequestMessagePayload, onlyCreateBuildDescription: Bool, retainBuildDescription: Bool) {
+        self.sessionHandle = sessionHandle
+        self.responseChannel = responseChannel
+        self.request = request
+        self.onlyCreateBuildDescription = onlyCreateBuildDescription
+        self.retainBuildDescription = retainBuildDescription
     }
 }
 

Sources/SWBTaskExecution/BuildDescriptionManager.swift

@@ -90,6 +90,9 @@ package final class BuildDescriptionManager: Sendable {
     /// The in-memory cache of build descriptions.
     private let inMemoryCachedBuildDescriptions: HeavyCache<BuildDescriptionSignature, BuildDescription>
 
+    /// Build descriptions explicitly retained by clients.
+    private let retainedBuildDescriptions: Registry<BuildDescriptionSignature, (BuildDescription, UInt)> = .init()
+
     /// The last build plan request. Used to generate a diff of the current plan for debugging purposes.
     private let lastBuildPlanRequest: SWBMutex<BuildPlanRequest?> = .init(nil)
 
@@ -254,35 +257,44 @@ package final class BuildDescriptionManager: Sendable {
     /// During normal operation (outside of tests), this should always be called on `queue`.
     package enum BuildDescriptionRequest {
         /// Retrieve or create a build description based on a build plan.
-        case newOrCached(BuildPlanRequest, bypassActualTasks: Bool, useSynchronousBuildDescriptionSerialization: Bool)
+        case newOrCached(BuildPlanRequest, bypassActualTasks: Bool, useSynchronousBuildDescriptionSerialization: Bool, retain: Bool)
         /// Retrieve an existing build description, build planning has been avoided. If the build description is not available then `getNewOrCachedBuildDescription` will fail.
-        case cachedOnly(BuildDescriptionID, request: BuildRequest, buildRequestContext: BuildRequestContext, workspaceContext: WorkspaceContext)
+        case cachedOnly(BuildDescriptionID, request: BuildRequest, buildRequestContext: BuildRequestContext, workspaceContext: WorkspaceContext, retain: Bool)
 
         var buildRequest: BuildRequest {
             switch self {
-            case .newOrCached(let planRequest, _, _): return planRequest.buildRequest
-            case .cachedOnly(_, let request, _, _): return request
+            case .newOrCached(let planRequest, _, _, _): return planRequest.buildRequest
+            case .cachedOnly(_, let request, _, _, _): return request
             }
         }
 
         var buildRequestContext: BuildRequestContext {
             switch self {
-            case .newOrCached(let planRequest, _, _): return planRequest.buildRequestContext
-            case .cachedOnly(_, _, let buildRequestContext, _): return buildRequestContext
+            case .newOrCached(let planRequest, _, _, _): return planRequest.buildRequestContext
+            case .cachedOnly(_, _, let buildRequestContext, _, _): return buildRequestContext
             }
         }
 
         var planRequest: BuildPlanRequest? {
             switch self {
-            case .newOrCached(let planRequest, _, _): return planRequest
+            case .newOrCached(let planRequest, _, _, _): return planRequest
             case .cachedOnly: return nil
             }
         }
 
         var workspaceContext: WorkspaceContext {
             switch self {
-            case .newOrCached(let planRequest, _, _): return planRequest.workspaceContext
-            case .cachedOnly(_, _, _, let workspaceContext): return workspaceContext
+            case .newOrCached(let planRequest, _, _, _): return planRequest.workspaceContext
+            case .cachedOnly(_, _, _, let workspaceContext, _): return workspaceContext
+            }
+        }
+
+        var retain: Bool {
+            switch self {
+            case .newOrCached(_, _, _, let retain):
+                retain
+            case .cachedOnly(_, _, _, _, let retain):
+                retain
             }
         }
 
@@ -305,8 +317,8 @@ package final class BuildDescriptionManager: Sendable {
 
         func signature(cacheDir: Path) throws -> BuildDescriptionSignature {
             switch self {
-            case .newOrCached(let planRequest, _, _): return try BuildDescriptionSignature.buildDescriptionSignature(planRequest, cacheDir: cacheDir)
-            case .cachedOnly(let buildDescriptionID, _, _, _):  return BuildDescriptionSignature.buildDescriptionSignature(buildDescriptionID)
+            case .newOrCached(let planRequest, _, _, _): return try BuildDescriptionSignature.buildDescriptionSignature(planRequest, cacheDir: cacheDir)
+            case .cachedOnly(let buildDescriptionID, _, _, _, _):  return BuildDescriptionSignature.buildDescriptionSignature(buildDescriptionID)
             }
         }
     }
@@ -317,6 +329,8 @@ package final class BuildDescriptionManager: Sendable {
             description = lastDescription
         } else if let inMemoryDescription = inMemoryCachedBuildDescriptions[signature] {
             description = inMemoryDescription
+        } else if let retainedDescription = retainedBuildDescriptions[signature] {
+            description = retainedDescription.0
         } else {
             description = nil
         }
@@ -397,18 +411,35 @@ package final class BuildDescriptionManager: Sendable {
             }
         }
 
+        if request.retain {
+            retainedBuildDescriptions.update(signature, update: { ($0.0, $0.1 + 1) }, default: { (buildDescription, 0) })
+        }
+
         return BuildDescriptionRetrievalInfo(buildDescription: buildDescription, source: source, inMemoryCacheSize: inMemoryCachedBuildDescriptions.count, onDiskCachePath: buildDescriptionPath)
     }
 
     /// Returns a build description for a particular workspace and request.
     ///
     /// - Returns: A build description, or nil if cancelled.
-    package func getBuildDescription(_ request: BuildPlanRequest, bypassActualTasks: Bool = false, useSynchronousBuildDescriptionSerialization: Bool = false, clientDelegate: any TaskPlanningClientDelegate, constructionDelegate: any BuildDescriptionConstructionDelegate) async throws -> BuildDescription? {
-        let descRequest = BuildDescriptionRequest.newOrCached(request, bypassActualTasks: bypassActualTasks, useSynchronousBuildDescriptionSerialization: useSynchronousBuildDescriptionSerialization)
+    package func getBuildDescription(_ request: BuildPlanRequest, bypassActualTasks: Bool = false, useSynchronousBuildDescriptionSerialization: Bool = false, retained: Bool, clientDelegate: any TaskPlanningClientDelegate, constructionDelegate: any BuildDescriptionConstructionDelegate) async throws -> BuildDescription? {
+        let descRequest = BuildDescriptionRequest.newOrCached(request, bypassActualTasks: bypassActualTasks, useSynchronousBuildDescriptionSerialization: useSynchronousBuildDescriptionSerialization, retain: retained)
         let retrievalInfo = try await getNewOrCachedBuildDescription(descRequest, clientDelegate: clientDelegate, constructionDelegate: constructionDelegate)
         return retrievalInfo?.buildDescription
     }
 
+    package func releaseBuildDescription(id: BuildDescriptionID) {
+        self.retainedBuildDescriptions.update(BuildDescriptionSignature.buildDescriptionSignature(id), update: {
+            let newCount = $0.1 - 1
+            if newCount == 0 {
+                return nil
+            } else {
+                return ($0.0, newCount)
+            }
+        }, default: {
+            nil
+        })
+    }
+
     /// Returns the path in which the`XCBuildData` directory will live. That location is uses to cache build descriptions for a particular workspace and request, the manifest, and the `build.db` database for llbuild.
     package static func cacheDirectory(_ request: BuildPlanRequest) throws -> Path {
         return try cacheDirectory(request.buildRequest, buildRequestContext: request.buildRequestContext, workspaceContext: request.workspaceContext)
@@ -514,7 +545,7 @@ package final class BuildDescriptionManager: Sendable {
         }
 
         // Unable to load from disk, create a new description
-        guard case let .newOrCached(request, bypassActualTasks, useSynchronousBuildDescriptionSerialization) = request else {
+        guard case let .newOrCached(request, bypassActualTasks, useSynchronousBuildDescriptionSerialization, _) = request else {
             preconditionFailure("entered build construction path but request was for existing cached description")
         }
 

Sources/SWBTestSupport/TaskExecutionTestSupport.swift

@@ -101,7 +101,7 @@ extension BuildDescription {
 
 extension BuildDescriptionManager {
     package func getNewOrCachedBuildDescription(_ request: BuildPlanRequest, bypassActualTasks: Bool = false, clientDelegate: any TaskPlanningClientDelegate, constructionDelegate: any BuildDescriptionConstructionDelegate) async throws -> BuildDescriptionRetrievalInfo? {
-        let descRequest = BuildDescriptionRequest.newOrCached(request, bypassActualTasks: bypassActualTasks, useSynchronousBuildDescriptionSerialization: true)
+        let descRequest = BuildDescriptionRequest.newOrCached(request, bypassActualTasks: bypassActualTasks, useSynchronousBuildDescriptionSerialization: true, retain: false)
         return try await getNewOrCachedBuildDescription(descRequest, clientDelegate: clientDelegate, constructionDelegate: constructionDelegate)
     }
 }