Use allocate with session to automatically clean up IPs. (apple#1544)

jglogan · web-flow · commit c56a659d337a · 2026-05-11T23:08:48.000-07:00
- Part of apple#1318. - Part of apple#1378. - Removes network plugin `deallocate()`, and allocate takes an `XPCServerSession` that registers an `onDisconnect` handler that performs deallocation. - ContainerService now tracks `networkSessions` for allocations.
diff --git a/Sources/Plugins/NetworkVmnet/NetworkVmnetHelper+Start.swift b/Sources/Plugins/NetworkVmnet/NetworkVmnetHelper+Start.swift
@@ -104,8 +104,7 @@ extension NetworkVmnetHelper {
                     identifier: serviceIdentifier,
                     routes: [
                         NetworkRoutes.state.rawValue: XPCServer.route(server.state),
-                        NetworkRoutes.allocate.rawValue: XPCServer.route(server.allocate),
-                        NetworkRoutes.deallocate.rawValue: XPCServer.route(server.deallocate),
+                        NetworkRoutes.allocate.rawValue: server.allocate,
                         NetworkRoutes.lookup.rawValue: XPCServer.route(server.lookup),
                         NetworkRoutes.disableAllocator.rawValue: XPCServer.route(server.disableAllocator),
                     ],
diff --git a/Sources/Services/ContainerAPIService/Server/Containers/ContainersService.swift b/Sources/Services/ContainerAPIService/Server/Containers/ContainersService.swift
@@ -36,6 +36,7 @@ public actor ContainersService {
         var snapshot: ContainerSnapshot
         var client: SandboxClient?
         var allocatedAttachments: [AllocatedAttachment]
+        var networkSessions: [XPCClientSession]
 
         func getClient() throws -> SandboxClient {
             guard let client else {
@@ -132,7 +133,8 @@ public actor ContainersService {
                         networks: [],
                         startedDate: nil
                     ),
-                    allocatedAttachments: []
+                    allocatedAttachments: [],
+                    networkSessions: []
                 )
                 results[config.id] = state
                 guard runtimePlugins.first(where: { $0.name == config.runtimeHandler }) != nil else {
@@ -394,7 +396,7 @@ public actor ContainersService {
                     networks: [],
                     startedDate: nil
                 )
-                await self.setContainerState(configuration.id, ContainerState(snapshot: snapshot, allocatedAttachments: []), context: context)
+                await self.setContainerState(configuration.id, ContainerState(snapshot: snapshot, allocatedAttachments: [], networkSessions: []), context: context)
             } catch {
                 throw error
             }
@@ -435,20 +437,23 @@ public actor ContainersService {
             let (config, _) = try Self.getContainerConfiguration(at: path)
 
             var allocatedAttachments = [AllocatedAttachment]()
+            var networkSessions = [XPCClientSession]()
             do {
                 for n in config.networks {
-                    let allocatedAttach = try await self.networksService?.allocate(
-                        id: n.network,
-                        hostname: n.options.hostname,
-                        macAddress: n.options.macAddress
-                    )
-                    guard var allocatedAttach = allocatedAttach else {
+                    guard
+                        let (allocatedAttach, session) = try await self.networksService?.allocate(
+                            id: n.network,
+                            hostname: n.options.hostname,
+                            macAddress: n.options.macAddress
+                        )
+                    else {
                         throw ContainerizationError(.internalError, message: "failed to allocate a network")
                     }
 
+                    var finalAttach = allocatedAttach
                     if let mtu = n.options.mtu {
                         let a = allocatedAttach.attachment
-                        allocatedAttach = AllocatedAttachment(
+                        finalAttach = AllocatedAttachment(
                             attachment: Attachment(
                                 network: a.network,
                                 hostname: a.hostname,
@@ -462,7 +467,8 @@ public actor ContainersService {
                             pluginInfo: allocatedAttach.pluginInfo
                         )
                     }
-                    allocatedAttachments.append(allocatedAttach)
+                    allocatedAttachments.append(finalAttach)
+                    networkSessions.append(session)
                 }
 
                 try Self.registerService(
@@ -487,20 +493,11 @@ public actor ContainersService {
 
                 state.client = sandboxClient
                 state.allocatedAttachments = allocatedAttachments
+                state.networkSessions = networkSessions
                 await self.setContainerState(id, state, context: context)
             } catch {
-                for allocatedAttach in allocatedAttachments {
-                    do {
-                        try await self.networksService?.deallocate(attachment: allocatedAttach.attachment)
-                    } catch {
-                        self.log.error(
-                            "failed to deallocate network attachment",
-                            metadata: [
-                                "id": "\(id)",
-                                "network": "\(allocatedAttach.attachment.network)",
-                                "error": "\(error)",
-                            ])
-                    }
+                for session in networkSessions {
+                    session.close()
                 }
 
                 let label = Self.fullLaunchdServiceLabel(
@@ -997,27 +994,17 @@ public actor ContainersService {
                 ])
         }
 
-        // Best effort deallocate network attachments for the container. Don't throw on
-        // failure so we can continue with state cleanup.
-        self.log.info("deallocating network attachments", metadata: ["id": "\(id)"])
-        for allocatedAttach in state.allocatedAttachments {
-            do {
-                try await self.networksService?.deallocate(attachment: allocatedAttach.attachment)
-            } catch {
-                self.log.error(
-                    "failed to deallocate network attachment",
-                    metadata: [
-                        "id": "\(id)",
-                        "network": "\(allocatedAttach.attachment.network)",
-                        "error": "\(error)",
-                    ])
-            }
+        // Close network sessions — the network helper auto-releases allocations on disconnect.
+        self.log.info("closing network sessions", metadata: ["id": "\(id)"])
+        for session in state.networkSessions {
+            session.close()
         }
 
         state.snapshot.status = .stopped
         state.snapshot.networks = []
         state.client = nil
         state.allocatedAttachments = []
+        state.networkSessions = []
         await self.setContainerState(id, state, context: context)
 
         let options = try getContainerCreationOptions(id: id)
diff --git a/Sources/Services/ContainerAPIService/Server/Networks/NetworksService.swift b/Sources/Services/ContainerAPIService/Server/Networks/NetworksService.swift
@@ -380,26 +380,26 @@ public actor NetworksService {
         }
     }
 
-    public func allocate(id: String, hostname: String, macAddress: MACAddress?) async throws -> AllocatedAttachment {
+    public func allocate(id: String, hostname: String, macAddress: MACAddress?) async throws -> (AllocatedAttachment, XPCClientSession) {
         guard let serviceState = serviceStates[id] else {
             throw ContainerizationError(.notFound, message: "no network for id \(id)")
         }
         guard let pluginInfo = serviceState.networkState.pluginInfo else {
             throw ContainerizationError(.internalError, message: "network \(id) missing plugin information")
         }
-        let (attach, additionalData) = try await serviceState.client.allocate(hostname: hostname, macAddress: macAddress)
-        return AllocatedAttachment(
-            attachment: attach,
-            additionalData: additionalData,
-            pluginInfo: pluginInfo
-        )
-    }
-
-    public func deallocate(attachment: Attachment) async throws {
-        guard let serviceState = serviceStates[attachment.network] else {
-            throw ContainerizationError(.notFound, message: "no network for id \(attachment.network)")
+        let session = serviceState.client.connect()
+        do {
+            let (attach, additionalData) = try await serviceState.client.allocate(hostname: hostname, macAddress: macAddress, on: session)
+            let alloc = AllocatedAttachment(
+                attachment: attach,
+                additionalData: additionalData,
+                pluginInfo: pluginInfo
+            )
+            return (alloc, session)
+        } catch {
+            session.close()
+            throw error
         }
-        return try await serviceState.client.deallocate(hostname: attachment.hostname)
     }
 
     private static func getClient(configuration: NetworkConfiguration) throws -> ContainerNetworkServiceClient.NetworkClient {
diff --git a/Sources/Services/ContainerNetworkService/Client/NetworkClient.swift b/Sources/Services/ContainerNetworkService/Client/NetworkClient.swift
@@ -71,12 +71,34 @@ extension NetworkClient {
         return (attachment, additionalData)
     }
 
-    public func deallocate(hostname: String) async throws {
-        let request = XPCMessage(route: NetworkRoutes.deallocate.rawValue)
-        request.set(key: NetworkKeys.hostname.rawValue, value: hostname)
+    /// Open a persistent connection to the network helper.
+    ///
+    /// The returned session should be reused for `allocate(on:)` calls. The
+    /// network helper automatically releases all allocations made over this
+    /// session when it closes.
+    public func connect() -> XPCClientSession {
+        createClient().openSession()
+    }
 
-        let client = createClient()
-        try await client.send(request)
+    /// Allocate a network attachment over an existing session.
+    ///
+    /// Use `connect()` to obtain a session, then pass it here. The session
+    /// must remain open for the lifetime of the allocation; closing it
+    /// releases the allocation on the network helper automatically.
+    public func allocate(
+        hostname: String,
+        macAddress: MACAddress? = nil,
+        on session: XPCClientSession
+    ) async throws -> (attachment: Attachment, additionalData: XPCMessage?) {
+        let request = XPCMessage(route: NetworkRoutes.allocate.rawValue)
+        request.set(key: NetworkKeys.hostname.rawValue, value: hostname)
+        if let macAddress = macAddress {
+            request.set(key: NetworkKeys.macAddress.rawValue, value: macAddress.description)
+        }
+        let response = try await session.send(request)
+        let attachment = try response.attachment()
+        let additionalData = response.additionalData()
+        return (attachment, additionalData)
     }
 
     public func lookup(hostname: String) async throws -> Attachment? {
diff --git a/Sources/Services/ContainerNetworkService/Client/NetworkRoutes.swift b/Sources/Services/ContainerNetworkService/Client/NetworkRoutes.swift
@@ -19,8 +19,6 @@ public enum NetworkRoutes: String {
     case state = "com.apple.container.network/state"
     /// Allocates parameters for attaching a sandbox to the network.
     case allocate = "com.apple.container.network/allocate"
-    /// Deallocates parameters for attaching a sandbox to the network.
-    case deallocate = "com.apple.container.network/deallocate"
     /// Disables the allocator if no sandboxes are attached.
     case disableAllocator = "com.apple.container.network/disableAllocator"
     /// Retrieves the allocation for a hostname.
diff --git a/Sources/Services/ContainerNetworkService/Server/NetworkService.swift b/Sources/Services/ContainerNetworkService/Server/NetworkService.swift
@@ -27,6 +27,7 @@ public actor NetworkService: Sendable {
     private let log: Logger
     private var allocator: AttachmentAllocator
     private var macAddresses: [UInt32: MACAddress]
+    private var allocationsBySession: [XPCServerSession: [(hostname: String, index: UInt32)]] = [:]
 
     /// Set up a network service for the specified network.
     public init(
@@ -56,7 +57,7 @@ public actor NetworkService: Sendable {
     }
 
     @Sendable
-    public func allocate(_ message: XPCMessage) async throws -> XPCMessage {
+    public func allocate(_ message: XPCMessage, _ session: XPCServerSession) async throws -> XPCMessage {
         log.debug("enter", metadata: ["func": "\(#function)"])
         defer { log.debug("exit", metadata: ["func": "\(#function)"]) }
 
@@ -99,20 +100,27 @@ public actor NetworkService: Sendable {
             }
         }
         macAddresses[index] = macAddress
+
+        if allocationsBySession[session] == nil {
+            allocationsBySession[session] = []
+            await session.onDisconnect { [weak self] in
+                await self?.releaseSession(session)
+            }
+        }
+        allocationsBySession[session]!.append((hostname: hostname, index: index))
+
         return reply
     }
 
-    @Sendable
-    public func deallocate(_ message: XPCMessage) async throws -> XPCMessage {
-        log.debug("enter", metadata: ["func": "\(#function)"])
-        defer { log.debug("exit", metadata: ["func": "\(#function)"]) }
-
-        let hostname = try message.hostname()
-        if let index = try await allocator.deallocate(hostname: hostname) {
-            macAddresses.removeValue(forKey: index)
+    private func releaseSession(_ session: XPCServerSession) async {
+        guard let allocations = allocationsBySession.removeValue(forKey: session) else {
+            return
+        }
+        for allocation in allocations {
+            _ = try? await allocator.deallocate(hostname: allocation.hostname)
+            macAddresses.removeValue(forKey: allocation.index)
         }
-        log.info("released attachments", metadata: ["hostname": "\(hostname)"])
-        return message.reply()
+        log.info("released session", metadata: ["allocations": "\(allocations.count)"])
     }
 
     @Sendable
