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