Merge pull request #104 from A3S-Lab/feat/e2b-runtime-readiness

fix(compat): gate runtime envd publication on readiness

Author: ZhiXiao-Lin

.github/workflows/e2b-runtime-image.yml

@@ -45,12 +45,12 @@ jobs:
             type=sha,format=long
             type=raw,value=edge,enable={{is_default_branch}}
 
-      - name: Build and push the multi-architecture image
+      - name: Build and push the runtime image
         uses: docker/build-push-action@v6
         with:
           context: .
           file: deploy/e2b/Dockerfile
-          platforms: linux/amd64,linux/arm64
+          platforms: ${{ github.ref == 'refs/heads/main' && 'linux/amd64,linux/arm64' || 'linux/amd64' }}
           push: true
           tags: ${{ steps.metadata.outputs.tags }}
           labels: ${{ steps.metadata.outputs.labels }}

deploy/e2b/README.md

@@ -10,9 +10,10 @@ The build pins:
 - Code Interpreter commit `5aeca43fe3fae2df260b1fb17c71fed5b5dac852`;
 - the JavaScript kernel commit declared in the Dockerfile.
 
-The image starts envd on port `49983`, initializes its default user and working
-directory, and starts the Code Interpreter service on port `49999`. A production
-ACL template selects the in-Sandbox daemon explicitly:
+The image waits for Jupyter and the Code Interpreter service on port `49999`
+before it starts envd on port `49983` and initializes envd's default user and
+working directory. A production ACL template selects the in-Sandbox daemon
+explicitly:
 
 ```acl
 template_policy "code-interpreter-v1" {

deploy/e2b/entrypoint.sh

@@ -11,10 +11,30 @@ terminate() {
 }
 trap terminate EXIT INT TERM
 
-/usr/local/bin/envd -isnotfc -no-cgroups &
-children+=("$!")
+wait_for_service() {
+  local name="$1"
+  local url="$2"
+  local pid="$3"
 
-/usr/bin/python3 /usr/local/lib/a3s-box-e2b/init-envd.py
+  for attempt in {1..150}; do
+    if curl --fail --silent --output /dev/null "${url}"; then
+      return 0
+    fi
+    if ! kill -0 "${pid}" 2>/dev/null; then
+      local status=0
+      wait "${pid}" || status=$?
+      echo "${name} exited before becoming healthy with status ${status}" >&2
+      if ((status == 0)); then
+        status=1
+      fi
+      return "${status}"
+    fi
+    sleep 0.2
+  done
+
+  echo "${name} did not become healthy within 30 seconds" >&2
+  return 1
+}
 
 runuser -u user -- env \
   HOME=/home/user \
@@ -23,17 +43,9 @@ runuser -u user -- env \
   --IdentityProvider.token= \
   --ServerApp.root_dir=/home/user &
 children+=("$!")
+jupyter_pid="$!"
 
-for attempt in {1..150}; do
-  if curl --fail --silent --output /dev/null http://127.0.0.1:8888/api/status; then
-    break
-  fi
-  if ((attempt == 150)); then
-    echo "Jupyter did not become healthy within 30 seconds" >&2
-    exit 1
-  fi
-  sleep 0.2
-done
+wait_for_service "Jupyter" "http://127.0.0.1:8888/api/status" "${jupyter_pid}"
 
 runuser -u user -- env \
   HOME=/home/user \
@@ -48,6 +60,17 @@ runuser -u user -- env \
   --no-use-colors \
   --timeout-keep-alive 640 &
 children+=("$!")
+code_interpreter_pid="$!"
+
+wait_for_service \
+  "Code Interpreter" \
+  "http://127.0.0.1:49999/health" \
+  "${code_interpreter_pid}"
+
+/usr/local/bin/envd -isnotfc -no-cgroups &
+children+=("$!")
+
+/usr/bin/python3 /usr/local/lib/a3s-box-e2b/init-envd.py
 
 set +e
 wait -n "${children[@]}"

src/compat/src/bin/a3s-box-e2b-fixture-server.rs

@@ -1,7 +1,9 @@
 use std::collections::BTreeMap;
+use std::num::NonZeroU16;
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex, MutexGuard};
+use std::time::Duration;
 
 use a3s_box_compat::control::{
     Clock, ControlService, ControlServiceDependencies, IdentityProviderResult, IssuedToken,
@@ -17,8 +19,8 @@ use a3s_box_compat::http::{
 use a3s_box_core::{
     resolve_execution, BoxConfig, CreateExecutionRequest, ExecutionGeneration, ExecutionId,
     ExecutionIsolation, ExecutionLease, ExecutionManager, ExecutionManagerError,
-    ExecutionManagerResult, ExecutionReservation, ExecutionState, ExecutionStatus, KillOutcome,
-    OperationId, ReconcileOutcome, ResourceConfig,
+    ExecutionManagerResult, ExecutionPortConnector, ExecutionPortStream, ExecutionReservation,
+    ExecutionState, ExecutionStatus, KillOutcome, OperationId, ReconcileOutcome, ResourceConfig,
 };
 use anyhow::{bail, Context, Result};
 use async_trait::async_trait;
@@ -39,9 +41,11 @@ async fn run() -> Result<()> {
     let port_file = parse_port_file()?;
     let clock = Arc::new(FixedClock(fixture_time()?));
     let tokens = Arc::new(FixtureTokens);
+    let executions = Arc::new(FixtureExecutionManager::new(clock.clone()));
     let service = Arc::new(ControlService::new(ControlServiceDependencies {
         repository: Arc::new(MemorySandboxRepository::default()),
-        executions: Arc::new(FixtureExecutionManager::new(clock.clone())),
+        executions: executions.clone(),
+        ports: executions,
         clock,
         identities: Arc::new(FixtureIdentities::default()),
         templates: Arc::new(FixtureTemplates),
@@ -462,3 +466,16 @@ impl ExecutionManager for FixtureExecutionManager {
         })
     }
 }
+
+#[async_trait]
+impl ExecutionPortConnector for FixtureExecutionManager {
+    async fn connect_port(
+        &self,
+        execution_id: &ExecutionId,
+        _generation: ExecutionGeneration,
+        _port: NonZeroU16,
+        _timeout: Duration,
+    ) -> ExecutionManagerResult<ExecutionPortStream> {
+        Err(ExecutionManagerError::NotFound(execution_id.clone()))
+    }
+}

src/compat/src/control/service.rs

@@ -1,17 +1,26 @@
 use std::collections::BTreeMap;
+use std::num::NonZeroU16;
 use std::sync::Arc;
 
-use a3s_box_core::{resolve_execution, CreateExecutionRequest, ExecutionManager, NetworkMode};
+use a3s_box_core::{
+    resolve_execution, CreateExecutionRequest, ExecutionLease, ExecutionManager,
+    ExecutionManagerError, ExecutionPortConnector, NetworkMode,
+};
 use chrono::{DateTime, Duration, Utc};
 use thiserror::Error;
 
 use super::{
-    Clock, CompareAndSwapResult, IdentityProviderError, LifecycleError, LifecycleFailure,
+    Clock, CompareAndSwapResult, EnvdMode, IdentityProviderError, LifecycleError, LifecycleFailure,
     LifecyclePolicy, LifecycleState, NewSandboxRecord, RepositoryError, SandboxCredentials,
     SandboxIdentityProvider, SandboxListFilter, SandboxPage, SandboxRecord, SandboxRepository,
     SecretToken, TemplateProvider, TemplateProviderError, TokenIssuer, TokenIssuerError,
     TokenResolver, TokenScope,
 };
+use crate::routing::ENVD_PORT;
+
+const RUNTIME_ENVD_READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
+const RUNTIME_ENVD_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(250);
+const RUNTIME_ENVD_RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(25);
 
 #[derive(Debug, Clone)]
 pub struct CreateSandboxRequest {
@@ -79,6 +88,7 @@ pub type ControlServiceResult<T> = std::result::Result<T, ControlServiceError>;
 pub struct ControlService {
     repository: Arc<dyn SandboxRepository>,
     executions: Arc<dyn ExecutionManager>,
+    ports: Arc<dyn ExecutionPortConnector>,
     clock: Arc<dyn Clock>,
     identities: Arc<dyn SandboxIdentityProvider>,
     templates: Arc<dyn TemplateProvider>,
@@ -89,6 +99,7 @@ pub struct ControlService {
 pub struct ControlServiceDependencies {
     pub repository: Arc<dyn SandboxRepository>,
     pub executions: Arc<dyn ExecutionManager>,
+    pub ports: Arc<dyn ExecutionPortConnector>,
     pub clock: Arc<dyn Clock>,
     pub identities: Arc<dyn SandboxIdentityProvider>,
     pub templates: Arc<dyn TemplateProvider>,
@@ -101,6 +112,7 @@ impl ControlService {
         Self {
             repository: dependencies.repository,
             executions: dependencies.executions,
+            ports: dependencies.ports,
             clock: dependencies.clock,
             identities: dependencies.identities,
             templates: dependencies.templates,
@@ -180,6 +192,24 @@ impl ControlService {
                 return Err(error.into());
             }
         };
+        if template.envd_mode == EnvdMode::Runtime {
+            if let Err(readiness_error) = self.wait_for_runtime_envd(&lease).await {
+                let cleanup = self
+                    .executions
+                    .kill(&lease.execution_id, lease.generation)
+                    .await;
+                let expected = record.generation();
+                record.mark_failed(LifecycleFailure::RuntimeFailed)?;
+                self.replace(expected, record).await?;
+                return Err(match cleanup {
+                    Ok(_) => readiness_error,
+                    Err(cleanup_error) => ExecutionManagerError::Internal(format!(
+                        "{readiness_error}; runtime cleanup failed: {cleanup_error}"
+                    )),
+                }
+                .into());
+            }
+        }
 
         let expected = record.generation();
         if let Err(error) = record.mark_running(lease) {
@@ -197,6 +227,44 @@ impl ControlService {
         })
     }
 
+    async fn wait_for_runtime_envd(
+        &self,
+        lease: &ExecutionLease,
+    ) -> Result<(), ExecutionManagerError> {
+        let port = NonZeroU16::new(ENVD_PORT).ok_or_else(|| {
+            ExecutionManagerError::Internal("envd port must be non-zero".to_string())
+        })?;
+        let deadline = tokio::time::Instant::now() + RUNTIME_ENVD_READY_TIMEOUT;
+        loop {
+            let last_error = match self
+                .ports
+                .connect_port(
+                    &lease.execution_id,
+                    lease.generation,
+                    port,
+                    RUNTIME_ENVD_CONNECT_TIMEOUT,
+                )
+                .await
+            {
+                Ok(stream) => {
+                    drop(stream);
+                    return Ok(());
+                }
+                Err(error @ ExecutionManagerError::InvalidRequest(_))
+                | Err(error @ ExecutionManagerError::Internal(_)) => return Err(error),
+                Err(error) => error,
+            };
+            if tokio::time::Instant::now() >= deadline {
+                return Err(ExecutionManagerError::Unavailable(format!(
+                    "runtime envd did not become ready within {} seconds: {}",
+                    RUNTIME_ENVD_READY_TIMEOUT.as_secs(),
+                    last_error
+                )));
+            }
+            tokio::time::sleep(RUNTIME_ENVD_RETRY_INTERVAL).await;
+        }
+    }
+
     pub async fn connect(
         &self,
         owner_id: &str,

src/compat/src/control/service_tests.rs

@@ -1,6 +1,7 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::num::NonZeroU32;
 
+use a3s_box_core::{ExecutionManagerError, ExecutionState};
 use chrono::Duration;
 
 use super::test_support::{assert_sandbox_request, create_request, test_time, TestHarness};
@@ -126,3 +127,55 @@ async fn failed_runtime_create_is_not_published() {
         Err(ControlServiceError::NotFound(_))
     ));
 }
+
+#[tokio::test]
+async fn runtime_envd_is_ready_before_the_sandbox_is_published() {
+    let harness = TestHarness::new();
+    let mut request = create_request("owner-1");
+    request.template_id = "runtime-envd-template".to_string();
+
+    let created = harness.service.create(request).await.unwrap();
+
+    assert_eq!(created.record.state(), LifecycleState::Running);
+    assert_eq!(created.record.envd_mode(), EnvdMode::Runtime);
+    assert_eq!(
+        harness.executions.port_requests(),
+        vec![(
+            "execution-operation-1".to_string(),
+            1,
+            crate::routing::ENVD_PORT,
+        )]
+    );
+}
+
+#[tokio::test]
+async fn permanent_runtime_envd_failure_stops_and_hides_the_execution() {
+    let harness = TestHarness::new();
+    harness.executions.fail_ports();
+    let mut request = create_request("owner-1");
+    request.template_id = "runtime-envd-template".to_string();
+
+    assert!(matches!(
+        harness.service.create(request).await,
+        Err(ControlServiceError::Execution(
+            ExecutionManagerError::InvalidRequest(_)
+        ))
+    ));
+    assert_eq!(
+        harness.executions.port_requests(),
+        vec![(
+            "execution-operation-1".to_string(),
+            1,
+            crate::routing::ENVD_PORT,
+        )]
+    );
+    assert_eq!(
+        harness.executions.execution_state("execution-operation-1"),
+        Some(ExecutionState::Stopped)
+    );
+    let sandbox_id = SandboxId::new("sandbox-1").unwrap();
+    assert!(matches!(
+        harness.service.get("owner-1", &sandbox_id).await,
+        Err(ControlServiceError::NotFound(_))
+    ));
+}