Add session API for registering additional toolchains

owenv · owenv · commit a65ecf6c163d · 2025-09-05T16:09:30.000-07:00
diff --git a/Sources/SWBBuildService/Messages.swift b/Sources/SWBBuildService/Messages.swift
@@ -302,6 +302,13 @@ private struct SetSessionUserPreferencesMsg: MessageHandler {
     }
 }
 
+private struct RegisterToolchainHandler: MessageHandler {
+    func handle(request: Request, message: RegisterToolchainRequest) async throws -> StringResponse {
+        let session = try request.session(for: message)
+        return StringResponse(try await session.core.registerToolchain(at: message.path))
+    }
+}
+
 /// Start a PIF transfer from the client.
 ///
 /// This will establish a workspace context in the relevant session by exchanging a PIF from the client to the service incrementally, only transferring subobjects as necessary.
@@ -1542,6 +1549,7 @@ public struct ServiceSessionMessageHandlers: ServiceExtension {
         service.registerMessageHandler(SetSessionSystemInfoMsg.self)
         service.registerMessageHandler(SetSessionUserInfoMsg.self)
         service.registerMessageHandler(SetSessionUserPreferencesMsg.self)
+        service.registerMessageHandler(RegisterToolchainHandler.self)
         service.registerMessageHandler(DeveloperPathHandler.self)
     }
 }
diff --git a/Sources/SWBCore/Core.swift b/Sources/SWBCore/Core.swift
@@ -533,6 +533,10 @@ public final class Core: Sendable {
         specRegistry.freeze()
     }
 
+    public func registerToolchain(at toolchainPath: Path) async throws -> String {
+        return try await toolchainRegistry.registerToolchain(at: toolchainPath, operatingSystem: hostOperatingSystem, delegate: registryDelegate, diagnoseAlreadyRegisteredToolchain: false, aliases: [])
+    }
+
     /// Dump information on the registered platforms.
     public func getPlatformsDump() -> String {
         var result = ""
@@ -615,7 +619,7 @@ public final class Core: Sendable {
     /// Dump information on the registered toolchains.
     public func getToolchainsDump() async -> String {
         var result = ""
-        for (_,toolchain) in toolchainRegistry.toolchainsByIdentifier.sorted(byKey: <) {
+        for toolchain in toolchainRegistry.toolchains.sorted(by: \.identifier) {
             result += "\(toolchain)\n"
         }
         return result
diff --git a/Sources/SWBCore/ToolchainRegistry.swift b/Sources/SWBCore/ToolchainRegistry.swift
@@ -428,6 +428,15 @@ extension Toolchain {
 
 /// The ToolchainRegistry manages the set of registered toolchains.
 public final class ToolchainRegistry: @unchecked Sendable {
+    enum Error: Swift.Error, CustomStringConvertible {
+        case toolchainAlreadyRegistered(String, Path)
+
+        var description: String {
+            switch self {
+            case .toolchainAlreadyRegistered(let identifier, let path): "toolchain '\(identifier)' already registered from \(path.str)"
+            }
+        }
+    }
     @_spi(Testing) public struct SearchPath: Sendable {
         public var path: Path
         public var strict: Bool
@@ -443,11 +452,14 @@ public final class ToolchainRegistry: @unchecked Sendable {
     let fs: any FSProxy
     let hostOperatingSystem: OperatingSystem
 
-    /// The map of toolchains by identifier.
-    @_spi(Testing) public private(set) var toolchainsByIdentifier = Dictionary<String, Toolchain>()
+    struct State {
+        /// The map of toolchains by identifier.
+        @_spi(Testing) public fileprivate(set) var toolchainsByIdentifier = Dictionary<String, Toolchain>()
 
-    /// Lower-cased alias -> toolchain (alias lookup is case-insensitive)
-    @_spi(Testing) public private(set) var toolchainsByAlias = Dictionary<String, Toolchain>()
+        /// Lower-cased alias -> toolchain (alias lookup is case-insensitive)
+        @_spi(Testing) public fileprivate(set) var toolchainsByAlias = Dictionary<String, Toolchain>()
+    }
+    private let state: SWBMutex<State> = .init(State())
 
     public static let defaultToolchainIdentifier: String = "com.apple.dt.toolchain.XcodeDefault"
 
@@ -503,40 +515,53 @@ public final class ToolchainRegistry: @unchecked Sendable {
             guard toolchainPath.basenameWithoutSuffix != "swift-latest" else { continue }
 
             do {
-                let toolchain = try await Toolchain(path: toolchainPath, operatingSystem: operatingSystem, aliases: aliases, fs: fs, pluginManager: delegate.pluginManager, platformRegistry: delegate.platformRegistry)
-                try register(toolchain)
+                _ = try await registerToolchain(at: toolchainPath, operatingSystem: operatingSystem, delegate: delegate, diagnoseAlreadyRegisteredToolchain: true, aliases: aliases)
             } catch let err {
                 delegate.issue(strict: strict, toolchainPath, "failed to load toolchain: \(err)")
             }
         }
     }
 
-    private func register(_ toolchain: Toolchain) throws {
-        if let duplicateToolchain = toolchainsByIdentifier[toolchain.identifier] {
-            throw StubError.error("toolchain '\(toolchain.identifier)' already registered from \(duplicateToolchain.path.str)")
+    func registerToolchain(at toolchainPath: Path, operatingSystem: OperatingSystem, delegate: any ToolchainRegistryDelegate, diagnoseAlreadyRegisteredToolchain: Bool, aliases: Set<String>) async throws -> String {
+        do {
+            let toolchain = try await Toolchain(path: toolchainPath, operatingSystem: operatingSystem, aliases: aliases, fs: fs, pluginManager: delegate.pluginManager, platformRegistry: delegate.platformRegistry)
+            try register(toolchain)
+            return toolchain.identifier
+        } catch Error.toolchainAlreadyRegistered(let identifier, _) where !diagnoseAlreadyRegisteredToolchain {
+            return identifier
         }
-        toolchainsByIdentifier[toolchain.identifier] = toolchain
-
-        for alias in toolchain.aliases {
-            guard !alias.isEmpty else { continue }
-            assert(alias.lowercased() == alias)
+    }
 
-            // When two toolchains have conflicting aliases, the highest-versioned toolchain wins (regardless of identifier)
-            if let existingToolchain = toolchainsByAlias[alias], existingToolchain.version >= toolchain.version {
-                continue
+    private func register(_ toolchain: Toolchain) throws {
+        try state.withLock { state in
+            if let duplicateToolchain = state.toolchainsByIdentifier[toolchain.identifier] {
+                throw Error.toolchainAlreadyRegistered(toolchain.identifier, duplicateToolchain.path)
             }
+            state.toolchainsByIdentifier[toolchain.identifier] = toolchain
+
+            for alias in toolchain.aliases {
+                guard !alias.isEmpty else { continue }
+                assert(alias.lowercased() == alias)
+
+                // When two toolchains have conflicting aliases, the highest-versioned toolchain wins (regardless of identifier)
+                if let existingToolchain = state.toolchainsByAlias[alias], existingToolchain.version >= toolchain.version {
+                    continue
+                }
 
-            toolchainsByAlias[alias] = toolchain
+                state.toolchainsByAlias[alias] = toolchain
+            }
         }
     }
 
     /// Look up the toolchain with the given identifier.
     public func lookup(_ identifier: String) -> Toolchain? {
-        let lowercasedIdentifier = identifier.lowercased()
-        if ["default", "xcode"].contains(lowercasedIdentifier) {
-            return toolchainsByIdentifier[ToolchainRegistry.defaultToolchainIdentifier] ?? toolchainsByAlias[lowercasedIdentifier]
-        } else {
-            return toolchainsByIdentifier[identifier] ?? toolchainsByAlias[lowercasedIdentifier]
+        state.withLock { state in
+            let lowercasedIdentifier = identifier.lowercased()
+            if ["default", "xcode"].contains(lowercasedIdentifier) {
+                return state.toolchainsByIdentifier[ToolchainRegistry.defaultToolchainIdentifier] ?? state.toolchainsByAlias[lowercasedIdentifier]
+            } else {
+                return state.toolchainsByIdentifier[identifier] ?? state.toolchainsByAlias[lowercasedIdentifier]
+            }
         }
     }
 
@@ -545,6 +570,8 @@ public final class ToolchainRegistry: @unchecked Sendable {
     }
 
     public var toolchains: Set<Toolchain> {
-        return Set(self.toolchainsByIdentifier.values)
+        state.withLock { state in
+            return Set(state.toolchainsByIdentifier.values)
+        }
     }
 }
diff --git a/Sources/SWBProtocol/Message.swift b/Sources/SWBProtocol/Message.swift
@@ -356,6 +356,20 @@ public struct MacCatalystUnavailableFrameworkNamesRequest: RequestMessage, Equat
     }
 }
 
+public struct RegisterToolchainRequest: SessionMessage, RequestMessage, Equatable, SerializableCodable {
+    public typealias ResponseMessage = StringResponse
+
+    public static let name = "REGISTER_TOOLCHAIN_REQUEST"
+
+    public let sessionHandle: String
+    public let path: Path
+
+    public init(sessionHandle: String, path: Path) {
+        self.sessionHandle = sessionHandle
+        self.path = path
+    }
+}
+
 public struct AppleSystemFrameworkNamesRequest: RequestMessage, Equatable, PendingSerializableCodable {
     public typealias ResponseMessage = StringListResponse
 
@@ -1168,6 +1182,7 @@ public struct IPCMessage: Serializable, Sendable {
         GetSpecsRequest.self,
         GetStatisticsRequest.self,
         GetToolchainsRequest.self,
+        RegisterToolchainRequest.self,
         GetBuildSettingsDescriptionRequest.self,
         ExecuteCommandLineToolRequest.self,
 
diff --git a/Sources/SwiftBuild/SWBBuildServiceSession.swift b/Sources/SwiftBuild/SWBBuildServiceSession.swift
@@ -629,6 +629,10 @@ public final class SWBBuildServiceSession: Sendable {
     public func setUserPreferences(enableDebugActivityLogs: Bool, enableBuildDebugging: Bool, enableBuildSystemCaching: Bool, activityTextShorteningLevel: Int, usePerConfigurationBuildLocations: Bool?, allowsExternalToolExecution: Bool) async throws {
         _ = try await service.send(request: SetSessionUserPreferencesRequest(sessionHandle: self.uid, enableDebugActivityLogs: enableDebugActivityLogs, enableBuildDebugging: enableBuildDebugging, enableBuildSystemCaching: enableBuildSystemCaching, activityTextShorteningLevel: ActivityTextShorteningLevel(rawValue: activityTextShorteningLevel) ?? .default, usePerConfigurationBuildLocations: usePerConfigurationBuildLocations, allowsExternalToolExecution: allowsExternalToolExecution))
     }
+
+    public func registerToolchain(at path: String) async throws -> String {
+        return try await service.send(request: RegisterToolchainRequest(sessionHandle: self.uid, path: Path(path))).value
+    }
 }
 
 extension SWBBuildServiceSession {
diff --git a/Tests/SWBCoreTests/ToolchainRegistryTests.swift b/Tests/SWBCoreTests/ToolchainRegistryTests.swift
@@ -124,7 +124,7 @@ import SWBServiceCore
             ("swift.xctoolchain", ["CFBundleIdentifier": "org.swift.3020161115a", "Aliases": ["swift"]]),
             ("swift-latest.xctoolchain", ["CFBundleIdentifier": "org.swift.latest"]),
         ]) { registry, warnings, errors in
-            #expect(registry.toolchainsByIdentifier.keys.sorted(by: <) == [ToolchainRegistry.defaultToolchainIdentifier, "d", "org.swift.3020161115a"])
+            #expect(registry.toolchains.map(\.identifier).sorted(by: <) == [ToolchainRegistry.defaultToolchainIdentifier, "d", "org.swift.3020161115a"])
 
             if strict {
                 #expect(warnings.isEmpty)
@@ -175,7 +175,7 @@ import SWBServiceCore
             ("swift-older.xctoolchain", ["CFBundleIdentifier": "org.swift.3020161114a", "Version": "3.0.220161211141", "Aliases": ["swift"]]),
         ], infoPlistName: "Info.plist") { registry, _, errors in
 
-            #expect(Set(registry.toolchainsByIdentifier.keys) == Set(["org.swift.3020161114a", "org.swift.3020161115a"] + additionalToolchains))
+            #expect(Set(registry.toolchains.map(\.identifier)) == Set(["org.swift.3020161114a", "org.swift.3020161115a"] + additionalToolchains))
             #expect(errors.count == 0, "\(errors)")
             #expect(registry.lookup("org.swift.3020161115a")?.identifier == "org.swift.3020161115a")
             #expect(registry.lookup("org.swift.3020161114a")?.identifier == "org.swift.3020161114a")
diff --git a/Tests/SwiftBuildTests/ToolchainTests.swift b/Tests/SwiftBuildTests/ToolchainTests.swift
@@ -0,0 +1,47 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+import Testing
+import SwiftBuild
+import SwiftBuildTestSupport
+import SWBTestSupport
+@_spi(Testing) import SWBUtil
+
+@Suite
+fileprivate struct ToolchainTests {
+    @Test
+    func lateRegistration() async throws {
+        try await withTemporaryDirectory { temporaryDirectory in
+            try await withAsyncDeferrable { deferrable in
+                let tmpDirPath = temporaryDirectory.path
+                let testSession = try await TestSWBSession(temporaryDirectory: temporaryDirectory)
+                await deferrable.addBlock {
+                    await #expect(throws: Never.self) {
+                        try await testSession.close()
+                    }
+                }
+                try localFS.createDirectory(tmpDirPath.join("Foo.xctoolchain"))
+                try await localFS.writePlist(tmpDirPath.join("Foo.xctoolchain/Info.plist"), ["Identifier" : "org.swift.foo"])
+                do {
+                    let identifier = try await testSession.session.registerToolchain(at: tmpDirPath.join("Foo.xctoolchain").str)
+                    #expect(identifier == "org.swift.foo")
+                }
+                // Late registration should be idempotent
+                do {
+                    let identifier = try await testSession.session.registerToolchain(at: tmpDirPath.join("Foo.xctoolchain").str)
+                    #expect(identifier == "org.swift.foo")
+                }
+            }
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -302,6 +302,13 @@ private struct SetSessionUserPreferencesMsg: MessageHandler {`
`302`	`302`	`}`
`303`	`303`	`}`
`304`	`304`
	`305`	`+private struct RegisterToolchainHandler: MessageHandler {`
	`306`	`+ func handle(request: Request, message: RegisterToolchainRequest) async throws -> StringResponse {`
	`307`	`+ let session = try request.session(for: message)`
	`308`	`+ return StringResponse(try await session.core.registerToolchain(at: message.path))`
	`309`	`+ }`
	`310`	`+}`
	`311`	`+`
`305`	`312`	`/// Start a PIF transfer from the client.`
`306`	`313`	`///`
`307`	`314`	`/// This will establish a workspace context in the relevant session by exchanging a PIF from the client to the service incrementally, only transferring subobjects as necessary.`
`@@ -1542,6 +1549,7 @@ public struct ServiceSessionMessageHandlers: ServiceExtension {`
`1542`	`1549`	`service.registerMessageHandler(SetSessionSystemInfoMsg.self)`
`1543`	`1550`	`service.registerMessageHandler(SetSessionUserInfoMsg.self)`
`1544`	`1551`	`service.registerMessageHandler(SetSessionUserPreferencesMsg.self)`
	`1552`	`+ service.registerMessageHandler(RegisterToolchainHandler.self)`
`1545`	`1553`	`service.registerMessageHandler(DeveloperPathHandler.self)`
`1546`	`1554`	`}`
`1547`	`1555`	`}`
Original file line number	Diff line number	Diff line change
`@@ -629,6 +629,10 @@ public final class SWBBuildServiceSession: Sendable {`
`629`	`629`	`public func setUserPreferences(enableDebugActivityLogs: Bool, enableBuildDebugging: Bool, enableBuildSystemCaching: Bool, activityTextShorteningLevel: Int, usePerConfigurationBuildLocations: Bool?, allowsExternalToolExecution: Bool) async throws {`
`630`	`630`	`_ = try await service.send(request: SetSessionUserPreferencesRequest(sessionHandle: self.uid, enableDebugActivityLogs: enableDebugActivityLogs, enableBuildDebugging: enableBuildDebugging, enableBuildSystemCaching: enableBuildSystemCaching, activityTextShorteningLevel: ActivityTextShorteningLevel(rawValue: activityTextShorteningLevel) ?? .default, usePerConfigurationBuildLocations: usePerConfigurationBuildLocations, allowsExternalToolExecution: allowsExternalToolExecution))`
`631`	`631`	`}`
	`632`	`+`
	`633`	`+ public func registerToolchain(at path: String) async throws -> String {`
	`634`	`+ return try await service.send(request: RegisterToolchainRequest(sessionHandle: self.uid, path: Path(path))).value`
	`635`	`+ }`
`632`	`636`	`}`
`633`	`637`
`634`	`638`	`extension SWBBuildServiceSession {`