Skip to main content

little_durable_objects/
sandbox.rs

1use std::{
2    collections::HashMap,
3    time::{Duration, Instant},
4};
5
6use anyhow::{Context, Result, ensure};
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9
10use crate::host::HostId;
11
12mod command_process;
13
14const PROVIDER_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
15const MAX_PROVIDER_OUTPUT_BYTES: usize = 1024 * 1024;
16
17#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
18#[serde(rename_all = "camelCase")]
19pub struct EnsureHostRequest {
20    pub namespace_id: String,
21    pub code_revision: String,
22    pub canonical_region: String,
23    pub host_id: HostId,
24    pub session_id: String,
25    pub host_token: String,
26    pub jwt_public_keys: String,
27    pub control_plane_url: String,
28    pub jwt_issuer: String,
29    pub invocation_jwt_audience: String,
30    pub image_ref: String,
31    pub working_directory: String,
32    pub actor_entrypoint: Option<String>,
33    pub actor_idle_timeout_ms: u64,
34    pub host_idle_timeout_ms: u64,
35}
36
37#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
38#[serde(rename_all = "camelCase")]
39pub struct ActorHostHandle {
40    pub host_id: HostId,
41    pub route: String,
42    pub canonical_region: String,
43    pub provisioning: Option<ActorHostProvisioning>,
44}
45
46#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
47#[serde(rename_all = "camelCase")]
48pub struct ActorHostProvisioning {
49    pub provider: String,
50    pub resource_id: String,
51    pub reused: bool,
52    pub started_at_ms: u64,
53    pub input_parsed_at_ms: Option<u64>,
54    pub sdk_loaded_at_ms: Option<u64>,
55    pub resources_resolved_at_ms: Option<u64>,
56    pub existing_host_checked_at_ms: Option<u64>,
57    pub sandbox_scheduled_at_ms: Option<u64>,
58    pub host_ready_observed_at_ms: Option<u64>,
59    pub route_read_at_ms: Option<u64>,
60    pub metadata_written_at_ms: Option<u64>,
61    pub completed_at_ms: u64,
62    #[serde(default)]
63    pub command_spawned_at_ms: Option<u64>,
64    #[serde(default)]
65    pub request_written_at_ms: Option<u64>,
66    #[serde(default)]
67    pub process_completed_at_ms: Option<u64>,
68    #[serde(default)]
69    pub response_decoded_at_ms: Option<u64>,
70}
71
72#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
73#[serde(rename_all = "camelCase")]
74pub struct WarmImageRequest {
75    pub namespace_id: String,
76    pub code_revision: String,
77    pub canonical_region: String,
78    pub image_ref: String,
79}
80
81#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
82#[serde(rename_all = "camelCase")]
83pub struct ImageWarmup {
84    pub provider: String,
85    pub resource_id: String,
86    pub total_ms: u64,
87}
88
89#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
90#[serde(rename_all = "camelCase")]
91pub struct TerminateHostsRequest {
92    pub namespace_id: String,
93    pub code_revision: String,
94    pub canonical_regions: Vec<String>,
95}
96
97#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
98#[serde(rename_all = "camelCase")]
99pub struct HostTermination {
100    pub provider: String,
101    pub resource_ids: Vec<String>,
102}
103
104#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
105#[serde(rename_all = "camelCase")]
106pub struct PublicHostRouteRequest {
107    pub namespace_id: String,
108    pub code_revision: String,
109    pub canonical_region: String,
110}
111
112#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
113#[serde(rename_all = "camelCase")]
114pub struct PublicHostRoute {
115    pub route: String,
116}
117
118#[async_trait]
119pub trait SandboxProvider: Send + Sync {
120    async fn ensure_host(&self, request: &EnsureHostRequest) -> Result<ActorHostHandle>;
121    async fn public_host_route(&self, request: &PublicHostRouteRequest) -> Result<PublicHostRoute>;
122    async fn warm_image(&self, request: &WarmImageRequest) -> Result<ImageWarmup>;
123    async fn terminate_hosts(&self, request: &TerminateHostsRequest) -> Result<HostTermination>;
124}
125
126#[derive(Clone)]
127pub struct HostSandboxRuntimeConfig {
128    pub control_plane_url: String,
129    pub jwt_issuer: String,
130    pub invocation_jwt_audience: String,
131    pub actor_idle_timeout_ms: u64,
132    pub host_idle_timeout_ms: u64,
133}
134
135pub struct CommandSandboxProvider {
136    provider_name: String,
137    processes: deadpool::managed::Pool<command_process::ProviderProcessManager>,
138}
139
140impl CommandSandboxProvider {
141    pub fn new(
142        provider_name: String,
143        command: String,
144        mut environment: HashMap<String, String>,
145    ) -> Result<Self> {
146        ensure!(
147            !provider_name.is_empty() && provider_name.trim() == provider_name,
148            "sandbox provider name must be non-empty without surrounding whitespace"
149        );
150        ensure!(
151            !command.is_empty() && command.trim() == command,
152            "DURABLE_OBJECT_SANDBOX_COMMAND must be non-empty without surrounding whitespace"
153        );
154        if let Ok(path) = std::env::var("PATH") {
155            environment.entry("PATH".into()).or_insert(path);
156        }
157        Ok(Self {
158            provider_name,
159            processes: command_process::pool(command, environment)?,
160        })
161    }
162}
163
164#[async_trait]
165impl SandboxProvider for CommandSandboxProvider {
166    async fn ensure_host(&self, request: &EnsureHostRequest) -> Result<ActorHostHandle> {
167        let (mut response, command): (ActorHostHandle, _) =
168            self.execute_timed("ensure_host", request).await?;
169        if let Some(provisioning) = &mut response.provisioning {
170            provisioning.command_spawned_at_ms = command.spawned_at_ms;
171            provisioning.request_written_at_ms = command.request_written_at_ms;
172            provisioning.process_completed_at_ms = command.process_completed_at_ms;
173            provisioning.response_decoded_at_ms = command.response_decoded_at_ms;
174        }
175        ensure!(
176            response.canonical_region == request.canonical_region,
177            "{} sandbox command returned a host in the wrong canonical region",
178            self.provider_name
179        );
180        ensure!(
181            !response.host_id.as_str().is_empty() && !response.route.is_empty(),
182            "{} sandbox command returned an invalid host",
183            self.provider_name
184        );
185        Ok(response)
186    }
187
188    async fn public_host_route(&self, request: &PublicHostRouteRequest) -> Result<PublicHostRoute> {
189        let response: PublicHostRoute = self.execute("public_host_route", request).await?;
190        ensure!(
191            !response.route.is_empty(),
192            "{} sandbox command returned an invalid public host route",
193            self.provider_name
194        );
195        Ok(response)
196    }
197
198    async fn warm_image(&self, request: &WarmImageRequest) -> Result<ImageWarmup> {
199        self.execute("warm_image", request).await
200    }
201
202    async fn terminate_hosts(&self, request: &TerminateHostsRequest) -> Result<HostTermination> {
203        self.execute("terminate_hosts", request).await
204    }
205}
206
207impl CommandSandboxProvider {
208    async fn execute<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
209        &self,
210        operation: &str,
211        request: &Request,
212    ) -> Result<Reply> {
213        Ok(self.execute_timed(operation, request).await?.0)
214    }
215
216    async fn execute_timed<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
217        &self,
218        operation: &str,
219        request: &Request,
220    ) -> Result<(Reply, ProviderCommandTimings)> {
221        let started_at = Instant::now();
222        let mut timings = ProviderCommandTimings::default();
223        match self
224            .execute_timed_inner(operation, request, started_at, &mut timings)
225            .await
226        {
227            Ok(response) => Ok((response, timings)),
228            Err(source) => Err(ProviderCommandFailure { source, timings }.into()),
229        }
230    }
231
232    async fn execute_timed_inner<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
233        &self,
234        operation: &str,
235        request: &Request,
236        started_at: Instant,
237        timings: &mut ProviderCommandTimings,
238    ) -> Result<Reply> {
239        let command = ProviderCommand { operation, request };
240        let execution = command_process::exchange(&self.processes, &command, started_at, timings);
241        tokio::time::timeout(PROVIDER_REQUEST_TIMEOUT, execution)
242            .await
243            .context("sandbox provider command timed out; outcome may be unknown")?
244    }
245}
246
247#[derive(Debug, Default)]
248struct ProviderCommandTimings {
249    spawned_at_ms: Option<u64>,
250    request_written_at_ms: Option<u64>,
251    process_completed_at_ms: Option<u64>,
252    response_decoded_at_ms: Option<u64>,
253}
254
255#[derive(Debug)]
256pub(crate) struct ProviderCommandFailure {
257    source: anyhow::Error,
258    timings: ProviderCommandTimings,
259}
260
261impl ProviderCommandFailure {
262    pub(crate) fn spawned_at_ms(&self) -> Option<u64> {
263        self.timings.spawned_at_ms
264    }
265
266    pub(crate) fn request_written_at_ms(&self) -> Option<u64> {
267        self.timings.request_written_at_ms
268    }
269
270    pub(crate) fn process_completed_at_ms(&self) -> Option<u64> {
271        self.timings.process_completed_at_ms
272    }
273
274    pub(crate) fn response_decoded_at_ms(&self) -> Option<u64> {
275        self.timings.response_decoded_at_ms
276    }
277}
278
279impl std::fmt::Display for ProviderCommandFailure {
280    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
281        self.source.fmt(formatter)
282    }
283}
284
285impl std::error::Error for ProviderCommandFailure {
286    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
287        self.source.source()
288    }
289}
290
291fn elapsed_ms(started_at: Instant) -> u64 {
292    u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
293}
294
295#[derive(Serialize)]
296struct ProviderCommand<'a, Request> {
297    operation: &'a str,
298    request: &'a Request,
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304
305    #[test]
306    fn decodes_provider_provisioning_timings() {
307        let handle: ActorHostHandle = serde_json::from_value(serde_json::json!({
308            "hostId": "host.v1.namespace.revision.session",
309            "route": "https://host.example.com",
310            "canonicalRegion": "north-america-east",
311            "provisioning": {
312                "provider": "modal",
313                "resourceId": "sb-actor",
314                "reused": false,
315                "startedAtMs": 0,
316                "resourcesResolvedAtMs": 12,
317                "existingHostCheckedAtMs": 34,
318                "sandboxScheduledAtMs": 56,
319                "hostReadyObservedAtMs": 123,
320                "routeReadAtMs": 125,
321                "metadataWrittenAtMs": 129,
322                "completedAtMs": 130
323            }
324        }))
325        .expect("actor host handle");
326
327        let provisioning = handle.provisioning.expect("provisioning timings");
328        assert_eq!(provisioning.resource_id, "sb-actor");
329        assert_eq!(provisioning.sandbox_scheduled_at_ms, Some(56));
330        assert_eq!(provisioning.completed_at_ms, 130);
331    }
332
333    #[test]
334    fn rejects_ambiguous_command_configuration() {
335        assert!(CommandSandboxProvider::new("".into(), "modal".into(), HashMap::new()).is_err());
336        assert!(CommandSandboxProvider::new("modal".into(), "".into(), HashMap::new()).is_err());
337    }
338
339    #[tokio::test]
340    async fn provider_reuses_arbitrarily_named_executables_and_discards_cancelled_exchanges()
341    -> Result<()> {
342        use std::os::unix::fs::PermissionsExt;
343        let directory = tempfile::tempdir()?;
344        let path = directory.path().join("custom-provider.js");
345        std::fs::write(
346            &path,
347            r#"#!/usr/bin/env node
348const readline = require('node:readline');
349let sequence = 0;
350async function reply(command) {
351  if (command.request.oversized) { process.stdout.write('x'.repeat(1024 * 1024 + 1)); return; }
352  if (command.request.fail) { process.stdout.write(JSON.stringify({status: 'failure', error: 'test failure'}) + '\n'); return; }
353  if (command.request.delay) await new Promise(resolve => setTimeout(resolve, command.request.delay));
354  const result = { pid: process.pid, sequence: ++sequence };
355  process.stdout.write(JSON.stringify({ status: 'success', result }) + '\n');
356}
357(async () => { for await (const line of readline.createInterface({input: process.stdin})) await reply(JSON.parse(line)); })();
358"#,
359        )?;
360        std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700))?;
361        let provider = CommandSandboxProvider::new(
362            "modal".into(),
363            path.display().to_string(),
364            HashMap::new(),
365        )?;
366        let first: serde_json::Value = provider.execute("test", &serde_json::json!({})).await?;
367        let second: serde_json::Value = provider.execute("test", &serde_json::json!({})).await?;
368        assert_eq!(first["pid"], second["pid"]);
369        assert_eq!(second["sequence"], 2);
370        let error = provider
371            .execute::<_, serde_json::Value>("test", &serde_json::json!({"fail": true}))
372            .await
373            .expect_err("provider failure must propagate");
374        assert!(error.to_string().contains("test failure"));
375        let recovered: serde_json::Value = provider.execute("test", &serde_json::json!({})).await?;
376        assert_eq!(first["pid"], recovered["pid"]);
377        assert!(
378            tokio::time::timeout(
379                Duration::from_millis(50),
380                provider
381                    .execute::<_, serde_json::Value>("test", &serde_json::json!({"delay": 500}))
382            )
383            .await
384            .is_err()
385        );
386        let next: serde_json::Value = provider.execute("test", &serde_json::json!({})).await?;
387        assert_ne!(first["pid"], next["pid"]);
388        assert_eq!(next["sequence"], 1);
389        let error = provider
390            .execute::<_, serde_json::Value>("test", &serde_json::json!({"oversized": true}))
391            .await
392            .expect_err("unbounded provider output must fail");
393        assert!(error.to_string().contains("stdout exceeds"));
394        let recovered: serde_json::Value = provider.execute("test", &serde_json::json!({})).await?;
395        assert_ne!(next["pid"], recovered["pid"]);
396        let request = serde_json::json!({"delay": 50});
397        let (one, two, three) = tokio::try_join!(
398            provider.execute::<_, serde_json::Value>("test", &request),
399            provider.execute::<_, serde_json::Value>("test", &request),
400            provider.execute::<_, serde_json::Value>("test", &request),
401        )?;
402        let pids: std::collections::HashSet<_> = [one, two, three]
403            .into_iter()
404            .map(|value| value["pid"].as_u64().unwrap())
405            .collect();
406        assert_eq!(
407            pids.len(),
408            2,
409            "provider concurrency stays within the process limit"
410        );
411        Ok(())
412    }
413}