From 4a5a2f79eb640efd8d115e19783ffa152a85040b Mon Sep 17 00:00:00 2001 From: Drew Newberry Date: Fri, 18 Sep 2026 00:08:33 -0700 Subject: [PATCH] refactor(vm): move managed launch into driver Signed-off-by: Drew Newberry --- Cargo.lock | 11 +- architecture/compute-runtimes.md | 15 +- crates/openshell-driver-vm/Cargo.toml | 7 +- crates/openshell-driver-vm/README.md | 5 +- crates/openshell-driver-vm/src/lib.rs | 4 + .../src/managed.rs} | 168 ++++-------------- crates/openshell-gateway/Cargo.toml | 19 +- crates/openshell-gateway/src/lib.rs | 20 +-- crates/openshell-server/src/compute/mod.rs | 121 ++++++++++++- crates/openshell-server/src/lib.rs | 1 + 10 files changed, 189 insertions(+), 182 deletions(-) rename crates/{openshell-gateway/src/vm.rs => openshell-driver-vm/src/managed.rs} (88%) diff --git a/Cargo.lock b/Cargo.lock index 109e6df8d..c981e9a40 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4313,6 +4313,7 @@ dependencies = [ "tempfile", "tokio", "tokio-stream", + "toml", "tonic", "tower-http", "tracing", @@ -4344,25 +4345,17 @@ name = "openshell-gateway" version = "0.0.0" dependencies = [ "async-trait", - "hyper-util", "miette", - "nix 0.29.0", "openshell-core", "openshell-driver-docker", "openshell-driver-kubernetes", "openshell-driver-mxc", "openshell-driver-podman", + "openshell-driver-vm", "openshell-otel", - "openshell-policy", "openshell-server", - "rustix 1.1.4", - "serde", "tempfile", "tokio", - "toml", - "tonic", - "tower", - "tracing", ] [[package]] diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 0077c8668..6edb57a15 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -153,6 +153,13 @@ compose that gateway with Docker, Podman, Kubernetes, and VM driver executables over the public UDS gRPC contract so an in-tree driver cannot silently depend on a server-only API. +Driver crates own their configuration defaults and backend-specific startup +contract. The standalone VM driver exposes a lightweight `managed` feature for +its configuration and subprocess arguments; this does not link libkrun or the +VM runtime into the gateway. The gateway composition crate only adapts that +launcher to `ComputeDriverFactory`. The server owns the generic managed-child +readiness probe, UDS connection, supervision, and socket cleanup. + ## Stop and Start Lifecycle The gateway persists lifecycle intent before mutating compute: @@ -328,10 +335,10 @@ can request a specific number of GPUs or the driver-specific default behaviour. For all in-tree drivers, this is equivalent to selecting a single GPU. VM runtime state paths are derived only from driver-validated sandbox IDs -matching `[A-Za-z0-9._-]{1,128}`. The gateway-owned VM driver socket uses a -private `run/` directory plus Unix peer UID/PID checks. Standalone -unauthenticated TCP mode is disabled unless explicitly enabled for local -development. +matching `[A-Za-z0-9._-]{1,128}`. The gateway-managed VM driver socket uses a +driver-configured private `run/` directory plus Unix peer UID/PID checks. +Standalone unauthenticated TCP mode is disabled unless explicitly enabled for +local development. Runtime-specific implementation notes belong in the driver crate README: diff --git a/crates/openshell-driver-vm/Cargo.toml b/crates/openshell-driver-vm/Cargo.toml index 9c2ef8c81..8e6f6a5f8 100644 --- a/crates/openshell-driver-vm/Cargo.toml +++ b/crates/openshell-driver-vm/Cargo.toml @@ -65,8 +65,12 @@ zstd = { version = "0.13", optional = true } [features] default = ["compute-driver", "telemetry"] +## Expose VM-specific managed launch configuration without linking the VM +## runtime implementation into the gateway process. +managed = ["dep:openshell-policy", "dep:rustix"] ## Build the standalone compute driver and its host runtime implementation. compute-driver = [ + "managed", "dep:base64", "dep:bollard", "dep:clap", @@ -78,7 +82,6 @@ compute-driver = [ "dep:nix", "dep:oci-client", "dep:openshell-otel", - "dep:openshell-policy", "dep:openshell-driver-podman", "dep:openshell-vfio", "dep:opentelemetry", @@ -87,7 +90,6 @@ compute-driver = [ "dep:prost", "dep:prost-types", "dep:rand", - "dep:rustix", "dep:sha2", "dep:tar", "dep:tokio-stream", @@ -117,6 +119,7 @@ defaults-without-telemetry = ["compute-driver"] openshell-otel-test-support = { path = "../openshell-otel-test-support" } temp-env = "0.3" tempfile = "3" +toml = { workspace = true } opentelemetry_sdk = { workspace = true, features = ["testing"] } # smol-rs/polling drives the BSD/macOS parent-death detection in diff --git a/crates/openshell-driver-vm/README.md b/crates/openshell-driver-vm/README.md index 2cce3b1c0..7a6e94336 100644 --- a/crates/openshell-driver-vm/README.md +++ b/crates/openshell-driver-vm/README.md @@ -11,7 +11,7 @@ The driver embeds libkrun, libkrunfw, the guest OCI unpacker, the portable guest ```mermaid flowchart LR subgraph host["Host"] - gateway["openshell-gateway
(vm::spawn)"] + gateway["openshell-gateway
(VM factory adapter)"] driver["openshell-driver-vm
libkrun"] supervisor["openshell-supervisor
host policy supervisor"] gateway <-->|"gRPC over UDS
compute-driver.sock"| driver @@ -305,5 +305,6 @@ so driver choice remains automatic unless the user explicitly overrides it. ## TODOs -- The gateway still configures the driver via CLI args; this will move to a gRPC bootstrap call so the driver interface is uniform across backends. See the `TODO(driver-abstraction)` note in `crates/openshell-gateway/src/vm.rs`. +- Managed launch still configures the driver via CLI args; a future gRPC + bootstrap call can make configuration uniform across standalone backends. - macOS local builds are codesigned by `tasks/scripts/gateway-vm.sh`; the generated Homebrew formula signs the release tarball driver for local installs. diff --git a/crates/openshell-driver-vm/src/lib.rs b/crates/openshell-driver-vm/src/lib.rs index 8e3f41a15..1e2cd077e 100644 --- a/crates/openshell-driver-vm/src/lib.rs +++ b/crates/openshell-driver-vm/src/lib.rs @@ -23,6 +23,8 @@ pub mod gpu; mod isolation; #[cfg(feature = "compute-driver")] pub mod lifecycle; +#[cfg(feature = "managed")] +mod managed; #[cfg(feature = "compute-driver")] pub mod otel_tracing; #[cfg(feature = "compute-driver")] @@ -40,6 +42,8 @@ pub use lifecycle::{ LaunchPlan, LifecycleError, LifecycleExtension, LifecycleExtensionRegistry, LifecycleResult, RestoreContext, }; +#[cfg(feature = "managed")] +pub use managed::{ManagedVmDriverProcess, VmComputeConfig, spawn_managed_vm_driver}; #[cfg(feature = "compute-driver")] pub use runtime::{ VM_RUNTIME_DIR_ENV, VmBackend, VmLaunchConfig, VsockPortMap, configured_runtime_dir, run_vm, diff --git a/crates/openshell-gateway/src/vm.rs b/crates/openshell-driver-vm/src/managed.rs similarity index 88% rename from crates/openshell-gateway/src/vm.rs rename to crates/openshell-driver-vm/src/managed.rs index 413d5a579..41acb7632 100644 --- a/crates/openshell-gateway/src/vm.rs +++ b/crates/openshell-driver-vm/src/managed.rs @@ -1,59 +1,31 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! VM compute driver plumbing. +//! Managed VM compute driver launch support. //! -//! This module owns everything needed to hand the gateway a `Channel` speaking -//! the `openshell.compute.v1.ComputeDriver` RPC surface against an -//! `openshell-driver-vm` subprocess over a Unix domain socket: +//! This module owns the VM-specific configuration and process launch contract +//! for an `openshell-driver-vm` subprocess: //! -//! - [`VmComputeConfig`]: gateway-local configuration (state dir, driver binary, +//! - [`VmComputeConfig`]: launch configuration (state dir, driver binary, //! VM shape, guest TLS material). -//! - [`spawn`]: spawn the driver subprocess, wait for its UDS to be ready, -//! and return a live gRPC channel plus a [`ManagedDriverProcess`] handle -//! that will reap the subprocess and clean up the socket on drop. +//! - [`spawn_managed_vm_driver`]: spawn the driver subprocess and return its +//! child handle and Unix domain socket path. //! - Helpers to resolve the driver binary, compute the socket path, and -//! validate guest TLS material when the gateway runs an `https://` control -//! plane. +//! validate guest TLS material for an `https://` control-plane endpoint. //! -//! The VM-driver fields deliberately live here rather than in -//! [`openshell_core::Config`] so the shared core stays free of driver-specific -//! plumbing. -//! -//! Process launch remains deliberately VM-specific at this binary composition -//! boundary: it translates gateway configuration into the standalone driver's -//! argv and then connects through the same public compute-driver RPC interface -//! used by operator-managed external drivers. +//! Socket readiness, RPC connection, process supervision, and cleanup remain +//! generic server concerns. -#[cfg(unix)] -use hyper_util::rt::TokioIo; -#[cfg(unix)] -use openshell_core::proto::compute::v1::{ - GetCapabilitiesRequest, compute_driver_client::ComputeDriverClient, -}; use openshell_core::{Error, Result, UpstreamProxyConfig}; #[cfg(unix)] -use openshell_otel::TraceContextInterceptor; -use openshell_server::AcquiredRemoteDriverEndpoint; -#[cfg(unix)] -use openshell_server::ManagedDriverProcess; -use openshell_server::config_file::OtlpConfig; -#[cfg(unix)] use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt}; #[cfg(unix)] use std::path::Path; use std::path::PathBuf; #[cfg(unix)] -use std::{io::ErrorKind, process::Stdio, sync::Arc, time::Duration}; -#[cfg(unix)] -use tokio::net::UnixStream; +use std::{io::ErrorKind, process::Stdio}; #[cfg(unix)] use tokio::process::Command; -use tonic::transport::Channel; -#[cfg(unix)] -use tonic::transport::Endpoint; -#[cfg(unix)] -use tower::service_fn; const DRIVER_BIN_NAME: &str = "openshell-driver-vm"; const COMPUTE_DRIVER_SOCKET_RUN_DIR: &str = "run"; @@ -533,16 +505,28 @@ pub fn compute_driver_guest_tls_paths( Ok(Some(VmGuestTlsPaths { ca, cert, key })) } -/// Launch the VM compute-driver subprocess, wait for its UDS to come up, -/// and return a gRPC `Channel` connected to it plus a process handle that -/// kills the subprocess and removes the socket on drop. +/// A launched VM compute-driver process and the socket it will listen on. +pub struct ManagedVmDriverProcess { + child: tokio::process::Child, + socket_path: PathBuf, +} + +impl ManagedVmDriverProcess { + /// Consume the launch result into the generic server-owned process parts. + #[must_use] + pub fn into_parts(self) -> (tokio::process::Child, PathBuf) { + (self.child, self.socket_path) + } +} + +/// Launch the VM compute-driver subprocess. #[cfg(unix)] -pub async fn spawn( +pub fn spawn_managed_vm_driver( gateway_log_level: &str, gateway_name: &str, vm_config: &VmComputeConfig, - otlp_config: Option<&OtlpConfig>, -) -> Result { + otlp_endpoint: Option<&str>, +) -> Result { vm_config.validate_configuration()?; let driver_bin = resolve_compute_driver_bin(vm_config)?; let socket_path = compute_driver_socket_path(vm_config); @@ -559,7 +543,7 @@ pub async fn spawn( .arg("--expected-peer-pid") .arg(std::process::id().to_string()); command.arg("--log-level").arg(gateway_log_level); - append_otlp_args(&mut command, otlp_config, gateway_name); + append_otlp_args(&mut command, otlp_endpoint, gateway_name); command.arg("--grpc-endpoint").arg(&vm_config.grpc_endpoint); command.arg("--state-dir").arg(&vm_config.state_dir); if !vm_config.default_image.trim().is_empty() { @@ -587,17 +571,13 @@ pub async fn spawn( } append_vm_proxy_and_spiffe_args(&mut command, vm_config); - let mut child = command.spawn().map_err(|e| { + let child = command.spawn().map_err(|e| { Error::execution(format!( "failed to launch vm compute driver '{}': {e}", driver_bin.display() )) })?; - let channel = wait_for_compute_driver(&socket_path, &mut child).await?; - let process = Arc::new(ManagedDriverProcess::new(child, socket_path)); - Ok(AcquiredRemoteDriverEndpoint::managed( - "vm", channel, process, - )) + Ok(ManagedVmDriverProcess { child, socket_path }) } fn validate_vm_sandbox_identity(config: &VmComputeConfig) -> Result<()> { @@ -670,93 +650,25 @@ fn append_vm_proxy_and_spiffe_args(command: &mut Command, config: &VmComputeConf } } -fn append_otlp_args(command: &mut Command, otlp_config: Option<&OtlpConfig>, gateway_name: &str) { - if let Some(config) = otlp_config { - command.arg("--otlp-endpoint").arg(&config.endpoint); +fn append_otlp_args(command: &mut Command, otlp_endpoint: Option<&str>, gateway_name: &str) { + if let Some(endpoint) = otlp_endpoint { + command.arg("--otlp-endpoint").arg(endpoint); command.arg("--gateway-name").arg(gateway_name); } } #[cfg(not(unix))] -pub async fn spawn( +pub fn spawn_managed_vm_driver( _gateway_log_level: &str, _gateway_name: &str, _vm_config: &VmComputeConfig, - _otlp_config: Option<&OtlpConfig>, -) -> Result { + _otlp_endpoint: Option<&str>, +) -> Result { Err(Error::config( "the vm compute driver requires unix domain socket support", )) } -#[cfg(unix)] -#[tracing::instrument( - name = "driver.wait_for_ready", - skip_all, - fields( - otel.name = "driver.wait_for_ready", - otel.status_code = tracing::field::Empty, - driver.name = "vm", - ) -)] -async fn wait_for_compute_driver( - socket_path: &Path, - child: &mut tokio::process::Child, -) -> Result { - let mut last_error: Option = None; - for _ in 0..100 { - let try_wait_result = child.try_wait().map_err(|e| { - Error::execution(format!("failed to poll vm compute driver process: {e}")) - })?; - if let Some(status) = try_wait_result { - return Err(Error::execution(format!( - "vm compute driver exited before becoming ready with status {status}" - ))); - } - - match connect_compute_driver(socket_path).await { - Ok(channel) => { - let mut client = - ComputeDriverClient::with_interceptor(channel.clone(), TraceContextInterceptor); - match client - .get_capabilities(tonic::Request::new(GetCapabilitiesRequest {})) - .await - { - Ok(_) => return Ok(channel), - Err(status) => last_error = Some(status.to_string()), - } - } - Err(err) => last_error = Some(err.to_string()), - } - - tokio::time::sleep(Duration::from_millis(100)).await; - } - - Err(Error::execution(format!( - "timed out waiting for vm compute driver socket '{}': {}", - socket_path.display(), - last_error.unwrap_or_else(|| "unknown error".to_string()) - ))) -} - -#[cfg(unix)] -async fn connect_compute_driver(socket_path: &Path) -> Result { - let socket_path = socket_path.to_path_buf(); - let display_path = socket_path.clone(); - Endpoint::from_static("http://[::]:50051") - .connect_with_connector(service_fn(move |_: tonic::transport::Uri| { - let socket_path = socket_path.clone(); - async move { UnixStream::connect(socket_path).await.map(TokioIo::new) } - })) - .await - .map_err(|e| { - Error::execution(format!( - "failed to connect to vm compute driver socket '{}': {e}", - display_path.display() - )) - }) -} - #[cfg(all(test, unix))] mod tests { use super::{ @@ -767,7 +679,6 @@ mod tests { validate_vm_sandbox_identity, }; use openshell_core::UpstreamProxyConfig; - use openshell_server::config_file::OtlpConfig; use std::os::unix::fs::PermissionsExt; use std::os::unix::net::UnixListener as StdUnixListener; use std::path::PathBuf; @@ -786,10 +697,7 @@ mod tests { let mut command = tokio::process::Command::new("openshell-driver-vm"); append_otlp_args( &mut command, - Some(&OtlpConfig { - endpoint: "http://collector.internal:4317".to_string(), - service_name: Some("custom-gateway".to_string()), - }), + Some("http://collector.internal:4317"), "production-us-west", ); diff --git a/crates/openshell-gateway/Cargo.toml b/crates/openshell-gateway/Cargo.toml index 32dcaeb24..fa6f5d9be 100644 --- a/crates/openshell-gateway/Cargo.toml +++ b/crates/openshell-gateway/Cargo.toml @@ -26,14 +26,7 @@ tokio = { workspace = true } openshell-driver-docker = { path = "../openshell-driver-docker", optional = true } openshell-driver-kubernetes = { path = "../openshell-driver-kubernetes", optional = true } openshell-driver-podman = { path = "../openshell-driver-podman", optional = true } -openshell-policy = { path = "../openshell-policy", optional = true } -hyper-util = { workspace = true, optional = true } -nix = { workspace = true, optional = true } -serde = { workspace = true, optional = true } -rustix = { workspace = true, optional = true } -tonic = { workspace = true, optional = true } -tower = { workspace = true, optional = true } -tracing = { workspace = true, optional = true } +openshell-driver-vm = { path = "../openshell-driver-vm", default-features = false, features = ["managed"], optional = true } [target.'cfg(target_os = "windows")'.dependencies] openshell-driver-mxc = { path = "../openshell-driver-mxc", optional = true } @@ -52,15 +45,8 @@ compute-driver-docker = ["dep:openshell-driver-docker", "dep:openshell-otel"] compute-driver-kubernetes = ["dep:openshell-driver-kubernetes", "dep:openshell-otel"] compute-driver-podman = ["dep:openshell-driver-podman", "dep:openshell-otel"] compute-driver-vm = [ - "dep:openshell-policy", + "dep:openshell-driver-vm", "dep:openshell-otel", - "dep:hyper-util", - "dep:nix", - "dep:serde", - "dep:rustix", - "dep:tonic", - "dep:tower", - "dep:tracing", ] telemetry = ["openshell-core/telemetry", "openshell-server/telemetry"] ## Convenience alias: every default feature except `telemetry`. Build a @@ -79,4 +65,3 @@ workspace = true [dev-dependencies] tempfile = "3" -toml = { workspace = true } diff --git a/crates/openshell-gateway/src/lib.rs b/crates/openshell-gateway/src/lib.rs index b67e93c10..260928dba 100644 --- a/crates/openshell-gateway/src/lib.rs +++ b/crates/openshell-gateway/src/lib.rs @@ -16,9 +16,6 @@ compile_error!( build a telemetry-free gateway with `--no-default-features --features defaults-without-telemetry`" ); -#[cfg(all(not(target_os = "windows"), feature = "compute-driver-vm"))] -mod vm; - #[cfg(any( all(target_os = "windows", feature = "compute-driver-mxc"), all( @@ -436,13 +433,16 @@ impl openshell_server::ComputeDriverFactory for VmFactory { &mut config.guest_tls_key, context.guest_tls_paths(), ); - let endpoint = vm::spawn( + let launch = openshell_driver_vm::spawn_managed_vm_driver( context.gateway_log_level(), context.gateway_name(), &config, - context.otlp_config(), - ) - .await?; + context.otlp_config().map(|config| config.endpoint.as_str()), + )?; + let (child, socket_path) = launch.into_parts(); + let endpoint = openshell_server::connect_managed_compute_driver("vm", socket_path, child) + .await + .map_err(|error| openshell_core::Error::execution(error.to_string()))?; Ok(openshell_server::ComputeDriverInstance::ManagedRemote( endpoint, )) @@ -452,10 +452,10 @@ impl openshell_server::ComputeDriverFactory for VmFactory { #[cfg(all(not(target_os = "windows"), feature = "compute-driver-vm"))] fn vm_config( context: openshell_server::ComputeDriverConfigContext<'_>, -) -> openshell_core::Result { - let mut config: vm::VmComputeConfig = context.driver_config()?; +) -> openshell_core::Result { + let mut config: openshell_driver_vm::VmComputeConfig = context.driver_config()?; if config.state_dir.as_os_str().is_empty() { - config.state_dir = vm::VmComputeConfig::default_state_dir(); + config.state_dir = openshell_driver_vm::VmComputeConfig::default_state_dir(); } Ok(config) } diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 9732999c3..d6f53a2a4 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -451,6 +451,28 @@ async fn managed_driver_shutdown_sends_sigterm_before_forcing_exit() { assert_eq!(std::fs::read_to_string(terminated).unwrap(), "terminated"); } +#[cfg(all(test, unix))] +#[tokio::test] +async fn managed_driver_connection_reports_early_process_exit() { + let dir = tempfile::tempdir().unwrap(); + let child = tokio::process::Command::new("sh") + .arg("-c") + .arg("exit 7") + .spawn() + .unwrap(); + + let error = + connect_managed_compute_driver("test", dir.path().join("missing-driver.sock"), child) + .await + .expect_err("an exited driver must fail readiness"); + + assert!( + error + .to_string() + .contains("test compute driver exited before becoming ready") + ); +} + #[derive(Debug)] pub struct AcquiredRemoteDriverEndpoint { pub(crate) name: String, @@ -4237,6 +4259,96 @@ fn is_failed_main_process_result(sandbox: &Sandbox) -> bool { }) } +/// Connect to a compute driver subprocess owned by the gateway. +/// +/// The caller owns driver-specific process construction. The server owns the +/// generic readiness probe, process supervision, and socket cleanup after the +/// child and socket path cross this boundary. +#[cfg(unix)] +#[tracing::instrument( + name = "driver.wait_for_ready", + skip_all, + fields( + otel.name = "driver.wait_for_ready", + otel.status_code = tracing::field::Empty, + driver.name = tracing::field::Empty, + ) +)] +pub async fn connect_managed_compute_driver( + name: impl Into, + socket_path: PathBuf, + mut child: tokio::process::Child, +) -> Result { + let name = name.into(); + tracing::Span::current().record("driver.name", &name); + let mut last_error: Option = None; + + for _ in 0..100 { + let status = child.try_wait().map_err(|error| { + ComputeError::Message(format!( + "failed to poll {name} compute driver process: {error}" + )) + })?; + if let Some(status) = status { + return Err(ComputeError::Message(format!( + "{name} compute driver exited before becoming ready with status {status}" + ))); + } + + match connect_compute_driver_socket(&socket_path).await { + Ok(channel) => { + let mut client = + ComputeDriverClient::with_interceptor(channel.clone(), TraceContextInterceptor); + match client + .get_capabilities(Request::new(GetCapabilitiesRequest {})) + .await + { + Ok(_) => { + let process = Arc::new(ManagedDriverProcess::new(child, socket_path)); + return Ok(AcquiredRemoteDriverEndpoint::managed( + name, channel, process, + )); + } + Err(status) => last_error = Some(status.to_string()), + } + } + Err(error) => last_error = Some(error.to_string()), + } + + tokio::time::sleep(Duration::from_millis(100)).await; + } + + Err(ComputeError::Message(format!( + "timed out waiting for {name} compute driver socket '{}': {}", + socket_path.display(), + last_error.unwrap_or_else(|| "unknown error".to_string()) + ))) +} + +#[cfg(not(unix))] +pub async fn connect_managed_compute_driver( + _name: impl Into, + _socket_path: PathBuf, + _child: tokio::process::Child, +) -> Result { + Err(ComputeError::Message( + "managed compute driver endpoints require unix domain socket support".to_string(), + )) +} + +#[cfg(unix)] +async fn connect_compute_driver_socket( + socket_path: &Path, +) -> Result { + let connector_path = socket_path.to_path_buf(); + Endpoint::from_static("http://[::]:50051") + .connect_with_connector(service_fn(move |_: tonic::transport::Uri| { + let connector_path = connector_path.clone(); + async move { UnixStream::connect(connector_path).await.map(TokioIo::new) } + })) + .await +} + /// Connect to an unmanaged remote compute driver that is already listening on /// `socket_path` and return the acquired endpoint. /// @@ -4252,14 +4364,7 @@ pub async fn connect_remote_compute_driver( let socket_path = socket_path.to_path_buf(); let deadline = tokio::time::Instant::now() + Duration::from_secs(30); let channel = loop { - let connector_path = socket_path.clone(); - match Endpoint::from_static("http://[::]:50051") - .connect_with_connector(service_fn(move |_: tonic::transport::Uri| { - let connector_path = connector_path.clone(); - async move { UnixStream::connect(connector_path).await.map(TokioIo::new) } - })) - .await - { + match connect_compute_driver_socket(&socket_path).await { Ok(channel) => break channel, Err(error) if tokio::time::Instant::now() < deadline => { tracing::debug!( diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 80cc2acd5..69b051e86 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -1106,6 +1106,7 @@ async fn terminate_signal() { pub use compute::{ AcquiredRemoteDriverEndpoint, DriverWatchStream, ManagedDriverProcess, SharedComputeDriver, + connect_managed_compute_driver, }; /// Driver instance returned by a compiled compute-driver factory.