Skip to main content

little_durable_objects/
sandbox.rs

1use std::{collections::HashMap, process::Stdio, time::Duration};
2
3use anyhow::{Context, Result, ensure};
4use async_trait::async_trait;
5use serde::{Deserialize, Serialize};
6use tokio::{io::AsyncWriteExt, process::Command};
7
8use crate::host::HostId;
9
10const PROVIDER_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
11
12#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
13#[serde(rename_all = "camelCase")]
14pub struct EnsureHostRequest {
15    pub namespace_id: String,
16    pub code_revision: String,
17    pub canonical_region: String,
18    pub host_id: HostId,
19    pub session_id: String,
20    pub host_token: String,
21    pub jwt_public_keys: String,
22    pub control_plane_url: String,
23    pub jwt_issuer: String,
24    pub invocation_jwt_audience: String,
25    pub image_ref: String,
26    pub working_directory: String,
27    pub actor_entrypoint: Option<String>,
28    pub actor_idle_timeout_ms: u64,
29    pub host_idle_timeout_ms: u64,
30}
31
32#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
33#[serde(rename_all = "camelCase")]
34pub struct ActorHostHandle {
35    pub host_id: HostId,
36    pub route: String,
37    pub canonical_region: String,
38    pub provisioning: Option<ActorHostProvisioning>,
39}
40
41#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
42#[serde(rename_all = "camelCase")]
43pub struct ActorHostProvisioning {
44    pub provider: String,
45    pub resource_id: String,
46    pub reused: bool,
47    pub resource_lookup_ms: u64,
48    pub existing_lookup_ms: u64,
49    pub create_ms: u64,
50    pub placement_ms: u64,
51    pub tunnel_ms: u64,
52    pub ready_ms: u64,
53    pub metadata_ms: u64,
54    pub total_ms: u64,
55}
56
57#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
58#[serde(rename_all = "camelCase")]
59pub struct WarmImageRequest {
60    pub namespace_id: String,
61    pub code_revision: String,
62    pub canonical_region: String,
63    pub image_ref: String,
64}
65
66#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
67#[serde(rename_all = "camelCase")]
68pub struct ImageWarmup {
69    pub provider: String,
70    pub resource_id: String,
71    pub total_ms: u64,
72}
73
74#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
75#[serde(rename_all = "camelCase")]
76pub struct TerminateHostsRequest {
77    pub namespace_id: String,
78    pub code_revision: String,
79    pub canonical_regions: Vec<String>,
80}
81
82#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
83#[serde(rename_all = "camelCase")]
84pub struct HostTermination {
85    pub provider: String,
86    pub resource_ids: Vec<String>,
87}
88
89#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
90#[serde(rename_all = "camelCase")]
91pub struct PublicHostRouteRequest {
92    pub namespace_id: String,
93    pub code_revision: String,
94    pub canonical_region: String,
95}
96
97#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
98#[serde(rename_all = "camelCase")]
99pub struct PublicHostRoute {
100    pub route: String,
101}
102
103#[async_trait]
104pub trait SandboxProvider: Send + Sync {
105    async fn ensure_host(&self, request: &EnsureHostRequest) -> Result<ActorHostHandle>;
106    async fn public_host_route(&self, request: &PublicHostRouteRequest) -> Result<PublicHostRoute>;
107    async fn warm_image(&self, request: &WarmImageRequest) -> Result<ImageWarmup>;
108    async fn terminate_hosts(&self, request: &TerminateHostsRequest) -> Result<HostTermination>;
109}
110
111#[derive(Clone)]
112pub struct HostSandboxRuntimeConfig {
113    pub control_plane_url: String,
114    pub jwt_issuer: String,
115    pub invocation_jwt_audience: String,
116    pub actor_idle_timeout_ms: u64,
117    pub host_idle_timeout_ms: u64,
118}
119
120pub struct CommandSandboxProvider {
121    provider_name: String,
122    command: String,
123    environment: HashMap<String, String>,
124}
125
126impl CommandSandboxProvider {
127    pub fn new(
128        provider_name: String,
129        command: String,
130        mut environment: HashMap<String, String>,
131    ) -> Result<Self> {
132        ensure!(
133            !provider_name.is_empty() && provider_name.trim() == provider_name,
134            "sandbox provider name must be non-empty without surrounding whitespace"
135        );
136        ensure!(
137            !command.is_empty() && command.trim() == command,
138            "DURABLE_OBJECT_SANDBOX_COMMAND must be non-empty without surrounding whitespace"
139        );
140        if let Ok(path) = std::env::var("PATH") {
141            environment.entry("PATH".into()).or_insert(path);
142        }
143        Ok(Self {
144            provider_name,
145            command,
146            environment,
147        })
148    }
149}
150
151#[async_trait]
152impl SandboxProvider for CommandSandboxProvider {
153    async fn ensure_host(&self, request: &EnsureHostRequest) -> Result<ActorHostHandle> {
154        let response: ActorHostHandle = self.execute("ensure_host", request).await?;
155        ensure!(
156            response.canonical_region == request.canonical_region,
157            "{} sandbox command returned a host in the wrong canonical region",
158            self.provider_name
159        );
160        ensure!(
161            !response.host_id.as_str().is_empty() && !response.route.is_empty(),
162            "{} sandbox command returned an invalid host",
163            self.provider_name
164        );
165        Ok(response)
166    }
167
168    async fn public_host_route(&self, request: &PublicHostRouteRequest) -> Result<PublicHostRoute> {
169        let response: PublicHostRoute = self.execute("public_host_route", request).await?;
170        ensure!(
171            !response.route.is_empty(),
172            "{} sandbox command returned an invalid public host route",
173            self.provider_name
174        );
175        Ok(response)
176    }
177
178    async fn warm_image(&self, request: &WarmImageRequest) -> Result<ImageWarmup> {
179        self.execute("warm_image", request).await
180    }
181
182    async fn terminate_hosts(&self, request: &TerminateHostsRequest) -> Result<HostTermination> {
183        self.execute("terminate_hosts", request).await
184    }
185}
186
187impl CommandSandboxProvider {
188    async fn execute<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
189        &self,
190        operation: &str,
191        request: &Request,
192    ) -> Result<Reply> {
193        let mut child = Command::new(&self.command)
194            .env_clear()
195            .envs(&self.environment)
196            .stdin(Stdio::piped())
197            .stdout(Stdio::piped())
198            .stderr(Stdio::piped())
199            .kill_on_drop(true)
200            .spawn()
201            .with_context(|| {
202                format!(
203                    "start {} sandbox command {:?}",
204                    self.provider_name, self.command
205                )
206            })?;
207        let document = serde_json::to_vec(&ProviderCommand { operation, request })?;
208        let mut stdin = child
209            .stdin
210            .take()
211            .with_context(|| format!("open {} sandbox command stdin", self.provider_name))?;
212        stdin
213            .write_all(&document)
214            .await
215            .with_context(|| format!("write {} sandbox command request", self.provider_name))?;
216        stdin
217            .shutdown()
218            .await
219            .with_context(|| format!("close {} sandbox command stdin", self.provider_name))?;
220        drop(stdin);
221        let output = tokio::time::timeout(PROVIDER_REQUEST_TIMEOUT, child.wait_with_output())
222            .await
223            .with_context(|| format!("{} sandbox command timed out", self.provider_name))??;
224        ensure!(
225            output.status.success(),
226            "{} sandbox command failed with {}: {}",
227            self.provider_name,
228            output.status,
229            String::from_utf8_lossy(&output.stderr).trim()
230        );
231        serde_json::from_slice(&output.stdout)
232            .with_context(|| format!("decode {} sandbox command response", self.provider_name))
233    }
234}
235
236#[derive(Serialize)]
237struct ProviderCommand<'a, Request> {
238    operation: &'a str,
239    request: &'a Request,
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    #[test]
247    fn decodes_provider_provisioning_timings() {
248        let handle: ActorHostHandle = serde_json::from_value(serde_json::json!({
249            "hostId": "host.v1.namespace.revision.session",
250            "route": "https://host.example.com",
251            "canonicalRegion": "north-america-east",
252            "provisioning": {
253                "provider": "modal",
254                "resourceId": "sb-actor",
255                "reused": false,
256                "resourceLookupMs": 12,
257                "existingLookupMs": 34,
258                "createMs": 56,
259                "placementMs": 78,
260                "tunnelMs": 90,
261                "readyMs": 123,
262                "metadataMs": 4,
263                "totalMs": 397
264            }
265        }))
266        .expect("actor host handle");
267
268        let provisioning = handle.provisioning.expect("provisioning timings");
269        assert_eq!(provisioning.resource_id, "sb-actor");
270        assert_eq!(provisioning.create_ms, 56);
271        assert_eq!(provisioning.total_ms, 397);
272    }
273
274    #[test]
275    fn rejects_ambiguous_command_configuration() {
276        assert!(CommandSandboxProvider::new("".into(), "modal".into(), HashMap::new()).is_err());
277        assert!(CommandSandboxProvider::new("modal".into(), "".into(), HashMap::new()).is_err());
278    }
279}