diff --git a/Sources/tart/Commands/Exec.swift b/Sources/tart/Commands/Exec.swift index 88f1d13..daa9504 100644 --- a/Sources/tart/Commands/Exec.swift +++ b/Sources/tart/Commands/Exec.swift @@ -1,6 +1,5 @@ import ArgumentParser import Foundation -import NIOPosix import GRPC import Cirruslabs_TartGuestAgent_Grpc_Swift @@ -41,12 +40,6 @@ struct Exec: AsyncParsableCommand { throw RuntimeError.VMNotRunning(name) } - // Create a gRPC channel connected to the VM's control socket - let group = MultiThreadedEventLoopGroup(numberOfThreads: 1) - defer { - try! group.syncShutdownGracefully() - } - // Change the current working directory to a VM's base directory // to work around Unix domain socket 104 byte limitation [1] // @@ -55,15 +48,6 @@ struct Exec: AsyncParsableCommand { FileManager.default.changeCurrentDirectoryPath(baseURL.path()) } - let channel = try GRPCChannelPool.with( - target: .unixDomainSocket(vmDir.controlSocketURL.relativePath), - transportSecurity: .plaintext, - eventLoopGroup: group, - ) - defer { - try! channel.close().wait() - } - // Switch controlling terminal into raw mode when remote pseudo-terminal is requested var state: State? = nil @@ -79,7 +63,10 @@ struct Exec: AsyncParsableCommand { // Execute a command in a running VM do { - try await execute(channel) + let controlSocketPath = vmDir.controlSocketURL.relativePath + try await withGuestAgentChannel(unixDomainSocketPath: controlSocketPath) { channel in + try await execute(channel) + } } catch let error as GRPCConnectionPoolError { throw RuntimeError.Generic("Failed to connect to the VM using its control socket: \(error.localizedDescription), is the Tart Guest Agent running?") } diff --git a/Sources/tart/GuestAgentChannel.swift b/Sources/tart/GuestAgentChannel.swift new file mode 100644 index 0000000..f754f2f --- /dev/null +++ b/Sources/tart/GuestAgentChannel.swift @@ -0,0 +1,28 @@ +import GRPC +import NIOPosix + +/// Connects to a guest agent's gRPC endpoint over a VM's control socket, runs +/// `body` with the resulting channel, and closes the channel afterwards on both +/// the success and error paths. +/// +/// The connection uses the process-wide singleton event loop group, which must +/// not be shut down, so there is no group lifecycle to manage here. +func withGuestAgentChannel( + unixDomainSocketPath socketPath: String, + _ body: (GRPCChannel) async throws -> T +) async throws -> T { + let channel = try GRPCChannelPool.with( + target: .unixDomainSocket(socketPath), + transportSecurity: .plaintext, + eventLoopGroup: .singletonMultiThreadedEventLoopGroup, + ) + + do { + let result = try await body(channel) + try await channel.close().get() + return result + } catch { + try? await channel.close().get() + throw error + } +} diff --git a/Sources/tart/MACAddressResolver/AgentResolver.swift b/Sources/tart/MACAddressResolver/AgentResolver.swift index 809ad06..a45739b 100644 --- a/Sources/tart/MACAddressResolver/AgentResolver.swift +++ b/Sources/tart/MACAddressResolver/AgentResolver.swift @@ -1,6 +1,5 @@ import Foundation import Network -import NIOPosix import GRPC import Cirruslabs_TartGuestAgent_Apple_Swift import Cirruslabs_TartGuestAgent_Grpc_Swift @@ -15,28 +14,15 @@ class AgentResolver { } private static func resolveIP(_ controlSocketPath: String) async throws -> IPv4Address? { - // Create a gRPC channel connected to the VM's control socket - let group = MultiThreadedEventLoopGroup(numberOfThreads: 1) - defer { - try! group.syncShutdownGracefully() + try await withGuestAgentChannel(unixDomainSocketPath: controlSocketPath) { channel in + // Invoke ResolveIP() gRPC method + let callOptions = CallOptions(timeLimit: .timeout(.seconds(1))) + let agentAsyncClient = AgentAsyncClient(channel: channel) + let resolveIPCall = agentAsyncClient.makeResolveIpCall(ResolveIPRequest(), callOptions: callOptions) + + let response = try await resolveIPCall.response + + return IPv4Address(response.ip) } - - let channel = try GRPCChannelPool.with( - target: .unixDomainSocket(controlSocketPath), - transportSecurity: .plaintext, - eventLoopGroup: group, - ) - defer { - try! channel.close().wait() - } - - // Invoke ResolveIP() gRPC method - let callOptions = CallOptions(timeLimit: .timeout(.seconds(1))) - let agentAsyncClient = AgentAsyncClient(channel: channel) - let resolveIPCall = agentAsyncClient.makeResolveIpCall(ResolveIPRequest(), callOptions: callOptions) - - let response = try await resolveIPCall.response - - return IPv4Address(response.ip) } }