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    command: String,
138    environment: HashMap<String, String>,
139}
140
141impl CommandSandboxProvider {
142    pub fn new(
143        provider_name: String,
144        command: String,
145        mut environment: HashMap<String, String>,
146    ) -> Result<Self> {
147        ensure!(
148            !provider_name.is_empty() && provider_name.trim() == provider_name,
149            "sandbox provider name must be non-empty without surrounding whitespace"
150        );
151        ensure!(
152            !command.is_empty() && command.trim() == command,
153            "DURABLE_OBJECT_SANDBOX_COMMAND must be non-empty without surrounding whitespace"
154        );
155        if let Ok(path) = std::env::var("PATH") {
156            environment.entry("PATH".into()).or_insert(path);
157        }
158        Ok(Self {
159            provider_name,
160            command,
161            environment,
162        })
163    }
164}
165
166#[async_trait]
167impl SandboxProvider for CommandSandboxProvider {
168    async fn ensure_host(&self, request: &EnsureHostRequest) -> Result<ActorHostHandle> {
169        let (mut response, command): (ActorHostHandle, _) =
170            self.execute_timed("ensure_host", request).await?;
171        if let Some(provisioning) = &mut response.provisioning {
172            provisioning.command_spawned_at_ms = command.spawned_at_ms;
173            provisioning.request_written_at_ms = command.request_written_at_ms;
174            provisioning.process_completed_at_ms = command.process_completed_at_ms;
175            provisioning.response_decoded_at_ms = command.response_decoded_at_ms;
176        }
177        ensure!(
178            response.canonical_region == request.canonical_region,
179            "{} sandbox command returned a host in the wrong canonical region",
180            self.provider_name
181        );
182        ensure!(
183            !response.host_id.as_str().is_empty() && !response.route.is_empty(),
184            "{} sandbox command returned an invalid host",
185            self.provider_name
186        );
187        Ok(response)
188    }
189
190    async fn public_host_route(&self, request: &PublicHostRouteRequest) -> Result<PublicHostRoute> {
191        let response: PublicHostRoute = self.execute("public_host_route", request).await?;
192        ensure!(
193            !response.route.is_empty(),
194            "{} sandbox command returned an invalid public host route",
195            self.provider_name
196        );
197        Ok(response)
198    }
199
200    async fn warm_image(&self, request: &WarmImageRequest) -> Result<ImageWarmup> {
201        self.execute("warm_image", request).await
202    }
203
204    async fn terminate_hosts(&self, request: &TerminateHostsRequest) -> Result<HostTermination> {
205        self.execute("terminate_hosts", request).await
206    }
207}
208
209impl CommandSandboxProvider {
210    async fn execute<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
211        &self,
212        operation: &str,
213        request: &Request,
214    ) -> Result<Reply> {
215        Ok(self.execute_timed(operation, request).await?.0)
216    }
217
218    async fn execute_timed<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
219        &self,
220        operation: &str,
221        request: &Request,
222    ) -> Result<(Reply, ProviderCommandTimings)> {
223        let started_at = Instant::now();
224        let mut timings = ProviderCommandTimings::default();
225        match self
226            .execute_timed_inner(operation, request, started_at, &mut timings)
227            .await
228        {
229            Ok(response) => Ok((response, timings)),
230            Err(source) => Err(ProviderCommandFailure { source, timings }.into()),
231        }
232    }
233
234    async fn execute_timed_inner<Request: Serialize, Reply: for<'de> Deserialize<'de>>(
235        &self,
236        operation: &str,
237        request: &Request,
238        started_at: Instant,
239        timings: &mut ProviderCommandTimings,
240    ) -> Result<Reply> {
241        let command = ProviderCommand { operation, request };
242        let execution = command_process::exchange(
243            &self.command,
244            &self.environment,
245            &command,
246            started_at,
247            timings,
248        );
249        tokio::time::timeout(PROVIDER_REQUEST_TIMEOUT, execution)
250            .await
251            .context("sandbox provider command timed out; outcome may be unknown")?
252    }
253}
254
255#[derive(Debug, Default)]
256struct ProviderCommandTimings {
257    spawned_at_ms: Option<u64>,
258    request_written_at_ms: Option<u64>,
259    process_completed_at_ms: Option<u64>,
260    response_decoded_at_ms: Option<u64>,
261}
262
263#[derive(Debug)]
264pub(crate) struct ProviderCommandFailure {
265    source: anyhow::Error,
266    timings: ProviderCommandTimings,
267}
268
269impl ProviderCommandFailure {
270    pub(crate) fn spawned_at_ms(&self) -> Option<u64> {
271        self.timings.spawned_at_ms
272    }
273
274    pub(crate) fn request_written_at_ms(&self) -> Option<u64> {
275        self.timings.request_written_at_ms
276    }
277
278    pub(crate) fn process_completed_at_ms(&self) -> Option<u64> {
279        self.timings.process_completed_at_ms
280    }
281
282    pub(crate) fn response_decoded_at_ms(&self) -> Option<u64> {
283        self.timings.response_decoded_at_ms
284    }
285}
286
287impl std::fmt::Display for ProviderCommandFailure {
288    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289        self.source.fmt(formatter)
290    }
291}
292
293impl std::error::Error for ProviderCommandFailure {
294    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
295        self.source.source()
296    }
297}
298
299fn elapsed_ms(started_at: Instant) -> u64 {
300    u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
301}
302
303#[derive(Serialize)]
304struct ProviderCommand<'a, Request> {
305    operation: &'a str,
306    request: &'a Request,
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312
313    #[test]
314    fn decodes_provider_provisioning_timings() {
315        let handle: ActorHostHandle = serde_json::from_value(serde_json::json!({
316            "hostId": "host.v1.namespace.revision.session",
317            "route": "https://host.example.com",
318            "canonicalRegion": "north-america-east",
319            "provisioning": {
320                "provider": "modal",
321                "resourceId": "sb-actor",
322                "reused": false,
323                "startedAtMs": 0,
324                "resourcesResolvedAtMs": 12,
325                "existingHostCheckedAtMs": 34,
326                "sandboxScheduledAtMs": 56,
327                "hostReadyObservedAtMs": 123,
328                "routeReadAtMs": 125,
329                "metadataWrittenAtMs": 129,
330                "completedAtMs": 130
331            }
332        }))
333        .expect("actor host handle");
334
335        let provisioning = handle.provisioning.expect("provisioning timings");
336        assert_eq!(provisioning.resource_id, "sb-actor");
337        assert_eq!(provisioning.sandbox_scheduled_at_ms, Some(56));
338        assert_eq!(provisioning.completed_at_ms, 130);
339    }
340
341    #[test]
342    fn rejects_ambiguous_command_configuration() {
343        assert!(CommandSandboxProvider::new("".into(), "modal".into(), HashMap::new()).is_err());
344        assert!(CommandSandboxProvider::new("modal".into(), "".into(), HashMap::new()).is_err());
345    }
346
347    #[tokio::test]
348    async fn provider_calls_use_independent_processes() -> Result<()> {
349        let (directory, provider) = test_provider()?;
350        let requests = (0..5)
351            .map(|index| {
352                serde_json::json!({
353                    "index": index,
354                    "barrier": directory.path(),
355                })
356            })
357            .collect::<Vec<_>>();
358        let replies = tokio::time::timeout(
359            Duration::from_secs(3),
360            futures_util::future::try_join_all(
361                requests
362                    .iter()
363                    .map(|request| provider.execute::<_, serde_json::Value>("test", request)),
364            ),
365        )
366        .await
367        .expect("all five independent processes must start before any replies")?;
368        let pids = replies
369            .iter()
370            .map(|reply| reply["pid"].as_u64().unwrap())
371            .collect::<std::collections::HashSet<_>>();
372        assert_eq!(pids.len(), 5);
373        for (index, reply) in replies.iter().enumerate() {
374            assert_eq!(reply["index"], index);
375        }
376        Ok(())
377    }
378
379    #[tokio::test]
380    async fn provider_failures_do_not_affect_other_calls() -> Result<()> {
381        let (_directory, provider) = test_provider()?;
382        for (request, message) in [
383            (serde_json::json!({"fail": true}), "test failure"),
384            (serde_json::json!({"oversized": true}), "stdout exceeds"),
385            (
386                serde_json::json!({"malformed": true}),
387                "decode provider response",
388            ),
389            (serde_json::json!({"exit": true}), "exited"),
390        ] {
391            let healthy_request = serde_json::json!({"index": 42});
392            let (failed, healthy) = tokio::join!(
393                provider.execute::<_, serde_json::Value>("test", &request),
394                provider.execute::<_, serde_json::Value>("test", &healthy_request),
395            );
396            assert!(failed.unwrap_err().to_string().contains(message));
397            assert_eq!(healthy?["index"], 42);
398        }
399        Ok(())
400    }
401
402    #[tokio::test]
403    async fn cancelling_one_provider_call_terminates_only_its_process() -> Result<()> {
404        let (directory, provider) = test_provider()?;
405        let marker = directory.path().join("cancelled.pid");
406        let request = serde_json::json!({"marker": marker, "delay": 30_000});
407        let mut cancelled = Box::pin(provider.execute::<_, serde_json::Value>("test", &request));
408        let wait_for_start = async {
409            while !marker.exists() {
410                tokio::time::sleep(Duration::from_millis(5)).await;
411            }
412        };
413        tokio::select! {
414            result = &mut cancelled => panic!("provider should still be waiting: {result:?}"),
415            started = tokio::time::timeout(Duration::from_secs(3), wait_for_start) => started?,
416        }
417        let pid: u32 = std::fs::read_to_string(marker)?.parse()?;
418        drop(cancelled);
419        let healthy: serde_json::Value = provider
420            .execute("test", &serde_json::json!({"index": 42}))
421            .await?;
422        assert_eq!(healthy["index"], 42);
423        tokio::time::timeout(Duration::from_secs(3), async {
424            while tokio::process::Command::new("kill")
425                .args(["-0", &pid.to_string()])
426                .stderr(std::process::Stdio::null())
427                .status()
428                .await?
429                .success()
430            {
431                tokio::time::sleep(Duration::from_millis(5)).await;
432            }
433            anyhow::Ok(())
434        })
435        .await??;
436        Ok(())
437    }
438
439    fn test_provider() -> Result<(tempfile::TempDir, CommandSandboxProvider)> {
440        use std::os::unix::fs::PermissionsExt;
441        let directory = tempfile::tempdir()?;
442        let path = directory.path().join("custom-provider");
443        std::fs::write(
444            &path,
445            r#"#!/usr/bin/env node
446const fs = require('node:fs');
447const command = JSON.parse(fs.readFileSync(0, 'utf8'));
448const request = command.request;
449const reply = result => process.stdout.write(JSON.stringify({status: 'success', result}) + '\n');
450if (request.barrier) {
451  fs.writeFileSync(request.barrier + '/' + process.pid + '.started', '');
452  const timer = setInterval(() => {
453    if (fs.readdirSync(request.barrier).filter(name => name.endsWith('.started')).length === 5) {
454      clearInterval(timer);
455      reply({pid: process.pid, index: request.index});
456    }
457  }, 5);
458} else if (request.fail) {
459  process.stdout.write(JSON.stringify({status: 'failure', error: 'test failure'}) + '\n');
460} else if (request.oversized) {
461  process.stdout.write('x'.repeat(1024 * 1024 + 1));
462} else if (request.malformed) {
463  process.stdout.write('not json\n');
464} else if (request.exit) {
465  process.exitCode = 1;
466} else {
467  if (request.marker) fs.writeFileSync(request.marker, String(process.pid));
468  setTimeout(() => reply({pid: process.pid, index: request.index}), request.delay ?? 0);
469}
470"#,
471        )?;
472        std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700))?;
473        let provider = CommandSandboxProvider::new(
474            "modal".into(),
475            path.display().to_string(),
476            HashMap::new(),
477        )?;
478        Ok((directory, provider))
479    }
480}