Skip to content

Commit c945253

Browse files
committed
feat: Add systemd socket activation for dstack-guest-agent
Replace systemd-socket-proxyd with native socket activation support. This eliminates the proxy layer and allows the sockets to survive service restarts, maintaining compatibility with containers that mount /var/run/dstack.sock or /var/run/tappd.sock directly. Changes: - Add dstack-guest-agent.socket with both sockets (ordered) - Remove dstack-socket.{socket,service} and tappd-socket.{socket,service} - Add listenfd dependency for receiving sockets from systemd - Add ActivatedUnixListener wrapper implementing rocket's Listener trait - Modify run_internal/run_internal_v0 to use activated sockets when available - Update dstack-guest-agent.service to require the socket unit
1 parent 8e6fae8 commit c945253

11 files changed

Lines changed: 172 additions & 92 deletions

Cargo.lock

Lines changed: 12 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
105105
safe-write = "0.1.2"
106106
nix = "0.29.0"
107107
sd-notify = "0.4.5"
108+
listenfd = "1.0"
108109
jemallocator = "0.5.4"
109110

110111
# Serialization/Parsing

basefiles/dstack-guest-agent.service

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
[Unit]
22
Description=dstack Guest Agent Service
3-
After=network.target tboot.service
3+
Requires=dstack-guest-agent.socket
4+
After=network.target tboot.service dstack-guest-agent.socket
45
Before=docker.service
56

67
[Service]
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# SPDX-FileCopyrightText: 2025 Phala Network <dstack@phala.network>
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
# Socket activation for dstack-guest-agent.
6+
# Provides backward compatibility for containers that mount sockets directly.
7+
# Socket order: dstack.sock (index 0), tappd.sock (index 1)
8+
9+
[Unit]
10+
Description=dstack guest agent sockets
11+
12+
[Socket]
13+
ListenStream=/var/run/dstack.sock
14+
ListenStream=/var/run/tappd.sock
15+
SocketMode=0777
16+
17+
[Install]
18+
WantedBy=sockets.target

basefiles/dstack-socket.service

Lines changed: 0 additions & 15 deletions
This file was deleted.

basefiles/dstack-socket.socket

Lines changed: 0 additions & 16 deletions
This file was deleted.

basefiles/tappd-socket.service

Lines changed: 0 additions & 15 deletions
This file was deleted.

basefiles/tappd-socket.socket

Lines changed: 0 additions & 16 deletions
This file was deleted.

guest-agent/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,3 +53,4 @@ tempfile.workspace = true
5353
rand.workspace = true
5454
or-panic.workspace = true
5555
cc-eventlog.workspace = true
56+
listenfd.workspace = true

guest-agent/src/main.rs

Lines changed: 58 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
//
33
// SPDX-License-Identifier: Apache-2.0
44

5-
use std::{fs::Permissions, future::pending, os::unix::fs::PermissionsExt};
5+
use std::{future::pending, os::unix::net::UnixListener as StdUnixListener};
66

77
use anyhow::{anyhow, Context, Result};
88
use clap::Parser;
@@ -16,6 +16,7 @@ use rocket::{
1616
use rocket_vsock_listener::VsockListener;
1717
use rpc_service::{AppState, ExternalRpcHandler, InternalRpcHandler, InternalRpcHandlerV0};
1818
use sd_notify::{notify as sd_notify, NotifyState};
19+
use socket_activation::{ActivatedSockets, ActivatedUnixListener};
1920
use std::time::Duration;
2021
use tokio::sync::oneshot;
2122
use tracing::{error, info};
@@ -25,6 +26,7 @@ mod guest_api_service;
2526
mod http_routes;
2627
mod models;
2728
mod rpc_service;
29+
mod socket_activation;
2830

2931
const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
3032
const GIT_REV: &str = git_version::git_version!(
@@ -52,6 +54,7 @@ struct Args {
5254
async fn run_internal_v0(
5355
state: AppState,
5456
figment: Figment,
57+
activated_socket: Option<StdUnixListener>,
5558
sock_ready_tx: oneshot::Sender<()>,
5659
) -> Result<()> {
5760
let rocket = rocket::custom(figment)
@@ -64,26 +67,36 @@ async fn run_internal_v0(
6467
.ignite()
6568
.await
6669
.map_err(|err| anyhow!("Failed to ignite rocket: {err}"))?;
67-
let endpoint = DefaultListener::bind_endpoint(&ignite)
68-
.map_err(|err| anyhow!("Failed to get endpoint: {err}"))?;
69-
let listener = DefaultListener::bind(&ignite)
70-
.await
71-
.map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
72-
if let Some(path) = endpoint.unix() {
73-
// Allow any user to connect to the socket
74-
fs_err::set_permissions(path, Permissions::from_mode(0o777))?;
70+
71+
if let Some(std_listener) = activated_socket {
72+
// Use systemd-activated socket
73+
info!("Using systemd-activated socket for tappd.sock");
74+
let listener = ActivatedUnixListener::new(std_listener)?;
75+
sock_ready_tx.send(()).ok();
76+
ignite
77+
.launch_on(listener)
78+
.await
79+
.map_err(|err: rocket::Error| anyhow!(err.to_string()))?;
80+
} else {
81+
// Fall back to binding our own socket
82+
let endpoint = DefaultListener::bind_endpoint(&ignite)
83+
.map_err(|err| anyhow!("Failed to get endpoint: {err}"))?;
84+
let listener = DefaultListener::bind(&ignite)
85+
.await
86+
.map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
87+
sock_ready_tx.send(()).ok();
88+
ignite
89+
.launch_on(listener)
90+
.await
91+
.map_err(|err| anyhow!(err.to_string()))?;
7592
}
76-
sock_ready_tx.send(()).ok();
77-
ignite
78-
.launch_on(listener)
79-
.await
80-
.map_err(|err| anyhow!(err.to_string()))?;
8193
Ok(())
8294
}
8395

8496
async fn run_internal(
8597
state: AppState,
8698
figment: Figment,
99+
activated_socket: Option<StdUnixListener>,
87100
sock_ready_tx: oneshot::Sender<()>,
88101
) -> Result<()> {
89102
let rocket = rocket::custom(figment)
@@ -93,20 +106,29 @@ async fn run_internal(
93106
.ignite()
94107
.await
95108
.map_err(|err| anyhow!("Failed to ignite rocket: {err}"))?;
96-
let endpoint = DefaultListener::bind_endpoint(&ignite)
97-
.map_err(|err| anyhow!("Failed to get endpoint: {err}"))?;
98-
let listener = DefaultListener::bind(&ignite)
99-
.await
100-
.map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
101-
if let Some(path) = endpoint.unix() {
102-
// Allow any user to connect to the socket
103-
fs_err::set_permissions(path, Permissions::from_mode(0o777))?;
109+
110+
if let Some(std_listener) = activated_socket {
111+
// Use systemd-activated socket
112+
info!("Using systemd-activated socket for dstack.sock");
113+
let listener = ActivatedUnixListener::new(std_listener)?;
114+
sock_ready_tx.send(()).ok();
115+
ignite
116+
.launch_on(listener)
117+
.await
118+
.map_err(|err: rocket::Error| anyhow!(err.to_string()))?;
119+
} else {
120+
// Fall back to binding our own socket
121+
let endpoint = DefaultListener::bind_endpoint(&ignite)
122+
.map_err(|err| anyhow!("Failed to get endpoint: {err}"))?;
123+
let listener = DefaultListener::bind(&ignite)
124+
.await
125+
.map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
126+
sock_ready_tx.send(()).ok();
127+
ignite
128+
.launch_on(listener)
129+
.await
130+
.map_err(|err| anyhow!(err.to_string()))?;
104131
}
105-
sock_ready_tx.send(()).ok();
106-
ignite
107-
.launch_on(listener)
108-
.await
109-
.map_err(|err| anyhow!(err.to_string()))?;
110132
Ok(())
111133
}
112134

@@ -219,11 +241,18 @@ async fn main() -> Result<()> {
219241
.extract()
220242
.context("Failed to extract bind address")?;
221243
let guest_api_figment = figment.select("guest-api");
244+
245+
// Get systemd-activated sockets if available
246+
let activated = ActivatedSockets::from_env();
247+
if activated.any_activated() {
248+
info!("Systemd socket activation detected");
249+
}
250+
222251
let (tappd_ready_tx, tappd_ready_rx) = oneshot::channel();
223252
let (sock_ready_tx, sock_ready_rx) = oneshot::channel();
224253
tokio::select!(
225-
res = run_internal_v0(state.clone(), internal_v0_figment, tappd_ready_tx) => res?,
226-
res = run_internal(state.clone(), internal_figment, sock_ready_tx) => res?,
254+
res = run_internal_v0(state.clone(), internal_v0_figment, activated.tappd, tappd_ready_tx) => res?,
255+
res = run_internal(state.clone(), internal_figment, activated.dstack, sock_ready_tx) => res?,
227256
res = run_external(state.clone(), external_figment) => res?,
228257
res = run_guest_api(state.clone(), guest_api_figment) => res?,
229258
_ = async {

0 commit comments

Comments
 (0)