diff --git a/.agents/design/burst-worker-cold-ready.md b/.agents/design/burst-worker-cold-ready.md new file mode 100644 index 00000000..55168109 --- /dev/null +++ b/.agents/design/burst-worker-cold-ready.md @@ -0,0 +1,131 @@ +# Burst Worker Cold-Ready Startup Plan + +## Context + +The burst worker launch test on June 10, 2026 showed two different timing +segments: + +- EC2 instance creation to worker service start was roughly 90 seconds. +- Worker service start to control-plane registration was much longer because + startup blocked on `PrepareGoldenSnapshot`. + +The important observation is that the worker can be useful for cold boots +before the golden snapshot is ready. The current startup path does not expose +that intermediate state because the worker prepares the golden snapshot before +starting its servers and heartbeat. + +## Goal + +Make a newly launched burst worker register as soon as it is cold-boot capable, +while preparing the golden snapshot in the background. + +Target behavior: + +- Worker becomes schedulable for cold boots as soon as networking, env, shared + mounts, gRPC, HTTP, and Redis heartbeat are ready. +- Golden snapshot preparation continues asynchronously. +- Once the golden snapshot is ready, the worker heartbeat advertises the golden + version and the control plane can prefer it for fast creates. + +This does not remove EC2 launch latency. It removes golden snapshot creation +from the critical path for worker registration. + +## Proposed Changes + +1. Move golden snapshot preparation out of the blocking worker startup path. + + Today `cmd/worker/main.go` calls `PrepareGoldenSnapshot()` before starting + metadata, HTTP/gRPC, and Redis heartbeat. Move this after server startup and + heartbeat setup, running in a background goroutine. + +2. Register the worker as cold-ready first. + + Heartbeat should be published with no `golden_version` until the snapshot is + ready. The control plane already treats empty `golden_version` as "no golden + snapshot available"; keep that meaning. + +3. Update heartbeat when golden prep completes. + + After background `PrepareGoldenSnapshot()` succeeds, call + `hb.SetGoldenVersion(qemuMgr.GoldenVersion())`. The next heartbeat should + update the registry. + +4. Add explicit logs for readiness phases. + + Suggested log points: + + - `worker cold-ready: starting heartbeat before golden snapshot` + - `worker golden snapshot preparation started in background` + - `worker golden-ready: version=` + - `worker golden preparation failed: ; continuing cold-ready` + +5. Fix AMI/systemd ordering for burst workers. + + The burst AMI currently enables `opensandbox-worker.service`, so systemd can + start it before user-data writes `/etc/opensandbox/worker.env`. That caused + repeated `Failed to load environment files` messages during boot. + + Change the burst Packer file to install the worker unit but leave it + disabled. User-data should start the worker exactly once after: + + - instance identity is known + - shared volumes are attached/mounted + - `/etc/opensandbox/worker.env` has been written and patched + +6. Keep user-data minimal. + + User-data should only do runtime-specific work: + + - fetch instance identity + - attach/mount shared volumes + - write env + - start worker + + Dependency installation, binaries, OCFS2 tools, AWS CLI, QEMU, kernel + modules, and rootfs assets should stay baked into the AMI. + +## Non-Goals + +- Do not change Spot instance type fallback strategy yet. +- Do not try to guarantee sub-10-second readiness from a brand-new EC2 launch. +- Do not implement downloaded/prebuilt QEMU memory snapshots in this pass. +- Do not change public API behavior. + +## Expected Impact + +Based on the June 10 test: + +- Current EC2-created-to-registered time was about 6 minutes 24 seconds. +- Worker service started about 91 seconds after EC2 creation. +- Moving golden prep to the background could make cold-ready registration close + to that worker-service-start time, likely around 90-100 seconds from EC2 + creation before further AMI cleanup. + +With AMI/systemd cleanup, a realistic next target is roughly 45-70 seconds from +EC2 creation to cold-ready in favorable cases. + +## Risks + +- Cold-ready workers may serve slower first sandboxes until golden prep + completes. +- Some scheduling paths may implicitly assume a non-empty `golden_version`. + Those paths need review before allowing all workloads onto cold-ready workers. +- Migration/checkpoint paths that require a known source golden version should + continue to require it. + +## Validation Plan + +1. Build and deploy a worker with background golden prep. +2. Launch a fresh burst worker and capture timestamps: + - scaler launch decision + - EC2 instance created + - user-data start + - worker service start + - first Redis heartbeat / CP registration + - golden snapshot ready +3. Confirm the CP sees the worker before golden snapshot readiness. +4. Create a sandbox on the cold-ready worker and verify it succeeds via cold + boot. +5. Wait for golden-ready heartbeat and verify subsequent creates use the golden + path. +6. Terminate the extra worker after the test to avoid unnecessary cost. diff --git a/.github/workflows/build-aws-worker-ami.yml b/.github/workflows/build-aws-worker-ami.yml new file mode 100644 index 00000000..b995bc17 --- /dev/null +++ b/.github/workflows/build-aws-worker-ami.yml @@ -0,0 +1,197 @@ +name: Build AWS Worker AMI + +on: + push: + branches: + - main + - feat/aws-poc-worker-support-clean-rebased + workflow_dispatch: + inputs: + region: + description: AWS region to build the AMI in + required: true + default: us-east-2 + builder_instance_type: + description: EC2 instance type used by Packer for the build + required: true + default: c5.4xlarge + cleanup_old_amis: + description: Deregister old OpenComputer AWS worker AMIs after a successful build + required: true + default: "false" + type: choice + options: + - "false" + - "true" + ssm_parameter_prefix: + description: SSM prefix to update after publishing the AMI + required: true + default: /opencomputer/aws-us-east-2-burst-prod + +env: + AWS_REGION: ${{ inputs.region || vars.AWS_REGION || 'us-east-2' }} + BUILDER_INSTANCE_TYPE: ${{ inputs.builder_instance_type || 'c5.4xlarge' }} + SSM_PARAMETER_PREFIX: ${{ inputs.ssm_parameter_prefix || '/opencomputer/aws-us-east-2-burst-prod' }} + +jobs: + build-ami: + name: Build AWS Worker AMI + runs-on: ubuntu-latest + environment: aws-prod + permissions: + id-token: write + contents: read + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: actions/setup-go@v5 + with: + go-version: "1.23" + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.AWS_ROLE_ARN }} + aws-region: ${{ env.AWS_REGION }} + + - name: Setup Packer + uses: hashicorp/setup-packer@main + + - name: Build binaries (amd64) + run: | + VERSION=$(git rev-parse --short HEAD) + echo "VERSION=$VERSION" >> "$GITHUB_ENV" + + AGENT_VERSION=$(git log -1 --pretty=format:%h -- cmd/agent internal/agent proto/agent) + if [ -z "$AGENT_VERSION" ]; then + AGENT_VERSION=$VERSION + fi + echo "AGENT_VERSION=$AGENT_VERSION" >> "$GITHUB_ENV" + + CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \ + -ldflags "-X main.WorkerVersion=$VERSION -X main.AgentVersion=$AGENT_VERSION" \ + -o bin/opensandbox-worker ./cmd/worker/ + + CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \ + -ldflags "-X main.Version=$AGENT_VERSION" \ + -o bin/osb-agent ./cmd/agent/ + + - name: Package rootfs context + run: | + tar czf /tmp/packer-rootfs-ctx.tar.gz \ + deploy/firecracker/rootfs/ \ + deploy/ec2/build-rootfs-docker.sh \ + scripts/claude-agent-wrapper/ + + - name: Package Vector configs + run: tar czf /tmp/packer-vector-ctx.tar.gz -C deploy vector + + - name: Packer init + run: packer init deploy/packer/worker-ami-aws-burst.pkr.hcl + + - name: Build and publish AMI + run: | + packer build \ + -var "worker_version=$VERSION" \ + -var "agent_version=$AGENT_VERSION" \ + -var "region=$AWS_REGION" \ + -var "instance_type=$BUILDER_INSTANCE_TYPE" \ + -var "tigris_endpoint=${{ secrets.TIGRIS_ENDPOINT }}" \ + -var "tigris_access_key_id=${{ secrets.TIGRIS_ACCESS_KEY_ID }}" \ + -var "tigris_secret_access_key=${{ secrets.TIGRIS_SECRET_ACCESS_KEY }}" \ + -var "tigris_goldens_bucket=${{ secrets.TIGRIS_GOLDENS_BUCKET }}" \ + deploy/packer/worker-ami-aws-burst.pkr.hcl | tee /tmp/packer-output.txt + + - name: Read AMI manifest + id: ami + run: | + AMI_ID=$(jq -r '.builds[-1].artifact_id | split(":")[-1]' packer-manifest-aws.json) + if [ -z "$AMI_ID" ] || [ "$AMI_ID" = "null" ]; then + echo "Could not read AMI ID from packer-manifest-aws.json" + cat packer-manifest-aws.json + exit 1 + fi + + GOLDEN_VERSION=$(grep -a 'Golden version:' /tmp/packer-output.txt | tail -1 | awk '{print $NF}' | tr -d '\r') + + echo "ami_id=$AMI_ID" >> "$GITHUB_OUTPUT" + echo "golden_version=$GOLDEN_VERSION" >> "$GITHUB_OUTPUT" + echo "AMI_ID=$AMI_ID" >> "$GITHUB_ENV" + echo "GOLDEN_VERSION=$GOLDEN_VERSION" >> "$GITHUB_ENV" + + - name: Verify AMI tags + run: | + aws ec2 describe-images \ + --region "$AWS_REGION" \ + --image-ids "$AMI_ID" \ + --query 'Images[0].{ImageId:ImageId,Name:Name,State:State,Role:Tags[?Key==`opensandbox-role`].Value|[0],Cloud:Tags[?Key==`opensandbox-cloud`].Value|[0],Version:Tags[?Key==`opensandbox-version`].Value|[0]}' \ + --output table + + - name: Update worker AMI pointer + run: | + aws ssm put-parameter \ + --region "$AWS_REGION" \ + --name "$SSM_PARAMETER_PREFIX/worker-ami-id" \ + --type String \ + --value "$AMI_ID" \ + --overwrite + + aws ssm put-parameter \ + --region "$AWS_REGION" \ + --name "$SSM_PARAMETER_PREFIX/worker-ami-version" \ + --type String \ + --value "$VERSION" \ + --overwrite + + echo "Updated SSM worker AMI pointer: $SSM_PARAMETER_PREFIX/worker-ami-id -> $AMI_ID" + + - name: Cleanup old AMIs + if: ${{ inputs.cleanup_old_amis == 'true' }} + run: | + set -euo pipefail + + mapfile -t OLD_AMIS < <( + aws ec2 describe-images \ + --region "$AWS_REGION" \ + --owners self \ + --filters \ + "Name=tag:opensandbox-role,Values=worker" \ + "Name=tag:opensandbox-cloud,Values=aws" \ + "Name=state,Values=available" \ + --query 'reverse(sort_by(Images,&CreationDate))[10:].ImageId' \ + --output text | tr '\t' '\n' | awk 'NF' + ) + + if [ "${#OLD_AMIS[@]}" -eq 0 ]; then + echo "No old AMIs to clean up" + exit 0 + fi + + for AMI in "${OLD_AMIS[@]}"; do + echo "Deregistering old AMI: $AMI" + SNAPSHOTS=$(aws ec2 describe-images \ + --region "$AWS_REGION" \ + --image-ids "$AMI" \ + --query 'Images[0].BlockDeviceMappings[].Ebs.SnapshotId' \ + --output text) + aws ec2 deregister-image --region "$AWS_REGION" --image-id "$AMI" + for SNAPSHOT in $SNAPSHOTS; do + echo "Deleting old AMI snapshot: $SNAPSHOT" + aws ec2 delete-snapshot --region "$AWS_REGION" --snapshot-id "$SNAPSHOT" || true + done + done + + - name: Summary + run: | + echo "## AWS Worker AMI Build Complete" >> "$GITHUB_STEP_SUMMARY" + echo "- **AMI ID:** \`$AMI_ID\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Worker version:** \`$VERSION\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Agent version:** \`$AGENT_VERSION\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Golden version:** \`${GOLDEN_VERSION:-unknown}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Golden cache:** \`${{ secrets.TIGRIS_GOLDENS_BUCKET != '' && 'Tigris enabled' || 'disabled' }}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Region:** \`$AWS_REGION\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **SSM pointer:** \`$SSM_PARAMETER_PREFIX/worker-ami-id\`" >> "$GITHUB_STEP_SUMMARY" + echo "" >> "$GITHUB_STEP_SUMMARY" + echo "The CP scaler reads this AMI through \`OPENSANDBOX_EC2_SSM_AMI_PARAM\`. Terraform creates the parameter but ignores value drift so later applies do not overwrite CI-published AMIs." >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/deploy-aws-burst-cp.yml b/.github/workflows/deploy-aws-burst-cp.yml new file mode 100644 index 00000000..9ee25958 --- /dev/null +++ b/.github/workflows/deploy-aws-burst-cp.yml @@ -0,0 +1,283 @@ +name: Deploy AWS Burst Control Plane + +on: + push: + branches: + - main + - feat/aws-poc-worker-support-clean-rebased + paths: + - "cmd/server/**" + - "internal/**" + - "web/**" + - "deploy/vector/**" + - "go.mod" + - "go.sum" + - ".github/workflows/deploy-aws-burst-cp.yml" + workflow_dispatch: + inputs: + region: + description: AWS region + required: true + default: us-east-2 + cell_id: + description: Cell ID + required: true + default: aws-us-east-2-burst-prod + cp_domain: + description: Public control-plane domain + required: true + default: cp-aws-us-east-2-burst-prod.opencomputer.dev + +env: + AWS_REGION: ${{ inputs.region || vars.AWS_REGION || 'us-east-2' }} + CELL_ID: ${{ inputs.cell_id || 'aws-us-east-2-burst-prod' }} + CP_DOMAIN: ${{ inputs.cp_domain || 'cp-aws-us-east-2-burst-prod.opencomputer.dev' }} + GOARCH: amd64 + +jobs: + deploy: + name: Deploy AWS Burst CP + runs-on: ubuntu-latest + timeout-minutes: 30 + permissions: + id-token: write + contents: read + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: "1.23" + + - uses: actions/setup-node@v4 + with: + node-version: "20" + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.AWS_ROLE_ARN }} + aws-region: ${{ env.AWS_REGION }} + + - name: Discover control plane + run: | + set -euo pipefail + + CP_INSTANCE_ID=$(aws ec2 describe-instances \ + --region "$AWS_REGION" \ + --filters \ + "Name=tag:Name,Values=oc-cp-$CELL_ID" \ + "Name=instance-state-name,Values=running" \ + --query 'Reservations[0].Instances[0].InstanceId' \ + --output text) + + if [ -z "$CP_INSTANCE_ID" ] || [ "$CP_INSTANCE_ID" = "None" ]; then + echo "No running CP instance found for $CELL_ID" + exit 1 + fi + + CP_PUBLIC_IP=$(aws ec2 describe-instances \ + --region "$AWS_REGION" \ + --instance-ids "$CP_INSTANCE_ID" \ + --query 'Reservations[0].Instances[0].PublicIpAddress' \ + --output text) + + CP_AZ=$(aws ec2 describe-instances \ + --region "$AWS_REGION" \ + --instance-ids "$CP_INSTANCE_ID" \ + --query 'Reservations[0].Instances[0].Placement.AvailabilityZone' \ + --output text) + + CP_SG_ID=$(aws ec2 describe-security-groups \ + --region "$AWS_REGION" \ + --filters "Name=group-name,Values=oc-sg-cp-$CELL_ID" \ + --query 'SecurityGroups[0].GroupId' \ + --output text) + + if [ -z "$CP_PUBLIC_IP" ] || [ "$CP_PUBLIC_IP" = "None" ]; then + echo "CP instance $CP_INSTANCE_ID has no public IP" + exit 1 + fi + + if [ -z "$CP_SG_ID" ] || [ "$CP_SG_ID" = "None" ]; then + echo "CP security group not found for $CELL_ID" + exit 1 + fi + + echo "CP_INSTANCE_ID=$CP_INSTANCE_ID" >> "$GITHUB_ENV" + echo "CP_PUBLIC_IP=$CP_PUBLIC_IP" >> "$GITHUB_ENV" + echo "CP_AZ=$CP_AZ" >> "$GITHUB_ENV" + echo "CP_SG_ID=$CP_SG_ID" >> "$GITHUB_ENV" + + - name: Build server binary + run: | + set -euo pipefail + VERSION=$(git rev-parse --short HEAD) + echo "VERSION=$VERSION" >> "$GITHUB_ENV" + CGO_ENABLED=0 GOOS=linux GOARCH="$GOARCH" go build \ + -ldflags "-X main.Version=$VERSION" \ + -o bin/opensandbox-server ./cmd/server/ + + - name: Build web dashboard + run: cd web && npm ci && npm run build + + - name: Package deploy artifacts + run: | + set -euo pipefail + tar czf bin/web-dist.tar.gz -C web dist + tar czf bin/vector-deploy.tar.gz -C deploy vector + + - name: Prepare SSH key + run: | + set -euo pipefail + mkdir -p ~/.ssh + printf '%s\n' "${{ secrets.AWS_CP_SSH_PRIVATE_KEY }}" > ~/.ssh/aws-burst-cp + chmod 600 ~/.ssh/aws-burst-cp + ssh-keygen -y -f ~/.ssh/aws-burst-cp > ~/.ssh/aws-burst-cp.pub + + - name: Authorize runner SSH + run: | + set -euo pipefail + RUNNER_IP=$(curl -fsSL https://checkip.amazonaws.com | tr -d '[:space:]') + echo "RUNNER_CIDR=${RUNNER_IP}/32" >> "$GITHUB_ENV" + + set +e + aws ec2 authorize-security-group-ingress \ + --region "$AWS_REGION" \ + --group-id "$CP_SG_ID" \ + --ip-permissions "IpProtocol=tcp,FromPort=22,ToPort=22,IpRanges=[{CidrIp=${RUNNER_IP}/32,Description=github-actions-cp-deploy-${GITHUB_RUN_ID}}]" + STATUS=$? + set -e + + if [ "$STATUS" -ne 0 ]; then + echo "SSH ingress may already exist; continuing so deploy can attempt SSH." + fi + + - name: Push SSH key with EC2 Instance Connect + run: | + set -euo pipefail + aws ec2-instance-connect send-ssh-public-key \ + --region "$AWS_REGION" \ + --instance-id "$CP_INSTANCE_ID" \ + --availability-zone "$CP_AZ" \ + --instance-os-user ubuntu \ + --ssh-public-key "file://$HOME/.ssh/aws-burst-cp.pub" + + - name: Wait for SSH + run: | + set -euo pipefail + for i in $(seq 1 30); do + if ssh -i ~/.ssh/aws-burst-cp \ + -o StrictHostKeyChecking=no \ + -o UserKnownHostsFile=/dev/null \ + -o ConnectTimeout=5 \ + ubuntu@"$CP_PUBLIC_IP" 'echo ssh-ready'; then + exit 0 + fi + sleep 5 + done + echo "SSH did not become ready" + exit 1 + + - name: Upload artifacts + run: | + set -euo pipefail + scp -i ~/.ssh/aws-burst-cp \ + -o StrictHostKeyChecking=no \ + -o UserKnownHostsFile=/dev/null \ + bin/opensandbox-server bin/web-dist.tar.gz bin/vector-deploy.tar.gz \ + ubuntu@"$CP_PUBLIC_IP":/tmp/ + + - name: Install and restart control plane + run: | + set -euo pipefail + ssh -i ~/.ssh/aws-burst-cp \ + -o StrictHostKeyChecking=no \ + -o UserKnownHostsFile=/dev/null \ + ubuntu@"$CP_PUBLIC_IP" 'bash -s' <<'REMOTE' + set -euo pipefail + + sudo mkdir -p /opt/opensandbox/web + sudo install -m 0755 /tmp/opensandbox-server /usr/local/bin/opensandbox-server + sudo tar xzf /tmp/web-dist.tar.gz -C /opt/opensandbox/web + + sudo tee /etc/systemd/system/opensandbox-server.service >/dev/null <<'UNIT' + [Unit] + Description=OpenComputer Control Plane + After=network-online.target cloud-init.target + Wants=network-online.target + + [Service] + Type=simple + WorkingDirectory=/opt/opensandbox + EnvironmentFile=/etc/opensandbox/server.env + ExecStart=/usr/local/bin/opensandbox-server + Restart=always + RestartSec=5 + + [Install] + WantedBy=multi-user.target + UNIT + + sudo systemctl daemon-reload + sudo systemctl enable opensandbox-server + sudo systemctl restart opensandbox-server + + mkdir -p /tmp/vector-deploy + tar xzf /tmp/vector-deploy.tar.gz -C /tmp/vector-deploy + sudo bash /tmp/vector-deploy/vector/install.sh control-plane || true + rm -rf /tmp/vector-deploy /tmp/opensandbox-server /tmp/web-dist.tar.gz /tmp/vector-deploy.tar.gz + + sudo systemctl is-active opensandbox-server + REMOTE + + - name: Health check + run: | + set -euo pipefail + + echo "Checking direct CP health..." + for i in $(seq 1 30); do + if curl -fsS --max-time 5 "http://${CP_PUBLIC_IP}:8080/health"; then + echo + break + fi + sleep 5 + done + + echo "Checking public domain health..." + for i in $(seq 1 30); do + if curl -fsS --max-time 10 "https://${CP_DOMAIN}/health"; then + echo + exit 0 + fi + sleep 5 + done + + echo "Public health check failed" + exit 1 + + - name: Revoke runner SSH + if: always() + run: | + set -euo pipefail + if [ -z "${RUNNER_CIDR:-}" ] || [ -z "${CP_SG_ID:-}" ]; then + echo "No runner SSH rule recorded; skipping revoke" + exit 0 + fi + + aws ec2 revoke-security-group-ingress \ + --region "$AWS_REGION" \ + --group-id "$CP_SG_ID" \ + --protocol tcp \ + --port 22 \ + --cidr "$RUNNER_CIDR" || true + + - name: Summary + if: always() + run: | + echo "## AWS Burst CP Deploy" >> "$GITHUB_STEP_SUMMARY" + echo "- **Cell:** \`$CELL_ID\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **CP instance:** \`${CP_INSTANCE_ID:-unknown}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **CP public IP:** \`${CP_PUBLIC_IP:-unknown}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Version:** \`${VERSION:-unknown}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- **Domain:** \`https://$CP_DOMAIN\`" >> "$GITHUB_STEP_SUMMARY" diff --git a/cloudflare-workers/api-edge/src/index.ts b/cloudflare-workers/api-edge/src/index.ts index d3193daa..ae6cbfc5 100644 --- a/cloudflare-workers/api-edge/src/index.ts +++ b/cloudflare-workers/api-edge/src/index.ts @@ -33,6 +33,9 @@ export interface Env extends DashboardEnv { // CF_API_TOKEN and CF_ZONE_ID are optional in DashboardEnv (custom domain // feature gates on them). Inherited. ASSETS?: Fetcher; + // Optional alpha spot-only cell. When unset, sandboxFamily="spot" creates + // fail closed rather than silently landing on on-demand capacity. + SPOT_CELL_ID?: string; } // ── small helpers ──────────────────────────────────────────────────────── @@ -440,18 +443,54 @@ async function createSandbox(req: Request, env: Env): Promise { // Read body once — used for size-gating, the hard-pin cell peek, and the // verbatim forward to the CP. - const bodyText = await req.text(); + let bodyText = await req.text(); let requestedCellID: string | null = null; let bodyCpuCount = 0; let bodyMemoryMB = 0; let bodyDiskMB = 0; + let sandboxFamily = ""; try { if (bodyText) { - const parsed = JSON.parse(bodyText) as { cellId?: unknown; cpuCount?: unknown; memoryMB?: unknown; diskMB?: unknown }; + const parsed = JSON.parse(bodyText) as { + cellId?: unknown; + cpuCount?: unknown; + memoryMB?: unknown; + diskMB?: unknown; + sandboxFamily?: unknown; + image?: unknown; + snapshot?: unknown; + }; if (typeof parsed.cellId === "string") requestedCellID = parsed.cellId; if (typeof parsed.cpuCount === "number") bodyCpuCount = parsed.cpuCount; if (typeof parsed.memoryMB === "number") bodyMemoryMB = parsed.memoryMB; if (typeof parsed.diskMB === "number") bodyDiskMB = parsed.diskMB; + if (typeof parsed.sandboxFamily === "string") sandboxFamily = parsed.sandboxFamily; + + if (sandboxFamily === "default") sandboxFamily = ""; + if (sandboxFamily && sandboxFamily !== "spot") { + return json({ error: `unsupported sandboxFamily ${sandboxFamily}` }, 400); + } + if (sandboxFamily === "spot") { + if (!env.SPOT_CELL_ID) { + return json({ error: "spot sandbox alpha is not configured" }, 503); + } + if (requestedCellID && requestedCellID !== env.SPOT_CELL_ID) { + return json({ error: "sandboxFamily spot cannot be combined with a different cellId" }, 400); + } + if (parsed.image != null || parsed.snapshot != null) { + return json({ error: "sandboxFamily spot does not support image or snapshot creates in alpha" }, 400); + } + if ((bodyCpuCount && bodyCpuCount !== 1) || (bodyMemoryMB && bodyMemoryMB !== 1024)) { + return json({ error: "sandboxFamily spot is currently limited to 1 vCPU and 1024 MB memory" }, 400); + } + parsed.cpuCount = 1; + parsed.memoryMB = 1024; + parsed.sandboxFamily = "spot"; + requestedCellID = env.SPOT_CELL_ID; + bodyCpuCount = 1; + bodyMemoryMB = 1024; + bodyText = JSON.stringify(parsed); + } } } catch { /* malformed JSON — let the CP reject with a proper 400 */ diff --git a/cloudflare-workers/api-edge/wrangler.prod.toml b/cloudflare-workers/api-edge/wrangler.prod.toml index 32bfc6dc..39316060 100644 --- a/cloudflare-workers/api-edge/wrangler.prod.toml +++ b/cloudflare-workers/api-edge/wrangler.prod.toml @@ -62,6 +62,8 @@ WORKER_ENV = "prod" # No prod cells registered yet. Populate when prod CPs come online and # their rows are written to the D1 `cells` table. CELLS = "" +# Optional alpha route for sandboxFamily="spot". Leave empty to disable. +SPOT_CELL_ID = "" # Ship Worker logs to opencomputer-log-tail-prod → Axiom dataset cf-prod. [[tail_consumers]] diff --git a/cloudflare-workers/api-edge/wrangler.toml b/cloudflare-workers/api-edge/wrangler.toml index f5d32907..e5524569 100644 --- a/cloudflare-workers/api-edge/wrangler.toml +++ b/cloudflare-workers/api-edge/wrangler.toml @@ -84,6 +84,8 @@ WORKER_ENV = "dev" # the D1 `cells` table (see schema.sql + scripts/seed_dev_xcell.sh). # Cross-cell testbed: dev2 (Azure westus2) + dev3 (AWS us-east-1). CELLS = "azure-us-west-2-b,aws-us-east-1-a" +# Optional alpha route for sandboxFamily="spot". Leave empty to disable. +SPOT_CELL_ID = "" # Ship console output + exceptions to opensandbox-log-tail, which forwards # to Axiom dataset cf-dev. Keeps CF Worker logs durable in the same store diff --git a/cmd/agent/listen_linux.go b/cmd/agent/listen_linux.go index 06f52a8e..9f324188 100644 --- a/cmd/agent/listen_linux.go +++ b/cmd/agent/listen_linux.go @@ -111,13 +111,8 @@ func (l *virtioSerialListener) Accept() (net.Conn, error) { } l.mu.Unlock() - // Drain any stale data from the previous gRPC session. - // After a gRPC connection drops, the serial port may have leftover - // bytes (GOAWAY frames, partial HTTP/2 frames). If we hand these to - // the new gRPC session, it gets a protocol error and closes immediately. - if instrumentVirtioSerial { - log.Printf("virtio-serial: Accept iter: active=false, draining stale...") - } + // Drain stale frames from the previous gRPC session before accepting a + // fresh host connection after restore. drainStaleData(l.f) // Wait for the host to connect (port becomes readable with fresh data) @@ -181,7 +176,14 @@ func waitForReadable(f *os.File, timeout time.Duration) bool { if err != nil || n <= 0 { return false } - return fds[0].Revents&(unix.POLLIN|unix.POLLHUP) != 0 + return readableForAccept(fds[0].Revents) +} + +func readableForAccept(revents int16) bool { + if revents&unix.POLLHUP != 0 { + return false + } + return revents&unix.POLLIN != 0 } // PrepareHibernate resets the active flag so that after migration restore, @@ -219,8 +221,7 @@ type virtioSerialConn struct { // instrumentVirtioSerial controls per-Read logging on the conn. Set true via // OSB_AGENT_TRACE_VIRTIO=1 env var to debug post-loadvm protocol confusion. -// HARDCODED TO TRUE FOR INVESTIGATION BUILD — revert before merging. -var instrumentVirtioSerial = true || os.Getenv("OSB_AGENT_TRACE_VIRTIO") == "1" +var instrumentVirtioSerial = os.Getenv("OSB_AGENT_TRACE_VIRTIO") == "1" func (c *virtioSerialConn) Read(b []byte) (int, error) { n, err := c.f.Read(b) diff --git a/cmd/agent/listen_linux_test.go b/cmd/agent/listen_linux_test.go new file mode 100644 index 00000000..eb264afb --- /dev/null +++ b/cmd/agent/listen_linux_test.go @@ -0,0 +1,28 @@ +package main + +import ( + "testing" + + "golang.org/x/sys/unix" +) + +func TestReadableForAccept(t *testing.T) { + tests := []struct { + name string + revents int16 + want bool + }{ + {name: "pollin", revents: unix.POLLIN, want: true}, + {name: "pollhup", revents: unix.POLLHUP, want: false}, + {name: "pollin pollhup", revents: unix.POLLIN | unix.POLLHUP, want: false}, + {name: "pollerr", revents: unix.POLLERR, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := readableForAccept(tt.revents); got != tt.want { + t.Fatalf("readableForAccept(%v) = %v, want %v", tt.revents, got, tt.want) + } + }) + } +} diff --git a/cmd/server/main.go b/cmd/server/main.go index efc2d557..0891f6eb 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -26,8 +26,8 @@ import ( "github.com/opensandbox/opensandbox/internal/db" "github.com/opensandbox/opensandbox/internal/edgeclient" "github.com/opensandbox/opensandbox/internal/metrics" - "github.com/opensandbox/opensandbox/internal/obslog" "github.com/opensandbox/opensandbox/internal/observability" + "github.com/opensandbox/opensandbox/internal/obslog" "github.com/opensandbox/opensandbox/internal/proxy" "github.com/opensandbox/opensandbox/internal/sandbox" "github.com/opensandbox/opensandbox/internal/storage" @@ -37,16 +37,11 @@ import ( var ServerVersion = "dev" func main() { - // Load secrets from the appropriate cloud vault (before config.Load reads - // env vars). Each loader is a no-op if its env trigger is unset: - // SECRETS_VAULT_NAME → Azure Key Vault (Azure cells) - // OPENSANDBOX_SECRETS_ARN → AWS Secrets Manager (AWS cells) - // One or the other should be set at a time; neither = env file authoritative. - if err := config.LoadSecretsFromKeyVault(); err != nil { - log.Fatalf("failed to load secrets from Key Vault: %v", err) - } - if err := config.LoadSecretsFromSecretsManager(); err != nil { - log.Fatalf("failed to load secrets from Secrets Manager: %v", err) + // Load secrets from the configured cloud secret store before config.Load + // reads env vars. No-op if neither Azure Key Vault nor AWS Secrets Manager + // is configured. + if err := config.LoadSecrets(); err != nil { + log.Fatalf("failed to load secrets: %v", err) } cfg, err := config.Load() @@ -289,6 +284,7 @@ func main() { spec := compute.WorkerSpec{ CellID: cfg.CellID, Region: cfg.Region, + WorkerPool: cfg.WorkerPool, DatabaseURL: workerDBURL, RedisURL: workerRedisURL, JWTSecret: cfg.JWTSecret, @@ -446,17 +442,25 @@ func main() { case "aws": ec2Pool, err := compute.NewEC2Pool(compute.EC2PoolConfig{ - Region: cfg.S3Region, - AccessKeyID: cfg.S3AccessKeyID, - SecretAccessKey: cfg.S3SecretAccessKey, - AMI: cfg.EC2AMI, - InstanceType: cfg.EC2InstanceType, - SubnetID: cfg.EC2SubnetID, - SecurityGroupID: cfg.EC2SecurityGroupID, - KeyName: cfg.EC2KeyName, - IAMInstanceProfile: cfg.EC2IAMInstanceProfile, - SecretsARN: cfg.SecretsARN, - SSMParameterName: cfg.EC2SSMParameterName, + Region: cfg.Region, + AccessKeyID: cfg.EC2AccessKeyID, + SecretAccessKey: cfg.EC2SecretAccessKey, + AMI: cfg.EC2AMI, + InstanceType: cfg.EC2InstanceType, + SubnetID: cfg.EC2SubnetID, + SecurityGroupID: cfg.EC2SecurityGroupID, + KeyName: cfg.EC2KeyName, + IAMInstanceProfile: cfg.EC2IAMInstanceProfile, + SecretsARN: cfg.SecretsARN, + SSMParameterName: cfg.EC2SSMParameterName, + MarketType: cfg.EC2MarketType, + CellID: cfg.CellID, + SharedSandboxDataVolumeID: cfg.EC2SharedSandboxDataVolumeID, + SharedGoldensVolumeID: cfg.EC2SharedGoldensVolumeID, + OCFS2ClusterName: cfg.EC2OCFS2ClusterName, + OCFS2ExpectedNodes: cfg.EC2OCFS2ExpectedNodes, + OCFS2MaxNodes: cfg.EC2OCFS2MaxNodes, + OCFS2NodeIPs: cfg.EC2OCFS2NodeIPs, }) if err != nil { log.Fatalf("opensandbox: failed to create EC2 pool: %v", err) @@ -482,11 +486,14 @@ func main() { // Pick the per-provider ranked size list. Empty → scaler defers to // the pool's single configured default (cfg.AzureVMSize / cfg.EC2InstanceType). var machineSizes []string + var workerImage string switch { case len(cfg.AzureVMSizes) > 0 && cfg.AzureSubscriptionID != "": machineSizes = cfg.AzureVMSizes + workerImage = cfg.AzureImageID case len(cfg.EC2InstanceTypes) > 0 && (cfg.EC2AMI != "" || cfg.EC2SSMParameterName != ""): machineSizes = cfg.EC2InstanceTypes + workerImage = cfg.EC2AMI } if len(machineSizes) > 0 { log.Printf("opensandbox: scaler size fallback ranked: %v", machineSizes) @@ -494,16 +501,20 @@ func main() { scalerState := controlplane.NewRedisScalerState(redisRegistry.RedisClient()) scaler := controlplane.NewScaler(controlplane.ScalerConfig{ - Pool: pool, - Registry: redisRegistry, - Store: opts.Store, - StateStore: scalerState, - WorkerImage: cfg.EC2WorkerImage, - Cooldown: time.Duration(cfg.ScaleCooldownSec) * time.Second, - MinWorkers: cfg.MinWorkersPerRegion, - MaxWorkers: cfg.MaxWorkersPerRegion, - IdleReserve: cfg.IdleReserveWorkers, - MachineSizes: machineSizes, + Pool: pool, + Registry: redisRegistry, + Store: opts.Store, + StateStore: scalerState, + WorkerImage: workerImage, + Cooldown: time.Duration(cfg.ScaleCooldownSec) * time.Second, + MinWorkers: cfg.MinWorkersPerRegion, + MaxWorkers: cfg.MaxWorkersPerRegion, + IdleReserve: cfg.IdleReserveWorkers, + MinIdleCapacity: cfg.MinIdleCapacity, + MinIdleCPUs: cfg.MinIdleCPUs, + DefaultSandboxCPUs: cfg.DefaultSandboxCPUs, + WorkerPool: cfg.WorkerPool, + MachineSizes: machineSizes, // For "migrated" event emit after scaler-driven migrations // (rolling replace, evacuation) — keeps D1 sandboxes_index // worker_id in sync with cell-PG truth. Without this, the @@ -512,6 +523,7 @@ func main() { RedisClient: redisRegistry.RedisClient(), CellID: cfg.CellID, }) + opts.WorkerEvacuator = scaler defer scaler.Stop() // Leader election: only the leader runs the scaler. The @@ -531,6 +543,19 @@ func main() { defer leaderElector.Stop() log.Printf("opensandbox: leader election started (instance=%s)", leaderElector.InstanceID()) } + + if opts.WorkerEvacuator == nil && redisRegistry != nil { + scalerState := controlplane.NewRedisScalerState(redisRegistry.RedisClient()) + opts.WorkerEvacuator = controlplane.NewScaler(controlplane.ScalerConfig{ + Registry: redisRegistry, + Store: opts.Store, + StateStore: scalerState, + MaxWorkers: cfg.MaxWorkersPerRegion, + RedisClient: redisRegistry.RedisClient(), + CellID: cfg.CellID, + }) + log.Printf("opensandbox: admin worker evacuator configured without compute pool") + } } // Background maintenance tasks diff --git a/cmd/worker/golden_upload.go b/cmd/worker/golden_upload.go index 636a2179..c0fc04c8 100644 --- a/cmd/worker/golden_upload.go +++ b/cmd/worker/golden_upload.go @@ -30,14 +30,13 @@ import ( // Reads OPENSANDBOX_GLOBAL_BLOB_* env vars for endpoint + creds. Fails loud // if blobstore isn't configured. func uploadGolden(path string) error { - // Bootstrap from the appropriate cloud vault — same pattern as main(). One - // of these picks up worker-global-blob-* secrets so the upload knows where - // to push; the other is a no-op when its trigger env var isn't set. - if err := config.LoadSecretsFromKeyVault(); err != nil { - return fmt.Errorf("keyvault bootstrap: %w", err) - } - if err := config.LoadSecretsFromSecretsManager(); err != nil { - return fmt.Errorf("secretsmanager bootstrap: %w", err) + // Bootstrap from the configured secret store the same way the main + // worker does — so this subcommand picks up worker-global-blob-* secrets + // when run on a host whose worker.env points at a KV / SM prefix. + // Without this, the user would have to manually export every Tigris env + // var inline. + if err := config.LoadSecrets(); err != nil { + return fmt.Errorf("secrets bootstrap: %w", err) } cfg, err := config.Load() if err != nil { diff --git a/cmd/worker/main.go b/cmd/worker/main.go index aed3e79b..fb680e75 100644 --- a/cmd/worker/main.go +++ b/cmd/worker/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "encoding/json" "fmt" "io" "log" @@ -26,6 +27,7 @@ import ( "github.com/opensandbox/opensandbox/internal/metrics" "github.com/opensandbox/opensandbox/internal/observability" "github.com/opensandbox/opensandbox/internal/obslog" + "github.com/opensandbox/opensandbox/internal/preemption" "github.com/opensandbox/opensandbox/internal/proxy" qm "github.com/opensandbox/opensandbox/internal/qemu" "github.com/opensandbox/opensandbox/internal/sandbox" @@ -44,7 +46,7 @@ var WorkerVersion = "dev" func main() { // Subcommands that don't need config/secrets. Must short-circuit before - // LoadSecretsFromKeyVault, which is slow and would fail outside Azure. + // LoadSecrets, which is slow and can fail without cloud credentials. // // "golden-version " prints the full-file hash used for golden-image // archive keys. Packer invokes this so the archive key matches what @@ -82,16 +84,11 @@ func main() { return } - // Load secrets from the appropriate cloud vault (before config.Load reads - // env vars). Each loader is a no-op if its env trigger is unset: - // SECRETS_VAULT_NAME → Azure Key Vault (Azure cells) - // OPENSANDBOX_SECRETS_ARN → AWS Secrets Manager (AWS cells) - // One or the other should be set at a time; neither = env file authoritative. - if err := config.LoadSecretsFromKeyVault(); err != nil { - log.Fatalf("failed to load secrets from Key Vault: %v", err) - } - if err := config.LoadSecretsFromSecretsManager(); err != nil { - log.Fatalf("failed to load secrets from Secrets Manager: %v", err) + // Load secrets from the configured cloud secret store before config.Load + // reads env vars. No-op if neither Azure Key Vault nor AWS Secrets Manager + // is configured. + if err := config.LoadSecrets(); err != nil { + log.Fatalf("failed to load secrets: %v", err) } cfg, err := config.Load() @@ -240,6 +237,7 @@ func main() { DataDir: cfg.DataDir, KernelPath: cfg.KernelPath, ImagesDir: cfg.ImagesDir, + GoldenDir: cfg.GoldenDir, QEMUBin: cfg.QEMUBin, AgentBinaryPath: "/usr/local/bin/osb-agent", AgentVersion: AgentVersion, @@ -311,9 +309,15 @@ func main() { if len(vms) == 0 { return } - log.Printf("opensandbox-worker: hibernating %d sandboxes...", len(vms)) + skip := shutdownHibernateSkipSet(context.Background(), store, cfg.WorkerID, vms) + hibernateCount := len(vms) - len(skip) + if hibernateCount == 0 { + log.Printf("opensandbox-worker: shutdown skipping hibernation for %d sandbox(es)", len(skip)) + return + } + log.Printf("opensandbox-worker: hibernating %d sandboxes (skipping %d)...", hibernateCount, len(skip)) shutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) - results := qmMgr.HibernateAll(shutCtx, checkpointStore) + results := qmMgr.HibernateAllExcept(shutCtx, checkpointStore, skip) cancel() // Log which VMs were NOT hibernated @@ -327,7 +331,7 @@ func main() { log.Printf("opensandbox-worker: %d VMs failed to hibernate: %v", len(failed), failed) } - processHibernateResults(results, store, checkpointStore, sandboxDBMgr, func(r interface{}) (string, string, error) { + processHibernateResults(results, store, checkpointStore, sandboxDBMgr, cfg.WorkerID, func(r interface{}) (string, string, error) { hr := r.(qm.HibernateAllResult) return hr.SandboxID, hr.HibernationKey, hr.Err }) @@ -603,7 +607,7 @@ func main() { } // HTTP server - httpServer := worker.NewHTTPServer(mgr, ptyMgr, execMgr, jwtIssuer, sandboxDBMgr, sbProxy, sbRouter, cfg.SandboxDomain) + httpServer := worker.NewHTTPServer(mgr, ptyMgr, execMgr, jwtIssuer, sandboxDBMgr, sbProxy, sbRouter, cfg.SandboxDomain, store) httpAddr := fmt.Sprintf(":%d", cfg.Port) log.Printf("opensandbox-worker: starting HTTP server on %s", httpAddr) go func() { @@ -623,6 +627,7 @@ func main() { if err != nil { log.Printf("opensandbox-worker: Redis heartbeat not available: %v", err) } else { + hb.SetPool(cfg.WorkerPool) hb.SetWorkerVersion(WorkerVersion) if qemuMgr != nil { hb.SetGoldenVersion(qemuMgr.GoldenVersion()) @@ -709,6 +714,18 @@ func main() { } defer hb.Stop() log.Println("opensandbox-worker: Redis heartbeat started") + + preemptMon := preemption.NewMonitor() + go func() { + notices := preemptMon.Watch(ctx) + for notice := range notices { + log.Printf("opensandbox-worker: PREEMPTION notice from %s - action=%s eta=%s, notifying resumable sandboxes", + preemptMon.Name(), notice.Action, notice.ETA.Format(time.RFC3339)) + worker.NotifyResumableSandboxesBeforeRestart(context.Background(), mgr, store, sandboxDBMgr, 25*time.Second, notice.ETA) + hb.Stop() + return + } + }() } } @@ -1125,6 +1142,48 @@ func deleteOldHibernation(store *storage.CheckpointStore, key string) { } } +func shutdownHibernateSkipSet(ctx context.Context, store *db.Store, workerID string, sandboxes []types.Sandbox) map[string]struct{} { + skip := make(map[string]struct{}) + if store == nil { + return skip + } + for _, sb := range sandboxes { + if sb.ID == "" { + continue + } + session, err := store.GetSandboxSession(ctx, sb.ID) + if err != nil { + log.Printf("opensandbox-worker: shutdown hibernate: could not load session %s: %v", sb.ID, err) + continue + } + if session.WorkerID != workerID { + log.Printf("opensandbox-worker: shutdown hibernate: skipping %s; DB owner is %s, not %s", sb.ID, session.WorkerID, workerID) + skip[sb.ID] = struct{}{} + continue + } + if sandboxSessionIsResumable(session) { + log.Printf("opensandbox-worker: shutdown hibernate: skipping resumable sandbox %s", sb.ID) + skip[sb.ID] = struct{}{} + } + } + return skip +} + +func sandboxSessionIsResumable(session *db.SandboxSession) bool { + if session == nil || len(session.Config) == 0 { + return false + } + var cfg types.SandboxConfig + if err := json.Unmarshal(session.Config, &cfg); err != nil { + return false + } + return cfg.IsResumable() +} + +func sandboxSessionOwnedByWorker(session *db.SandboxSession, workerID string) bool { + return session != nil && session.WorkerID == workerID +} + // processHibernateResults handles results from HibernateAll for both backends. // // In addition to updating cell-local PG, we LogEvent("hibernated") into the @@ -1134,10 +1193,32 @@ func deleteOldHibernation(store *storage.CheckpointStore, key string) { // sync. Without this, the bulk-shutdown path silently skipped the lifecycle // event the gRPC HibernateSandbox handler emits per call, and D1 stayed // "running" until something else nudged it. -func processHibernateResults(results interface{}, store *db.Store, checkpointStore *storage.CheckpointStore, sandboxDBs *sandbox.SandboxDBManager, extract func(interface{}) (string, string, error)) { +func processHibernateResults(results interface{}, store *db.Store, checkpointStore *storage.CheckpointStore, sandboxDBs *sandbox.SandboxDBManager, workerID string, extract func(interface{}) (string, string, error)) { switch rs := results.(type) { case []qm.HibernateAllResult: for _, r := range rs { + var session *db.SandboxSession + if store != nil { + var err error + session, err = store.GetSandboxSession(context.Background(), r.SandboxID) + if err != nil { + log.Printf("opensandbox-worker: shutdown hibernate: could not load session %s: %v", r.SandboxID, err) + } + } + if store != nil && session != nil && !sandboxSessionOwnedByWorker(session, workerID) { + log.Printf("opensandbox-worker: shutdown hibernate: ignoring result for %s; DB owner is %s, not %s", r.SandboxID, session.WorkerID, workerID) + if sandboxDBs != nil { + _ = sandboxDBs.Remove(r.SandboxID) + } + continue + } + if sandboxSessionIsResumable(session) { + log.Printf("opensandbox-worker: shutdown hibernate: ignoring result for resumable sandbox %s", r.SandboxID) + if sandboxDBs != nil { + _ = sandboxDBs.Remove(r.SandboxID) + } + continue + } if r.Err != nil { log.Printf("opensandbox-worker: hibernate failed for %s: %v", r.SandboxID, r.Err) if store != nil { @@ -1153,14 +1234,11 @@ func processHibernateResults(results interface{}, store *db.Store, checkpointSto continue } log.Printf("opensandbox-worker: hibernated %s (key=%s)", r.SandboxID, r.HibernationKey) - if store != nil { - session, err := store.GetSandboxSession(context.Background(), r.SandboxID) - if err == nil { - _, superseded, _ := store.CreateHibernation(context.Background(), r.SandboxID, session.OrgID, - r.HibernationKey, 0, session.Region, session.Template, session.Config) - deleteOldHibernation(checkpointStore, superseded) - _ = store.UpdateSandboxSessionStatus(context.Background(), r.SandboxID, "hibernated", nil) - } + if store != nil && session != nil { + _, superseded, _ := store.CreateHibernation(context.Background(), r.SandboxID, session.OrgID, + r.HibernationKey, 0, session.Region, session.Template, session.Config) + deleteOldHibernation(checkpointStore, superseded) + _ = store.UpdateSandboxSessionStatus(context.Background(), r.SandboxID, "hibernated", nil) } if sandboxDBs != nil { if sdb, err := sandboxDBs.Get(r.SandboxID); err == nil { diff --git a/deploy/packer/worker-ami-aws-burst.pkr.hcl b/deploy/packer/worker-ami-aws-burst.pkr.hcl new file mode 100644 index 00000000..9f5b019f --- /dev/null +++ b/deploy/packer/worker-ami-aws-burst.pkr.hcl @@ -0,0 +1,345 @@ +# worker-ami-aws-burst.pkr.hcl — Build an immutable AMI for OpenSandbox Burst workers on AWS. +# +# Mirrors deploy/packer/worker-ami.pkr.hcl (Azure variant) but targets the +# amazon-ebs builder. The setup script (`deploy/azure/setup-azure-host.sh`) +# is cloud-agnostic in practice — it installs QEMU + kernel modules + systemd +# units + Vector and never talks to Azure-specific APIs. We reuse it as-is. +# +# Differences from the Azure file: +# - amazon-ebs source on Ubuntu 24.04 LTS x86_64 instead of azure-arm. +# - Optional Tigris/S3-compatible rootfs blob caching. Same rootfs inputs +# reuse the same cached default.ext4, which keeps AMI builds fast and +# golden versions stable when the guest image did not change. +# - Installs awscli (needed by deploy/vector/populate-vector-env.sh AWS path +# and by the worker user-data shared-disk attach). +# - Tags the AMI for the terraform `aws_ami` data source lookup +# (opensandbox-role=worker, opensandbox-cloud=aws). +# +# Usage: +# # 1. Build binaries for linux/amd64: +# CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags "-X main.WorkerVersion=$(git rev-parse --short HEAD)" \ +# -o bin/opensandbox-worker ./cmd/worker/ +# CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o bin/osb-agent ./cmd/agent/ +# +# # 2. Build the rootfs context tarball: +# tar czf /tmp/packer-rootfs-ctx.tar.gz deploy/firecracker/rootfs/ deploy/ec2/build-rootfs-docker.sh scripts/claude-agent-wrapper/ +# +# # 3. Run packer: +# packer init deploy/packer/worker-ami-aws-burst.pkr.hcl +# packer build -var "worker_version=$(git rev-parse --short HEAD)" deploy/packer/worker-ami-aws-burst.pkr.hcl +# +# # 4. The data source in opencomputer-infra/terraform/aws/us-east-2-poc/ami.tf +# # picks up the new AMI on the next `tofu apply`. + +packer { + required_plugins { + amazon = { + version = ">= 1.3.0" + source = "github.com/hashicorp/amazon" + } + } +} + +# --------------------------------------------------------------------- +# Variables +# --------------------------------------------------------------------- + +variable "worker_version" { + type = string + description = "Worker version (git SHA). Baked into AMI name and tags." +} + +variable "agent_version" { + type = string + default = "" +} + +variable "region" { + type = string + default = "us-east-2" +} + +variable "instance_type" { + type = string + default = "c5.4xlarge" + description = "Builder instance type. Needs enough memory for Docker rootfs build (~8GB) but doesn't need to run guest VMs, so non-metal is fine and saves ~10× vs c5.metal." +} + +variable "worker_binary" { + type = string + default = "bin/opensandbox-worker" +} + +variable "agent_binary" { + type = string + default = "bin/osb-agent" +} + +variable "rootfs_context" { + type = string + default = "/tmp/packer-rootfs-ctx.tar.gz" + description = "Pre-built tarball of rootfs + agent wrapper sources." +} + +variable "vector_context" { + type = string + default = "/tmp/packer-vector-ctx.tar.gz" + description = "Pre-built tarball of deploy/vector/ (config + populator + units). Pre-create with: tar czf /tmp/packer-vector-ctx.tar.gz deploy/vector/" +} + +variable "tigris_endpoint" { + type = string + default = "" + description = "Optional S3-compatible endpoint for Tigris rootfs/golden cache." +} + +variable "tigris_access_key_id" { + type = string + default = "" + sensitive = true + description = "Optional Tigris access key for rootfs/golden cache." +} + +variable "tigris_secret_access_key" { + type = string + default = "" + sensitive = true + description = "Optional Tigris secret key for rootfs/golden cache." +} + +variable "tigris_goldens_bucket" { + type = string + default = "" + description = "Optional Tigris bucket for content-addressed rootfs cache and golden uploads. Empty = skip cache." +} + +# --------------------------------------------------------------------- +# Source +# --------------------------------------------------------------------- + +source "amazon-ebs" "worker" { + region = var.region + instance_type = var.instance_type + ssh_username = "ubuntu" + ssh_pty = true + + ami_name = "opensandbox-burst-worker-${var.worker_version}-${formatdate("YYYYMMDD-hhmm", timestamp())}" + ami_description = "OpenSandbox Burst worker AMI (Ubuntu 24.04, QEMU/KVM nested-virt). Built from git ${var.worker_version}." + + source_ami_filter { + filters = { + name = "ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*" + architecture = "x86_64" + virtualization-type = "hvm" + root-device-type = "ebs" + } + most_recent = true + owners = ["099720109477"] # Canonical + } + + ena_support = true + sriov_support = true + + launch_block_device_mappings { + device_name = "/dev/sda1" + volume_size = 50 + volume_type = "gp3" + delete_on_termination = true + } + + # AMI tags — the terraform `aws_ami` data source in the AWS leaf filters + # on these to pick the most-recent worker AMI for this cloud. + tags = { + Name = "opensandbox-burst-worker-${var.worker_version}" + "opensandbox-role" = "worker" + "opensandbox-cloud" = "aws" + "opensandbox-version" = var.worker_version + } + + # Volume snapshot tag — propagates so the EBS snapshot underlying the AMI + # has the same provenance metadata as the AMI itself. + snapshot_tags = { + "opensandbox-role" = "worker" + "opensandbox-cloud" = "aws" + "opensandbox-version" = var.worker_version + } + + run_tags = { + Name = "packer-opensandbox-worker-build" + } +} + +# --------------------------------------------------------------------- +# Build +# --------------------------------------------------------------------- + +build { + sources = ["source.amazon-ebs.worker"] + + # 1. Upload pre-built binaries. + provisioner "file" { + source = var.worker_binary + destination = "/tmp/opensandbox-worker" + } + provisioner "file" { + source = var.agent_binary + destination = "/tmp/osb-agent" + } + + # 2. Upload rootfs build context. + provisioner "file" { + source = var.rootfs_context + destination = "/tmp/rootfs-ctx.tar.gz" + } + + # 3. Upload the EC2 worker systemd unit (the Azure variant uses a different + # unit; the EC2 one was already drafted at deploy/ec2/opensandbox-worker.service). + provisioner "file" { + source = "deploy/ec2/opensandbox-worker.service" + destination = "/tmp/opensandbox-worker.service" + } + + # 4. Upload Vector config + populator. Packer's file provisioner doesn't + # do recursive directory upload reliably across SSH clients, so we + # tar/extract the same way we do the rootfs context above. See + # var.vector_context for the pre-build command. + provisioner "file" { + source = var.vector_context + destination = "/tmp/vector-ctx.tar.gz" + } + provisioner "shell" { + inline = [ + "mkdir -p /tmp/vector", + "tar xzf /tmp/vector-ctx.tar.gz -C /tmp/vector --strip-components=2", # strip deploy/vector/ prefix + "rm /tmp/vector-ctx.tar.gz", + ] + } + + # 5. Run the (misleadingly-named-but-cloud-agnostic) setup script. Installs + # QEMU, kernel modules, Docker for rootfs build, Vector, systemd units. + provisioner "shell" { + execute_command = "chmod +x {{ .Path }}; {{ .Vars }} sudo -E bash '{{ .Path }}'" + script = "deploy/azure/setup-azure-host.sh" + } + + # 6. AWS-specific: install awscli (used by populate-vector-env.sh and by + # the worker user-data's shared-disk attach), bake OCFS2 dependencies for + # the shared data volume, then install binaries and build the golden rootfs. + provisioner "shell" { + execute_command = "chmod +x {{ .Path }}; {{ .Vars }} sudo -E bash '{{ .Path }}'" + environment_vars = [ + "TIGRIS_ENDPOINT=${var.tigris_endpoint}", + "TIGRIS_ACCESS_KEY_ID=${var.tigris_access_key_id}", + "TIGRIS_SECRET_ACCESS_KEY=${var.tigris_secret_access_key}", + "TIGRIS_GOLDENS_BUCKET=${var.tigris_goldens_bucket}", + "AWS_DEFAULT_REGION=auto", + ] + inline = [ + # awscli v2 — apt's `awscli` is v1 and missing some commands we use. + "apt-get update -qq", + "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq unzip ocfs2-tools \"linux-modules-extra-$(uname -r)\"", + "curl -fsSL 'https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip' -o /tmp/awscliv2.zip", + "cd /tmp && unzip -q awscliv2.zip && ./aws/install --update", + "rm -rf /tmp/awscliv2.zip /tmp/aws", + "aws --version", + "modprobe ocfs2", + "modprobe ocfs2_dlmfs", + "modprobe ocfs2_stack_o2cb", + "command -v mount.ocfs2", + "systemctl disable --now apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service 2>/dev/null || true", + + # Install worker + agent binaries. + "mv /tmp/opensandbox-worker /usr/local/bin/opensandbox-worker", + "chmod +x /usr/local/bin/opensandbox-worker", + "mv /tmp/osb-agent /usr/local/bin/osb-agent", + "chmod +x /usr/local/bin/osb-agent", + + # Install systemd unit. + "mv /tmp/opensandbox-worker.service /etc/systemd/system/opensandbox-worker.service", + "systemctl daemon-reload", + "systemctl enable opensandbox-worker.service", + + # Build or restore the golden rootfs. The cache key is content-addressed + # from the guest agent, rootfs sources, and guest kernel modules. + "mkdir -p /tmp/rootfs-ctx", + "cd /tmp/rootfs-ctx && tar xzf /tmp/rootfs-ctx.tar.gz", + "INPUT_HASH=$({ sha256sum /usr/local/bin/osb-agent; find /tmp/rootfs-ctx -type f | sort | xargs sha256sum; sha256sum /opt/opensandbox/guest-modules/*.ko* 2>/dev/null; } | sha256sum | awk '{print $1}')", + "echo \"Rootfs input hash: $INPUT_HASH\"", + "ROOTFS_UUID=$(echo \"$INPUT_HASH\" | head -c 32 | sed 's/\\(........\\)\\(....\\)\\(....\\)\\(....\\)\\(............\\)/\\1-\\2-\\3-\\4-\\5/')", + "export ROOTFS_UUID", + "INPUT_HASH_SHORT=$(echo \"$INPUT_HASH\" | cut -c1-16)", + "CACHE_KEY=\"rootfs-cache/$INPUT_HASH_SHORT/default.ext4\"", + "CACHE_HIT=0", + "mkdir -p /data/firecracker/images /opt/opensandbox/images", + "if [ -n \"$TIGRIS_ENDPOINT\" ] && [ -n \"$TIGRIS_ACCESS_KEY_ID\" ] && [ -n \"$TIGRIS_SECRET_ACCESS_KEY\" ] && [ -n \"$TIGRIS_GOLDENS_BUCKET\" ]; then", + " export AWS_ACCESS_KEY_ID=\"$TIGRIS_ACCESS_KEY_ID\" AWS_SECRET_ACCESS_KEY=\"$TIGRIS_SECRET_ACCESS_KEY\"", + " echo \"Checking rootfs cache: s3://$TIGRIS_GOLDENS_BUCKET/$CACHE_KEY\"", + " if aws s3 cp --endpoint-url \"$TIGRIS_ENDPOINT\" \"s3://$TIGRIS_GOLDENS_BUCKET/$CACHE_KEY\" /data/firecracker/images/default.ext4; then", + " CACHE_HIT=1", + " echo 'Rootfs restored from cache — skipping Docker build'", + " else", + " echo 'Rootfs cache miss — building from source'", + " fi", + "else", + " echo 'Tigris cache credentials incomplete; rootfs cache disabled'", + "fi", + "if [ \"$CACHE_HIT\" != \"1\" ]; then", + " cd /tmp/rootfs-ctx && ROOTFS_UUID=\"$ROOTFS_UUID\" bash deploy/ec2/build-rootfs-docker.sh /usr/local/bin/osb-agent /data/firecracker/images default", + "fi", + "cp /data/firecracker/images/default.ext4 /opt/opensandbox/images/default.ext4", + + # Inject guest kernel modules into rootfs. + "GUEST_MODDIR=/opt/opensandbox/guest-modules", + "if [ -d \"$GUEST_MODDIR\" ] && [ -f /opt/opensandbox/images/default.ext4 ]; then", + " MNTDIR=$(mktemp -d)", + " mount -o loop /opt/opensandbox/images/default.ext4 $MNTDIR", + " mkdir -p $MNTDIR/lib/modules/extra", + " cp $GUEST_MODDIR/*.ko* $MNTDIR/lib/modules/extra/ 2>/dev/null || true", + " umount $MNTDIR", + " rmdir $MNTDIR", + "fi", + + # Stamp the golden version (hash of the final ext4) — workers read this + # at boot to decide whether to fetch a newer golden from S3. + "GOLDEN_VERSION=$(/usr/local/bin/opensandbox-worker golden-version /opt/opensandbox/images/default.ext4 2>/dev/null || sha256sum /opt/opensandbox/images/default.ext4 | awk '{print $1}')", + "echo \"$GOLDEN_VERSION\" > /opt/opensandbox/images/golden-version", + "echo \"Golden version: $GOLDEN_VERSION\"", + "if [ \"$CACHE_HIT\" != \"1\" ] && [ -n \"$TIGRIS_ENDPOINT\" ] && [ -n \"$TIGRIS_ACCESS_KEY_ID\" ] && [ -n \"$TIGRIS_SECRET_ACCESS_KEY\" ] && [ -n \"$TIGRIS_GOLDENS_BUCKET\" ]; then", + " export AWS_ACCESS_KEY_ID=\"$TIGRIS_ACCESS_KEY_ID\" AWS_SECRET_ACCESS_KEY=\"$TIGRIS_SECRET_ACCESS_KEY\"", + " echo \"Uploading rootfs cache: s3://$TIGRIS_GOLDENS_BUCKET/$CACHE_KEY\"", + " aws s3 cp --endpoint-url \"$TIGRIS_ENDPOINT\" /opt/opensandbox/images/default.ext4 \"s3://$TIGRIS_GOLDENS_BUCKET/$CACHE_KEY\" || echo 'rootfs cache upload failed — continuing'", + "fi", + ] + } + + # 7. Optional: upload the golden to Tigris so future hydration paths + # + future per-instance prefetch path can fetch it without rebuilding. + provisioner "shell" { + execute_command = "chmod +x {{ .Path }}; {{ .Vars }} sudo -E bash '{{ .Path }}'" + environment_vars = [ + "TIGRIS_ENDPOINT=${var.tigris_endpoint}", + "TIGRIS_ACCESS_KEY_ID=${var.tigris_access_key_id}", + "TIGRIS_SECRET_ACCESS_KEY=${var.tigris_secret_access_key}", + "TIGRIS_GOLDENS_BUCKET=${var.tigris_goldens_bucket}", + "AWS_DEFAULT_REGION=auto", + ] + inline = [ + "set -e", + "if [ -z \"$TIGRIS_ENDPOINT\" ] || [ -z \"$TIGRIS_ACCESS_KEY_ID\" ] || [ -z \"$TIGRIS_SECRET_ACCESS_KEY\" ] || [ -z \"$TIGRIS_GOLDENS_BUCKET\" ]; then", + " echo 'Tigris cache credentials incomplete; skipping golden upload (worker AMI still includes the baked golden)'", + " exit 0", + "fi", + "export AWS_ACCESS_KEY_ID=\"$TIGRIS_ACCESS_KEY_ID\" AWS_SECRET_ACCESS_KEY=\"$TIGRIS_SECRET_ACCESS_KEY\"", + "GOLDEN_VERSION=$(cat /opt/opensandbox/images/golden-version)", + "S3_KEY=\"bases/$GOLDEN_VERSION/default.ext4\"", + "echo \"Uploading default.ext4 -> s3://$TIGRIS_GOLDENS_BUCKET/$S3_KEY (~4GB, will take a moment)\"", + "aws s3 cp --endpoint-url \"$TIGRIS_ENDPOINT\" /opt/opensandbox/images/default.ext4 \"s3://$TIGRIS_GOLDENS_BUCKET/$S3_KEY\" || echo 'Tigris upload failed — continuing (AMI golden is the only copy)'", + ] + } + + # 8. Write a manifest so external tooling can pin to the resulting AMI ID. + post-processor "manifest" { + output = "packer-manifest-aws.json" + strip_path = true + } +} diff --git a/deploy/packer/worker-ami-aws.pkr.hcl b/deploy/packer/worker-ami-aws.pkr.hcl deleted file mode 100644 index 765ece05..00000000 --- a/deploy/packer/worker-ami-aws.pkr.hcl +++ /dev/null @@ -1,215 +0,0 @@ -# worker-ami-aws.pkr.hcl — Packer template for AWS EC2 worker AMIs. -# -# Counterpart to worker-ami.pkr.hcl (which builds Azure managed images + -# publishes them to a Compute Gallery). This builds an AMI in EC2 and -# optionally publishes the AMI ID to SSM Parameter Store so EC2Pool's -# RefreshAMI loop picks it up automatically. -# -# Build: packer build -var "aws_region=us-east-1" deploy/packer/worker-ami-aws.pkr.hcl -# -# Inputs (most overridable via -var or env vars): -# - aws_region AWS region to build in -# - source_ami_filter Filter for the base Ubuntu AMI (default: 24.04 LTS) -# - instance_type Builder instance type (small is fine) -# - ssm_param_name If set, AMI ID is written here on success - -packer { - required_plugins { - amazon = { - version = ">= 1.3.0" - source = "github.com/hashicorp/amazon" - } - } -} - -variable "aws_region" { - type = string - default = "us-east-1" -} - -variable "instance_type" { - type = string - default = "t3.medium" -} - -variable "ami_name_prefix" { - type = string - default = "opensandbox-worker" -} - -variable "worker_version" { - type = string - description = "Worker version tag (e.g. git short SHA). Used in AMI name + tags." -} - -variable "ssm_param_name" { - type = string - description = "If set, the built AMI ID is written to this SSM parameter (e.g. /opensandbox/dev/worker-ami-id)." - default = "" -} - -variable "guest_image_version" { - type = string - description = "Hash of the canonical guest rootfs (default.ext4). Pulled from R2/S3 during AMI build." - default = "" -} - -variable "guest_image_url" { - type = string - description = "Canonical URL for the guest rootfs blob (e.g. https://r2.opensandbox.dev/golden-store/bases/{hash}/default.ext4)." - default = "" -} - -variable "worker_binary" { - type = string - default = "bin/opensandbox-worker" - description = "Path to the cross-compiled worker binary (linux/amd64)." -} - -variable "agent_binary" { - type = string - default = "bin/osb-agent" - description = "Path to the cross-compiled agent binary (linux/amd64)." -} - -variable "kernel_path" { - type = string - default = "deploy/firecracker/vmlinux" - description = "Path to the Linux kernel image baked into the AMI." -} - -# Latest Ubuntu 24.04 LTS amd64 from Canonical's account. -data "amazon-ami" "ubuntu" { - region = var.aws_region - most_recent = true - owners = ["099720109477"] - filters = { - name = "ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*" - root-device-type = "ebs" - virtualization-type = "hvm" - architecture = "x86_64" - } -} - -source "amazon-ebs" "worker" { - region = var.aws_region - source_ami = data.amazon-ami.ubuntu.id - instance_type = var.instance_type - ssh_username = "ubuntu" - - ami_name = "${var.ami_name_prefix}-${var.worker_version}-${formatdate("YYYYMMDD-hhmmss", timestamp())}" - ami_description = "OpenSandbox worker AMI — version ${var.worker_version}" - - tags = { - Name = "${var.ami_name_prefix}-${var.worker_version}" - "opensandbox:version" = var.worker_version - "opensandbox:role" = "worker" - "opensandbox:guest_version" = var.guest_image_version - } - - # Bigger root volume so we can stage rootfs images and tools. - launch_block_device_mappings { - device_name = "/dev/sda1" - volume_size = 30 - volume_type = "gp3" - delete_on_termination = true - } -} - -build { - sources = ["source.amazon-ebs.worker"] - - # Install system packages (KVM, QEMU, mdadm, xfs tools, etc.) - provisioner "shell" { - inline = [ - "set -euo pipefail", - "export DEBIAN_FRONTEND=noninteractive", - "sudo apt-get update -qq", - "sudo apt-get install -y -qq qemu-system-x86 qemu-utils mdadm xfsprogs e2fsprogs jq curl ca-certificates", - "sudo systemctl disable --now systemd-resolved 2>/dev/null || true", - "sudo mkdir -p /opt/opensandbox/images /usr/local/bin /etc/opensandbox", - ] - } - - # Upload worker + agent binaries - provisioner "file" { - source = var.worker_binary - destination = "/tmp/opensandbox-worker" - } - provisioner "file" { - source = var.agent_binary - destination = "/tmp/osb-agent" - } - provisioner "file" { - source = var.kernel_path - destination = "/tmp/vmlinux" - } - - provisioner "shell" { - inline = [ - "sudo install -m 0755 /tmp/opensandbox-worker /usr/local/bin/opensandbox-worker", - "sudo install -m 0755 /tmp/osb-agent /usr/local/bin/osb-agent", - "sudo install -m 0644 /tmp/vmlinux /opt/opensandbox/vmlinux", - ] - } - - # Pull the canonical guest rootfs from R2/S3 if configured. - provisioner "shell" { - inline = [ - "set -euo pipefail", - "if [ -n '${var.guest_image_url}' ] && [ -n '${var.guest_image_version}' ]; then", - " echo 'Pulling guest rootfs ${var.guest_image_version} from ${var.guest_image_url}...'", - " sudo mkdir -p /opt/opensandbox/images/bases/${var.guest_image_version}", - " sudo curl -fsSL -o /opt/opensandbox/images/bases/${var.guest_image_version}/default.ext4 '${var.guest_image_url}'", - " sudo cp /opt/opensandbox/images/bases/${var.guest_image_version}/default.ext4 /opt/opensandbox/images/default.ext4", - "fi", - ] - } - - # systemd unit (mirrors the one written by Azure setup-host.sh) - provisioner "shell" { - inline = [ - "set -euo pipefail", - "sudo tee /etc/systemd/system/opensandbox-worker.service > /dev/null <<'UNIT'", - "[Unit]", - "Description=OpenSandbox Worker (QEMU backend)", - "After=network-online.target", - "Wants=network-online.target", - "", - "[Service]", - "Type=simple", - "ExecStartPre=/sbin/modprobe vhost_vsock", - "EnvironmentFile=/etc/opensandbox/worker.env", - "ExecStart=/usr/local/bin/opensandbox-worker", - "Restart=on-failure", - "RestartSec=5", - "LimitNOFILE=1000000", - "LimitNPROC=infinity", - "KillMode=process", - "TimeoutStopSec=300", - "", - "[Install]", - "WantedBy=multi-user.target", - "UNIT", - "sudo systemctl daemon-reload", - "sudo systemctl enable opensandbox-worker", - ] - } - - # Publish AMI ID to SSM if configured - post-processor "shell-local" { - only = ["amazon-ebs.worker"] - inline = [ - "AMI_ID=$(jq -r '.builds[-1].artifact_id' manifest.json | cut -d: -f2)", - "if [ -n '${var.ssm_param_name}' ] && [ -n \"$AMI_ID\" ]; then", - " aws ssm put-parameter --region ${var.aws_region} --name '${var.ssm_param_name}' --type String --overwrite --value \"$AMI_ID\"", - " aws ssm put-parameter --region ${var.aws_region} --name \"$(dirname '${var.ssm_param_name}')/worker-ami-version\" --type String --overwrite --value '${var.worker_version}'", - " echo \"Published $AMI_ID to ${var.ssm_param_name}\"", - "fi", - ] - } - - post-processor "manifest" { - output = "manifest.json" - } -} diff --git a/deploy/server.env.example b/deploy/server.env.example index f441193e..06b62187 100644 --- a/deploy/server.env.example +++ b/deploy/server.env.example @@ -52,3 +52,4 @@ OPENSANDBOX_SENTRY_ENVIRONMENT=production # server-stripe-secret-key → STRIPE_SECRET_KEY # server-stripe-webhook-secret → STRIPE_WEBHOOK_SECRET # server-sentry-dsn → OPENSANDBOX_SENTRY_DSN +# server-ocfs2-node-ips → OPENSANDBOX_OCFS2_NODE_IPS diff --git a/deploy/vector/populate-vector-env.sh b/deploy/vector/populate-vector-env.sh index 96e0064c..80c05b63 100755 --- a/deploy/vector/populate-vector-env.sh +++ b/deploy/vector/populate-vector-env.sh @@ -114,42 +114,102 @@ fi [ -f /etc/opensandbox/worker.env ] && . /etc/opensandbox/worker.env # shellcheck disable=SC1091 [ -f /etc/opensandbox/server.env ] && . /etc/opensandbox/server.env -VAULT_NAME="${OPENSANDBOX_AZURE_KEY_VAULT_NAME:-}" -if [ -z "$VAULT_NAME" ]; then - log "OPENSANDBOX_AZURE_KEY_VAULT_NAME unset — host has no KV configured (e.g. dev VM without managed identity); skipping (Vector will use whatever vector.env is on disk)" - exit 0 -fi +# Cloud detection — explicit OPENSANDBOX_CLOUD wins, else probe IMDS endpoints. +detect_cloud() { + if [ -n "${OPENSANDBOX_CLOUD:-}" ]; then + echo "$OPENSANDBOX_CLOUD"; return + fi + # AWS IMDSv2 token PUT — succeeds only on AWS. + if curl -fsS -X PUT --max-time 2 \ + -H "X-aws-ec2-metadata-token-ttl-seconds: 60" \ + http://169.254.169.254/latest/api/token >/dev/null 2>&1; then + echo "aws"; return + fi + # Azure IMDS — requires Metadata: true header and api-version. + if curl -fsS -H 'Metadata: true' --max-time 2 \ + "http://169.254.169.254/metadata/instance?api-version=2018-02-01" \ + >/dev/null 2>&1; then + echo "azure"; return + fi + echo "none" +} +CLOUD=$(detect_cloud) +log "detected cloud: $CLOUD" -# IMDS → AAD token for Key Vault -IMDS_RESP=$(curl -sf -H 'Metadata: true' \ - "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fvault.azure.net" \ - || true) -AAD_TOKEN=$(echo "$IMDS_RESP" | python3 -c 'import sys,json; print(json.load(sys.stdin).get("access_token",""))' 2>/dev/null) -if [ -z "$AAD_TOKEN" ]; then - log "failed to acquire IMDS token (managed identity not attached?); skipping" - exit 0 -fi +# kv_get is the cloud-agnostic secret-fetch front. The provider-specific +# setup happens once below (Azure: get an AAD token; AWS: nothing — aws CLI +# uses the EC2 instance profile via the SDK's default chain) and then each +# logical secret name is dereferenced the same way regardless of cloud. -# Helper: fetch one secret value or empty string. -kv_get() { - local name=$1 - local resp - resp=$(curl -sf -H "Authorization: Bearer $AAD_TOKEN" \ - "https://${VAULT_NAME}.vault.azure.net/secrets/${name}?api-version=7.4" \ +case "$CLOUD" in +azure) + VAULT_NAME="${OPENSANDBOX_AZURE_KEY_VAULT_NAME:-}" + if [ -z "$VAULT_NAME" ]; then + log "OPENSANDBOX_AZURE_KEY_VAULT_NAME unset — host has no KV configured (e.g. dev VM without managed identity); skipping (Vector will use whatever vector.env is on disk)" + exit 0 + fi + IMDS_RESP=$(curl -sf -H 'Metadata: true' \ + "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fvault.azure.net" \ || true) - echo "$resp" | python3 -c 'import sys,json; print(json.load(sys.stdin).get("value",""))' 2>/dev/null -} + AAD_TOKEN=$(echo "$IMDS_RESP" | python3 -c 'import sys,json; print(json.load(sys.stdin).get("access_token",""))' 2>/dev/null) + if [ -z "$AAD_TOKEN" ]; then + log "failed to acquire IMDS token (managed identity not attached?); skipping" + exit 0 + fi + kv_get() { + local name=$1 + local resp + resp=$(curl -sf -H "Authorization: Bearer $AAD_TOKEN" \ + "https://${VAULT_NAME}.vault.azure.net/secrets/${name}?api-version=7.4" \ + || true) + echo "$resp" | python3 -c 'import sys,json; print(json.load(sys.stdin).get("value",""))' 2>/dev/null + } + ;; + +aws) + SECRETS_PREFIX="${OPENSANDBOX_AWS_SECRETS_PREFIX:-}" + if [ -z "$SECRETS_PREFIX" ]; then + log "OPENSANDBOX_AWS_SECRETS_PREFIX unset — host has no Secrets Manager prefix configured; skipping (Vector will use whatever vector.env is on disk)" + exit 0 + fi + if ! command -v aws >/dev/null 2>&1; then + log "aws CLI not installed in AMI — populator can't fetch from Secrets Manager. Bake awscli into the worker image (see deploy/packer/worker-ami-aws-burst.pkr.hcl)." + exit 0 + fi + # Auto-detect region from IMDSv2 so we don't have to plumb it via env. + AWS_IMDS_TOKEN=$(curl -fsS -X PUT --max-time 2 \ + -H "X-aws-ec2-metadata-token-ttl-seconds: 60" \ + http://169.254.169.254/latest/api/token 2>/dev/null || true) + AWS_REGION_DETECTED=$(curl -fsS --max-time 2 \ + -H "X-aws-ec2-metadata-token: $AWS_IMDS_TOKEN" \ + http://169.254.169.254/latest/meta-data/placement/region 2>/dev/null || true) + export AWS_REGION="${OPENSANDBOX_REGION:-${AWS_REGION_DETECTED:-us-east-2}}" + + kv_get() { + local name=$1 + aws secretsmanager get-secret-value \ + --secret-id "${SECRETS_PREFIX}${name}" \ + --query SecretString \ + --output text 2>/dev/null || echo "" + } + ;; + +*) + log "no recognized cloud detected and OPENSANDBOX_CLOUD is empty — skipping. Vector will use whatever vector.env is on disk." + exit 0 + ;; +esac TOKEN_VALUE=$(kv_get "$TOKEN_SECRET") if [ -z "$TOKEN_VALUE" ]; then - log "secret $TOKEN_SECRET not found in $VAULT_NAME (or no access); skipping" + log "secret $TOKEN_SECRET not found in $CLOUD secret store (or no access); skipping" exit 0 fi DATASET_VALUE=$(kv_get "$DATASET_SECRET") if [ -z "$DATASET_VALUE" ]; then - log "secret $DATASET_SECRET not found in $VAULT_NAME (or no access); skipping — Vector won't have a dataset to ship to" + log "secret $DATASET_SECRET not found in $CLOUD secret store (or no access); skipping — Vector won't have a dataset to ship to" exit 0 fi @@ -169,7 +229,7 @@ fi CELL_ID_VALUE=$(kv_get "$CELL_ID_SECRET") if [ -z "$CELL_ID_VALUE" ]; then CELL_ID_VALUE="${OPENCOMPUTER_CELL_ID:-unknown}" - log "secret $CELL_ID_SECRET not found in $VAULT_NAME — falling back to OPENCOMPUTER_CELL_ID=$CELL_ID_VALUE" + log "secret $CELL_ID_SECRET not found in $CLOUD secret store — falling back to OPENCOMPUTER_CELL_ID=$CELL_ID_VALUE" fi # Auto-detect HOST_IP via the kernel's source-address selection (skips link-local). @@ -194,4 +254,4 @@ metrics_status="absent" if [ -n "$METRICS_TOKEN_VALUE" ] && [ -n "$METRICS_DATASET_VALUE" ]; then metrics_status="present" fi -log "populated $ENV_FILE (logs token+dataset from $VAULT_NAME, metrics=$metrics_status, cell_id=$CELL_ID_VALUE, host_ip=${HOST_IP:-unknown})" +log "populated $ENV_FILE (logs token+dataset from $CLOUD, metrics=$metrics_status, cell_id=$CELL_ID_VALUE, host_ip=${HOST_IP:-unknown})" diff --git a/deploy/worker.env.example b/deploy/worker.env.example index d4133699..6f3234e6 100644 --- a/deploy/worker.env.example +++ b/deploy/worker.env.example @@ -5,10 +5,32 @@ # Non-secret config values stay in this file. # ── Secrets Service ────────────────────────────────────────────── -# Set this to enable secrets loading from secrets service. -# When set, all values marked [KEY VAULT] below are fetched automatically. -# Remove those values from this file — they're managed in the vault. +# Set ONE of the following — whichever matches the cloud the worker runs in. +# When set, all values marked "managed in secrets service" below are fetched +# automatically and don't need to be set here. +# +# Azure Key Vault — vault name (e.g. opencomputer-prod-kv). SECRETS_VAULT_NAME=opencomputer-prod-kv +# OPENSANDBOX_AZURE_KEY_VAULT_NAME=opencomputer-prod-kv # legacy alias for SECRETS_VAULT_NAME +# +# AWS Secrets Manager — prefix under which the cell's secrets live, e.g. +# "opencomputer/aws-us-east-2-poc/". The provider lists secrets under this +# prefix and dereferences each name in secretMapping (internal/config/secrets.go). +# OPENSANDBOX_AWS_SECRETS_PREFIX=opencomputer/aws-us-east-2-poc/ + +# Cloud identifier — set by cloud-init from terraform on AWS workers. Used by +# deploy/vector/populate-vector-env.sh and the Go preemption monitor to pick +# the right backend. Leave unset for local dev (script auto-detects via IMDS). +# OPENSANDBOX_CLOUD=aws + +# CPU overcommit ratio — multiplies the physical-vCPU-derived MAX_CAPACITY +# before the worker advertises capacity to the CP via Redis heartbeat. Memory +# is never overcommitted because host OOM is a worse failure mode than CPU +# contention. Recommended values: +# 1 no overcommit (Azure default, predictable performance) +# 2 ~36% cheaper per sandbox on AWS r8i for agent workloads (typical default) +# 3 aggressive — only if sandboxes are dominantly idle (LLM-bound) +# OPENSANDBOX_CPU_OVERCOMMIT_RATIO=1 # ── Worker Config (not secret — stays in this file) ────────────── OPENSANDBOX_MODE=worker diff --git a/docs/api-reference/sandboxes/create.mdx b/docs/api-reference/sandboxes/create.mdx index bc16af7f..906da964 100644 --- a/docs/api-reference/sandboxes/create.mdx +++ b/docs/api-reference/sandboxes/create.mdx @@ -14,11 +14,15 @@ Create a new sandbox. - CPU cores. Must match an allowed tier: `1`, `2`, or `4`. If omitted but `memoryMB` is set, inferred automatically. + CPU cores. Must match an allowed tier: `1`, `2`, `4`, `8`, or `16`. If omitted but `memoryMB` is set, inferred automatically. - Memory in MB. Must match an allowed tier: `1024`, `4096`, `8192`, or `16384`. If omitted but `cpuCount` is set, inferred automatically. + Memory in MB. Must match an allowed tier: `1024`, `4096`, `8192`, `16384`, `32768`, or `65536`. If omitted but `cpuCount` is set, inferred automatically. + + + + Create a [Burst Sandbox](/sandboxes/burst-sandboxes). Disk is preserved across infrastructure restarts; processes may restart. The allowed CPU/memory combinations are: @@ -29,6 +33,8 @@ The allowed CPU/memory combinations are: | 4096 MB (4 GB) | 1 | | 8192 MB (8 GB) | 2 | | 16384 MB (16 GB) | 4 | +| 32768 MB (32 GB) | 8 | +| 65536 MB (64 GB) | 16 | The 1 GB tier provides 1 vCPU on a best-effort basis. For guaranteed CPU allocation, use the 4 GB tier or above. diff --git a/docs/docs.json b/docs/docs.json index da0d6124..365ca4de 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -46,6 +46,13 @@ "sandboxes/patches", "sandboxes/preview-urls", "sandboxes/elasticity", + { + "group": "Burst Sandboxes", + "tag": "Alpha", + "pages": [ + "sandboxes/burst-sandboxes" + ] + }, { "group": "Usage", "tag": "Preview", diff --git a/docs/reference/python-sdk/sandbox.mdx b/docs/reference/python-sdk/sandbox.mdx index 04c1cced..7a0855da 100644 --- a/docs/reference/python-sdk/sandbox.mdx +++ b/docs/reference/python-sdk/sandbox.mdx @@ -33,6 +33,10 @@ Create a new sandbox. [HTTP API →](/api-reference/sandboxes/create) Arbitrary metadata + + Create a [Burst Sandbox](/sandboxes/burst-sandboxes). Disk is preserved across infrastructure restarts; processes may restart. + + Declarative image definition (see [Image](/reference/python-sdk/image)) @@ -51,8 +55,19 @@ Create a new sandbox. [HTTP API →](/api-reference/sandboxes/create) sandbox = await Sandbox.create(template="my-stack", timeout=600) ``` +Create a Burst Sandbox: + +```python +sandbox = await Sandbox.create( + burst=True, + timeout=300, +) +``` + - `cpuCount` and `memoryMB` are not available in the Python SDK. Use the [HTTP API](/api-reference/sandboxes/create) for custom resources. + Burst Sandboxes are alpha. They preserve filesystem state across + infrastructure restarts, may restart running processes, and are priced roughly + 2x cheaper than on-demand sandboxes. --- diff --git a/docs/reference/typescript-sdk/sandbox.mdx b/docs/reference/typescript-sdk/sandbox.mdx index 45525539..ef130807 100644 --- a/docs/reference/typescript-sdk/sandbox.mdx +++ b/docs/reference/typescript-sdk/sandbox.mdx @@ -33,6 +33,10 @@ Create a new sandbox. [HTTP API →](/api-reference/sandboxes/create) Arbitrary metadata + + Create a [Burst Sandbox](/sandboxes/burst-sandboxes). Disk is preserved across infrastructure restarts; processes may restart. + + CPU cores @@ -63,6 +67,21 @@ Create a new sandbox. [HTTP API →](/api-reference/sandboxes/create) const sandbox = await Sandbox.create({ template: "my-stack", timeout: 600 }); ``` +Create a Burst Sandbox: + +```typescript +const sandbox = await Sandbox.create({ + burst: true, + timeout: 300, +}); +``` + + + Burst Sandboxes are alpha. They preserve filesystem state across + infrastructure restarts, may restart running processes, and are priced roughly + 2x cheaper than on-demand sandboxes. + + --- ### `Sandbox.connect(sandboxId, opts?)` diff --git a/docs/sandboxes/burst-sandboxes.mdx b/docs/sandboxes/burst-sandboxes.mdx new file mode 100644 index 00000000..90bc99fc --- /dev/null +++ b/docs/sandboxes/burst-sandboxes.mdx @@ -0,0 +1,235 @@ +--- +title: "Burst Sandboxes" +description: "Alpha lower-cost sandboxes that preserve disk across infrastructure restarts" +--- + + + Burst Sandboxes are in alpha. They preserve filesystem state across + infrastructure restarts, but running processes, in-memory state, terminal + sessions, and open network connections may restart. + + +Burst Sandboxes are lower-cost OpenComputer sandboxes designed for workloads that can recover from a process restart. They expose the same OpenComputer API for commands, files, PTY, preview URLs, images, snapshots, scaling, and lifecycle operations. + +If OpenComputer receives advance infrastructure interruption notice, the sandbox gets up to **25 seconds** to flush state before it restarts on healthy capacity. The sandbox's filesystem is preserved and made available after resume. + +## Pricing + +Burst Sandboxes are priced roughly **2x cheaper than on-demand sandboxes**. + +## Create a Burst Sandbox + +Set `burst: true` when creating the sandbox. + + + +```typescript TypeScript +import { Sandbox } from "@opencomputer/sdk"; + +const sandbox = await Sandbox.create({ + burst: true, + memoryMB: 4096, + timeout: 300, +}); + +const result = await sandbox.exec.run("echo hello from burst"); +console.log(result.stdout); + +await sandbox.kill(); +``` + +```python Python +from opencomputer import Sandbox + +async with await Sandbox.create( + burst=True, + memory_mb=4096, + timeout=300, +) as sandbox: + result = await sandbox.exec.run("echo hello from burst") + print(result.stdout) +``` + +```bash HTTP +curl -X POST https://app.opencomputer.dev/api/sandboxes \ + -H "X-API-Key: $OPENCOMPUTER_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "burst": true, + "memoryMB": 4096, + "timeout": 300 + }' +``` + + + +## What Happens on Interruption + +Infrastructure capacity can be reclaimed or restarted by the cloud provider. When OpenComputer receives advance notice, it notifies the sandbox and starts the Burst restart flow. + +The sandbox receives these environment variables: + +- `OPENSANDBOX_BURST=true` +- `OPENSANDBOX_RESUMABLE=true` +- `OPENSANDBOX_RESUME_NOTICE_SECONDS=25` + +During the notice window, write important state to disk. After resume, restart your process from durable state. + +## Good Fits + +Burst Sandboxes work well for cost-sensitive workloads that can restart from disk: + +- Batch code execution where failed jobs can be requeued. +- CI-style checks, test runners, linters, and formatters. +- AI agent tool calls with durable task state outside the sandbox. +- Parallel exploration jobs where one attempt can restart. +- Development, previews, and experiments where cost matters more than availability. + +## Poor Fits + +Use on-demand sandboxes when process continuity is required: + +- User-facing interactive sessions where a rare disconnect is still disruptive. +- Long-running stateful services that keep important state only in RAM. +- Databases, queues, or coordinators running inside the sandbox. +- Workloads with strict deadlines and no retry budget. +- Any task that cannot safely restart after a partial command failure. + +## Alpha Limitations + +During alpha, Burst Sandboxes have these restrictions: + +| Capability | Burst alpha behavior | +| --- | --- | +| Process state | May restart | +| Filesystem | Preserved across Burst restarts | +| Notice window | Up to 25 seconds when advance notice is available | +| Sudden host failure | May resume without advance notice | +| Capacity | Best effort; create may fail or wait when Burst capacity is full | + +## Reliability Pattern + +For agent or batch systems, run your work from a small process wrapper and install a restart-notice hook. + +1. Store task state on disk or outside the sandbox. +2. Write outputs to durable storage as the task progresses. +3. Install `/home/sandbox/.opencomputer/on-restart-notice` to flush or checkpoint before restart. +4. Make your startup command resume from the last durable state. +5. Use on-demand sandboxes for workloads with no process restart budget. + +Example setup: + + + +```typescript TypeScript +import { Sandbox } from "@opencomputer/sdk"; + +const sandbox = await Sandbox.create({ + burst: true, + timeout: 0, +}); + +await sandbox.exec.run("mkdir -p /home/sandbox/.opencomputer /home/sandbox/app"); + +await sandbox.files.write( + "/home/sandbox/.opencomputer/on-restart-notice", + `#!/bin/sh +set -eu +echo "restart notice: ${1:-25}s" >> /home/sandbox/app/events.log +touch /home/sandbox/app/restarting +sync +`, +); + +await sandbox.exec.run("chmod +x /home/sandbox/.opencomputer/on-restart-notice"); + +await sandbox.files.write( + "/home/sandbox/app/worker.sh", + `#!/bin/sh +set -eu +cd /home/sandbox/app +if [ -f restarting ]; then + echo "resumed after restart" >> events.log + rm -f restarting +else + echo "started" >> events.log +fi + +i=$(cat counter 2>/dev/null || echo 0) +while true; do + i=$((i + 1)) + echo "$i" > counter + sync + sleep 5 +done +`, +); + +await sandbox.exec.run("chmod +x /home/sandbox/app/worker.sh"); +await sandbox.exec.background("/home/sandbox/app/worker.sh", { + maxRunAfterDisconnect: 0, +}); +``` + +```python Python +from opencomputer import Sandbox + +sandbox = await Sandbox.create(burst=True, timeout=0) + +await sandbox.exec.run("mkdir -p /home/sandbox/.opencomputer /home/sandbox/app") + +await sandbox.files.write( + "/home/sandbox/.opencomputer/on-restart-notice", + """#!/bin/sh +set -eu +echo "restart notice: ${1:-25}s" >> /home/sandbox/app/events.log +touch /home/sandbox/app/restarting +sync +""", +) + +await sandbox.exec.run("chmod +x /home/sandbox/.opencomputer/on-restart-notice") + +await sandbox.files.write( + "/home/sandbox/app/worker.sh", + """#!/bin/sh +set -eu +cd /home/sandbox/app +if [ -f restarting ]; then + echo "resumed after restart" >> events.log + rm -f restarting +else + echo "started" >> events.log +fi + +i=$(cat counter 2>/dev/null || echo 0) +while true; do + i=$((i + 1)) + echo "$i" > counter + sync + sleep 5 +done +""", +) + +await sandbox.exec.run("chmod +x /home/sandbox/app/worker.sh") +await sandbox.exec.background( + "/home/sandbox/app/worker.sh", + max_run_after_disconnect=0, +) +``` + + + +When OpenComputer receives notice, it runs the hook with the notice window in seconds. If the sandbox resumes on another worker, start your process again and read the state you wrote under `/home/sandbox/app`. + +## Related + + + + Resource tiers and autoscaling + + + Save and resume sandbox state + + diff --git a/internal/api/admin_drain_visualizer.go b/internal/api/admin_drain_visualizer.go new file mode 100644 index 00000000..813dc507 --- /dev/null +++ b/internal/api/admin_drain_visualizer.go @@ -0,0 +1,437 @@ +package api + +import ( + "net/http" + + "github.com/labstack/echo/v4" +) + +// adminDrainVisualizerPage serves an operator-only view for watching live +// migration drains. It intentionally lives under /admin so the existing API-key +// middleware guards the drain/evacuate controls. +func (s *Server) adminDrainVisualizerPage(c echo.Context) error { + return c.HTML(http.StatusOK, adminDrainVisualizerHTML) +} + +const adminDrainVisualizerHTML = ` + + + + + OpenComputer drain visualizer + + + +
+

Drain visualizer

+
+ connecting + + + +
+
+
+
+
Workers
0
+
Sandboxes
0
+
Draining
0
+
Moves
0
+
+
+
+

Events

+
+
+
+ + +` diff --git a/internal/api/admin_workers.go b/internal/api/admin_workers.go index 2dd5e206..a9297dc7 100644 --- a/internal/api/admin_workers.go +++ b/internal/api/admin_workers.go @@ -1,9 +1,16 @@ package api import ( + "bytes" + "encoding/json" + "fmt" "net/http" + "strings" + "time" "github.com/labstack/echo/v4" + "github.com/opensandbox/opensandbox/internal/controlplane" + "github.com/opensandbox/opensandbox/pkg/types" ) // adminSetWorkerDraining toggles the in-memory `Draining` flag on a worker so @@ -49,3 +56,144 @@ func (s *Server) adminSetWorkerDraining(c echo.Context) error { "draining": drain, }) } + +// adminEvacuateWorker starts the scaler's live-migration drain loop for a +// worker. Unlike scaler scale-down, this does not terminate the machine after +// it becomes empty; it is an operator/test hook for spot evacuation drills. +// +// POST /admin/workers/:id/evacuate +func (s *Server) adminEvacuateWorker(c echo.Context) error { + if s.workerEvacuator == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{ + "error": "worker evacuator not configured", + }) + } + + workerID := c.Param("id") + if workerID == "" { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "worker id required"}) + } + + if err := s.workerEvacuator.EvacuateWorker(c.Request().Context(), workerID); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": err.Error()}) + } + + return c.JSON(http.StatusAccepted, map[string]any{ + "workerID": workerID, + "evacuating": true, + }) +} + +type adminRecreateSandboxRequest struct { + TargetWorkerID string `json:"targetWorkerId,omitempty"` +} + +func (s *Server) adminRecreateSandbox(c echo.Context) error { + if s.store == nil || s.workerRegistry == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "store and worker registry are required"}) + } + sandboxID := c.Param("id") + if sandboxID == "" { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "sandbox id required"}) + } + + var req adminRecreateSandboxRequest + _ = c.Bind(&req) + + session, err := s.store.GetSandboxSession(c.Request().Context(), sandboxID) + if err != nil { + return c.JSON(http.StatusNotFound, map[string]string{"error": "sandbox not found"}) + } + var cfg types.SandboxConfig + if len(session.Config) > 0 { + if err := json.Unmarshal(session.Config, &cfg); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "invalid sandbox config: " + err.Error()}) + } + } + if !cfg.IsResumable() { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "sandbox is not resumable"}) + } + cfg.SandboxID = sandboxID + + target := s.findRecreateTarget(session.WorkerID, session.Region, req.TargetWorkerID) + if target == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "no eligible target worker"}) + } + if target.HTTPAddr == "" { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "target worker has no HTTP address"}) + } + + body, _ := json.Marshal(map[string]any{ + "sandboxId": sandboxID, + "config": cfg, + }) + url := strings.TrimRight(target.HTTPAddr, "/") + "/admin/resumable/recreate" + httpReq, err := http.NewRequestWithContext(c.Request().Context(), http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()}) + } + httpReq.Header.Set("Content-Type", "application/json") + client := &http.Client{Timeout: 120 * time.Second} + resp, err := client.Do(httpReq) + if err != nil { + return c.JSON(http.StatusBadGateway, map[string]string{"error": "target worker recreate failed: " + err.Error()}) + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + var workerErr map[string]string + _ = json.NewDecoder(resp.Body).Decode(&workerErr) + msg := workerErr["error"] + if msg == "" { + msg = fmt.Sprintf("target worker returned HTTP %d", resp.StatusCode) + } + return c.JSON(http.StatusBadGateway, map[string]string{"error": msg}) + } + var sb types.Sandbox + if err := json.NewDecoder(resp.Body).Decode(&sb); err != nil { + return c.JSON(http.StatusBadGateway, map[string]string{"error": "decode target worker response: " + err.Error()}) + } + + if err := s.store.UpdateSandboxSessionForRecreate(c.Request().Context(), sandboxID, target.ID); err != nil { + return c.JSON(http.StatusInternalServerError, map[string]string{"error": "update sandbox session: " + err.Error()}) + } + if target.GoldenVersion != "" { + _ = s.store.SetSandboxGoldenVersion(c.Request().Context(), sandboxID, target.GoldenVersion) + } + if s.sandboxAPIProxy != nil { + s.sandboxAPIProxy.InvalidateRouteCache(sandboxID) + } + s.emitEvent("resumed", sandboxID, target.ID, "resumable recreate") + + return c.JSON(http.StatusOK, map[string]any{ + "sandboxID": sandboxID, + "previousWorker": session.WorkerID, + "workerID": target.ID, + "status": sb.Status, + }) +} + +func (s *Server) findRecreateTarget(sourceWorkerID, region, targetWorkerID string) *controlplane.WorkerEntry { + if targetWorkerID != "" { + w := s.workerRegistry.GetWorker(targetWorkerID) + if w == nil || w.Draining { + return nil + } + return w + } + var best *controlplane.WorkerEntry + bestScore := 1e18 + for _, w := range s.workerRegistry.GetAllWorkers() { + if w == nil || w.ID == sourceWorkerID || w.Draining { + continue + } + if region != "" && w.Region != region { + continue + } + score := float64(w.Current)*1000 + w.CPUPct + w.MemPct + w.DiskPct + if best == nil || score < bestScore { + best = w + bestScore = score + } + } + return best +} diff --git a/internal/api/capacity.go b/internal/api/capacity.go index d38bb2d8..6233c5e1 100644 --- a/internal/api/capacity.go +++ b/internal/api/capacity.go @@ -249,11 +249,11 @@ func (s *Server) getCapacityBillableEvents(c echo.Context) error { eventType := c.QueryParam("eventType") switch eventType { - case "", db.BillableEventReservedUsage, db.BillableEventOverageUsage, db.BillableEventDiskOverageUsage: + case "", db.BillableEventReservedUsage, db.BillableEventOverageUsage, db.BillableEventBurstUsage, db.BillableEventDiskOverageUsage: // ok default: return c.JSON(http.StatusBadRequest, map[string]string{ - "error": "eventType must be one of reserved_usage, overage_usage, disk_overage_usage (or omitted for all)", + "error": "eventType must be one of reserved_usage, overage_usage, burst_usage, disk_overage_usage (or omitted for all)", }) } diff --git a/internal/api/internal_sandbox.go b/internal/api/internal_sandbox.go index 407e2314..70cd1ffe 100644 --- a/internal/api/internal_sandbox.go +++ b/internal/api/internal_sandbox.go @@ -113,7 +113,12 @@ func (s *Server) internalCreateSandbox(c echo.Context) error { // so forks of that sandbox inherited a no-network config. cfg.EnsureNetworkEnabledDefault() + if err := types.ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": err.Error()}) + } + // Same memory/cpu defaults the public POST /api/sandboxes applies. + // Sandbox-family defaults are applied above and should not be overwritten. if cfg.MemoryMB == 0 { cfg.MemoryMB = 4096 cfg.CpuCount = 1 diff --git a/internal/api/router.go b/internal/api/router.go index c492d1ea..054aafa7 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -35,35 +35,40 @@ var errSandboxNotAvailable = map[string]string{ "error": "sandbox execution not available in server-only mode", } +type workerEvacuator interface { + EvacuateWorker(ctx context.Context, workerID string) error +} + // Server holds the API server dependencies. type Server struct { - echo *echo.Echo - manager sandbox.Manager - router *sandbox.SandboxRouter // routes all sandbox interactions (state machine, auto-wake, rolling timeout) - ptyManager *sandbox.PTYManager - store *db.Store // nil in combined/dev mode without PG - jwtIssuer *auth.JWTIssuer // nil if JWT not configured - capTokenIssuer *auth.JWTIssuer // verifies edge→CP capability tokens; nil if SESSION_JWT_SECRET unset - cfAdminSecret string // HMAC shared with CreditAccount DO for /admin/halt-org and /admin/resume-org; empty disables auth (dev only) - cfEventSecret string // HMAC shared with the api-edge Worker for /internal/secret-refresh and other edge-→cell push paths - cellID string // this control plane's cell_id (for the cap-token cell check) - mode string // "server", "worker", "combined" - workerID string // this worker's ID - region string // this worker's region - httpAddr string // public HTTP address for direct access - execSessionManager *sandbox.ExecSessionManager // nil if not configured - sandboxDBs *sandbox.SandboxDBManager // per-sandbox SQLite manager - workos *auth.WorkOSMiddleware // nil if WorkOS not configured - workerRegistry *controlplane.RedisWorkerRegistry // nil in combined/worker mode - checkpointStore *storage.CheckpointStore // nil if hibernation not configured - sandboxDomain string // base domain for sandbox subdomains - cfClient *cloudflare.Client // nil if Cloudflare not configured - pendingCreates sync.Map // map[sandboxID]*pendingCreate — async sandbox creation tracking - sandboxAPIProxy *proxy.SandboxAPIProxy // nil except in server mode (proxies data-plane to workers) - stripeClient *billing.StripeClient // nil if Stripe not configured - redisClient *redis.Client // nil if Redis not configured (for health checks) - adminEvents *AdminEventBus // real-time event bus for admin dashboard - ready int32 // atomic: 1 = ready, 0 = not ready + echo *echo.Echo + manager sandbox.Manager + router *sandbox.SandboxRouter // routes all sandbox interactions (state machine, auto-wake, rolling timeout) + ptyManager *sandbox.PTYManager + store *db.Store // nil in combined/dev mode without PG + jwtIssuer *auth.JWTIssuer // nil if JWT not configured + capTokenIssuer *auth.JWTIssuer // verifies edge→CP capability tokens; nil if SESSION_JWT_SECRET unset + cfAdminSecret string // HMAC shared with CreditAccount DO for /admin/halt-org and /admin/resume-org; empty disables auth (dev only) + cfEventSecret string // HMAC shared with the api-edge Worker for /internal/secret-refresh and other edge-→cell push paths + cellID string // this control plane's cell_id (for the cap-token cell check) + mode string // "server", "worker", "combined" + workerID string // this worker's ID + region string // this worker's region + httpAddr string // public HTTP address for direct access + execSessionManager *sandbox.ExecSessionManager // nil if not configured + sandboxDBs *sandbox.SandboxDBManager // per-sandbox SQLite manager + workos *auth.WorkOSMiddleware // nil if WorkOS not configured + workerRegistry *controlplane.RedisWorkerRegistry // nil in combined/worker mode + workerEvacuator workerEvacuator // nil when autoscaler is disabled + checkpointStore *storage.CheckpointStore // nil if hibernation not configured + sandboxDomain string // base domain for sandbox subdomains + cfClient *cloudflare.Client // nil if Cloudflare not configured + pendingCreates sync.Map // map[sandboxID]*pendingCreate — async sandbox creation tracking + sandboxAPIProxy *proxy.SandboxAPIProxy // nil except in server mode (proxies data-plane to workers) + stripeClient *billing.StripeClient // nil if Stripe not configured + redisClient *redis.Client // nil if Redis not configured (for health checks) + adminEvents *AdminEventBus // real-time event bus for admin dashboard + ready int32 // atomic: 1 = ready, 0 = not ready // Axiom log query (sandbox session logs read API). // Empty token = endpoint returns 503. @@ -94,34 +99,35 @@ func (s *Server) SetAxiomQueryConfig(queryToken, dataset string) { // pendingCreate tracks an async sandbox creation. type pendingCreate struct { ready chan struct{} // closed when creation completes - err error // set before closing ready + err error // set before closing ready } // ServerOpts holds optional dependencies for the API server. type ServerOpts struct { - Store *db.Store - JWTIssuer *auth.JWTIssuer - SessionJWTSecret string // shared edge↔CP HMAC secret; enables /internal/sandboxes/create - CFAdminSecret string // HMAC shared with CF CreditAccount DO; enables /admin/halt-org and /admin/resume-org - CFEventSecret string // HMAC shared with the api-edge Worker; enables /internal/secret-refresh and other edge-→cell push paths - CellID string // this control plane's cell_id - Mode string // "server", "worker", "combined" - WorkerID string - Region string - HTTPAddr string + Store *db.Store + JWTIssuer *auth.JWTIssuer + SessionJWTSecret string // shared edge↔CP HMAC secret; enables /internal/sandboxes/create + CFAdminSecret string // HMAC shared with CF CreditAccount DO; enables /admin/halt-org and /admin/resume-org + CFEventSecret string // HMAC shared with the api-edge Worker; enables /internal/secret-refresh and other edge-→cell push paths + CellID string // this control plane's cell_id + Mode string // "server", "worker", "combined" + WorkerID string + Region string + HTTPAddr string ExecSessionManager *sandbox.ExecSessionManager - SandboxDBs *sandbox.SandboxDBManager - Router *sandbox.SandboxRouter // nil in server-only mode - SandboxProxy *proxy.SandboxProxy // nil if subdomain routing not configured - ControlPlaneProxy *proxy.ControlPlaneProxy // nil except in server mode (routes subdomains to workers) - SandboxDomain string // base domain for sandbox subdomains - WorkOSConfig *auth.WorkOSConfig // nil if WorkOS not configured - WorkerRegistry *controlplane.RedisWorkerRegistry // nil in combined/worker mode - CheckpointStore *storage.CheckpointStore // nil if hibernation not configured - CFClient *cloudflare.Client // nil if Cloudflare not configured - SandboxAPIProxy *proxy.SandboxAPIProxy // nil except in server mode (proxies data-plane to workers) - StripeClient *billing.StripeClient // nil if Stripe not configured - RedisClient *redis.Client // nil if Redis not configured (for health checks) + SandboxDBs *sandbox.SandboxDBManager + Router *sandbox.SandboxRouter // nil in server-only mode + SandboxProxy *proxy.SandboxProxy // nil if subdomain routing not configured + ControlPlaneProxy *proxy.ControlPlaneProxy // nil except in server mode (routes subdomains to workers) + SandboxDomain string // base domain for sandbox subdomains + WorkOSConfig *auth.WorkOSConfig // nil if WorkOS not configured + WorkerRegistry *controlplane.RedisWorkerRegistry // nil in combined/worker mode + WorkerEvacuator *controlplane.Scaler // nil when autoscaler is disabled + CheckpointStore *storage.CheckpointStore // nil if hibernation not configured + CFClient *cloudflare.Client // nil if Cloudflare not configured + SandboxAPIProxy *proxy.SandboxAPIProxy // nil except in server mode (proxies data-plane to workers) + StripeClient *billing.StripeClient // nil if Stripe not configured + RedisClient *redis.Client // nil if Redis not configured (for health checks) } // NewServer creates a new API server with all routes configured. @@ -153,6 +159,7 @@ func NewServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, apiKey string, o s.sandboxDBs = opts.SandboxDBs s.router = opts.Router s.workerRegistry = opts.WorkerRegistry + s.workerEvacuator = opts.WorkerEvacuator s.checkpointStore = opts.CheckpointStore s.sandboxDomain = opts.SandboxDomain s.cfClient = opts.CFClient @@ -267,8 +274,11 @@ func NewServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, apiKey string, o admin.GET("/report", s.adminReport) admin.POST("/events/clear", s.adminClearEvents) admin.POST("/workers/:id/drain", s.adminSetWorkerDraining) + admin.POST("/workers/:id/evacuate", s.adminEvacuateWorker) + admin.POST("/sandboxes/:id/recreate", s.adminRecreateSandbox) admin.GET("/demo/migration", s.demoPingPongPage) admin.GET("/demo/chaos", s.demoChaosPage) + admin.GET("/demo/drain", s.adminDrainVisualizerPage) // Signed URL endpoints (self-authenticated via HMAC, no API key required) e.GET("/api/sandboxes/:id/files/download", s.signedDownload) diff --git a/internal/api/sandbox.go b/internal/api/sandbox.go index dbd4b400..62e3f257 100644 --- a/internal/api/sandbox.go +++ b/internal/api/sandbox.go @@ -32,6 +32,13 @@ func (s *Server) createSandbox(c echo.Context) error { // to sandbox_sessions.config_json is explicit and forks inherit it correctly. cfg.EnsureNetworkEnabledDefault() + if err := types.ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{ + "error": err.Error(), + }) + } + cfg = withBurstSandboxEnv(cfg) + // Validate CPU/memory against allowed tiers. // Allowed tiers (memoryMB → vCPU): 1024→1, 4096→1, 8192→2, 16384→4, 32768→8, 65536→16. if err := types.ValidateResourceTier(&cfg); err != nil { @@ -71,7 +78,8 @@ func (s *Server) createSandbox(c echo.Context) error { } } - // Default to 4GB/1vCPU if not specified (all plans) + // Default to 4GB/1vCPU if not specified (all plans). Sandbox-family + // defaults are applied above and should not be overwritten here. if cfg.MemoryMB == 0 { cfg.MemoryMB = 4096 cfg.CpuCount = 1 @@ -177,6 +185,8 @@ func (s *Server) createSandbox(c echo.Context) error { "error": err.Error(), }) } + sb.SandboxFamily = cfg.SandboxFamily + sb.Burst = cfg.IsResumable() // Register with sandbox router for rolling timeout tracking. // timeout == 0 means "persistent" (no auto-hibernate). Negative values are @@ -493,7 +503,11 @@ func (s *Server) createSandboxRemote(c echo.Context, ctx context.Context, cfg ty region = "iad" } - worker, grpcClient, err := s.workerRegistry.GetLeastLoadedWorker(region) + workerPool := controlplane.WorkerPoolOnDemand + if cfg.IsResumable() { + workerPool = controlplane.WorkerPoolBurst + } + worker, grpcClient, err := s.workerRegistry.GetLeastLoadedWorkerForPool(region, workerPool) if err != nil { // No worker immediately available — poll for up to 30s // (scaler may be launching a new worker) @@ -504,17 +518,17 @@ func (s *Server) createSandboxRemote(c echo.Context, ctx context.Context, cfg ty select { case <-deadline: return c.JSON(http.StatusServiceUnavailable, map[string]string{ - "error": "no workers available in region " + region + " (waited 30s)", + "error": fmt.Sprintf("no %s workers available in region %s (waited 30s)", workerPool, region), }) case <-ctx.Done(): return c.JSON(http.StatusServiceUnavailable, map[string]string{ "error": "request cancelled while waiting for capacity", }) case <-ticker.C: - worker, grpcClient, err = s.workerRegistry.GetLeastLoadedWorker(region) + worker, grpcClient, err = s.workerRegistry.GetLeastLoadedWorkerForPool(region, workerPool) } } - log.Printf("sandbox: worker became available after queuing (region=%s)", region) + log.Printf("sandbox: %s worker became available after queuing (region=%s)", workerPool, region) } // Resolve template (org-scoped lookup with public fallback). @@ -669,11 +683,15 @@ func (s *Server) createSandboxRemote(c echo.Context, ctx context.Context, cfg ty s.emitEvent("create", grpcResp.SandboxId, worker.ID, fmt.Sprintf("created on %s", worker.ID[len(worker.ID)-8:])) resp := map[string]interface{}{ - "sandboxID": grpcResp.SandboxId, - "token": token, - "status": grpcResp.Status, - "region": region, - "workerID": worker.ID, + "sandboxID": grpcResp.SandboxId, + "token": token, + "status": grpcResp.Status, + "region": region, + "workerID": worker.ID, + "cpuCount": cfg.CpuCount, + "memoryMB": cfg.MemoryMB, + "sandboxFamily": cfg.SandboxFamily, + "burst": cfg.IsResumable(), } if s.sandboxDomain != "" { resp["sandboxDomain"] = s.sandboxDomain @@ -1184,6 +1202,36 @@ func (s *Server) effectivePlan(c echo.Context, orgID uuid.UUID) string { return "" } +func sandboxSessionFamily(session *db.SandboxSession) string { + if session == nil || len(session.Config) == 0 { + return types.SandboxFamilyDefault + } + var cfg types.SandboxConfig + if err := json.Unmarshal(session.Config, &cfg); err != nil { + return types.SandboxFamilyDefault + } + return cfg.SandboxFamily +} + +func isSpotSandboxSession(session *db.SandboxSession) bool { + return sandboxSessionFamily(session) == types.SandboxFamilySpot +} + +func withBurstSandboxEnv(cfg types.SandboxConfig) types.SandboxConfig { + if !cfg.IsResumable() { + return cfg + } + envs := make(map[string]string, len(cfg.Envs)+3) + for k, v := range cfg.Envs { + envs[k] = v + } + envs["OPENSANDBOX_BURST"] = "true" + envs["OPENSANDBOX_RESUMABLE"] = "true" + envs["OPENSANDBOX_RESUME_NOTICE_SECONDS"] = "25" + cfg.Envs = envs + return cfg +} + func (s *Server) setLimits(c echo.Context) error { id := c.Param("id") ctx := c.Request().Context() diff --git a/internal/billing/billable_events_sender.go b/internal/billing/billable_events_sender.go index 0dd17c5e..9c2ea3af 100644 --- a/internal/billing/billable_events_sender.go +++ b/internal/billing/billable_events_sender.go @@ -138,7 +138,7 @@ func (s *BillableEventsSender) shipOne(ctx context.Context, p db.PendingBillable } // meterEventNameFor maps an outbox event_type → Stripe meter event_name. -// All three event types route to a single meter each — overage is flat +// All event types route to a single meter each — overage and burst are flat // (the per-tier memory_mb on the outbox row is preserved for analytics // but ignored when shipping; Stripe sums GB-seconds across rows). func (s *BillableEventsSender) meterEventNameFor(eventType string, memoryMB int) (string, error) { @@ -153,6 +153,11 @@ func (s *BillableEventsSender) meterEventNameFor(eventType string, memoryMB int) return "", fmt.Errorf("overage meter not provisioned (run EnsureProducts)") } return s.stripe.OverageMeterEventName, nil + case db.BillableEventBurstUsage: + if s.stripe.BurstMeterEventName == "" { + return "", fmt.Errorf("burst meter not provisioned (run EnsureProducts)") + } + return s.stripe.BurstMeterEventName, nil case db.BillableEventDiskOverageUsage: if s.stripe.DiskOverageMeterEventName == "" { return "", fmt.Errorf("disk overage meter not provisioned") diff --git a/internal/billing/billable_events_sender_test.go b/internal/billing/billable_events_sender_test.go index 5d7d9d96..9991f3df 100644 --- a/internal/billing/billable_events_sender_test.go +++ b/internal/billing/billable_events_sender_test.go @@ -10,12 +10,13 @@ import ( // meterEventNameFor pins the (event_type, memory_mb) → Stripe meter // routing. Overage is flat — memory_mb is preserved on the outbox row // for analytics but ignored at ship time so a 1 GB and 64 GB sandbox -// hit the same meter. +// hit the same meter. Burst is also flat, but routed to its own meter. func newSenderForTest() *BillableEventsSender { stripe := &StripeClient{ ReservedMeterEventName: "sandbox_compute_sandbox_reserved", OverageMeterEventName: "sandbox_compute_sandbox_overage", + BurstMeterEventName: "sandbox_compute_sandbox_burst", DiskOverageMeterEventName: "sandbox_compute_sandbox_disk_overage", } return &BillableEventsSender{stripe: stripe} @@ -56,6 +57,17 @@ func TestSender_meterEventName_diskOverage(t *testing.T) { } } +func TestSender_meterEventName_burstFlat(t *testing.T) { + s := newSenderForTest() + got, err := s.meterEventNameFor(db.BillableEventBurstUsage, 0) + if err != nil { + t.Fatalf("err: %v", err) + } + if got != "sandbox_compute_sandbox_burst" { + t.Errorf("got %q", got) + } +} + func TestSender_meterEventName_unknownTypeRejected(t *testing.T) { s := newSenderForTest() _, err := s.meterEventNameFor("totally_made_up", 0) @@ -74,6 +86,7 @@ func TestSender_meterEventName_missingProvisionRejected(t *testing.T) { }{ {db.BillableEventReservedUsage, "reserved meter not provisioned"}, {db.BillableEventOverageUsage, "overage meter not provisioned"}, + {db.BillableEventBurstUsage, "burst meter not provisioned"}, {db.BillableEventDiskOverageUsage, "disk overage meter not provisioned"}, } for _, tc := range cases { diff --git a/internal/billing/capacity_reconciler.go b/internal/billing/capacity_reconciler.go index 5b99dd5e..32b9f964 100644 --- a/internal/billing/capacity_reconciler.go +++ b/internal/billing/capacity_reconciler.go @@ -14,7 +14,7 @@ import ( // it scans for closed 15-min buckets that are at least `settle` past // their end and haven't been processed yet, runs the per-second // integration walk for each (org, bucket), and emits the resulting -// reserved_usage / overage_usage / disk_overage_usage rows to the +// reserved_usage / overage_usage / burst_usage / disk_overage_usage rows to the // `billable_events` outbox. // // Phase 2 runs in shadow: the rows are written but not delivered to @@ -152,6 +152,7 @@ func (r *CapacityReconciler) processBucket(ctx context.Context, orgID uuid.UUID, // is org-level, summed across all running events at each segment. type BucketTotals struct { OverageGBSecondsByTier map[int]float64 + BurstGBSeconds float64 DiskOverageGBSeconds float64 ReservedFloorGBSeconds float64 // reservedGb × secs accumulated across segments — for shadow validation only } @@ -176,6 +177,7 @@ type clippedEvent struct { From, To time.Time TierMB int DiskMB int + Burst bool } // clipEvent restricts a ScaleEvent's lifetime to the bucket window. @@ -192,7 +194,7 @@ func clipEvent(e db.ScaleEvent, bucketStart, bucketEnd time.Time) (clippedEvent, if !to.After(from) { return clippedEvent{}, false } - return clippedEvent{From: from, To: to, TierMB: e.MemoryMB, DiskMB: e.DiskMB}, true + return clippedEvent{From: from, To: to, TierMB: e.MemoryMB, DiskMB: e.DiskMB, Burst: e.Burst}, true } func collectBoundaries(events []clippedEvent, bucketStart, bucketEnd time.Time) []time.Time { @@ -216,9 +218,10 @@ func collectBoundaries(events []clippedEvent, bucketStart, bucketEnd time.Time) } type segment struct { - From, To time.Time - RunningByTier map[int]int // tier_mb → running GB at this tier - DiskOverageMB int // sum of (disk_mb − 20480) over running events + From, To time.Time + RunningByTier map[int]int // non-burst tier_mb → running GB at this tier + BurstGB int + DiskOverageMB int // sum of (disk_mb − 20480) over running events } func walkSegments(boundaries []time.Time, events []clippedEvent) []segment { @@ -229,6 +232,7 @@ func walkSegments(boundaries []time.Time, events []clippedEvent) []segment { continue } tiers := map[int]int{} + burstGB := 0 diskOver := 0 for _, e := range events { // Event is "running" in [from, to) if it covers the segment @@ -236,13 +240,18 @@ func walkSegments(boundaries []time.Time, events []clippedEvent) []segment { // boundary set includes every event endpoint, every segment // is fully contained in any event whose span covers `from`. if !e.From.After(from) && !e.To.Before(to) { - tiers[e.TierMB] += e.TierMB / 1024 + gb := e.TierMB / 1024 + if e.Burst { + burstGB += gb + } else { + tiers[e.TierMB] += gb + } if e.DiskMB > 20480 { diskOver += e.DiskMB - 20480 } } } - segs = append(segs, segment{From: from, To: to, RunningByTier: tiers, DiskOverageMB: diskOver}) + segs = append(segs, segment{From: from, To: to, RunningByTier: tiers, BurstGB: burstGB, DiskOverageMB: diskOver}) } return segs } @@ -251,6 +260,10 @@ func integrateSegments(segs []segment, reservedGB int) BucketTotals { out := BucketTotals{OverageGBSecondsByTier: map[int]float64{}} for _, s := range segs { secs := s.To.Sub(s.From).Seconds() + if s.BurstGB > 0 { + out.BurstGBSeconds += float64(s.BurstGB) * secs + } + usage := 0 for _, gb := range s.RunningByTier { usage += gb @@ -281,7 +294,9 @@ func integrateSegments(segs []segment, reservedGB int) BucketTotals { // actual usage — the customer paid for the floor whether or not // they used it). // - overage_usage — one row per (org, sandbox_tier, bucket) where -// the tier's overage contribution is non-zero. +// the non-burst tier's overage contribution is non-zero. +// - burst_usage — one row per (org, bucket) when burst sandboxes ran. +// Burst is billed independently and does not consume reserved floors. // - disk_overage_usage — one row per (org, bucket) when any sandbox // in the bucket exceeded the 20 GB allowance. // @@ -319,6 +334,20 @@ func emitBucket(ctx context.Context, store *db.Store, orgID uuid.UUID, bucketSta } } + if totals.BurstGBSeconds > 0 { + ev := db.BillableEvent{ + OrgID: orgID, + EventType: db.BillableEventBurstUsage, + MemoryMB: 0, + GBSeconds: totals.BurstGBSeconds, + BucketStart: bucketStart, + BucketEnd: bucketEnd, + } + if _, err := store.UpsertBillableEvent(ctx, ev); err != nil { + return err + } + } + if totals.DiskOverageGBSeconds > 0 { ev := db.BillableEvent{ OrgID: orgID, diff --git a/internal/billing/capacity_reconciler_test.go b/internal/billing/capacity_reconciler_test.go index 75671d11..67383a45 100644 --- a/internal/billing/capacity_reconciler_test.go +++ b/internal/billing/capacity_reconciler_test.go @@ -66,6 +66,19 @@ func TestIntegrateBucket_zeroReservation_replaysLegacy(t *testing.T) { } } +func TestIntegrateBucket_burstUsageSeparateFromOverage(t *testing.T) { + bs, be := canonicalBucket() + totals := IntegrateBucket(bs, be, 8, []db.ScaleEvent{ + {MemoryMB: 8192, DiskMB: 20480, Burst: true, StartedAt: bs, EndedAt: &be}, + {MemoryMB: 8192, DiskMB: 20480, StartedAt: bs, EndedAt: &be}, + }) + eq(t, "burst", totals.BurstGBSeconds, 8*900) + eq(t, "reserved floor", totals.ReservedFloorGBSeconds, 8*900) + if len(totals.OverageGBSecondsByTier) != 0 { + t.Errorf("expected burst to stay out of overage, got %v", totals.OverageGBSecondsByTier) + } +} + func TestIntegrateBucket_fullCoverage_noOverage(t *testing.T) { bs, be := canonicalBucket() // One 8 GB sandbox, reservation covers it. diff --git a/internal/billing/pricing.go b/internal/billing/pricing.go index 62db397b..8fed72a1 100644 --- a/internal/billing/pricing.go +++ b/internal/billing/pricing.go @@ -54,8 +54,8 @@ var TierPriceKey = map[int]string{ // full lifetime of the sandbox (running OR hibernated, since the workspace // qcow2 still occupies host disk). const ( - DiskFreeAllowanceMB = 20480 // 20GB included with every sandbox - DiskOveragePricePerGBPerSecond = 0.0000001 // ~$0.26 per GB-month + DiskFreeAllowanceMB = 20480 // 20GB included with every sandbox + DiskOveragePricePerGBPerSecond = 0.0000001 // ~$0.26 per GB-month DiskOverageMetadataKey = "sandbox_disk_overage" ) @@ -68,6 +68,8 @@ const ( // reserve capacity pay this. // - Overage: customer-facing name "instant." Flat per-GB rate // regardless of sandbox tier. +// - Burst: customer-facing name "burst." Flat per-GB rate for +// restartable burst sandboxes, independent of reserved capacity. // // **Dollar rates are deliberately not in code.** Stripe Prices are // configured in the Stripe Dashboard (or via Stripe API by the @@ -92,6 +94,10 @@ const ( // Instant (overage): single meter, flat across all sandbox sizes. OverageMeterKey = "sandbox_overage" OveragePriceKey = "sandbox_overage_v1" + + // Burst: single meter, flat across all sandbox sizes. + BurstMeterKey = "sandbox_burst" + BurstPriceKey = "sandbox_burst_v1" ) // DiskOverageGBSeconds returns the chargeable GB-seconds for one usage summary diff --git a/internal/billing/stripe.go b/internal/billing/stripe.go index a66b391d..0bb2749a 100644 --- a/internal/billing/stripe.go +++ b/internal/billing/stripe.go @@ -32,17 +32,19 @@ type StripeClient struct { MeterEventNames map[int]string // Disk overage meter / price (single dimension: GB-seconds above 20GB). - DiskOveragePriceID string + DiskOveragePriceID string DiskOverageMeterEventName string - // Phase-3 unified-pipeline meters and prices. Two flat meters - // (overage + reserved) used by new orgs (`billing_mode='unified'`). + // Phase-3 unified-pipeline meters and prices. Flat meters + // (overage + reserved + burst) used by new orgs (`billing_mode='unified'`). // Legacy per-tier meters above are untouched and continue to serve // existing orgs via UsageReporter. OveragePriceID string // flat overage Price for `overage_usage` events OverageMeterEventName string // "sandbox_compute_sandbox_overage" ReservedPriceID string // flat reserved Price for `reserved_usage` events ReservedMeterEventName string // "sandbox_compute_sandbox_reserved" + BurstPriceID string // flat burst Price for `burst_usage` events + BurstMeterEventName string // "sandbox_compute_sandbox_burst" // Per-agent paywalled-feature prices. Configured from env and // referenced by name from the dashboard subscribe handlers. Empty @@ -236,8 +238,8 @@ func (s *StripeClient) EnsureProducts() error { log.Printf("billing: created disk overage price (id=%s)", p.ID) } - // 5. Phase-3 unified-pipeline meters. Two flat meters (overage + - // reserved) at unit GB-seconds. Code creates the *meters* (their + // 5. Phase-3 unified-pipeline meters. Flat meters (overage, reserved, + // and burst) at unit GB-seconds. Code creates the *meters* (their // event names are stable wire-protocol coupling) but **does not // create Prices** — those are configured in the Stripe Dashboard // so pricing changes don't need a code deploy. The Price IDs @@ -246,6 +248,7 @@ func (s *StripeClient) EnsureProducts() error { // linked to the meter in Stripe. s.OverageMeterEventName = ensureMeter(existingMeters, "sandbox_compute_"+OverageMeterKey, "Sandbox Instant Compute (GB-seconds)") s.ReservedMeterEventName = ensureMeter(existingMeters, "sandbox_compute_"+ReservedMeterKey, "Sandbox Reserved Capacity (GB-seconds)") + s.BurstMeterEventName = ensureMeter(existingMeters, "sandbox_compute_"+BurstMeterKey, "Sandbox Burst Compute (GB-seconds)") if id, ok := existingPrices[OveragePriceKey]; ok { s.OveragePriceID = id @@ -259,6 +262,12 @@ func (s *StripeClient) EnsureProducts() error { } else { log.Printf("billing: no reserved price configured for meter %s — meter events will flow but won't appear on invoices until a Price is created in Stripe", s.ReservedMeterEventName) } + if id, ok := existingPrices[BurstPriceKey]; ok { + s.BurstPriceID = id + log.Printf("billing: found existing burst price (id=%s)", id) + } else { + log.Printf("billing: no burst price configured for meter %s — meter events will flow but won't appear on invoices until a Price is created in Stripe", s.BurstMeterEventName) + } return nil } diff --git a/internal/compute/ec2.go b/internal/compute/ec2.go index 4d3673a4..0d11500c 100644 --- a/internal/compute/ec2.go +++ b/internal/compute/ec2.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "log" + "strconv" "strings" "sync" @@ -50,27 +51,55 @@ func wrapEC2CreateErr(err error, format string, args ...any) error { return wrapped } +func supportsEC2NestedVirtualization(instanceType string) bool { + family, _, _ := strings.Cut(strings.ToLower(instanceType), ".") + switch family { + case "c8i", "m8i", "r8i": + return true + default: + return false + } +} + const ( // AWS tag keys (kept consistent with the Azure pool's azure-prefixed tags). awsTagRole = "opensandbox:role" + awsTagCell = "opensandbox:cell" awsTagInstanceType = "opensandbox:instance-type" awsTagDraining = "opensandbox:draining" + awsTagOCFS2Slot = "opensandbox:ocfs2-slot" + awsTagOCFS2IP = "opensandbox:ocfs2-ip" awsTagWorker = "worker" ) // EC2PoolConfig configures the EC2 compute pool. type EC2PoolConfig struct { - Region string - AccessKeyID string // empty = use default credential chain (IAM role preferred) - SecretAccessKey string - AMI string // static AMI ID; empty if SSMParameterName is set - InstanceType string // e.g. "c7gd.metal", "r7gd.xlarge", "m7i.large" - SubnetID string - SecurityGroupID string - KeyName string // optional SSH key pair (debug use only) - IAMInstanceProfile string // attached to instances; gives them Secrets Manager + S3 read - SecretsARN string // Secrets Manager ARN; passed to worker via WorkerSpec.SecretsRef - SSMParameterName string // SSM parameter for dynamic AMI ID (e.g. /opensandbox/dev/worker-ami-id) + Region string + AccessKeyID string // empty = use default credential chain (IAM role preferred) + SecretAccessKey string + AMI string // static AMI ID; empty if SSMParameterName is set + InstanceType string // e.g. "c7gd.metal", "r7gd.xlarge", "m7i.large" + SubnetID string + SecurityGroupID string + KeyName string // optional SSH key pair (debug use only) + IAMInstanceProfile string // attached to instances; gives them Secrets Manager + S3 read + SecretsARN string // Secrets Manager ARN; passed to worker via WorkerSpec.SecretsRef + SSMParameterName string // SSM parameter for dynamic AMI ID (e.g. /opensandbox/dev/worker-ami-id) + MarketType string // empty/on-demand or spot + CellID string + SharedSandboxDataVolumeID string // optional io2 Multi-Attach volume mounted at /data/sandboxes via OCFS2 + SharedGoldensVolumeID string // optional io2 Multi-Attach volume for golden image cache + OCFS2ClusterName string + OCFS2ExpectedNodes int + OCFS2MaxNodes int + OCFS2NodeIPs []string // fixed private IPs, one per OCFS2 node slot +} + +type ocfs2Assignment struct { + Enabled bool + Slot int + IP string + NodeIPs []string } // EC2Pool implements compute.Pool using AWS EC2 instances. @@ -150,8 +179,38 @@ func (p *EC2Pool) CreateMachine(ctx context.Context, opts MachineOpts) (*Machine return nil, fmt.Errorf("ec2: no AMI set (configure AMI or SSMParameterName)") } - userData := p.buildUserData(opts) + ocfs2, err := p.allocateOCFS2Slot(ctx) + if err != nil { + return nil, err + } + + userData := p.buildUserData(opts, ocfs2) machineName := fmt.Sprintf("osb-worker-%s", randomSuffix()) + instanceTags := []ec2types.Tag{ + {Key: aws.String("Name"), Value: aws.String(machineName)}, + {Key: aws.String("Role"), Value: aws.String("worker")}, + {Key: aws.String(awsTagRole), Value: aws.String(awsTagWorker)}, + {Key: aws.String(awsTagInstanceType), Value: aws.String(instanceType)}, + } + if ocfs2.Enabled { + instanceTags = append(instanceTags, + ec2types.Tag{Key: aws.String(awsTagOCFS2Slot), Value: aws.String(strconv.Itoa(ocfs2.Slot))}, + ec2types.Tag{Key: aws.String(awsTagOCFS2IP), Value: aws.String(ocfs2.IP)}, + ) + } + volumeTags := []ec2types.Tag{ + {Key: aws.String(awsTagRole), Value: aws.String(awsTagWorker)}, + } + if p.cfg.CellID != "" { + instanceTags = append(instanceTags, + ec2types.Tag{Key: aws.String("Cell"), Value: aws.String(p.cfg.CellID)}, + ec2types.Tag{Key: aws.String(awsTagCell), Value: aws.String(p.cfg.CellID)}, + ) + volumeTags = append(volumeTags, + ec2types.Tag{Key: aws.String("Cell"), Value: aws.String(p.cfg.CellID)}, + ec2types.Tag{Key: aws.String(awsTagCell), Value: aws.String(p.cfg.CellID)}, + ) + } input := &ec2.RunInstancesInput{ ImageId: aws.String(ami), @@ -162,20 +221,29 @@ func (p *EC2Pool) CreateMachine(ctx context.Context, opts MachineOpts) (*Machine TagSpecifications: []ec2types.TagSpecification{ { ResourceType: ec2types.ResourceTypeInstance, - Tags: []ec2types.Tag{ - {Key: aws.String("Name"), Value: aws.String(machineName)}, - {Key: aws.String(awsTagRole), Value: aws.String(awsTagWorker)}, - {Key: aws.String(awsTagInstanceType), Value: aws.String(instanceType)}, - }, + Tags: instanceTags, }, { ResourceType: ec2types.ResourceTypeVolume, - Tags: []ec2types.Tag{ - {Key: aws.String(awsTagRole), Value: aws.String(awsTagWorker)}, - }, + Tags: volumeTags, }, }, } + if strings.EqualFold(p.cfg.MarketType, "spot") { + input.InstanceMarketOptions = &ec2types.InstanceMarketOptionsRequest{ + MarketType: ec2types.MarketTypeSpot, + SpotOptions: &ec2types.SpotMarketOptions{ + InstanceInterruptionBehavior: ec2types.InstanceInterruptionBehaviorTerminate, + SpotInstanceType: ec2types.SpotInstanceTypeOneTime, + }, + } + } + + if supportsEC2NestedVirtualization(instanceType) { + input.CpuOptions = &ec2types.CpuOptionsRequest{ + NestedVirtualization: ec2types.NestedVirtualizationSpecificationEnabled, + } + } if p.cfg.SubnetID != "" { input.SubnetId = aws.String(p.cfg.SubnetID) @@ -183,6 +251,9 @@ func (p *EC2Pool) CreateMachine(ctx context.Context, opts MachineOpts) (*Machine if p.cfg.SecurityGroupID != "" { input.SecurityGroupIds = []string{p.cfg.SecurityGroupID} } + if ocfs2.Enabled { + input.PrivateIpAddress = aws.String(ocfs2.IP) + } if p.cfg.KeyName != "" { input.KeyName = aws.String(p.cfg.KeyName) } @@ -285,6 +356,67 @@ func (p *EC2Pool) DrainMachine(ctx context.Context, machineID string) error { return nil } +func (p *EC2Pool) allocateOCFS2Slot(ctx context.Context) (ocfs2Assignment, error) { + if p.cfg.SharedSandboxDataVolumeID == "" || len(p.cfg.OCFS2NodeIPs) == 0 { + return ocfs2Assignment{}, nil + } + + used := make(map[int]bool, len(p.cfg.OCFS2NodeIPs)) + ipToSlot := make(map[string]int, len(p.cfg.OCFS2NodeIPs)) + for i, ip := range p.cfg.OCFS2NodeIPs { + ip = strings.TrimSpace(ip) + if ip == "" { + return ocfs2Assignment{}, fmt.Errorf("ec2: OCFS2 node IP slot %d is empty", i) + } + ipToSlot[ip] = i + } + + filters := []ec2types.Filter{ + {Name: aws.String("tag:" + awsTagRole), Values: []string{awsTagWorker}}, + {Name: aws.String("instance-state-name"), Values: []string{"pending", "running", "stopping", "stopped"}}, + } + if p.cfg.CellID != "" { + filters = append(filters, ec2types.Filter{Name: aws.String("tag:" + awsTagCell), Values: []string{p.cfg.CellID}}) + } + + result, err := p.client.DescribeInstances(ctx, &ec2.DescribeInstancesInput{Filters: filters}) + if err != nil { + return ocfs2Assignment{}, fmt.Errorf("ec2: describe workers for OCFS2 slot allocation: %w", err) + } + for _, res := range result.Reservations { + for _, inst := range res.Instances { + if inst.PrivateIpAddress != nil { + if slot, ok := ipToSlot[aws.ToString(inst.PrivateIpAddress)]; ok { + used[slot] = true + } + } + for _, tag := range inst.Tags { + if aws.ToString(tag.Key) != awsTagOCFS2Slot { + continue + } + slot, convErr := strconv.Atoi(aws.ToString(tag.Value)) + if convErr == nil && slot >= 0 && slot < len(p.cfg.OCFS2NodeIPs) { + used[slot] = true + } + } + } + } + + for slot, ip := range p.cfg.OCFS2NodeIPs { + if used[slot] { + continue + } + return ocfs2Assignment{ + Enabled: true, + Slot: slot, + IP: strings.TrimSpace(ip), + NodeIPs: append([]string(nil), p.cfg.OCFS2NodeIPs...), + }, nil + } + + return ocfs2Assignment{}, fmt.Errorf("ec2: no free OCFS2 node slots available (%d configured)", len(p.cfg.OCFS2NodeIPs)) +} + // CleanupOrphanedResources reclaims ENIs and EBS volumes left by failed // VM creates. Mirrors the AzurePool's NIC/disk cleanup. // @@ -417,39 +549,63 @@ func (p *EC2Pool) instanceToMachine(inst *ec2types.Instance) *Machine { // buildUserData returns the EC2 instance user-data script. Combines the // CP-supplied WorkerSpec with EC2-specific cloud-init (NVMe instance-store // mount, AMI-baked rootfs copy, machine-id stamping). -func (p *EC2Pool) buildUserData(opts MachineOpts) string { +func (p *EC2Pool) buildUserData(opts MachineOpts, ocfs2 ocfs2Assignment) string { _ = opts // opts.Region/Size honored at instance launch; cloud-init is cell-uniform var sb strings.Builder sb.WriteString("#!/bin/bash\nset -euo pipefail\n\n") + sb.WriteString("oc_boot_log() { echo \"opensandbox-worker-bootstrap $(date -Is) $*\"; }\n") + sb.WriteString("oc_boot_log 'user-data start'\n\n") + sb.WriteString("systemctl stop opensandbox-worker.service 2>/dev/null || true\n") + sb.WriteString("systemctl disable opensandbox-worker.service 2>/dev/null || true\n") + sb.WriteString("systemctl reset-failed opensandbox-worker.service 2>/dev/null || true\n\n") + sb.WriteString("systemctl stop opensandbox-server.service 2>/dev/null || true\n") + sb.WriteString("systemctl disable opensandbox-server.service 2>/dev/null || true\n") + sb.WriteString("systemctl reset-failed opensandbox-server.service 2>/dev/null || true\n\n") + + sb.WriteString("# Instance identity from EC2 metadata (IMDSv2)\n") + sb.WriteString("TOKEN=$(curl -fsS -X PUT 'http://169.254.169.254/latest/api/token' -H 'X-aws-ec2-metadata-token-ttl-seconds: 300')\n") + sb.WriteString("MY_IP=$(curl -fsS -H \"X-aws-ec2-metadata-token: $TOKEN\" http://169.254.169.254/latest/meta-data/local-ipv4)\n") + sb.WriteString("INSTANCE_ID=$(curl -fsS -H \"X-aws-ec2-metadata-token: $TOKEN\" http://169.254.169.254/latest/meta-data/instance-id)\n") + sb.WriteString("WORKER_ID=\"w-aws-${INSTANCE_ID}\"\n\n") + sb.WriteString("oc_boot_log \"instance identity ready: $INSTANCE_ID $MY_IP\"\n\n") // NVMe instance store handling. Larger metal/x.gd instance families expose // multiple NVMe drives at /dev/nvme[1-N]n1; smaller instances rely on EBS // (the attached data volume). RAID 0 across instance store NVMe when present. - sb.WriteString("# Mount data: prefer NVMe instance store (RAID 0), else first EBS data volume\n") + sb.WriteString("# Mount data: prefer EC2 instance-store NVMe (RAID 0). Otherwise use root fs for /data.\n") sb.WriteString("if ! mountpoint -q /data 2>/dev/null; then\n") sb.WriteString(" mkdir -p /data\n") sb.WriteString(" ROOT_DEV=$(lsblk -no PKNAME $(findmnt -n -o SOURCE /) 2>/dev/null | head -1)\n") sb.WriteString(" NVME_DISKS=()\n") - sb.WriteString(" for d in /dev/nvme1n1 /dev/nvme2n1 /dev/nvme3n1 /dev/nvme4n1 /dev/nvme5n1; do\n") - sb.WriteString(" [ -b \"$d\" ] || continue\n") - sb.WriteString(" [ \"$(basename $d)\" = \"$ROOT_DEV\" ] && continue\n") - sb.WriteString(" NVME_DISKS+=(\"$d\")\n") - sb.WriteString(" done\n") + sb.WriteString(" while read -r name model; do\n") + sb.WriteString(" [ -n \"${name:-}\" ] || continue\n") + sb.WriteString(" [ \"$name\" = \"$ROOT_DEV\" ] && continue\n") + sb.WriteString(" [ \"$model\" = \"Amazon EC2 NVMe Instance Storage\" ] || continue\n") + sb.WriteString(" NVME_DISKS+=(\"/dev/$name\")\n") + sb.WriteString(" done < <(lsblk -dn -o NAME,MODEL)\n") + sb.WriteString(" if [ ${#NVME_DISKS[@]} -eq 0 ]; then\n") + sb.WriteString(" echo 'No EC2 instance-store NVMe found; using root filesystem for /data'\n") + sb.WriteString(" fi\n") sb.WriteString(" if [ ${#NVME_DISKS[@]} -gt 1 ]; then\n") sb.WriteString(" mdadm --create /dev/md0 --level=0 --raid-devices=${#NVME_DISKS[@]} \"${NVME_DISKS[@]}\" --run --force\n") sb.WriteString(" mkfs.xfs -f -m reflink=1 /dev/md0 && mount /dev/md0 /data\n") sb.WriteString(" elif [ ${#NVME_DISKS[@]} -eq 1 ]; then\n") sb.WriteString(" mkfs.xfs -f -m reflink=1 \"${NVME_DISKS[0]}\" && mount \"${NVME_DISKS[0]}\" /data\n") - sb.WriteString(" else\n") - sb.WriteString(" for d in /dev/nvme1n1 /dev/sdb /dev/xvdb; do\n") - sb.WriteString(" [ -b \"$d\" ] || continue\n") - sb.WriteString(" mkfs.xfs -f -m reflink=1 \"$d\" && mount \"$d\" /data && break\n") - sb.WriteString(" done\n") sb.WriteString(" fi\n") sb.WriteString("fi\n") sb.WriteString("mkdir -p /data/sandboxes /data/firecracker/images\n") + sb.WriteString("oc_boot_log 'base data mount ready'\n\n") + + if p.cfg.SharedSandboxDataVolumeID != "" { + sb.WriteString(p.sharedSandboxDataUserData(ocfs2)) + } + if p.cfg.SharedGoldensVolumeID != "" { + sb.WriteString(p.sharedGoldensUserData()) + } + sb.WriteString("# Copy AMI-baked rootfs images to data disk if not already present\n") sb.WriteString("if [ -d /opt/opensandbox/images ] && [ ! -f /data/firecracker/images/default.ext4 ]; then\n") + sb.WriteString(" oc_boot_log 'copying AMI-baked rootfs to data disk'\n") sb.WriteString(" cp /opt/opensandbox/images/*.ext4 /data/firecracker/images/ 2>/dev/null || true\n") sb.WriteString("fi\n") sb.WriteString("if [ -d /opt/opensandbox/images/bases ] && [ ! -d /data/firecracker/images/bases ]; then\n") @@ -467,10 +623,6 @@ func (p *EC2Pool) buildUserData(opts MachineOpts) string { sb.WriteString(fmt.Sprintf("echo '%s' | base64 -d > /etc/opensandbox/worker.env\n\n", envB64)) sb.WriteString("# Patch worker identity from EC2 instance metadata (IMDSv2)\n") - sb.WriteString("TOKEN=$(curl -s -X PUT 'http://169.254.169.254/latest/api/token' -H 'X-aws-ec2-metadata-token-ttl-seconds: 300')\n") - sb.WriteString("MY_IP=$(curl -s -H \"X-aws-ec2-metadata-token: $TOKEN\" http://169.254.169.254/latest/meta-data/local-ipv4)\n") - sb.WriteString("INSTANCE_ID=$(curl -s -H \"X-aws-ec2-metadata-token: $TOKEN\" http://169.254.169.254/latest/meta-data/instance-id)\n") - sb.WriteString("WORKER_ID=\"w-aws-${INSTANCE_ID}\"\n") sb.WriteString("sed -i \"s|OPENSANDBOX_GRPC_ADVERTISE=.*|OPENSANDBOX_GRPC_ADVERTISE=${MY_IP}:9090|\" /etc/opensandbox/worker.env\n") sb.WriteString("sed -i \"s|OPENSANDBOX_HTTP_ADDR=.*|OPENSANDBOX_HTTP_ADDR=http://${MY_IP}:8081|\" /etc/opensandbox/worker.env\n") sb.WriteString("sed -i \"s|OPENSANDBOX_WORKER_ID=.*|OPENSANDBOX_WORKER_ID=${WORKER_ID}|\" /etc/opensandbox/worker.env\n") @@ -481,8 +633,150 @@ func (p *EC2Pool) buildUserData(opts MachineOpts) string { sb.WriteString("rm -rf /data/sandboxes/golden-snapshot /data/sandboxes/golden\n\n") // Start worker + sb.WriteString("oc_boot_log 'starting opensandbox-worker service'\n") sb.WriteString("systemctl restart opensandbox-worker\n") + sb.WriteString("oc_boot_log 'user-data complete'\n") return sb.String() } +func (p *EC2Pool) sharedSandboxDataUserData(ocfs2 ocfs2Assignment) string { + clusterName := p.cfg.OCFS2ClusterName + if clusterName == "" { + clusterName = "opensandbox" + } + expectedNodes := p.cfg.OCFS2ExpectedNodes + if expectedNodes <= 0 { + expectedNodes = 1 + } + maxNodes := p.cfg.OCFS2MaxNodes + if maxNodes <= 0 { + maxNodes = expectedNodes + } + if ocfs2.Enabled && maxNodes < len(ocfs2.NodeIPs) { + maxNodes = len(ocfs2.NodeIPs) + } + if maxNodes < expectedNodes { + maxNodes = expectedNodes + } + + var sb strings.Builder + sb.WriteString("# Shared sandbox data: OCFS2 over io2 Multi-Attach\n") + sb.WriteString("oc_boot_log 'validating baked OCFS2 dependencies'\n") + sb.WriteString("command -v mount.ocfs2 >/dev/null 2>&1 || { echo 'ERROR: AMI missing ocfs2-tools; rebuild worker AMI'; exit 1; }\n") + sb.WriteString("modprobe ocfs2 || { echo 'ERROR: AMI missing ocfs2 kernel module; rebuild worker AMI with linux-modules-extra'; exit 1; }\n") + sb.WriteString("modprobe ocfs2_dlmfs || { echo 'ERROR: AMI missing ocfs2_dlmfs kernel module; rebuild worker AMI'; exit 1; }\n") + sb.WriteString("modprobe ocfs2_stack_o2cb || { echo 'ERROR: AMI missing ocfs2_stack_o2cb kernel module; rebuild worker AMI'; exit 1; }\n") + sb.WriteString(fmt.Sprintf("SANDBOX_VOLUME_ID=%q\n", p.cfg.SharedSandboxDataVolumeID)) + sb.WriteString(fmt.Sprintf("OCFS2_CLUSTER_NAME=%q\n", clusterName)) + sb.WriteString(fmt.Sprintf("OCFS2_EXPECTED_NODES=%d\n", expectedNodes)) + sb.WriteString(fmt.Sprintf("OCFS2_MAX_NODES=%d\n", maxNodes)) + if ocfs2.Enabled { + sb.WriteString(fmt.Sprintf("OCFS2_NODE_SLOT=%d\n", ocfs2.Slot)) + sb.WriteString(fmt.Sprintf("OCFS2_NODE_IP=%q\n", ocfs2.IP)) + } + sb.WriteString("oc_boot_log \"attaching shared sandbox data volume $SANDBOX_VOLUME_ID\"\n") + sb.WriteString("aws ec2 attach-volume --region " + shellQuote(p.cfg.Region) + " --volume-id \"$SANDBOX_VOLUME_ID\" --instance-id \"$INSTANCE_ID\" --device /dev/sdg || true\n") + sb.WriteString("SANDBOX_DEV=\"\"\n") + sb.WriteString("SANDBOX_VOL_NO_DASH=\"${SANDBOX_VOLUME_ID//-/}\"\n") + sb.WriteString("for i in $(seq 1 180); do\n") + sb.WriteString(" if [ -e \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${SANDBOX_VOL_NO_DASH}\" ]; then\n") + sb.WriteString(" SANDBOX_DEV=$(readlink -f \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${SANDBOX_VOL_NO_DASH}\")\n") + sb.WriteString(" elif [ -e \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${SANDBOX_VOL_NO_DASH}_1\" ]; then\n") + sb.WriteString(" SANDBOX_DEV=$(readlink -f \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${SANDBOX_VOL_NO_DASH}_1\")\n") + sb.WriteString(" else\n") + sb.WriteString(" SANDBOX_DEV=$(lsblk -dn -o NAME,SERIAL | awk -v v=\"$SANDBOX_VOL_NO_DASH\" '$2 == v {print \"/dev/\"$1; exit}')\n") + sb.WriteString(" fi\n") + sb.WriteString(" [ -n \"${SANDBOX_DEV:-}\" ] && break\n") + sb.WriteString(" sleep 1\n") + sb.WriteString("done\n") + sb.WriteString("if [ -z \"${SANDBOX_DEV:-}\" ]; then echo \"ERROR: shared sandbox data volume not attached\"; lsblk -o NAME,MODEL,SERIAL,SIZE,FSTYPE,MOUNTPOINT || true; exit 1; fi\n") + sb.WriteString("SANDBOX_SERIAL=$(lsblk -dn -o SERIAL \"$SANDBOX_DEV\" 2>/dev/null | head -1 || true)\n") + sb.WriteString("if [ \"$SANDBOX_SERIAL\" != \"$SANDBOX_VOL_NO_DASH\" ]; then echo \"ERROR: $SANDBOX_DEV serial $SANDBOX_SERIAL does not match sandbox volume $SANDBOX_VOLUME_ID\"; lsblk -o NAME,MODEL,SERIAL,SIZE,FSTYPE,MOUNTPOINT || true; exit 1; fi\n") + sb.WriteString("echo \"Using shared sandbox data volume $SANDBOX_VOLUME_ID at $SANDBOX_DEV\"\n") + sb.WriteString("oc_boot_log \"shared sandbox data volume visible at $SANDBOX_DEV\"\n") + sb.WriteString("install -d -m 0755 /etc/ocfs2 /etc/sysconfig\n") + if ocfs2.Enabled { + sb.WriteString("oc_boot_log \"using static OCFS2 slot $OCFS2_NODE_SLOT at $OCFS2_NODE_IP\"\n") + sb.WriteString("OCFS2_NODE_NAMES=()\n") + sb.WriteString("OCFS2_NODE_IPS=()\n") + for _, ip := range ocfs2.NodeIPs { + ip = strings.TrimSpace(ip) + sb.WriteString(fmt.Sprintf("OCFS2_NODE_NAMES+=(%q)\n", awsPrivateDNSShortName(ip))) + sb.WriteString(fmt.Sprintf("OCFS2_NODE_IPS+=(%q)\n", ip)) + } + sb.WriteString("{ echo \"cluster:\"; echo \" node_count = ${#OCFS2_NODE_IPS[@]}\"; echo \" name = $OCFS2_CLUSTER_NAME\"; echo \"\"; for i in \"${!OCFS2_NODE_IPS[@]}\"; do echo \"node:\"; echo \" ip_port = 7777\"; echo \" ip_address = ${OCFS2_NODE_IPS[$i]}\"; echo \" number = $i\"; echo \" name = ${OCFS2_NODE_NAMES[$i]}\"; echo \" cluster = $OCFS2_CLUSTER_NAME\"; echo \"\"; done; } > /etc/ocfs2/cluster.conf\n") + } else { + sb.WriteString("oc_boot_log 'discovering OCFS2 peer nodes'\n") + sb.WriteString("mapfile -t OCFS2_NODES < <(for i in $(seq 1 60); do aws ec2 describe-instances --region " + shellQuote(p.cfg.Region) + " --filters \"Name=tag:Cell,Values=" + shellEscapedDouble(p.cfg.CellID) + "\" \"Name=tag:Role,Values=worker\" \"Name=instance-state-name,Values=running\" --query 'Reservations[].Instances[].PrivateDnsName' --output text | tr '\\t' '\\n' | awk 'NF { sub(/\\..*/, \"\", $0); print }' | sort -u; break; done)\n") + sb.WriteString("for i in $(seq 1 60); do\n") + sb.WriteString(" [ \"${#OCFS2_NODES[@]}\" -ge \"$OCFS2_EXPECTED_NODES\" ] && break\n") + sb.WriteString(" sleep 2\n") + sb.WriteString(" mapfile -t OCFS2_NODES < <(aws ec2 describe-instances --region " + shellQuote(p.cfg.Region) + " --filters \"Name=tag:Cell,Values=" + shellEscapedDouble(p.cfg.CellID) + "\" \"Name=tag:Role,Values=worker\" \"Name=instance-state-name,Values=running\" --query 'Reservations[].Instances[].PrivateDnsName' --output text | tr '\\t' '\\n' | awk 'NF { sub(/\\..*/, \"\", $0); print }' | sort -u)\n") + sb.WriteString("done\n") + sb.WriteString("if [ \"${#OCFS2_NODES[@]}\" -lt \"$OCFS2_EXPECTED_NODES\" ]; then echo \"ERROR: found ${#OCFS2_NODES[@]} OCFS2 nodes, expected $OCFS2_EXPECTED_NODES\"; exit 1; fi\n") + sb.WriteString("oc_boot_log \"OCFS2 peer nodes: ${OCFS2_NODES[*]}\"\n") + sb.WriteString("{ echo \"cluster:\"; echo \" node_count = ${#OCFS2_NODES[@]}\"; echo \" name = $OCFS2_CLUSTER_NAME\"; echo \"\"; n=0; for node in \"${OCFS2_NODES[@]}\"; do ip=$(getent ahostsv4 \"$node\" | awk '{print $1; exit}'); [ -n \"${ip:-}\" ] || { echo \"ERROR: could not resolve OCFS2 node $node\"; exit 1; }; echo \"node:\"; echo \" ip_port = 7777\"; echo \" ip_address = $ip\"; echo \" number = $n\"; echo \" name = $node\"; echo \" cluster = $OCFS2_CLUSTER_NAME\"; echo \"\"; n=$((n + 1)); done; } > /etc/ocfs2/cluster.conf\n") + } + sb.WriteString("cat > /etc/default/o2cb </dev/null 2>&1 && o2cb register-cluster \"$OCFS2_CLUSTER_NAME\" || true\n") + sb.WriteString("[ -x /etc/init.d/o2cb ] && /etc/init.d/o2cb online \"$OCFS2_CLUSTER_NAME\" || true\n") + sb.WriteString("mkdir -p /data/sandboxes\n") + sb.WriteString("FSTYPE=$(blkid -s TYPE -o value \"$SANDBOX_DEV\" 2>/dev/null || true)\n") + sb.WriteString("if [ -z \"$FSTYPE\" ]; then mkfs.ocfs2 -F -N \"$OCFS2_MAX_NODES\" -L opensandbox-sandboxes -T vmstore \"$SANDBOX_DEV\"; fi\n") + sb.WriteString("FSTYPE=$(blkid -s TYPE -o value \"$SANDBOX_DEV\" 2>/dev/null || true)\n") + sb.WriteString("if [ \"$FSTYPE\" != \"ocfs2\" ]; then echo \"ERROR: shared sandbox data volume $SANDBOX_DEV has filesystem '$FSTYPE', expected ocfs2\"; lsblk -o NAME,MODEL,SERIAL,SIZE,FSTYPE,MOUNTPOINT || true; exit 1; fi\n") + sb.WriteString("if ! grep -q 'LABEL=opensandbox-sandboxes' /etc/fstab; then echo 'LABEL=opensandbox-sandboxes /data/sandboxes ocfs2 noauto,_netdev,noatime 0 0' >> /etc/fstab; fi\n") + sb.WriteString("oc_boot_log 'mounting OCFS2 shared sandbox data volume'\n") + sb.WriteString("timeout 90 mount -t ocfs2 -o noatime \"$SANDBOX_DEV\" /data/sandboxes\n") + sb.WriteString("oc_boot_log 'OCFS2 shared sandbox data mounted'\n") + sb.WriteString("chown root:root /data/sandboxes\n\n") + return sb.String() +} + +func (p *EC2Pool) sharedGoldensUserData() string { + var sb strings.Builder + sb.WriteString("# Shared golden image volume\n") + sb.WriteString("mkdir -p /opt/opensandbox/goldens-shared /var/lib/opensandbox/golden\n") + sb.WriteString(fmt.Sprintf("GOLDENS_VOLUME_ID=%q\n", p.cfg.SharedGoldensVolumeID)) + sb.WriteString("aws ec2 attach-volume --region " + shellQuote(p.cfg.Region) + " --volume-id \"$GOLDENS_VOLUME_ID\" --instance-id \"$INSTANCE_ID\" --device /dev/sdf || true\n") + sb.WriteString("GOLDENS_DEV=\"\"\n") + sb.WriteString("GOLDENS_VOL_NO_DASH=\"${GOLDENS_VOLUME_ID//-/}\"\n") + sb.WriteString("for i in $(seq 1 120); do\n") + sb.WriteString(" if [ -e \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${GOLDENS_VOL_NO_DASH}\" ]; then GOLDENS_DEV=$(readlink -f \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${GOLDENS_VOL_NO_DASH}\"); fi\n") + sb.WriteString(" if [ -z \"${GOLDENS_DEV:-}\" ] && [ -e \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${GOLDENS_VOL_NO_DASH}_1\" ]; then GOLDENS_DEV=$(readlink -f \"/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${GOLDENS_VOL_NO_DASH}_1\"); fi\n") + sb.WriteString(" if [ -z \"${GOLDENS_DEV:-}\" ]; then GOLDENS_DEV=$(lsblk -dn -o NAME,SERIAL | awk -v v=\"$GOLDENS_VOL_NO_DASH\" '$2 == v {print \"/dev/\"$1; exit}'); fi\n") + sb.WriteString(" [ -n \"${GOLDENS_DEV:-}\" ] && break\n") + sb.WriteString(" sleep 1\n") + sb.WriteString("done\n") + sb.WriteString("if [ -n \"${GOLDENS_DEV:-}\" ]; then\n") + sb.WriteString(" GOLDENS_SERIAL=$(lsblk -dn -o SERIAL \"$GOLDENS_DEV\" 2>/dev/null | head -1 || true)\n") + sb.WriteString(" if [ \"$GOLDENS_SERIAL\" != \"$GOLDENS_VOL_NO_DASH\" ]; then echo \"WARN: $GOLDENS_DEV serial $GOLDENS_SERIAL does not match golden volume $GOLDENS_VOLUME_ID\"; GOLDENS_DEV=\"\"; fi\n") + sb.WriteString("fi\n") + sb.WriteString("if [ -n \"${GOLDENS_DEV:-}\" ]; then\n") + sb.WriteString(" GOLDENS_FSTYPE=$(blkid -s TYPE -o value \"$GOLDENS_DEV\" 2>/dev/null || true)\n") + sb.WriteString(" case \"$GOLDENS_FSTYPE\" in\n") + sb.WriteString(" ext2|ext3|ext4) mount -t \"$GOLDENS_FSTYPE\" -o ro,noload,noatime \"$GOLDENS_DEV\" /opt/opensandbox/goldens-shared || true ;;\n") + sb.WriteString(" xfs) mount -t xfs -o ro,noatime \"$GOLDENS_DEV\" /opt/opensandbox/goldens-shared || true ;;\n") + sb.WriteString(" '') echo \"WARN: shared golden volume $GOLDENS_VOLUME_ID has no filesystem; continuing without it\" ;;\n") + sb.WriteString(" *) echo \"WARN: shared golden volume $GOLDENS_VOLUME_ID has unsupported filesystem '$GOLDENS_FSTYPE'; continuing without it\" ;;\n") + sb.WriteString(" esac\n") + sb.WriteString("fi\n") + sb.WriteString("if [ -d /opt/opensandbox/goldens-shared/golden ]; then ln -sfn /opt/opensandbox/goldens-shared/golden /var/lib/opensandbox/golden; fi\n\n") + return sb.String() +} + +func shellQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", "'\"'\"'") + "'" +} + +func shellEscapedDouble(s string) string { + return strings.ReplaceAll(s, `"`, `\"`) +} + +func awsPrivateDNSShortName(ip string) string { + return "ip-" + strings.ReplaceAll(strings.TrimSpace(ip), ".", "-") +} diff --git a/internal/compute/ec2_ocfs2_test.go b/internal/compute/ec2_ocfs2_test.go new file mode 100644 index 00000000..f8229967 --- /dev/null +++ b/internal/compute/ec2_ocfs2_test.go @@ -0,0 +1,48 @@ +package compute + +import ( + "strings" + "testing" +) + +func TestEC2UserDataStaticOCFS2Slots(t *testing.T) { + p := &EC2Pool{cfg: EC2PoolConfig{ + Region: "us-east-2", + CellID: "aws-us-east-2-burst-prod", + SharedSandboxDataVolumeID: "vol-123", + OCFS2ClusterName: "opensandbox", + OCFS2MaxNodes: 4, + }} + + userData := p.buildUserData(MachineOpts{}, ocfs2Assignment{ + Enabled: true, + Slot: 1, + IP: "10.60.1.11", + NodeIPs: []string{"10.60.1.10", "10.60.1.11", "10.60.1.12", "10.60.1.13"}, + }) + + for _, want := range []string{ + `OCFS2_NODE_SLOT=1`, + `OCFS2_NODE_IP="10.60.1.11"`, + `OCFS2_NODE_NAMES+=("ip-10-60-1-10")`, + `OCFS2_NODE_NAMES+=("ip-10-60-1-13")`, + `OCFS2_NODE_IPS+=("10.60.1.12")`, + `node_count = ${#OCFS2_NODE_IPS[@]}`, + `mkfs.ocfs2 -F -N "$OCFS2_MAX_NODES"`, + } { + if !strings.Contains(userData, want) { + t.Fatalf("user-data missing %q", want) + } + } + + if strings.Contains(userData, "discovering OCFS2 peer nodes") { + t.Fatal("static OCFS2 user-data should not discover live EC2 peers") + } +} + +func TestAWSPrivateDNSShortName(t *testing.T) { + got := awsPrivateDNSShortName("10.60.1.122") + if got != "ip-10-60-1-122" { + t.Fatalf("awsPrivateDNSShortName() = %q", got) + } +} diff --git a/internal/compute/pool.go b/internal/compute/pool.go index 901388cf..a974a859 100644 --- a/internal/compute/pool.go +++ b/internal/compute/pool.go @@ -40,8 +40,9 @@ type MachineOpts struct { // cloud means writing a new Pool; the CP doesn't change. type WorkerSpec struct { // Cell + region identity - CellID string // "{cloud}-{region}-{slot}", e.g. "azure-us-west-2-b" - Region string // AWS-style hyphenated, e.g. "us-west-2", "us-east-1" + CellID string // "{cloud}-{region}-{slot}", e.g. "azure-us-west-2-b" + Region string // AWS-style hyphenated, e.g. "us-west-2", "us-east-1" + WorkerPool string // "ondemand" or "burst"; empty means ondemand // Connectivity back to the control plane DatabaseURL string @@ -128,6 +129,7 @@ func BuildWorkerEnv(spec WorkerSpec) string { {"OPENSANDBOX_JWT_SECRET", spec.JWTSecret}, {"OPENSANDBOX_REGION", spec.Region}, {"OPENSANDBOX_CELL_ID", spec.CellID}, + {"OPENSANDBOX_WORKER_POOL", spec.WorkerPool}, {"OPENSANDBOX_MAX_CAPACITY", itoa(spec.MaxCapacity)}, {"OPENSANDBOX_DEFAULT_SANDBOX_MEMORY_MB", itoa(spec.DefaultMemoryMB)}, {"OPENSANDBOX_DEFAULT_SANDBOX_CPUS", itoa(spec.DefaultCPUs)}, diff --git a/internal/compute/quota_test.go b/internal/compute/quota_test.go index 422bd84c..f01c8ac8 100644 --- a/internal/compute/quota_test.go +++ b/internal/compute/quota_test.go @@ -60,6 +60,30 @@ func TestIsEC2QuotaErr(t *testing.T) { } } +func TestSupportsEC2NestedVirtualization(t *testing.T) { + tests := []struct { + instanceType string + want bool + }{ + {"r8i.48xlarge", true}, + {"m8i.4xlarge", true}, + {"c8i.2xlarge", true}, + {"R8I.48XLARGE", true}, + {"r5d.metal", false}, + {"c7gd.metal", false}, + {"m7i.large", false}, + {"", false}, + } + + for _, tc := range tests { + t.Run(tc.instanceType, func(t *testing.T) { + if got := supportsEC2NestedVirtualization(tc.instanceType); got != tc.want { + t.Fatalf("supportsEC2NestedVirtualization(%q) = %v, want %v", tc.instanceType, got, tc.want) + } + }) + } +} + func TestWrapAzureCreateErrTagsQuota(t *testing.T) { quotaSrc := errors.New("AllocationFailed: no capacity") wrapped := wrapAzureCreateErr(quotaSrc, "azure: create VM foo failed: %w", quotaSrc) diff --git a/internal/config/config.go b/internal/config/config.go index 8a314487..9a21293d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -102,28 +102,39 @@ type Config struct { // QEMU VM configuration (worker mode) KernelPath string // Path to vmlinux kernel ImagesDir string // Path to base rootfs images + GoldenDir string // Path to worker-local golden VM snapshot cache QEMUBin string // Path to qemu-system binary (default: "qemu-system-x86_64") + WorkerPool string // Worker placement pool advertised in heartbeats: "ondemand" (default) or "burst" // AWS EC2 compute pool (server mode only — for auto-scaling worker machines) - EC2AMI string // Custom AMI for worker instances - EC2InstanceType string // single fallback type; used only when EC2InstanceTypes is empty - EC2InstanceTypes []string // ranked list of instance types tried in order on quota/capacity errors - EC2SubnetID string // VPC subnet for worker instances - EC2SecurityGroupID string // Security group (allow 8080, 9090, 9091) - EC2KeyName string // SSH key pair name (for debugging) - EC2WorkerImage string // Docker image for containerized workers - EC2IAMInstanceProfile string // IAM instance profile for worker instances (Secrets Manager + S3) - EC2SSMParameterName string // SSM parameter name for dynamic AMI ID (e.g. /opensandbox/prod/worker-ami-id) + EC2AMI string // Custom AMI for worker instances + EC2InstanceType string // single fallback type; used only when EC2InstanceTypes is empty + EC2InstanceTypes []string // ranked list of instance types tried in order on quota/capacity errors + EC2SubnetID string // VPC subnet for worker instances + EC2SecurityGroupID string // Security group (allow 8080, 9090, 9091) + EC2KeyName string // SSH key pair name (for debugging) + EC2WorkerImage string // Docker image for containerized workers + EC2IAMInstanceProfile string // IAM instance profile for worker instances (Secrets Manager + S3) + EC2SSMParameterName string // SSM parameter name for dynamic AMI ID (e.g. /opensandbox/prod/worker-ami-id) + EC2MarketType string // empty/on-demand or spot + EC2AccessKeyID string // optional EC2-specific static credentials; empty = IAM/default chain + EC2SecretAccessKey string + EC2SharedSandboxDataVolumeID string // optional io2 Multi-Attach volume mounted at /data/sandboxes + EC2SharedGoldensVolumeID string // optional io2 Multi-Attach volume mounted read-only for golden images + EC2OCFS2ClusterName string + EC2OCFS2ExpectedNodes int + EC2OCFS2MaxNodes int + EC2OCFS2NodeIPs []string // Azure compute pool (server mode — for auto-scaling worker VMs) - AzureSubscriptionID string // Azure subscription ID - AzureResourceGroup string // Resource group for worker VMs - AzureVMSize string // single fallback size; used only when AzureVMSizes is empty + AzureSubscriptionID string // Azure subscription ID + AzureResourceGroup string // Resource group for worker VMs + AzureVMSize string // single fallback size; used only when AzureVMSizes is empty AzureVMSizes []string // ranked list of VM sizes tried in order on quota/capacity errors - AzureImageID string // Custom image ID or URN - AzureSubnetID string // Full resource ID of the VNet subnet - AzureSSHPublicKey string // SSH public key for worker VMs - AzureKeyVaultName string // Key Vault name for dynamic image ID refresh (e.g. "opensandbox-prod") + AzureImageID string // Custom image ID or URN + AzureSubnetID string // Full resource ID of the VNet subnet + AzureSSHPublicKey string // SSH public key for worker VMs + AzureKeyVaultName string // Key Vault name for dynamic image ID refresh (e.g. "opensandbox-prod") // AzureWorkerIdentityID is the full resource ID of a UserAssigned managed // identity to attach to every worker VM. The identity must already have // "Key Vault Secrets Officer" on the regional KV so workers can fetch @@ -143,6 +154,8 @@ type Config struct { MinWorkersPerRegion int // Minimum total workers per region, default 1 MaxWorkersPerRegion int // Maximum workers per region (hard cap), default 10 IdleReserveWorkers int // Target idle workers for burst absorption, default 1 + MinIdleCapacity int // Minimum spare sandbox capacity slots per region. When >0, overrides MinWorkers/IdleReserve. + MinIdleCPUs int // Minimum spare sandbox CPU units per region. When >0, overrides MinIdleCapacity. // Stripe billing StripeSecretKey string @@ -171,6 +184,25 @@ type Config struct { // Env vars take precedence over secret values (for local overrides). SecretsARN string + // AWSSecretsPrefix is the per-cell secret name prefix in AWS Secrets + // Manager, e.g. "opencomputer/aws-us-east-2-poc/". When set, LoadSecrets + // selects the AWS Secrets Manager provider and lists+loads everything + // under this prefix. Mirrors AzureKeyVaultName for the AWS path. + AWSSecretsPrefix string + + // Cloud is a short identifier for the cloud the binary is running on + // ("aws", "azure", or empty for local/dev). Set by cloud-init from + // terraform — preferred over runtime probing for fast startup. + Cloud string + + // CPUOvercommitRatio multiplies the worker's physical vCPU count when + // publishing capacity in the Redis heartbeat. 1 = no overcommit + // (matches Azure default). 2 = recommended for agent workloads on AWS + // — roughly halves $/sandbox without measurable UX impact for + // LLM-bound sandboxes. Memory is never overcommitted because host OOM + // is a worse failure mode than CPU contention. + CPUOvercommitRatio int + // Secret encryption key (hex-encoded 32 bytes / 64 hex chars) for encrypting // project secrets at rest in PostgreSQL. Required if using project secrets. SecretEncryptionKey string @@ -294,7 +326,7 @@ func Load() (*Config, error) { WorkOSCookieDomain: os.Getenv("WORKOS_COOKIE_DOMAIN"), WorkOSFrontendURL: os.Getenv("WORKOS_FRONTEND_URL"), - RedisURL: os.Getenv("OPENSANDBOX_REDIS_URL"), + RedisURL: os.Getenv("OPENSANDBOX_REDIS_URL"), MaxCapacity: envOrDefaultInt("OPENSANDBOX_MAX_CAPACITY", 50), @@ -320,28 +352,42 @@ func Load() (*Config, error) { DefaultSandboxCPUs: envOrDefaultInt("OPENSANDBOX_DEFAULT_SANDBOX_CPUS", 1), DefaultSandboxDiskMB: envOrDefaultInt("OPENSANDBOX_DEFAULT_SANDBOX_DISK_MB", 0), - KernelPath: os.Getenv("OPENSANDBOX_KERNEL_PATH"), // default derived from DataDir - ImagesDir: os.Getenv("OPENSANDBOX_IMAGES_DIR"), // default derived from DataDir - QEMUBin: envOrDefault("OPENSANDBOX_QEMU_BIN", "qemu-system-x86_64"), - - EC2AMI: os.Getenv("OPENSANDBOX_EC2_AMI"), - EC2InstanceType: envOrDefault("OPENSANDBOX_EC2_INSTANCE_TYPE", "c7gd.metal"), - EC2InstanceTypes: splitCSV(os.Getenv("OPENSANDBOX_EC2_INSTANCE_TYPES")), - EC2SubnetID: os.Getenv("OPENSANDBOX_EC2_SUBNET_ID"), - EC2SecurityGroupID: os.Getenv("OPENSANDBOX_EC2_SECURITY_GROUP_ID"), - EC2KeyName: os.Getenv("OPENSANDBOX_EC2_KEY_NAME"), - EC2WorkerImage: envOrDefault("OPENSANDBOX_EC2_WORKER_IMAGE", "opensandbox-worker:latest"), - EC2IAMInstanceProfile: os.Getenv("OPENSANDBOX_EC2_IAM_INSTANCE_PROFILE"), - EC2SSMParameterName: os.Getenv("OPENSANDBOX_EC2_SSM_AMI_PARAM"), - - AzureSubscriptionID: os.Getenv("OPENSANDBOX_AZURE_SUBSCRIPTION_ID"), - AzureResourceGroup: os.Getenv("OPENSANDBOX_AZURE_RESOURCE_GROUP"), - AzureVMSize: envOrDefault("OPENSANDBOX_AZURE_VM_SIZE", "Standard_D16s_v5"), - AzureVMSizes: splitCSV(os.Getenv("OPENSANDBOX_AZURE_VM_SIZES")), - AzureImageID: os.Getenv("OPENSANDBOX_AZURE_IMAGE_ID"), - AzureSubnetID: os.Getenv("OPENSANDBOX_AZURE_SUBNET_ID"), - AzureSSHPublicKey: os.Getenv("OPENSANDBOX_AZURE_SSH_PUBLIC_KEY"), - AzureKeyVaultName: os.Getenv("OPENSANDBOX_AZURE_KEY_VAULT_NAME"), + KernelPath: os.Getenv("OPENSANDBOX_KERNEL_PATH"), // default derived from DataDir + ImagesDir: os.Getenv("OPENSANDBOX_IMAGES_DIR"), // default derived from DataDir + GoldenDir: os.Getenv("OPENSANDBOX_GOLDEN_DIR"), // default derived from DataDir + QEMUBin: envOrDefault("OPENSANDBOX_QEMU_BIN", "qemu-system-x86_64"), + WorkerPool: envOrDefault("OPENSANDBOX_WORKER_POOL", "ondemand"), + + EC2AMI: os.Getenv("OPENSANDBOX_EC2_AMI"), + EC2InstanceType: envOrDefault("OPENSANDBOX_EC2_INSTANCE_TYPE", "c7gd.metal"), + EC2InstanceTypes: splitCSV(os.Getenv("OPENSANDBOX_EC2_INSTANCE_TYPES")), + EC2SubnetID: os.Getenv("OPENSANDBOX_EC2_SUBNET_ID"), + EC2SecurityGroupID: os.Getenv("OPENSANDBOX_EC2_SECURITY_GROUP_ID"), + EC2KeyName: os.Getenv("OPENSANDBOX_EC2_KEY_NAME"), + EC2WorkerImage: envOrDefault("OPENSANDBOX_EC2_WORKER_IMAGE", "opensandbox-worker:latest"), + EC2IAMInstanceProfile: os.Getenv("OPENSANDBOX_EC2_IAM_INSTANCE_PROFILE"), + EC2SSMParameterName: os.Getenv("OPENSANDBOX_EC2_SSM_AMI_PARAM"), + EC2MarketType: os.Getenv("OPENSANDBOX_EC2_MARKET_TYPE"), + EC2AccessKeyID: os.Getenv("OPENSANDBOX_EC2_ACCESS_KEY_ID"), + EC2SecretAccessKey: os.Getenv("OPENSANDBOX_EC2_SECRET_ACCESS_KEY"), + EC2SharedSandboxDataVolumeID: os.Getenv("OPENSANDBOX_SHARED_SANDBOX_DATA_VOLUME_ID"), + EC2SharedGoldensVolumeID: os.Getenv("OPENSANDBOX_SHARED_GOLDENS_VOLUME_ID"), + EC2OCFS2ClusterName: envOrDefault("OPENSANDBOX_OCFS2_CLUSTER_NAME", "opensandbox"), + EC2OCFS2ExpectedNodes: envOrDefaultInt("OPENSANDBOX_OCFS2_EXPECTED_NODES", 1), + EC2OCFS2MaxNodes: envOrDefaultInt("OPENSANDBOX_OCFS2_MAX_NODES", 4), + EC2OCFS2NodeIPs: splitCSV(os.Getenv("OPENSANDBOX_OCFS2_NODE_IPS")), + + AzureSubscriptionID: os.Getenv("OPENSANDBOX_AZURE_SUBSCRIPTION_ID"), + AzureResourceGroup: os.Getenv("OPENSANDBOX_AZURE_RESOURCE_GROUP"), + AzureVMSize: envOrDefault("OPENSANDBOX_AZURE_VM_SIZE", "Standard_D16s_v5"), + AzureVMSizes: splitCSV(os.Getenv("OPENSANDBOX_AZURE_VM_SIZES")), + AzureImageID: os.Getenv("OPENSANDBOX_AZURE_IMAGE_ID"), + AzureSubnetID: os.Getenv("OPENSANDBOX_AZURE_SUBNET_ID"), + AzureSSHPublicKey: os.Getenv("OPENSANDBOX_AZURE_SSH_PUBLIC_KEY"), + AzureKeyVaultName: os.Getenv("OPENSANDBOX_AZURE_KEY_VAULT_NAME"), + AWSSecretsPrefix: os.Getenv("OPENSANDBOX_AWS_SECRETS_PREFIX"), + Cloud: os.Getenv("OPENSANDBOX_CLOUD"), + CPUOvercommitRatio: envOrDefaultInt("OPENSANDBOX_CPU_OVERCOMMIT_RATIO", 1), AzureWorkerIdentityID: os.Getenv("OPENSANDBOX_AZURE_WORKER_IDENTITY_ID"), CFAPIToken: os.Getenv("OPENSANDBOX_CF_API_TOKEN"), @@ -351,12 +397,14 @@ func Load() (*Config, error) { MinWorkersPerRegion: envOrDefaultInt("OPENSANDBOX_MIN_WORKERS", 1), MaxWorkersPerRegion: envOrDefaultInt("OPENSANDBOX_MAX_WORKERS", 10), IdleReserveWorkers: envOrDefaultInt("OPENSANDBOX_IDLE_RESERVE", 1), + MinIdleCapacity: envOrDefaultInt("OPENSANDBOX_MIN_IDLE_CAPACITY", 0), + MinIdleCPUs: envOrDefaultInt("OPENSANDBOX_MIN_IDLE_CPUS", 0), StripeSecretKey: os.Getenv("STRIPE_SECRET_KEY"), StripeWebhookSecret: os.Getenv("STRIPE_WEBHOOK_SECRET"), StripeTelegramAgentPriceID: os.Getenv("STRIPE_TELEGRAM_AGENT_PRICE_ID"), - StripeSuccessURL: envOrDefault("STRIPE_SUCCESS_URL", "http://localhost:3000/billing?success=true"), - StripeCancelURL: envOrDefault("STRIPE_CANCEL_URL", "http://localhost:3000/billing?cancelled=true"), + StripeSuccessURL: envOrDefault("STRIPE_SUCCESS_URL", "http://localhost:3000/billing?success=true"), + StripeCancelURL: envOrDefault("STRIPE_CANCEL_URL", "http://localhost:3000/billing?cancelled=true"), SegmentWriteKey: os.Getenv("SEGMENT_WRITE_KEY"), @@ -418,6 +466,20 @@ func Load() (*Config, error) { cfg.CellID = cfg.Region + "-default" } + // CPU overcommit: multiply the physical-capacity ceiling by the + // configured ratio so the worker advertises (and the CP places onto) + // the inflated slot count. Ratio < 1 is treated as 1 (the operator + // cannot under-commit via this knob — set MaxCapacity directly for + // that). Applied after all loading so callers see one consistent + // number; the heartbeat, gRPC capacity guard, and any future consumer + // all agree. + if cfg.CPUOvercommitRatio < 1 { + cfg.CPUOvercommitRatio = 1 + } + if cfg.CPUOvercommitRatio > 1 && cfg.MaxCapacity > 0 { + cfg.MaxCapacity = cfg.MaxCapacity * cfg.CPUOvercommitRatio + } + if portStr := os.Getenv("OPENSANDBOX_PORT"); portStr != "" { port, err := strconv.Atoi(portStr) if err != nil { diff --git a/internal/config/keyvault.go b/internal/config/keyvault.go index 2d4b7aab..09dd5d11 100644 --- a/internal/config/keyvault.go +++ b/internal/config/keyvault.go @@ -1,188 +1,84 @@ -// Package config provides configuration loading from Azure Key Vault. -// -// If SECRETS_VAULT_NAME is set, LoadSecretsFromKeyVault fetches all secrets -// from the vault and maps them to environment variables. The mapping is: -// -// Key Vault secret name → Environment variable -// server-database-url → OPENSANDBOX_DATABASE_URL -// server-jwt-secret → OPENSANDBOX_JWT_SECRET -// worker-s3-secret-key → OPENSANDBOX_S3_SECRET_ACCESS_KEY -// ...etc -// -// Secrets already set in the environment are NOT overwritten — env vars take -// precedence over Key Vault. This allows local overrides for development. +// Package config — Azure Key Vault implementation of SecretsProvider. // // Authentication uses Azure Default Credential (Managed Identity on VMs, -// CLI credentials locally). No explicit credentials needed. +// CLI credentials locally). The trigger env var is OPENSANDBOX_AZURE_KEY_VAULT_NAME +// (legacy: SECRETS_VAULT_NAME); LoadSecrets() in secrets.go selects this +// provider when either is set. + package config import ( "context" + "fmt" + "log" "os" - "time" - "github.com/opensandbox/opensandbox/internal/secrets" + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azsecrets" ) -// kvMapping maps Key Vault entry names to environment variable names. Despite -// the historical "secret" terminology this includes both real secrets (DB -// passwords, JWT keys) and non-secret cell config (region, cell_id, capacity -// tuning). The principle: only bootstrap pointers + per-VM identity stay in -// the env file; everything else is in KV so a cell's configuration is one -// source-of-truth rather than scattered across N worker.env files. -// -// Entries not in this map are silently ignored — the allowlist is the -// safety guard preventing a stray vault entry from accidentally shadowing -// an unrelated env var (e.g., PATH). -var kvMapping = map[string]string{ - // Server secrets - "server-database-url": "OPENSANDBOX_DATABASE_URL", - "server-redis-url": "OPENSANDBOX_REDIS_URL", - "server-jwt-secret": "OPENSANDBOX_JWT_SECRET", - "server-api-key": "OPENSANDBOX_API_KEY", - "server-secret-encryption-key": "OPENSANDBOX_SECRET_ENCRYPTION_KEY", - "server-workos-api-key": "WORKOS_API_KEY", - "server-workos-client-id": "WORKOS_CLIENT_ID", - "server-cf-api-token": "OPENSANDBOX_CF_API_TOKEN", - "server-cf-zone-id": "OPENSANDBOX_CF_ZONE_ID", - "server-stripe-secret-key": "STRIPE_SECRET_KEY", - "server-stripe-webhook-secret": "STRIPE_WEBHOOK_SECRET", - "server-sentry-dsn": "OPENSANDBOX_SENTRY_DSN", - // Machine-size fallback lists (PR #209). Comma-separated ranked - // instance types the autoscaler tries in order on quota / capacity - // errors. Empty value = use the single VMSize / InstanceType - // configured on the pool (pre-fallback behavior). - "server-azure-vm-sizes": "OPENSANDBOX_AZURE_VM_SIZES", - "server-ec2-instance-types": "OPENSANDBOX_EC2_INSTANCE_TYPES", - // Legacy Axiom mappings — kept for backwards compat with existing prod - // KVs that pre-date the `shared-` prefix. New deploys should use - // `shared-axiom-*` instead. Safe to leave: in server mode only - // `server-axiom-*` is loaded; in worker mode only `worker-axiom-*`. New - // `shared-*` mappings below win for new envs that have only those. - "server-axiom-query-token": "AXIOM_QUERY_TOKEN", - "server-axiom-dataset": "AXIOM_DATASET", +// azureKeyVaultProvider fetches secrets by listing the vault and dereferencing +// every name that matches secretMapping for the current mode. +type azureKeyVaultProvider struct { + vaultName string +} - // Server-side cell config + shared secrets. These mirror the worker-* keys - // of the same name — both sides need them, and the prefix filter loads only - // secrets for the current mode. When dev1/prod consolidate to cell-* we can - // drop the duplicates; for now this keeps the layout symmetric and explicit. - "server-cell-id": "OPENSANDBOX_CELL_ID", - "server-region": "OPENSANDBOX_REGION", - "server-sandbox-domain": "OPENSANDBOX_SANDBOX_DOMAIN", - "server-cf-event-endpoint": "OPENSANDBOX_CF_EVENT_ENDPOINT", - "server-cf-event-secret": "OPENSANDBOX_CF_EVENT_SECRET", - "server-cf-admin-secret": "OPENSANDBOX_CF_ADMIN_SECRET", - "server-session-jwt-secret": "OPENSANDBOX_SESSION_JWT_SECRET", - "server-halt-list-url": "OPENSANDBOX_HALT_LIST_URL", +func (p *azureKeyVaultProvider) Name() string { return "azure-keyvault" } - // Worker secrets - "worker-jwt-secret": "OPENSANDBOX_JWT_SECRET", - "worker-database-url": "OPENSANDBOX_DATABASE_URL", - "worker-redis-url": "OPENSANDBOX_REDIS_URL", - "worker-s3-access-key": "OPENSANDBOX_S3_ACCESS_KEY_ID", - "worker-s3-secret-key": "OPENSANDBOX_S3_SECRET_ACCESS_KEY", - "worker-sentry-dsn": "OPENSANDBOX_SENTRY_DSN", - "worker-axiom-ingest-token": "AXIOM_INGEST_TOKEN", // legacy; superseded by shared-axiom-ingest-token - "worker-axiom-dataset": "AXIOM_DATASET", // legacy; superseded by shared-axiom-dataset +func (p *azureKeyVaultProvider) Load(ctx context.Context, mode string) (int, int, error) { + vaultURL := fmt.Sprintf("https://%s.vault.azure.net/", p.vaultName) - // Worker per-cell config (non-secret but cell-scoped — every worker in the - // cell shares these, so KV is the single source of truth) - "worker-region": "OPENSANDBOX_REGION", - "worker-cell-id": "OPENSANDBOX_CELL_ID", - "worker-max-capacity": "OPENSANDBOX_MAX_CAPACITY", - "worker-default-sandbox-memory-mb": "OPENSANDBOX_DEFAULT_SANDBOX_MEMORY_MB", - "worker-default-sandbox-cpus": "OPENSANDBOX_DEFAULT_SANDBOX_CPUS", - "worker-default-sandbox-disk-mb": "OPENSANDBOX_DEFAULT_SANDBOX_DISK_MB", - "worker-sandbox-domain": "OPENSANDBOX_SANDBOX_DOMAIN", - "worker-s3-bucket": "OPENSANDBOX_S3_BUCKET", - "worker-s3-region": "OPENSANDBOX_S3_REGION", - "worker-s3-endpoint": "OPENSANDBOX_S3_ENDPOINT", - "worker-s3-force-path-style": "OPENSANDBOX_S3_FORCE_PATH_STYLE", - "worker-cf-event-endpoint": "OPENSANDBOX_CF_EVENT_ENDPOINT", - "worker-halt-list-url": "OPENSANDBOX_HALT_LIST_URL", - "worker-segment-write-key": "SEGMENT_WRITE_KEY", + cred, err := azidentity.NewDefaultAzureCredential(nil) + if err != nil { + return 0, 0, fmt.Errorf("keyvault: azure credential: %w", err) + } - // CF-cutover event pipe (worker) - "worker-cf-event-secret": "OPENSANDBOX_CF_EVENT_SECRET", - "worker-cf-admin-secret": "OPENSANDBOX_CF_ADMIN_SECRET", - "worker-session-jwt-secret": "OPENSANDBOX_SESSION_JWT_SECRET", + client, err := azsecrets.NewClient(vaultURL, cred, nil) + if err != nil { + return 0, 0, fmt.Errorf("keyvault: client: %w", err) + } - // Phase 2 — global blob store (Tigris primary) - "worker-global-blob-name": "OPENSANDBOX_GLOBAL_BLOB_NAME", - "worker-global-blob-endpoint": "OPENSANDBOX_GLOBAL_BLOB_ENDPOINT", - "worker-global-blob-region": "OPENSANDBOX_GLOBAL_BLOB_REGION", - "worker-global-blob-access-key-id": "OPENSANDBOX_GLOBAL_BLOB_ACCESS_KEY_ID", - "worker-global-blob-secret-access-key": "OPENSANDBOX_GLOBAL_BLOB_SECRET_ACCESS_KEY", - "worker-global-blob-goldens-bucket": "OPENSANDBOX_GLOBAL_BLOB_GOLDENS_BUCKET", - "worker-global-blob-templates-bucket": "OPENSANDBOX_GLOBAL_BLOB_TEMPLATES_BUCKET", - "worker-global-blob-events-bucket": "OPENSANDBOX_GLOBAL_BLOB_EVENTS_BUCKET", + loaded, skipped := 0, 0 - // Phase 2 — global blob store (optional fallback) - "worker-global-blob-fallback-name": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_NAME", - "worker-global-blob-fallback-endpoint": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_ENDPOINT", - "worker-global-blob-fallback-region": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_REGION", - "worker-global-blob-fallback-access-key-id": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_ACCESS_KEY_ID", - "worker-global-blob-fallback-secret-access-key": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_SECRET_ACCESS_KEY", + pager := client.NewListSecretPropertiesPager(nil) + for pager.More() { + page, err := pager.NextPage(ctx) + if err != nil { + return loaded, skipped, fmt.Errorf("keyvault: list secrets: %w", err) + } - // Shared (mode-agnostic — loaded in both server and worker) - "pg-password": "OPENSANDBOX_PG_PASSWORD", - "shared-axiom-ingest-token": "AXIOM_INGEST_TOKEN", - "shared-axiom-query-token": "AXIOM_QUERY_TOKEN", - "shared-axiom-dataset": "AXIOM_DATASET", - // Platform-logs: Vector reads these from /etc/opensandbox/vector.env, - // populated by populate-vector-env.service via its own IMDS+KV REST call - // (not by this Go-side loader, because Vector starts as its own systemd - // unit before the Go binary). The entries here exist for two reasons: - // 1. Discoverability — kvMapping is the single source of truth for - // "what shared-* secrets does this deployment need in KV". - // 2. Side-effect: the Go binary ALSO loads them into its own env at - // startup; future Go code that wants to surface platform-stream - // config (e.g. an admin endpoint) gets them for free. - "shared-axiom-platform-ingest-token": "AXIOM_PLATFORM_TOKEN", - "shared-axiom-platform-dataset": "AXIOM_PLATFORM_DATASET", - // Cell identifier — stamped on every log + metric event so platform - // dashboards can filter per cell. Same dual-consumer pattern as the - // platform-* secrets above: Vector reads it from /etc/opensandbox/vector.env - // (written by populate-vector-env.sh) for its remap substitutions, and - // the Go binary reads it from cfg.CellID (which falls back to - // "-default" when this isn't in KV — see config.go). - "shared-cell-id": "OPENSANDBOX_CELL_ID", + for _, prop := range page.Value { + name := prop.ID.Name() + envVar, mapped := secretMapping[name] + if !mapped { + continue + } + if !shouldLoadForMode(name, mode) { + continue + } + // Skip the network round-trip when the env is already set — + // callers can override locally without paying for the GET. + if os.Getenv(envVar) != "" { + skipped++ + continue + } - // Edge integration (mode-agnostic — both server and worker need them). - // cf-edge-base-url is consumed by the CP's edgeclient for HMAC'd - // /internal/templates + /internal/secret-stores lookups; secret-encryption-key - // is the AES-256-GCM key shared with the api-edge Worker so the edge can - // encrypt secret-store entries and any cell can decrypt them. - "shared-cf-edge-base-url": "OPENSANDBOX_CF_EDGE_BASE_URL", - "shared-secret-encryption-key": "OPENSANDBOX_SECRET_ENCRYPTION_KEY", -} + resp, err := client.GetSecret(ctx, name, "", nil) + if err != nil { + log.Printf("keyvault: failed to get secret %s: %v (skipping)", name, err) + continue + } + if resp.Value == nil { + continue + } -// LoadSecretsFromKeyVault fetches secrets from Azure Key Vault and sets them -// as environment variables. Only loads secrets relevant to the current mode -// (server or worker), determined by the secret name prefix. -// -// Skips secrets that are already set in the environment — local env wins -// (the .env file is the "emergency override" path). -// -// Does nothing if SECRETS_VAULT_NAME is not set. Now delegates to the -// internal/secrets KeyVaultBackend so server, worker, and the AzurePool -// runtime image-refresh path share one implementation. -func LoadSecretsFromKeyVault() error { - vaultName := os.Getenv("SECRETS_VAULT_NAME") - if vaultName == "" { - return nil // Key Vault not configured — use env file as-is + if setIfUnset(envVar, *resp.Value) { + loaded++ + } else { + skipped++ + } + } } - mode := os.Getenv("OPENSANDBOX_MODE") // "server" or "worker" - be, err := secrets.NewKeyVaultBackend(vaultName, kvMapping, mode) - if err != nil { - return err - } - if be == nil { - return nil // shouldn't happen — vaultName non-empty was checked above - } - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - _, _, err = be.LoadAllToEnv(ctx) - return err + return loaded, skipped, nil } diff --git a/internal/config/secrets.go b/internal/config/secrets.go new file mode 100644 index 00000000..ce09df93 --- /dev/null +++ b/internal/config/secrets.go @@ -0,0 +1,222 @@ +// Package config — multi-cloud secret loading. +// +// A SecretsProvider knows how to fetch secrets from a cloud secret store and +// populate them as environment variables, honoring the same conventions: +// +// - secretMapping is the single source of truth for "what logical secret +// name maps to what env var". Both Azure Key Vault and AWS Secrets +// Manager providers use the same map; the difference is the transport. +// - Mode-prefix filtering: only `-*`, `pg-*`, and `shared-*` secrets +// are loaded, where mode is the value of OPENSANDBOX_MODE ("server" or +// "worker"). Without this a server would try to load worker-only +// secrets and vice versa. +// - Env-var precedence: secrets already set in os.Environ() are NOT +// overwritten. Lets local dev shells override. +// +// LoadSecrets() is the entrypoint — it picks a provider by environment: +// +// OPENSANDBOX_AZURE_KEY_VAULT_NAME or SECRETS_VAULT_NAME → Azure +// OPENSANDBOX_AWS_SECRETS_PREFIX → AWS Secrets Manager +// (neither set) → no-op +// +// Both env vars set is treated as "both providers run in order" — useful for +// migrations. + +package config + +import ( + "context" + "log" + "os" + "strings" + "time" +) + +// SecretsProvider is implemented by per-cloud secret loaders. Implementations +// must honor mode-prefix filtering and env-var precedence; see Load doc. +type SecretsProvider interface { + // Name returns a short identifier used only for log messages. + Name() string + + // Load fetches secrets and sets matching env vars. Returns (loaded, + // skipped, err). `mode` is the value of OPENSANDBOX_MODE; pass "" + // for mode-agnostic loaders. + Load(ctx context.Context, mode string) (loaded int, skipped int, err error) +} + +// secretMapping is the canonical map of "logical secret name → env var +// name". Implementations look up each fetched secret's logical name here +// to decide what env var to set. +// +// Logical secret names are prefixed: +// +// server-* loaded only when OPENSANDBOX_MODE=server +// worker-* loaded only when OPENSANDBOX_MODE=worker +// pg-* grandfathered shared (Postgres password) — loaded in both modes +// shared-* mode-agnostic; loaded in both modes +// +// Backends apply different transports for fetching these by name (KV list- +// then-get, SM list-by-prefix-then-get). The names themselves are the same +// across clouds — operators see the same logical inventory regardless of +// where the cell runs. +var secretMapping = map[string]string{ + // ----- Server secrets ----- + "server-database-url": "OPENSANDBOX_DATABASE_URL", + "server-redis-url": "OPENSANDBOX_REDIS_URL", + "server-jwt-secret": "OPENSANDBOX_JWT_SECRET", + "server-api-key": "OPENSANDBOX_API_KEY", + "server-secret-encryption-key": "OPENSANDBOX_SECRET_ENCRYPTION_KEY", + "server-workos-api-key": "WORKOS_API_KEY", + "server-workos-client-id": "WORKOS_CLIENT_ID", + "server-cf-api-token": "OPENSANDBOX_CF_API_TOKEN", + "server-cf-zone-id": "OPENSANDBOX_CF_ZONE_ID", + "server-stripe-secret-key": "STRIPE_SECRET_KEY", + "server-stripe-webhook-secret": "STRIPE_WEBHOOK_SECRET", + "server-sentry-dsn": "OPENSANDBOX_SENTRY_DSN", + "server-azure-vm-sizes": "OPENSANDBOX_AZURE_VM_SIZES", + "server-ec2-instance-types": "OPENSANDBOX_EC2_INSTANCE_TYPES", + "server-min-idle-capacity": "OPENSANDBOX_MIN_IDLE_CAPACITY", + "server-min-idle-cpus": "OPENSANDBOX_MIN_IDLE_CPUS", + "server-ocfs2-node-ips": "OPENSANDBOX_OCFS2_NODE_IPS", + "server-s3-access-key": "OPENSANDBOX_S3_ACCESS_KEY_ID", + "server-s3-secret-key": "OPENSANDBOX_S3_SECRET_ACCESS_KEY", + "server-s3-bucket": "OPENSANDBOX_S3_BUCKET", + "server-s3-region": "OPENSANDBOX_S3_REGION", + "server-s3-endpoint": "OPENSANDBOX_S3_ENDPOINT", + "server-s3-force-path-style": "OPENSANDBOX_S3_FORCE_PATH_STYLE", + + // Legacy Axiom mappings — kept for backwards compat with existing prod + // KVs that pre-date the `shared-` prefix. New deploys should use + // `shared-axiom-*` instead. + "server-axiom-query-token": "AXIOM_QUERY_TOKEN", + "server-axiom-dataset": "AXIOM_DATASET", + + // Server-side cell config + shared secrets. + "server-cell-id": "OPENSANDBOX_CELL_ID", + "server-region": "OPENSANDBOX_REGION", + "server-sandbox-domain": "OPENSANDBOX_SANDBOX_DOMAIN", + "server-cf-event-endpoint": "OPENSANDBOX_CF_EVENT_ENDPOINT", + "server-cf-event-secret": "OPENSANDBOX_CF_EVENT_SECRET", + "server-cf-admin-secret": "OPENSANDBOX_CF_ADMIN_SECRET", + "server-session-jwt-secret": "OPENSANDBOX_SESSION_JWT_SECRET", + "server-halt-list-url": "OPENSANDBOX_HALT_LIST_URL", + + // ----- Worker secrets ----- + "worker-jwt-secret": "OPENSANDBOX_JWT_SECRET", + "worker-database-url": "OPENSANDBOX_DATABASE_URL", + "worker-redis-url": "OPENSANDBOX_REDIS_URL", + "worker-s3-access-key": "OPENSANDBOX_S3_ACCESS_KEY_ID", + "worker-s3-secret-key": "OPENSANDBOX_S3_SECRET_ACCESS_KEY", + "worker-sentry-dsn": "OPENSANDBOX_SENTRY_DSN", + "worker-axiom-ingest-token": "AXIOM_INGEST_TOKEN", + "worker-axiom-dataset": "AXIOM_DATASET", + + // Worker per-cell config. + "worker-region": "OPENSANDBOX_REGION", + "worker-cell-id": "OPENSANDBOX_CELL_ID", + "worker-max-capacity": "OPENSANDBOX_MAX_CAPACITY", + "worker-default-sandbox-memory-mb": "OPENSANDBOX_DEFAULT_SANDBOX_MEMORY_MB", + "worker-default-sandbox-cpus": "OPENSANDBOX_DEFAULT_SANDBOX_CPUS", + "worker-default-sandbox-disk-mb": "OPENSANDBOX_DEFAULT_SANDBOX_DISK_MB", + "worker-sandbox-domain": "OPENSANDBOX_SANDBOX_DOMAIN", + "worker-s3-bucket": "OPENSANDBOX_S3_BUCKET", + "worker-s3-region": "OPENSANDBOX_S3_REGION", + "worker-s3-endpoint": "OPENSANDBOX_S3_ENDPOINT", + "worker-s3-force-path-style": "OPENSANDBOX_S3_FORCE_PATH_STYLE", + "worker-cf-event-endpoint": "OPENSANDBOX_CF_EVENT_ENDPOINT", + "worker-halt-list-url": "OPENSANDBOX_HALT_LIST_URL", + "worker-segment-write-key": "SEGMENT_WRITE_KEY", + "worker-cf-event-secret": "OPENSANDBOX_CF_EVENT_SECRET", + "worker-cf-admin-secret": "OPENSANDBOX_CF_ADMIN_SECRET", + "worker-session-jwt-secret": "OPENSANDBOX_SESSION_JWT_SECRET", + "worker-global-blob-name": "OPENSANDBOX_GLOBAL_BLOB_NAME", + "worker-global-blob-endpoint": "OPENSANDBOX_GLOBAL_BLOB_ENDPOINT", + "worker-global-blob-region": "OPENSANDBOX_GLOBAL_BLOB_REGION", + "worker-global-blob-access-key-id": "OPENSANDBOX_GLOBAL_BLOB_ACCESS_KEY_ID", + "worker-global-blob-secret-access-key": "OPENSANDBOX_GLOBAL_BLOB_SECRET_ACCESS_KEY", + "worker-global-blob-goldens-bucket": "OPENSANDBOX_GLOBAL_BLOB_GOLDENS_BUCKET", + "worker-global-blob-templates-bucket": "OPENSANDBOX_GLOBAL_BLOB_TEMPLATES_BUCKET", + "worker-global-blob-events-bucket": "OPENSANDBOX_GLOBAL_BLOB_EVENTS_BUCKET", + "worker-global-blob-fallback-name": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_NAME", + "worker-global-blob-fallback-endpoint": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_ENDPOINT", + "worker-global-blob-fallback-region": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_REGION", + "worker-global-blob-fallback-access-key-id": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_ACCESS_KEY_ID", + "worker-global-blob-fallback-secret-access-key": "OPENSANDBOX_GLOBAL_BLOB_FALLBACK_SECRET_ACCESS_KEY", + + // ----- Shared / mode-agnostic ----- + "pg-password": "OPENSANDBOX_PG_PASSWORD", + "shared-axiom-ingest-token": "AXIOM_INGEST_TOKEN", + "shared-axiom-query-token": "AXIOM_QUERY_TOKEN", + "shared-axiom-dataset": "AXIOM_DATASET", + "shared-axiom-platform-ingest-token": "AXIOM_PLATFORM_TOKEN", + "shared-axiom-platform-dataset": "AXIOM_PLATFORM_DATASET", + "shared-cell-id": "OPENSANDBOX_CELL_ID", + "shared-cf-edge-base-url": "OPENSANDBOX_CF_EDGE_BASE_URL", + "shared-secret-encryption-key": "OPENSANDBOX_SECRET_ENCRYPTION_KEY", +} + +// shouldLoadForMode returns true if the given logical secret name applies to +// the running mode. Empty mode loads everything; otherwise only matching +// prefixes plus `pg-*` and `shared-*` pass. +func shouldLoadForMode(name, mode string) bool { + if mode == "" { + return true + } + return strings.HasPrefix(name, mode+"-") || + strings.HasPrefix(name, "pg-") || + strings.HasPrefix(name, "shared-") +} + +// setIfUnset writes name=value to the process environment, but only when +// name is not already set. Returns true if the value was applied. +func setIfUnset(name, value string) bool { + if os.Getenv(name) != "" { + return false + } + _ = os.Setenv(name, value) + return true +} + +// LoadSecrets picks a provider by environment and loads secrets from it. +// Mode is read from OPENSANDBOX_MODE. Returns nil and does nothing if no +// provider is configured — local dev with a worker.env file works as-is. +// +// If both Azure and AWS providers are configured (unusual but legal during +// a cross-cloud migration), both run in declaration order. +func LoadSecrets() error { + mode := os.Getenv("OPENSANDBOX_MODE") + + var providers []SecretsProvider + if name := azureVaultName(); name != "" { + providers = append(providers, &azureKeyVaultProvider{vaultName: name}) + } + if prefix := os.Getenv("OPENSANDBOX_AWS_SECRETS_PREFIX"); prefix != "" { + providers = append(providers, &awsSecretsManagerProvider{prefix: prefix}) + } + + if len(providers) == 0 { + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + for _, p := range providers { + loaded, skipped, err := p.Load(ctx, mode) + if err != nil { + return err + } + log.Printf("secrets: %s loaded %d, skipped %d (already set)", p.Name(), loaded, skipped) + } + return nil +} + +// azureVaultName reads the Azure KV name from either of the two env vars +// historically used. Kept here so both `secrets.go` and `keyvault.go` see +// the same lookup precedence. +func azureVaultName() string { + if v := os.Getenv("OPENSANDBOX_AZURE_KEY_VAULT_NAME"); v != "" { + return v + } + return os.Getenv("SECRETS_VAULT_NAME") +} diff --git a/internal/config/secretsmanager.go b/internal/config/secretsmanager.go index 69ef2d53..71591686 100644 --- a/internal/config/secretsmanager.go +++ b/internal/config/secretsmanager.go @@ -1,58 +1,129 @@ +// Package config — AWS Secrets Manager implementation of SecretsProvider. +// +// Authentication uses the AWS Default Credential chain (EC2 IAM role on +// instances, AWS_PROFILE / SSO / env locally). The trigger env var is +// OPENSANDBOX_AWS_SECRETS_PREFIX; LoadSecrets() in secrets.go selects this +// provider when it is set. +// +// Layout: secrets are stored under a flat per-cell prefix, e.g. +// +// opencomputer/aws-us-east-2-poc/worker-jwt-secret +// opencomputer/aws-us-east-2-poc/worker-redis-url +// opencomputer/aws-us-east-2-poc/shared-axiom-ingest-token +// +// The provider lists everything under the prefix, strips it, looks up the +// remaining logical name in secretMapping (cloud-agnostic, defined in +// secrets.go), and sets the matching env var if not already populated. + package config import ( "context" + "fmt" + "log" "os" - "time" + "strings" - "github.com/opensandbox/opensandbox/internal/secrets" + "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/secretsmanager" + "github.com/aws/aws-sdk-go-v2/service/secretsmanager/types" ) -// LoadSecretsFromSecretsManager is the AWS analogue of LoadSecretsFromKeyVault. -// Maps every kebab-case key from kvMapping into the process environment via -// the same translation Azure KV uses (server-jwt-secret → OPENSANDBOX_JWT_SECRET). -// -// Two trigger env vars, picked one at a time: -// -// - OPENSANDBOX_SECRETS_ARN: single bundled SM secret with JSON SecretString. -// One GetSecretValue, one IAM scope. Cheap ($0.40/mo) but requires the -// producer (Infisical or otherwise) to write a JSON bundle. -// -// - OPENSANDBOX_SECRETS_AWS_REGION: per-key mode. Each kvMapping key is its -// own SM secret. BatchGetSecretValue fetches them in chunks of 20. This -// is what Infisical's "Sync each secret to its own secret" mode produces. -// Costs $0.40/mo per secret but matches Infisical's default sync shape. -// -// Existing env vars take precedence (emergency override path). -// -// Does nothing if neither trigger is set. -func LoadSecretsFromSecretsManager() error { - bundleARN := os.Getenv("OPENSANDBOX_SECRETS_ARN") - listRegion := os.Getenv("OPENSANDBOX_SECRETS_AWS_REGION") - if bundleARN == "" && listRegion == "" { - return nil // not configured; defer to env file / Azure KV / combined mode +// awsSecretsManagerProvider fetches secrets by listing the cell's prefix and +// dereferencing every name that matches secretMapping for the current mode. +type awsSecretsManagerProvider struct { + prefix string +} + +func (p *awsSecretsManagerProvider) Name() string { return "aws-secretsmanager" } + +func (p *awsSecretsManagerProvider) Load(ctx context.Context, mode string) (int, int, error) { + // Region: explicit env var wins (matches the cell config), else AWS_REGION, + // else fall through to the default chain (IMDS on EC2). Passing nothing + // when LoadDefaultConfig can't resolve a region anywhere makes the first + // API call return "Missing Region", which surfaces unclearly — prefer the + // explicit env-var path even on EC2. + region := os.Getenv("OPENSANDBOX_REGION") + if region == "" { + region = os.Getenv("AWS_REGION") } - mode := os.Getenv("OPENSANDBOX_MODE") - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() + var loadOpts []func(*awsconfig.LoadOptions) error + if region != "" { + loadOpts = append(loadOpts, awsconfig.WithRegion(region)) + } - be, err := secrets.NewSecretsManagerBackend(ctx, listRegion, bundleARN) + cfg, err := awsconfig.LoadDefaultConfig(ctx, loadOpts...) if err != nil { - return err + return 0, 0, fmt.Errorf("secretsmanager: load aws config: %w", err) } - be.NameMap = kvMapping - be.ModePrefixFilter = mode + client := secretsmanager.NewFromConfig(cfg) - if bundleARN != "" { - _, _, err = be.LoadAllToEnv(ctx) - return err - } - // list/per-key mode — fetch each kvMapping key as its own SM secret - names := make([]string, 0, len(kvMapping)) - for k := range kvMapping { - names = append(names, k) + loaded, skipped := 0, 0 + + var nextToken *string + for { + out, err := client.ListSecrets(ctx, &secretsmanager.ListSecretsInput{ + MaxResults: aws.Int32(100), + NextToken: nextToken, + Filters: []types.Filter{ + { + Key: types.FilterNameStringTypeName, + Values: []string{p.prefix}, + }, + }, + }) + if err != nil { + return loaded, skipped, fmt.Errorf("secretsmanager: list: %w", err) + } + + for _, entry := range out.SecretList { + if entry.Name == nil { + continue + } + fullName := *entry.Name + // `name` filter is a prefix-match; defensively strip and skip if not ours. + if !strings.HasPrefix(fullName, p.prefix) { + continue + } + logicalName := strings.TrimPrefix(fullName, p.prefix) + + envVar, mapped := secretMapping[logicalName] + if !mapped { + continue + } + if !shouldLoadForMode(logicalName, mode) { + continue + } + if os.Getenv(envVar) != "" { + skipped++ + continue + } + + val, err := client.GetSecretValue(ctx, &secretsmanager.GetSecretValueInput{ + SecretId: aws.String(fullName), + }) + if err != nil { + log.Printf("secretsmanager: failed to get secret %s: %v (skipping)", logicalName, err) + continue + } + if val.SecretString == nil { + continue + } + + if setIfUnset(envVar, *val.SecretString) { + loaded++ + } else { + skipped++ + } + } + + if out.NextToken == nil { + break + } + nextToken = out.NextToken } - _, _, err = be.LoadAllByNameList(ctx, names) - return err + + return loaded, skipped, nil } diff --git a/internal/controlplane/reconcile.go b/internal/controlplane/reconcile.go index c6d265ad..bbf952e4 100644 --- a/internal/controlplane/reconcile.go +++ b/internal/controlplane/reconcile.go @@ -3,6 +3,7 @@ package controlplane import ( "context" "encoding/json" + "fmt" "log" "strings" "time" @@ -11,6 +12,7 @@ import ( "github.com/redis/go-redis/v9" "github.com/opensandbox/opensandbox/internal/db" + "github.com/opensandbox/opensandbox/pkg/types" pb "github.com/opensandbox/opensandbox/proto/worker" ) @@ -18,15 +20,16 @@ import ( // believes is stopped on this worker but the worker may still be hosting. // // Why this exists: -// When the worker is unreachable, the customer's DELETE goes through the -// cell-side fallback at internal/api/sandbox.go's destroy handler — the -// cell marks the session stopped in PG and publishes a `stopped` lifecycle -// event "to close the drift" with D1. Worker never receives the gRPC -// Destroy. When the worker becomes reachable again, the cell already -// thinks the sandbox is dead, but the worker still has m.vms[id] alive, -// qemu still running, usage_ticker still emitting `usage_tick` events. -// That window has run for 74h+ in the wild before a worker restart finally -// cleared the m.vms map. +// +// When the worker is unreachable, the customer's DELETE goes through the +// cell-side fallback at internal/api/sandbox.go's destroy handler — the +// cell marks the session stopped in PG and publishes a `stopped` lifecycle +// event "to close the drift" with D1. Worker never receives the gRPC +// Destroy. When the worker becomes reachable again, the cell already +// thinks the sandbox is dead, but the worker still has m.vms[id] alive, +// qemu still running, usage_ticker still emitting `usage_tick` events. +// That window has run for 74h+ in the wild before a worker restart finally +// cleared the m.vms map. // // This reconcile closes that gap: on worker rejoin (RedisWorkerRegistry's // OnWorkerRejoined callback), the cell sweeps every sandbox it has marked @@ -102,30 +105,32 @@ func isSandboxNotFound(err error) bool { // ReconcileRunningOnWorker is the symmetric direction of ReconcileStoppedOnWorker: // -// forward — cell says STOPPED on this worker, worker may still be hosting → -// re-issue Destroy via RPC. ReconcileStoppedOnWorker above. -// reverse — cell says RUNNING on this worker, worker doesn't have it → -// close the row on the cell side. THIS function. +// forward — cell says STOPPED on this worker, worker may still be hosting → +// re-issue Destroy via RPC. ReconcileStoppedOnWorker above. +// reverse — cell says RUNNING on this worker, worker doesn't have it → +// close the row on the cell side. THIS function. // // Why both directions are needed: -// The cell-side fallback at internal/api/sandbox.go (worker-unreachable -// destroy path) covers the forward direction. There's no symmetric fallback -// for "worker died, never finished EndScaleEvent on its way out" — when a -// worker process crashes/OOMs/restarts, its m.vms is cleared but cell PG -// keeps the scale event open. usage-reporter sums (now - started_at) -// indefinitely; customer gets billed for compute that hasn't run for days. +// +// The cell-side fallback at internal/api/sandbox.go (worker-unreachable +// destroy path) covers the forward direction. There's no symmetric fallback +// for "worker died, never finished EndScaleEvent on its way out" — when a +// worker process crashes/OOMs/restarts, its m.vms is cleared but cell PG +// keeps the scale event open. usage-reporter sums (now - started_at) +// indefinitely; customer gets billed for compute that hasn't run for days. // // Empirical fingerprint from prod: 49 still-open scale events on two workers // that were known to have restarted ~3-5 days prior, accumulating ~$2k of // phantom Pro billing per restart event. // // Process: -// 1. Ask cell PG: what sandboxes are status='running' on this worker? -// 2. Ask worker (existing ListSandboxes RPC): what do you actually have? -// 3. For each cell-PG-running entry the worker doesn't claim: -// - UpdateSandboxSessionStatus → stopped -// - EndScaleEvent → closes the open billing row -// - publish stopped lifecycle event so events-ingest updates D1 +// 1. Ask cell PG: what sandboxes are status='running' on this worker? +// 2. Ask worker (existing ListSandboxes RPC): what do you actually have? +// 3. For each cell-PG-running entry the worker doesn't claim: +// - resumable: recreate from the shared cell disk on an eligible worker +// and move routing to that worker +// - non-resumable: UpdateSandboxSessionStatus → stopped, EndScaleEvent, +// publish stopped lifecycle event so events-ingest updates D1 func ReconcileRunningOnWorker(ctx context.Context, registry *RedisWorkerRegistry, store *db.Store, cellID, workerID string) { if store == nil || registry == nil { return @@ -162,12 +167,18 @@ func ReconcileRunningOnWorker(ctx context.Context, registry *RedisWorkerRegistry log.Printf("controlplane: reverse-reconcile %s: cell-running=%d worker-has=%d", workerID, len(cellRunning), len(workerHas)) reason := "reverse_reconcile_worker_lost_session" - var closed, alive int + var closed, alive, recreated int for _, ref := range cellRunning { if _, has := workerHas[ref.SandboxID]; has { alive++ continue } + if ok, err := recreateResumableSandbox(ctx, registry, store, cellID, ref.SandboxID, workerID); err != nil { + log.Printf("controlplane: reverse-reconcile %s: resumable recreate %s failed: %v", workerID, ref.SandboxID, err) + } else if ok { + recreated++ + continue + } // Close the cell-side state. Order matters: status first (so future // dashboard reads see the right thing immediately), then scale event // (so usage-reporter stops billing), then the lifecycle event (so D1 @@ -187,7 +198,82 @@ func ReconcileRunningOnWorker(ctx context.Context, registry *RedisWorkerRegistry publishStoppedLifecycleEvent(ctx, registry.RedisClient(), cellID, ref.SandboxID, ref.OrgID.String(), workerID, reason) closed++ } - log.Printf("controlplane: reverse-reconcile %s: closed=%d still-alive-on-worker=%d (of %d cell-running)", workerID, closed, alive, len(cellRunning)) + log.Printf("controlplane: reverse-reconcile %s: recreated=%d closed=%d still-alive-on-worker=%d (of %d cell-running)", workerID, recreated, closed, alive, len(cellRunning)) +} + +func recreateResumableSandbox(ctx context.Context, registry *RedisWorkerRegistry, store *db.Store, cellID, sandboxID, oldWorkerID string) (bool, error) { + session, err := store.GetSandboxSession(ctx, sandboxID) + if err != nil { + return false, err + } + if session.Status != "running" { + return false, nil + } + var cfg types.SandboxConfig + if len(session.Config) > 0 { + if err := json.Unmarshal(session.Config, &cfg); err != nil { + return false, fmt.Errorf("parse sandbox config: %w", err) + } + } + if !cfg.IsResumable() { + return false, nil + } + cfg.EnsureNetworkEnabledDefault() + cfg.SandboxID = sandboxID + if cfg.Envs == nil { + cfg.Envs = map[string]string{} + } + cfg.Envs["OPENSANDBOX_BURST"] = "true" + cfg.Envs["OPENSANDBOX_RESUMABLE"] = "true" + cfg.Envs["OPENSANDBOX_RESUME_NOTICE_SECONDS"] = "25" + + worker, client, err := registry.GetLeastLoadedWorkerForPool(session.Region, WorkerPoolBurst) + if err != nil { + return true, fmt.Errorf("pick worker: %w", err) + } + + rpcCtx, cancel := context.WithTimeout(ctx, 60*time.Second) + resp, err := client.CreateSandbox(rpcCtx, &pb.CreateSandboxRequest{ + SandboxId: sandboxID, + Template: cfg.Template, + Timeout: int32(cfg.Timeout), + Envs: cfg.Envs, + MemoryMb: int32(cfg.MemoryMB), + CpuCount: int32(cfg.CpuCount), + NetworkEnabled: cfg.IsNetworkEnabled(), + Port: int32(cfg.Port), + EgressAllowlist: cfg.EgressAllowlist, + SecretAllowedHosts: flattenSecretAllowedHostsForReconcile(cfg.SecretAllowedHosts), + SecretEnvs: cfg.SecretEnvs, + DiskMb: int32(cfg.DiskMB), + }) + cancel() + if err != nil { + return true, fmt.Errorf("worker CreateSandbox on %s: %w", worker.ID, err) + } + if resp == nil || resp.SandboxId == "" { + return true, fmt.Errorf("worker CreateSandbox on %s returned empty response", worker.ID) + } + if err := store.CompleteMigration(ctx, sandboxID, worker.ID); err != nil { + return true, fmt.Errorf("update session worker: %w", err) + } + if worker.GoldenVersion != "" { + _ = store.SetSandboxGoldenVersion(ctx, sandboxID, worker.GoldenVersion) + } + PublishLifecycle(ctx, registry.RedisClient(), cellID, "migrated", sandboxID, worker.ID, session.OrgID, "resumable_recreate") + log.Printf("controlplane: resumable recreate %s: %s -> %s", sandboxID, oldWorkerID, worker.ID) + return true, nil +} + +func flattenSecretAllowedHostsForReconcile(m map[string][]string) map[string]string { + if len(m) == 0 { + return nil + } + out := make(map[string]string, len(m)) + for k, hosts := range m { + out[k] = strings.Join(hosts, ",") + } + return out } // publishStoppedLifecycleEvent emits a `stopped` event onto this cell's events diff --git a/internal/controlplane/redis_registry.go b/internal/controlplane/redis_registry.go index 745fd765..d79f6b2b 100644 --- a/internal/controlplane/redis_registry.go +++ b/internal/controlplane/redis_registry.go @@ -52,15 +52,16 @@ const ( // WorkerEntry represents a worker in the Redis-backed registry. type WorkerEntry struct { - ID string `json:"worker_id"` - MachineID string `json:"machine_id,omitempty"` // EC2 instance ID - Region string `json:"region"` - GRPCAddr string `json:"grpc_addr"` - HTTPAddr string `json:"http_addr"` - Capacity int `json:"capacity"` - Current int `json:"current"` - CPUPct float64 `json:"cpu_pct"` - MemPct float64 `json:"mem_pct"` + ID string `json:"worker_id"` + MachineID string `json:"machine_id,omitempty"` // EC2 instance ID + Pool string `json:"pool,omitempty"` + Region string `json:"region"` + GRPCAddr string `json:"grpc_addr"` + HTTPAddr string `json:"http_addr"` + Capacity int `json:"capacity"` + Current int `json:"current"` + CPUPct float64 `json:"cpu_pct"` + MemPct float64 `json:"mem_pct"` DiskPct float64 `json:"disk_pct"` TotalMemoryMB int `json:"total_memory_mb,omitempty"` CommittedMemoryMB int `json:"committed_memory_mb,omitempty"` @@ -74,6 +75,20 @@ type WorkerEntry struct { Sandboxes map[string]SandboxStats `json:"sandboxes,omitempty"` } +const ( + WorkerPoolOnDemand = "ondemand" + WorkerPoolBurst = "burst" +) + +func NormalizeWorkerPool(pool string) string { + switch pool { + case WorkerPoolBurst: + return WorkerPoolBurst + default: + return WorkerPoolOnDemand + } +} + // SandboxStats is the per-sandbox snapshot ingested from worker heartbeats. // Mirrors internal/worker.SandboxStatsWire — kept separate to avoid an // import cycle (CP shouldn't depend on the worker package). @@ -88,13 +103,13 @@ type SandboxStats struct { // backed by Redis pub/sub for real-time updates and periodic SCAN for reconciliation. // It also maintains a persistent gRPC connection pool to workers. type RedisWorkerRegistry struct { - rdb *redis.Client - mu sync.RWMutex - workers map[string]*WorkerEntry // in-memory hot cache - conns map[string]*grpc.ClientConn // persistent gRPC connections - clients map[string]pb.SandboxWorkerClient // cached gRPC clients - rrCounter uint64 // round-robin counter for tie-breaking - stop chan struct{} + rdb *redis.Client + mu sync.RWMutex + workers map[string]*WorkerEntry // in-memory hot cache + conns map[string]*grpc.ClientConn // persistent gRPC connections + clients map[string]pb.SandboxWorkerClient // cached gRPC clients + rrCounter uint64 // round-robin counter for tie-breaking + stop chan struct{} // onWorkerRejoined fires when a worker registers — both genuinely new and // after being pruned for missed heartbeats. Used by the reconcile-on- @@ -330,6 +345,7 @@ func (r *RedisWorkerRegistry) handleHeartbeat(entry WorkerEntry) { existing, ok := r.workers[entry.ID] if ok { // Update existing entry + existing.Pool = NormalizeWorkerPool(entry.Pool) existing.Current = entry.Current existing.Capacity = entry.Capacity existing.CPUPct = entry.CPUPct @@ -338,6 +354,7 @@ func (r *RedisWorkerRegistry) handleHeartbeat(entry WorkerEntry) { existing.TotalMemoryMB = entry.TotalMemoryMB existing.CommittedMemoryMB = entry.CommittedMemoryMB existing.Draining = drainOverride + existing.Sandboxes = entry.Sandboxes if entry.GoldenVersion != "" { existing.GoldenVersion = entry.GoldenVersion } @@ -360,9 +377,10 @@ func (r *RedisWorkerRegistry) handleHeartbeat(entry WorkerEntry) { // on that ID so it's a cheap no-op; for rejoins, this is where we // re-issue Destroy for sandboxes the cell published "stopped" for // during the unreachable window. See internal/controlplane/reconcile.go. + entry.Pool = NormalizeWorkerPool(entry.Pool) entry.Draining = drainOverride r.workers[entry.ID] = &entry - log.Printf("redis_registry: new worker registered: %s (region=%s, grpc=%s, draining=%v)", entry.ID, entry.Region, entry.GRPCAddr, drainOverride) + log.Printf("redis_registry: new worker registered: %s (pool=%s, region=%s, grpc=%s, draining=%v)", entry.ID, entry.Pool, entry.Region, entry.GRPCAddr, drainOverride) if r.onWorkerRejoined != nil { // Fire in a goroutine — reconcile may take a few seconds (DB query // + a DestroySandbox RPC per stale entry) and we don't want @@ -489,18 +507,23 @@ func (r *RedisWorkerRegistry) dialWorkerLocked(workerID, grpcAddr string) { // score, since they're spikier and a transient burst shouldn't permanently // disprefer a worker. func (r *RedisWorkerRegistry) GetLeastLoadedWorker(region string) (*WorkerEntry, pb.SandboxWorkerClient, error) { + return r.GetLeastLoadedWorkerForPool(region, WorkerPoolOnDemand) +} + +func (r *RedisWorkerRegistry) GetLeastLoadedWorkerForPool(region string, pool string) (*WorkerEntry, pb.SandboxWorkerClient, error) { r.mu.Lock() defer r.mu.Unlock() + pool = NormalizeWorkerPool(pool) routingCtx, routingCancel := context.WithTimeout(context.Background(), time.Second) defer routingCancel() - eligible := r.collectEligibleLocked(region, false /* anyRegion */) + eligible := r.collectEligibleLocked(region, false /* anyRegion */, pool) if len(eligible) == 0 && region != "" { - eligible = r.collectEligibleLocked(region, true /* anyRegion */) + eligible = r.collectEligibleLocked(region, true /* anyRegion */, pool) } if len(eligible) == 0 { - return nil, nil, fmt.Errorf("no workers available") + return nil, nil, fmt.Errorf("no %s workers available", pool) } // Apply the cross-CP routing counter so an in-flight placement on the @@ -573,9 +596,12 @@ func (r *RedisWorkerRegistry) GetLeastLoadedWorker(region string) (*WorkerEntry, // collectEligibleLocked returns workers passing the routing eligibility gates. // If anyRegion is true the region filter is dropped (used for cross-region // fallback when a region is starved). Caller must hold r.mu. -func (r *RedisWorkerRegistry) collectEligibleLocked(region string, anyRegion bool) []*WorkerEntry { +func (r *RedisWorkerRegistry) collectEligibleLocked(region string, anyRegion bool, pool string) []*WorkerEntry { var out []*WorkerEntry for _, w := range r.workers { + if NormalizeWorkerPool(w.Pool) != pool { + continue + } if !anyRegion && region != "" && w.Region != region { continue } @@ -744,27 +770,39 @@ func (r *RedisWorkerRegistry) Regions() []string { // GetWorkersByRegion returns workers in a region (satisfies ScalerRegistry). func (r *RedisWorkerRegistry) GetWorkersByRegion(region string) []*WorkerInfo { + return r.GetWorkersByRegionAndPool(region, "") +} + +func (r *RedisWorkerRegistry) GetWorkersByRegionAndPool(region, pool string) []*WorkerInfo { r.mu.RLock() defer r.mu.RUnlock() + if pool != "" { + pool = NormalizeWorkerPool(pool) + } var result []*WorkerInfo for _, w := range r.workers { + if pool != "" && NormalizeWorkerPool(w.Pool) != pool { + continue + } if w.Region == region { result = append(result, &WorkerInfo{ - ID: w.ID, - MachineID: w.MachineID, - Region: w.Region, - GRPCAddr: w.GRPCAddr, - HTTPAddr: w.HTTPAddr, - Capacity: w.Capacity, - Current: w.Current, - CPUPct: w.CPUPct, - MemPct: w.MemPct, - DiskPct: w.DiskPct, + ID: w.ID, + MachineID: w.MachineID, + Pool: NormalizeWorkerPool(w.Pool), + Region: w.Region, + GRPCAddr: w.GRPCAddr, + HTTPAddr: w.HTTPAddr, + Capacity: w.Capacity, + Current: w.Current, + CPUPct: w.CPUPct, + MemPct: w.MemPct, + DiskPct: w.DiskPct, + Draining: w.Draining, TotalMemoryMB: w.TotalMemoryMB, CommittedMemoryMB: w.CommittedMemoryMB, - GoldenVersion: w.GoldenVersion, - WorkerVersion: w.WorkerVersion, + GoldenVersion: w.GoldenVersion, + WorkerVersion: w.WorkerVersion, }) } } diff --git a/internal/controlplane/scaler.go b/internal/controlplane/scaler.go index ed306b21..5ebaef46 100644 --- a/internal/controlplane/scaler.go +++ b/internal/controlplane/scaler.go @@ -20,10 +20,11 @@ import ( ) const ( - scaleUpThreshold = 0.50 // Scale up when utilization > 50% (gives ~3 min runway for new worker to boot) - scaleDownThreshold = 0.20 // Scale down when utilization < 20% - maxWorkersPerRegion = 10 // Hard cap to prevent runaway launches + scaleUpThreshold = 0.50 // Scale up when utilization > 50% (gives ~3 min runway for new worker to boot) + scaleDownThreshold = 0.20 // Scale down when utilization < 20% + maxWorkersPerRegion = 10 // Hard cap to prevent runaway launches pendingWorkerTTL = 10 * time.Minute // How long to wait for a launched worker to register + defaultWorkerCap = 50 // Capacity estimate for pending workers before first heartbeat. // Resource-based scaling thresholds (applied per-worker, trigger on ANY worker exceeding) resourceCPUThreshold = 70.0 // Scale up if any worker CPU > 70% @@ -40,9 +41,15 @@ const ( emergencyCPUThreshold = 95.0 emergencyMemThreshold = 95.0 emergencyDiskThreshold = 90.0 - evacuationBatchSize = 3 // sandboxes to migrate per eval cycle per worker - evacuationCooldown = 60 * time.Second // per-worker cooldown between evacuation batches - drainTimeout = 45 * time.Minute // max time to drain a worker via live migration (allows 30 sandboxes × 10min each in batches of 3) + evacuationBatchSize = 3 // sandboxes to migrate per eval cycle per worker + evacuationCooldown = 60 * time.Second // per-worker cooldown between evacuation batches + drainTimeout = 45 * time.Minute // max time to drain a worker via live migration (allows 30 sandboxes × 10min each in batches of 3) + drainBatchSuccessPause = 2 * time.Second // pause between successful drain batches + drainBatchFailurePause = 5 * time.Second // pause before retrying after a failed drain batch + spotDrainFailurePause = 1 * time.Second // spot/preemption drains should keep moving inside the 2min notice window + normalDrainBackoff = 5 * time.Minute // conservative retry backoff for ordinary scaler drains + spotDrainBackoff = 1 * time.Second // aggressive retry backoff for spot/preemption drains + spotDrainConcurrency = 10 // rolling parallelism for urgent spot/preemption drains creationFailureThreshold = 3 // consecutive failures before exponential backoff creationBackoffMin = 1 * time.Minute // initial backoff after threshold hit @@ -76,14 +83,24 @@ type OrphanCleaner interface { type ScalerConfig struct { Pool compute.Pool Registry ScalerRegistry - Store *db.Store // for updating session worker_id after migration + Store *db.Store // for updating session worker_id after migration StateStore ScalerStateStore // optional: persists scaler state to Redis (nil = in-memory) WorkerImage string Cooldown time.Duration // minimum time between scale-up actions per region Interval time.Duration // how often to evaluate scaling (0 = default 30s) - MinWorkers int // minimum total workers per region (0 = default 1). Always kept running. - MaxWorkers int // maximum workers per region (0 = default 10). Hard cap to prevent runaway launches. - IdleReserve int // target idle (0 sandbox) workers for burst absorption (0 = default 1). Separate from MinWorkers. + MinWorkers int // minimum total workers per region (0 = default 1). Always kept running. + MaxWorkers int // maximum workers per region (0 = default 10). Hard cap to prevent runaway launches. + IdleReserve int // target idle (0 sandbox) workers for burst absorption (0 = default 1). Separate from MinWorkers. + // MinIdleCapacity is the target spare sandbox slot capacity per region. + // When >0, this replaces the worker-count based MinWorkers/IdleReserve + // pre-provisioning logic while still respecting MaxWorkers. + MinIdleCapacity int + // MinIdleCPUs is the target spare sandbox CPU units per region. When >0, + // it takes precedence over MinIdleCapacity and is converted to slots via + // DefaultSandboxCPUs. + MinIdleCPUs int + DefaultSandboxCPUs int + WorkerPool string // optional placement pool filter: ondemand or burst // Event emit for D1 sandboxes_index sync. After a scaler-triggered // migration succeeds (rolling replace, evacuation), XADD a "migrated" @@ -112,6 +129,41 @@ type drainState struct { MachineID string `json:"machine_id"` Region string `json:"region"` StartedAt time.Time `json:"started_at"` + Reason string `json:"reason,omitempty"` +} + +const ( + drainReasonScaleDown = "scale_down" + drainReasonRollingReplace = "rolling_replace" + drainReasonSpotPreemption = "spot_preemption" + maxDrainMigrationFailures = 3 +) + +type drainPolicy struct { + reason string + backoff time.Duration + failurePause time.Duration + concurrency int +} + +func drainPolicyForReason(reason string) drainPolicy { + if reason == drainReasonSpotPreemption { + return drainPolicy{ + reason: reason, + backoff: spotDrainBackoff, + failurePause: spotDrainFailurePause, + concurrency: spotDrainConcurrency, + } + } + if reason == "" { + reason = drainReasonScaleDown + } + return drainPolicy{ + reason: reason, + backoff: normalDrainBackoff, + failurePause: drainBatchFailurePause, + concurrency: evacuationBatchSize, + } } // Scaler manages autoscaling of workers via the compute Pool. @@ -123,17 +175,21 @@ type Scaler struct { image string cooldown time.Duration interval time.Duration - minWorkers int - maxWorkers int - idleReserve int - - rdb *redis.Client - cellID string - - mu sync.Mutex // protects stop/cancel - cancel context.CancelFunc - wg sync.WaitGroup - running bool + minWorkers int + maxWorkers int + idleReserve int + minIdleCap int + minIdleCPUs int + sandboxCPUs int + workerPool string + + rdb *redis.Client + cellID string + + mu sync.Mutex // protects stop/cancel + cancel context.CancelFunc + wg sync.WaitGroup + running bool machineSizes []string // ranked list of provider-specific sizes for scale-up fallback @@ -164,6 +220,10 @@ func NewScaler(cfg ScalerConfig) *Scaler { if idleReserve < 0 { idleReserve = 0 } + sandboxCPUs := cfg.DefaultSandboxCPUs + if sandboxCPUs <= 0 { + sandboxCPUs = 1 + } stateStore := cfg.StateStore if stateStore == nil { stateStore = NewInMemoryScalerState() @@ -180,6 +240,10 @@ func NewScaler(cfg ScalerConfig) *Scaler { minWorkers: minWorkers, maxWorkers: maxWorkers, idleReserve: idleReserve, + minIdleCap: cfg.MinIdleCapacity, + minIdleCPUs: cfg.MinIdleCPUs, + sandboxCPUs: sandboxCPUs, + workerPool: cfg.WorkerPool, machineSizes: cfg.MachineSizes, rdb: cfg.RedisClient, cellID: cfg.CellID, @@ -227,7 +291,7 @@ func (s *Scaler) Start() { } } }() - log.Printf("scaler: autoscaling controller started (interval=%s, cooldown=%s)", s.interval, s.cooldown) + log.Printf("scaler: autoscaling controller started (pool=%s, interval=%s, cooldown=%s)", NormalizeWorkerPool(s.workerPool), s.interval, s.cooldown) } // Stop stops the autoscaling loop. Can be called multiple times (idempotent). @@ -243,6 +307,85 @@ func (s *Scaler) Stop() { s.wg.Wait() } +func (s *Scaler) workersByRegion(region string) []*WorkerInfo { + workers := s.registry.GetWorkersByRegion(region) + if s.workerPool == "" { + return workers + } + pool := NormalizeWorkerPool(s.workerPool) + filtered := workers[:0] + for _, w := range workers { + if NormalizeWorkerPool(w.Pool) == pool { + filtered = append(filtered, w) + } + } + return filtered +} + +func (s *Scaler) regionUtilization(workers []*WorkerInfo) float64 { + if len(workers) == 0 { + return 0 + } + var totalCapacity, totalCurrent int + for _, w := range workers { + totalCapacity += w.Capacity + totalCurrent += w.Current + } + if totalCapacity == 0 { + return 0 + } + return float64(totalCurrent) / float64(totalCapacity) +} + +func (s *Scaler) regionResourcePressure(workers []*WorkerInfo) (maxCPU, maxMem, maxDisk float64) { + for _, w := range workers { + if w.CPUPct > maxCPU { + maxCPU = w.CPUPct + } + if w.MemPct > maxMem { + maxMem = w.MemPct + } + if w.DiskPct > maxDisk { + maxDisk = w.DiskPct + } + } + return maxCPU, maxMem, maxDisk +} + +// EvacuateWorker starts the normal live-migration drain loop for a specific +// worker without terminating the machine when the worker becomes empty. This is +// intended for operator-triggered evacuation tests and spot-preemption drills. +func (s *Scaler) EvacuateWorker(_ context.Context, workerID string) error { + s.mu.Lock() + defer s.mu.Unlock() + + target := s.getWorkerInfo(workerID) + if target == nil { + return fmt.Errorf("worker %s not found", workerID) + } + if target.MachineID == "" { + return fmt.Errorf("worker %s has no machine id", workerID) + } + if !s.state.TryAcquireEvacuationLock() { + return fmt.Errorf("another evacuation is already running") + } + + s.state.SetDraining(target.MachineID, &drainState{ + WorkerID: target.ID, + MachineID: target.MachineID, + Region: target.Region, + StartedAt: time.Now(), + Reason: drainReasonSpotPreemption, + }) + + go func() { + defer s.state.ReleaseEvacuationLock() + s.drainWorker(target.ID, target.MachineID, target.Region, drainReasonSpotPreemption) + }() + + return nil +} + func (s *Scaler) evaluate() { ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() @@ -284,18 +427,53 @@ func (s *Scaler) evaluate() { regions = poolRegions } } + s.recreateResumableOnDeadWorkers(ctx) for _, region := range regions { s.evaluateRegion(ctx, region) } } +func (s *Scaler) recreateResumableOnDeadWorkers(ctx context.Context) { + if s.store == nil { + return + } + redisRegistry, ok := s.registry.(*RedisWorkerRegistry) + if !ok { + return + } + liveWorkers := make(map[string]bool) + for _, region := range s.registry.Regions() { + for _, w := range s.registry.GetWorkersByRegion(region) { + liveWorkers[w.ID] = true + } + } + if len(liveWorkers) == 0 { + return + } + deadRefs, err := s.store.ListRunningSandboxesOnDeadWorkers(ctx, liveWorkers) + if err != nil { + log.Printf("scaler: resumable dead-worker scan failed: %v", err) + return + } + for _, ref := range deadRefs { + ok, err := recreateResumableSandbox(ctx, redisRegistry, s.store, s.cellID, ref.SandboxID, ref.WorkerID) + if err != nil { + log.Printf("scaler: resumable recreate %s from dead worker %s failed: %v", ref.SandboxID, ref.WorkerID, err) + continue + } + if ok { + log.Printf("scaler: resumable recreate %s from dead worker %s succeeded", ref.SandboxID, ref.WorkerID) + } + } +} + func (s *Scaler) evaluateRegion(ctx context.Context, region string) { s.mu.Lock() defer s.mu.Unlock() - workers := s.registry.GetWorkersByRegion(region) - utilization := s.registry.RegionUtilization(region) - maxCPU, maxMem, maxDisk := s.registry.RegionResourcePressure(region) + workers := s.workersByRegion(region) + utilization := s.regionUtilization(workers) + maxCPU, maxMem, maxDisk := s.regionResourcePressure(workers) // Expire stale pending launches s.expirePending(region) @@ -358,10 +536,15 @@ func (s *Scaler) evaluateRegion(ctx context.Context, region string) { } } - // Ensure minimum workers are running (pre-provisioned capacity). - // Ignores cooldowns but respects creation failure backoff. totalWorkers := len(workers) + len(s.state.GetPendingLaunches(region)) - if totalWorkers < s.minWorkers { + + if s.hasMinIdleReserve() { + if s.ensureMinIdleReserve(ctx, region, workers, totalWorkers) { + return + } + } else if totalWorkers < s.minWorkers { + // Ensure minimum workers are running (pre-provisioned capacity). + // Ignores cooldowns but respects creation failure backoff. if until, ok := s.state.GetCreationBackoffUntil(region); ok { log.Printf("scaler: region %s below minimum workers (%d/%d) but creation backoff active until %s", region, totalWorkers, s.minWorkers, until.Format(time.RFC3339)) @@ -376,29 +559,31 @@ func (s *Scaler) evaluateRegion(ctx context.Context, region string) { return } - // Headroom: maintain a pool of idle workers for burst absorption. - // Uses minWorkers as the reserve target — this is separate from the - // minimum total workers check above. When bin-packing overflows into - // reserve workers, we launch replacements one at a time so there's - // always warm capacity without thrashing. - idleWorkers := 0 - for _, w := range workers { - if s.state.IsDraining(w.MachineID) { - continue + if !s.hasMinIdleReserve() { + // Headroom: maintain a pool of idle workers for burst absorption. + // Uses minWorkers as the reserve target — this is separate from the + // minimum total workers check above. When bin-packing overflows into + // reserve workers, we launch replacements one at a time so there's + // always warm capacity without thrashing. + idleWorkers := 0 + for _, w := range workers { + if s.state.IsDraining(w.MachineID) { + continue + } + if w.Current == 0 { + idleWorkers++ + } } - if w.Current == 0 { - idleWorkers++ + pendingCount := len(s.state.GetPendingLaunches(region)) + reserveTarget := s.idleReserve + idleOrPending := idleWorkers + pendingCount + if idleOrPending < reserveTarget && totalWorkers < s.maxWorkers { + // Launch 1 at a time to avoid over-provisioning + log.Printf("scaler: region %s reserve low (%d idle + %d pending < %d target), launching 1", + region, idleWorkers, pendingCount, reserveTarget) + s.scaleUp(ctx, region) } } - pendingCount := len(s.state.GetPendingLaunches(region)) - reserveTarget := s.idleReserve - idleOrPending := idleWorkers + pendingCount - if idleOrPending < reserveTarget && totalWorkers+pendingCount < s.maxWorkers { - // Launch 1 at a time to avoid over-provisioning - log.Printf("scaler: region %s reserve low (%d idle + %d pending < %d target), launching 1", - region, idleWorkers, pendingCount, reserveTarget) - s.scaleUp(ctx, region) - } if needsScaleUp { // Cascade of guards. Important: each guard logs and SKIPS scale-up but @@ -438,7 +623,7 @@ func (s *Scaler) evaluateRegion(ctx context.Context, region string) { region, reason, maxCPU, maxMem, maxDisk, utilization*100) s.scaleUp(ctx, region) } - } else if utilization < scaleDownThreshold && len(workers) > s.minWorkers { + } else if utilization < scaleDownThreshold && s.canScaleDown(region, workers) { // Phase 4: Scale down via smart drain (live-migrate sandboxes, then destroy) log.Printf("scaler: region %s utilization %.1f%% < %.0f%%, initiating smart drain", region, utilization*100, scaleDownThreshold*100) s.smartScaleDown(ctx, region, workers) @@ -448,6 +633,92 @@ func (s *Scaler) evaluateRegion(ctx context.Context, region string) { s.rollingReplace(ctx, region, workers) } +func (s *Scaler) hasMinIdleReserve() bool { + return s.minIdleCPUs > 0 || s.minIdleCap > 0 +} + +func (s *Scaler) minIdleReserveTarget() (target, multiplier int, label string) { + if s.minIdleCPUs > 0 { + return s.minIdleCPUs, s.sandboxCPUs, "cpu" + } + return s.minIdleCap, 1, "slot" +} + +func (s *Scaler) ensureMinIdleReserve(ctx context.Context, region string, workers []*WorkerInfo, totalWorkers int) bool { + target, multiplier, label := s.minIdleReserveTarget() + idleCapacity, pendingCapacity, estimate := s.idleCapacity(region, workers, multiplier) + if idleCapacity+pendingCapacity >= target { + return false + } + if until, ok := s.state.GetCreationBackoffUntil(region); ok { + log.Printf("scaler: region %s idle %s capacity low (%d live + %d pending < %d target) but creation backoff active until %s", + region, label, idleCapacity, pendingCapacity, target, until.Format(time.RFC3339)) + return true + } + remainingWorkers := s.maxWorkers - totalWorkers + if remainingWorkers <= 0 { + log.Printf("scaler: region %s idle %s capacity low (%d live + %d pending < %d target) but at max workers (%d/%d)", + region, label, idleCapacity, pendingCapacity, target, totalWorkers, s.maxWorkers) + return false + } + + deficit := target - idleCapacity - pendingCapacity + launches := (deficit + estimate - 1) / estimate + if launches > remainingWorkers { + launches = remainingWorkers + } + log.Printf("scaler: region %s idle %s capacity low (%d live + %d pending < %d target), launching %d worker(s) using ~%d %s/worker", + region, label, idleCapacity, pendingCapacity, target, launches, estimate, label) + for i := 0; i < launches; i++ { + s.scaleUp(ctx, region) + } + return true +} + +func (s *Scaler) idleCapacity(region string, workers []*WorkerInfo, multiplier int) (idleCapacity, pendingCapacity, estimate int) { + if multiplier <= 0 { + multiplier = 1 + } + estimate = defaultWorkerCap * multiplier + for _, w := range workers { + if w.Draining || s.state.IsDraining(w.MachineID) { + continue + } + workerCapacity := w.Capacity * multiplier + if workerCapacity > estimate { + estimate = workerCapacity + } + if w.Capacity > w.Current { + idleCapacity += (w.Capacity - w.Current) * multiplier + } + } + pendingCapacity = len(s.state.GetPendingLaunches(region)) * estimate + return idleCapacity, pendingCapacity, estimate +} + +func (s *Scaler) canScaleDown(region string, workers []*WorkerInfo) bool { + if len(workers) <= s.minWorkers { + return false + } + if !s.hasMinIdleReserve() { + return true + } + + target, multiplier, _ := s.minIdleReserveTarget() + idleCapacity, pendingCapacity, _ := s.idleCapacity(region, workers, multiplier) + bestRemovable := 0 + for _, w := range workers { + if w.Draining || s.state.IsDraining(w.MachineID) { + continue + } + spare := (w.Capacity - w.Current) * multiplier + if spare > bestRemovable { + bestRemovable = spare + } + } + return idleCapacity+pendingCapacity-bestRemovable >= target +} + func (s *Scaler) scaleUp(_ context.Context, region string) { // Check creation failure backoff if until, ok := s.state.GetCreationBackoffUntil(region); ok { @@ -550,7 +821,7 @@ func (s *Scaler) expirePending(region string) { // Get currently registered worker machine IDs registered := make(map[string]bool) - for _, w := range s.registry.GetWorkersByRegion(region) { + for _, w := range s.workersByRegion(region) { if w.MachineID != "" { registered[w.MachineID] = true } @@ -742,7 +1013,7 @@ func (s *Scaler) evacuateHotWorkers(_ context.Context, region string, workers [] // need 2x the workers of one at 50% actual — expensive dead weight for an // idle-heavy workload like sandboxes. func (s *Scaler) findMigrationTarget(region, excludeWorkerID string, requiredMemMB int32) *WorkerInfo { - workers := s.registry.GetWorkersByRegion(region) + workers := s.workersByRegion(region) var best *WorkerInfo bestScore := -1.0 @@ -750,9 +1021,13 @@ func (s *Scaler) findMigrationTarget(region, excludeWorkerID string, requiredMem if w.ID == excludeWorkerID { continue } - if s.state.IsDraining(w.MachineID) { + if w.Draining { continue } + if s.state.IsDraining(w.MachineID) { + log.Printf("scaler: migration target %s had stale drain state without live registry drain flag; clearing", w.ID) + s.state.RemoveDraining(w.MachineID) + } // Subtract in-flight migrations from remaining capacity pending := s.state.GetInFlight(w.ID) remaining := w.Capacity - w.Current - pending @@ -922,9 +1197,10 @@ func (s *Scaler) smartScaleDown(_ context.Context, region string, workers []*Wor MachineID: target.MachineID, Region: region, StartedAt: time.Now(), + Reason: drainReasonScaleDown, }) - go s.drainWorker(target.ID, target.MachineID, region) + go s.drainWorker(target.ID, target.MachineID, region, drainReasonScaleDown) } // rollingReplace executes a quota-aware rolling replacement of stale workers @@ -1034,6 +1310,7 @@ func (s *Scaler) rollingReplace(ctx context.Context, region string, workers []*W MachineID: target.MachineID, Region: region, StartedAt: time.Now(), + Reason: drainReasonRollingReplace, }) // Run the dance in a goroutine so we don't block the scaler tick. @@ -1058,7 +1335,7 @@ func (s *Scaler) replaceOneStale(ctx context.Context, region string, target *Wor // fires. drainWorker handles per-sandbox findMigrationTarget so each // sandbox lands on whichever current-version worker has the most room // at that exact moment. - s.drainWorker(target.ID, target.MachineID, region) + s.drainWorker(target.ID, target.MachineID, region, drainReasonRollingReplace) // 2. Terminate the (now-empty) stale worker. Frees the quota slot for the // replacement scaleUp below. We do this even on partial drain — the @@ -1091,9 +1368,13 @@ func (s *Scaler) replaceOneStale(ctx context.Context, region string, target *Wor // (e.g., S3 auth, no targets), falls back to waiting for natural expiry — // sandboxes will timeout or be destroyed by users. No new sandboxes are routed // to draining workers. -func (s *Scaler) drainWorker(workerID, machineID, region string) { +func (s *Scaler) drainWorker(workerID, machineID, region, reason string) { ctx, cancel := context.WithTimeout(context.Background(), drainTimeout) defer cancel() + policy := drainPolicyForReason(reason) + if machineID != "" { + defer s.state.RemoveDraining(machineID) + } // Mark worker as draining so routing skips it s.registry.SetDraining(workerID, true) @@ -1106,7 +1387,6 @@ func (s *Scaler) drainWorker(workerID, machineID, region string) { } migrationFailures := 0 - const maxMigrationFailures = 3 // after 3 failed attempts, stop trying migration for { select { @@ -1162,14 +1442,14 @@ func (s *Scaler) drainWorker(workerID, machineID, region string) { // and a periodic reset lets us succeed once the transient clears. If // migration genuinely can't complete, drainTimeout terminates the loop // and the next eval tick takes over. - if migrationFailures >= maxMigrationFailures { - log.Printf("scaler: drain: %d migration failures on %s, backing off 5min before retry (%d sandboxes remaining)", - migrationFailures, workerID, len(running)) + if migrationFailures >= maxDrainMigrationFailures { + log.Printf("scaler: drain: %d migration failures on %s, backing off %s before retry (reason=%s, %d sandboxes remaining)", + migrationFailures, workerID, policy.backoff, policy.reason, len(running)) select { case <-ctx.Done(): log.Printf("scaler: drain: timeout reached for worker %s", workerID) return - case <-time.After(5 * time.Minute): + case <-time.After(policy.backoff): } migrationFailures = 0 continue @@ -1187,13 +1467,13 @@ func (s *Scaler) drainWorker(workerID, machineID, region string) { // utilization-based scale-up path in Evaluate() won't trigger. if probe := s.findMigrationTarget(region, workerID, 0); probe == nil { effective := 0 - for _, w := range s.registry.GetWorkersByRegion(region) { + for _, w := range s.workersByRegion(region) { if !s.state.IsDraining(w.MachineID) { effective++ } } pending := len(s.state.GetPendingLaunches(region)) - if effective+pending < s.maxWorkers { + if s.pool != nil && effective+pending < s.maxWorkers { log.Printf("scaler: drain: no migration target for %s (%d sandboxes), triggering scale-up", workerID, len(running)) s.scaleUp(ctx, region) select { @@ -1203,46 +1483,73 @@ func (s *Scaler) drainWorker(workerID, machineID, region string) { } continue } + if s.pool == nil { + log.Printf("scaler: drain: no migration target for %s (%d sandboxes) and no compute pool configured, waiting for natural expiry", workerID, len(running)) + s.waitForNaturalDrain(ctx, workerID) + return + } log.Printf("scaler: drain: no migration target for %s (%d sandboxes) and at max workers (%d), waiting for natural expiry", workerID, len(running), s.maxWorkers) s.waitForNaturalDrain(ctx, workerID) return } - // Migrate a batch — bounded parallelism to avoid overwhelming - // network/disk on source and target workers. Each sandbox picks its - // own target (based on its own memory footprint) inside liveMigrateSandbox. - batch := running - if len(batch) > evacuationBatchSize { - batch = batch[:evacuationBatchSize] - } - - batchFailed := false - var wg sync.WaitGroup - var failCount int64 - for _, sandboxID := range batch { - wg.Add(1) - go func(sbID string) { - defer wg.Done() - if err := s.liveMigrateSandbox(ctx, sbID, workerID, ""); err != nil { - log.Printf("scaler: drain: migrate %s failed: %v", sbID, err) - atomic.AddInt64(&failCount, 1) - } - }(sandboxID) - } - wg.Wait() + // Migrate with bounded rolling parallelism. A slow sandbox occupies one + // worker slot, but does not block the remaining pending sandboxes from + // starting as other slots complete. + failCount := s.drainMigrateSandboxes(ctx, workerID, running, policy.concurrency) + batchFailed := failCount > 0 if failCount > 0 { migrationFailures += int(failCount) - batchFailed = true } if batchFailed { - time.Sleep(5 * time.Second) + time.Sleep(policy.failurePause) } else { - time.Sleep(2 * time.Second) + time.Sleep(drainBatchSuccessPause) } } } +func (s *Scaler) drainMigrateSandboxes(ctx context.Context, workerID string, sandboxIDs []string, concurrency int) int64 { + if concurrency <= 0 { + concurrency = evacuationBatchSize + } + if concurrency > len(sandboxIDs) { + concurrency = len(sandboxIDs) + } + + jobs := make(chan string) + var wg sync.WaitGroup + var failCount int64 + + for i := 0; i < concurrency; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for sandboxID := range jobs { + if err := s.liveMigrateSandbox(ctx, sandboxID, workerID, ""); err != nil { + log.Printf("scaler: drain: migrate %s failed: %v", sandboxID, err) + atomic.AddInt64(&failCount, 1) + } + } + }() + } + + for _, sandboxID := range sandboxIDs { + select { + case <-ctx.Done(): + close(jobs) + wg.Wait() + return failCount + case jobs <- sandboxID: + } + } + close(jobs) + wg.Wait() + + return failCount +} + // waitForNaturalDrain polls until the worker has 0 sandboxes or the context expires. // Sandboxes expire via timeout or user-initiated destroy. func (s *Scaler) waitForNaturalDrain(ctx context.Context, workerID string) { @@ -1288,7 +1595,7 @@ func (s *Scaler) checkDrainingWorkers(ctx context.Context, region string) { } // Check if worker has 0 sandboxes - workers := s.registry.GetWorkersByRegion(region) + workers := s.workersByRegion(region) for _, w := range workers { if w.MachineID == machineID && w.Current == 0 { log.Printf("scaler: worker %s fully drained (0 sandboxes), destroying machine %s", @@ -1541,7 +1848,7 @@ func (s *Scaler) liveMigrateSandbox(ctx context.Context, sandboxID, sourceWorker // no QEMU, drain's ListSandboxes returns 0 and exits cleanly leaving // the row stuck. Mark error instead so the sandbox is visibly broken. if qmpSucceeded { - if err := s.store.FailMigrationPostQMP(ctx, sandboxID, "migration failed after QMP transfer; source VM gone, target failed to complete"); err != nil { + if err := s.store.FailMigrationPostQMP(ctx, sandboxID, db.PostQMPMigrationFailureMessage); err != nil { log.Printf("scaler: migrate %s: FailMigrationPostQMP failed: %v", sandboxID, err) } // Source VM is gone, target failed. Emit a `stopped` event so diff --git a/internal/controlplane/scaler_test.go b/internal/controlplane/scaler_test.go index 663bed17..ab4f5852 100644 --- a/internal/controlplane/scaler_test.go +++ b/internal/controlplane/scaler_test.go @@ -173,12 +173,12 @@ func (p *mockPool) DestroyMachine(_ context.Context, machineID string) error { return nil } -func (p *mockPool) DrainMachine(_ context.Context, _ string) error { return nil } -func (p *mockPool) StartMachine(_ context.Context, _ string) error { return nil } -func (p *mockPool) StopMachine(_ context.Context, _ string) error { return nil } -func (p *mockPool) HealthCheck(_ context.Context, _ string) error { return nil } -func (p *mockPool) CleanupOrphanedResources(_ context.Context) (int, error) { return 0, nil } -func (p *mockPool) ListMachines(_ context.Context) ([]*compute.Machine, error) { return nil, nil } +func (p *mockPool) DrainMachine(_ context.Context, _ string) error { return nil } +func (p *mockPool) StartMachine(_ context.Context, _ string) error { return nil } +func (p *mockPool) StopMachine(_ context.Context, _ string) error { return nil } +func (p *mockPool) HealthCheck(_ context.Context, _ string) error { return nil } +func (p *mockPool) CleanupOrphanedResources(_ context.Context) (int, error) { return 0, nil } +func (p *mockPool) ListMachines(_ context.Context) ([]*compute.Machine, error) { return nil, nil } func (p *mockPool) SupportedRegions(_ context.Context) ([]string, error) { return []string{"us-east-1"}, nil } @@ -402,6 +402,211 @@ func TestMinWorkersEnforced(t *testing.T) { } } +func TestMinIdleCapacityOverridesMinWorkers(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 5, + MaxWorkers: 10, + MinIdleCapacity: 120, + }) + + s.evaluateRegion(context.Background(), "us-east-1") + time.Sleep(50 * time.Millisecond) + + if got := atomic.LoadInt32(&pool.created); got != 2 { + t.Fatalf("expected 2 launches to add ~100 spare slots, got %d", got) + } +} + +func TestMinIdleCapacityCountsPendingLaunches(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 1, + MaxWorkers: 10, + MinIdleCapacity: 100, + }) + s.state.AddPendingLaunch("us-east-1", pendingLaunch{ + MachineID: "osb-worker-pending-test", + LaunchedAt: time.Now(), + }) + + s.evaluateRegion(context.Background(), "us-east-1") + time.Sleep(50 * time.Millisecond) + + if got := atomic.LoadInt32(&pool.created); got != 0 { + t.Fatalf("expected no launches because pending capacity satisfies reserve, got %d", got) + } +} + +func TestMinIdleCapacityRespectsMaxWorkers(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + for i := 0; i < 2; i++ { + reg.addWorker(&WorkerInfo{ + ID: fmt.Sprintf("w%d", i), MachineID: fmt.Sprintf("osb-worker-w%d", i), Region: "us-east-1", + Capacity: 50, Current: 45, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + } + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 1, + MaxWorkers: 2, + MinIdleCapacity: 100, + }) + + s.evaluateRegion(context.Background(), "us-east-1") + time.Sleep(50 * time.Millisecond) + + if got := atomic.LoadInt32(&pool.created); got != 0 { + t.Fatalf("expected no launches at max workers, got %d", got) + } +} + +func TestMinIdleCapacityPreventsScaleDownBelowReserve(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + reg.addWorker(&WorkerInfo{ + ID: "w2", MachineID: "osb-worker-w2", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 1, + MaxWorkers: 10, + MinIdleCapacity: 80, + }) + + if s.canScaleDown("us-east-1", reg.GetWorkersByRegion("us-east-1")) { + t.Fatal("expected min idle capacity reserve to block scale-down") + } +} + +func TestMinIdleCPUsOverridesMinIdleCapacity(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 1, + MaxWorkers: 10, + MinIdleCapacity: 200, + MinIdleCPUs: 120, + DefaultSandboxCPUs: 2, + }) + + s.evaluateRegion(context.Background(), "us-east-1") + time.Sleep(50 * time.Millisecond) + + if got := atomic.LoadInt32(&pool.created); got != 1 { + t.Fatalf("expected 1 launch using cpu reserve target, got %d", got) + } +} + +func TestMinIdleCPUsCountsPendingLaunches(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 1, + MaxWorkers: 10, + MinIdleCPUs: 200, + DefaultSandboxCPUs: 2, + }) + s.state.AddPendingLaunch("us-east-1", pendingLaunch{ + MachineID: "osb-worker-pending-test", + LaunchedAt: time.Now(), + }) + + s.evaluateRegion(context.Background(), "us-east-1") + time.Sleep(50 * time.Millisecond) + + if got := atomic.LoadInt32(&pool.created); got != 0 { + t.Fatalf("expected no launches because pending CPU reserve satisfies target, got %d", got) + } +} + +func TestMinIdleCPUsPreventsScaleDownBelowReserve(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + reg.addWorker(&WorkerInfo{ + ID: "w2", MachineID: "osb-worker-w2", Region: "us-east-1", + Capacity: 50, Current: 0, CPUPct: 0, MemPct: 0, DiskPct: 0, + }) + + s := NewScaler(ScalerConfig{ + Pool: pool, + Registry: reg, + Cooldown: 1 * time.Second, + Interval: 100 * time.Millisecond, + MinWorkers: 1, + MaxWorkers: 10, + MinIdleCPUs: 160, + DefaultSandboxCPUs: 2, + }) + + if s.canScaleDown("us-east-1", reg.GetWorkersByRegion("us-east-1")) { + t.Fatal("expected min idle CPU reserve to block scale-down") + } +} + // ============================================================ // Test: Scale-down triggers // ============================================================ @@ -414,14 +619,17 @@ func TestSmartScaleDownTargetsLeastLoaded(t *testing.T) { reg.addWorker(&WorkerInfo{ ID: "w1", MachineID: "osb-worker-w1", Region: "us-east-1", Capacity: 50, Current: 5, CPUPct: 20, MemPct: 20, DiskPct: 20, + TotalMemoryMB: 64000, CommittedMemoryMB: 12800, }) reg.addWorker(&WorkerInfo{ ID: "w2", MachineID: "osb-worker-w2", Region: "us-east-1", Capacity: 50, Current: 2, CPUPct: 10, MemPct: 10, DiskPct: 10, + TotalMemoryMB: 64000, CommittedMemoryMB: 6400, }) reg.addWorker(&WorkerInfo{ ID: "w3", MachineID: "osb-worker-w3", Region: "us-east-1", Capacity: 50, Current: 5, CPUPct: 20, MemPct: 20, DiskPct: 20, + TotalMemoryMB: 64000, CommittedMemoryMB: 12800, }) s := newTestScaler(reg, pool) @@ -463,10 +671,17 @@ func TestScaleDownSkipsAlreadyDraining(t *testing.T) { reg.addWorker(&WorkerInfo{ ID: "w2", MachineID: "osb-worker-w2", Region: "us-east-1", Capacity: 50, Current: 2, CPUPct: 10, MemPct: 10, DiskPct: 10, + TotalMemoryMB: 64000, CommittedMemoryMB: 6400, }) reg.addWorker(&WorkerInfo{ ID: "w3", MachineID: "osb-worker-w3", Region: "us-east-1", Capacity: 50, Current: 5, CPUPct: 20, MemPct: 20, DiskPct: 20, + TotalMemoryMB: 64000, CommittedMemoryMB: 12800, + }) + reg.addWorker(&WorkerInfo{ + ID: "w4", MachineID: "osb-worker-w4", Region: "us-east-1", + Capacity: 50, Current: 10, CPUPct: 20, MemPct: 20, DiskPct: 20, + TotalMemoryMB: 64000, CommittedMemoryMB: 12800, }) s := newTestScaler(reg, pool) @@ -708,7 +923,7 @@ func TestFindMigrationTargetSkipsDraining(t *testing.T) { }) reg.addWorker(&WorkerInfo{ ID: "draining", MachineID: "osb-worker-draining", Region: "us-east-1", - Capacity: 50, Current: 5, CPUPct: 10, MemPct: 10, DiskPct: 10, + Capacity: 50, Current: 5, CPUPct: 10, MemPct: 10, DiskPct: 10, Draining: true, }) s := newTestScaler(reg, pool) @@ -720,6 +935,31 @@ func TestFindMigrationTargetSkipsDraining(t *testing.T) { } } +func TestFindMigrationTargetClearsStaleDrainState(t *testing.T) { + reg := newMockRegistry() + pool := newMockPool() + + reg.addWorker(&WorkerInfo{ + ID: "hot", MachineID: "osb-worker-hot", Region: "us-east-1", + Capacity: 50, Current: 45, CPUPct: 85, MemPct: 50, DiskPct: 30, + }) + reg.addWorker(&WorkerInfo{ + ID: "candidate", MachineID: "osb-worker-candidate", Region: "us-east-1", + Capacity: 50, Current: 5, CPUPct: 10, MemPct: 10, DiskPct: 10, + }) + + s := newTestScaler(reg, pool) + s.state.SetDraining("osb-worker-candidate", &drainState{WorkerID: "candidate"}) + + target := s.findMigrationTarget("us-east-1", "hot", 0) + if target == nil || target.ID != "candidate" { + t.Fatalf("expected stale-drain candidate to be usable, got %#v", target) + } + if s.state.IsDraining("osb-worker-candidate") { + t.Fatalf("expected stale drain state to be cleared") + } +} + func TestFindMigrationTargetAccountsForInFlight(t *testing.T) { reg := newMockRegistry() pool := newMockPool() @@ -857,7 +1097,7 @@ func TestDrainTimeoutCancelsDrainKeepsWorker(t *testing.T) { WorkerID: "w2", MachineID: "osb-worker-w2", Region: "us-east-1", - StartedAt: time.Now().Add(-20 * time.Minute), // well past drainTimeout + StartedAt: time.Now().Add(-(drainTimeout + time.Minute)), }) ctx := context.Background() @@ -1778,6 +2018,36 @@ func TestInMemoryScalerState_Draining(t *testing.T) { } } +func TestDrainPolicyForReason(t *testing.T) { + normal := drainPolicyForReason(drainReasonScaleDown) + if normal.backoff != normalDrainBackoff { + t.Fatalf("expected normal drain backoff %s, got %s", normalDrainBackoff, normal.backoff) + } + if normal.failurePause != drainBatchFailurePause { + t.Fatalf("expected normal drain failure pause %s, got %s", drainBatchFailurePause, normal.failurePause) + } + if normal.concurrency != evacuationBatchSize { + t.Fatalf("expected normal drain concurrency %d, got %d", evacuationBatchSize, normal.concurrency) + } + + spot := drainPolicyForReason(drainReasonSpotPreemption) + if spot.backoff != spotDrainBackoff { + t.Fatalf("expected spot drain backoff %s, got %s", spotDrainBackoff, spot.backoff) + } + if spot.failurePause != spotDrainFailurePause { + t.Fatalf("expected spot drain failure pause %s, got %s", spotDrainFailurePause, spot.failurePause) + } + if spot.concurrency != spotDrainConcurrency { + t.Fatalf("expected spot drain concurrency %d, got %d", spotDrainConcurrency, spot.concurrency) + } + if spot.backoff >= normal.backoff { + t.Fatalf("expected spot drain backoff to be shorter than normal drain backoff") + } + if spot.concurrency <= normal.concurrency { + t.Fatalf("expected spot drain concurrency to be higher than normal drain concurrency") + } +} + func TestInMemoryScalerState_AllDraining(t *testing.T) { m := NewInMemoryScalerState() diff --git a/internal/controlplane/spot_evacuation_sim_test.go b/internal/controlplane/spot_evacuation_sim_test.go new file mode 100644 index 00000000..2edf435f --- /dev/null +++ b/internal/controlplane/spot_evacuation_sim_test.go @@ -0,0 +1,141 @@ +package controlplane + +import ( + "math" + "testing" + "time" +) + +type spotEvacuationScenario struct { + name string + sourceSlots int + sourceUsedPct int + spareSlots int + spareUsed int + migrationP95 time.Duration + noticeWindow time.Duration +} + +type spotEvacuationResult struct { + workloads int + spareAvailable int + moved int + remaining int + waves int + duration time.Duration +} + +func simulateSpotEvacuation(sc spotEvacuationScenario) spotEvacuationResult { + workloads := sc.sourceSlots * sc.sourceUsedPct / 100 + spareAvailable := sc.spareSlots - sc.spareUsed + if spareAvailable < 0 { + spareAvailable = 0 + } + + moved := workloads + if moved > spareAvailable { + moved = spareAvailable + } + remaining := workloads - moved + waves := int(math.Ceil(float64(moved) / float64(evacuationBatchSize))) + + var duration time.Duration + if waves > 0 { + duration = time.Duration(waves)*sc.migrationP95 + time.Duration(waves-1)*drainBatchSuccessPause + } + + return spotEvacuationResult{ + workloads: workloads, + spareAvailable: spareAvailable, + moved: moved, + remaining: remaining, + waves: waves, + duration: duration, + } +} + +func TestSpotEvacuationScenariosWithOneSpareWorker(t *testing.T) { + const ( + sourceSlots = 100 + spareSlots = 100 + ) + + cases := []struct { + sc spotEvacuationScenario + wantWithinNotice bool + }{ + { + sc: spotEvacuationScenario{ + name: "source 25 percent full", + sourceSlots: sourceSlots, + sourceUsedPct: 25, + spareSlots: spareSlots, + migrationP95: 5 * time.Second, + noticeWindow: 2 * time.Minute, + }, + wantWithinNotice: true, + }, + { + sc: spotEvacuationScenario{ + name: "source 50 percent full", + sourceSlots: sourceSlots, + sourceUsedPct: 50, + spareSlots: spareSlots, + migrationP95: 5 * time.Second, + noticeWindow: 2 * time.Minute, + }, + wantWithinNotice: true, + }, + { + sc: spotEvacuationScenario{ + name: "source 75 percent full", + sourceSlots: sourceSlots, + sourceUsedPct: 75, + spareSlots: spareSlots, + migrationP95: 5 * time.Second, + noticeWindow: 2 * time.Minute, + }, + wantWithinNotice: false, + }, + } + + for _, tc := range cases { + t.Run(tc.sc.name, func(t *testing.T) { + got := simulateSpotEvacuation(tc.sc) + + if got.remaining != 0 { + t.Fatalf("expected spare worker to fit all workloads, moved=%d remaining=%d", got.moved, got.remaining) + } + + withinNotice := got.duration <= tc.sc.noticeWindow + if withinNotice != tc.wantWithinNotice { + t.Fatalf("duration=%s waves=%d workloads=%d, within notice=%v want %v", + got.duration, got.waves, got.workloads, withinNotice, tc.wantWithinNotice) + } + + t.Logf("%s: workloads=%d waves=%d estimated=%s notice=%s", + tc.sc.name, got.workloads, got.waves, got.duration, tc.sc.noticeWindow) + }) + } +} + +func TestSpotEvacuationRequiresEnoughSpareCapacity(t *testing.T) { + sc := spotEvacuationScenario{ + name: "50 percent source, half-used spare", + sourceSlots: 100, + sourceUsedPct: 50, + spareSlots: 80, + spareUsed: 40, + migrationP95: 5 * time.Second, + noticeWindow: 2 * time.Minute, + } + + got := simulateSpotEvacuation(sc) + + if got.spareAvailable != 40 { + t.Fatalf("expected 40 spare slots, got %d", got.spareAvailable) + } + if got.moved != 40 || got.remaining != 10 { + t.Fatalf("expected partial evacuation: moved=40 remaining=10, got moved=%d remaining=%d", got.moved, got.remaining) + } +} diff --git a/internal/controlplane/worker_registry.go b/internal/controlplane/worker_registry.go index 7dc9e77f..495ccfb1 100644 --- a/internal/controlplane/worker_registry.go +++ b/internal/controlplane/worker_registry.go @@ -13,22 +13,24 @@ import ( // WorkerInfo represents a registered worker. type WorkerInfo struct { - ID string `json:"worker_id"` - MachineID string `json:"machine_id,omitempty"` // EC2 instance ID - Region string `json:"region"` - GRPCAddr string `json:"grpc_addr"` - HTTPAddr string `json:"http_addr"` - Capacity int `json:"capacity"` - Current int `json:"current"` - CPUPct float64 `json:"cpu_pct"` - MemPct float64 `json:"mem_pct"` - DiskPct float64 `json:"disk_pct"` - TotalMemoryMB int `json:"total_memory_mb,omitempty"` - CommittedMemoryMB int `json:"committed_memory_mb,omitempty"` - GoldenVersion string `json:"golden_version,omitempty"` - WorkerVersion string `json:"worker_version,omitempty"` - LastSeen time.Time `json:"-"` - MissedBeats int `json:"-"` + ID string `json:"worker_id"` + MachineID string `json:"machine_id,omitempty"` // EC2 instance ID + Pool string `json:"pool,omitempty"` + Region string `json:"region"` + GRPCAddr string `json:"grpc_addr"` + HTTPAddr string `json:"http_addr"` + Capacity int `json:"capacity"` + Current int `json:"current"` + CPUPct float64 `json:"cpu_pct"` + MemPct float64 `json:"mem_pct"` + DiskPct float64 `json:"disk_pct"` + Draining bool `json:"draining,omitempty"` + TotalMemoryMB int `json:"total_memory_mb,omitempty"` + CommittedMemoryMB int `json:"committed_memory_mb,omitempty"` + GoldenVersion string `json:"golden_version,omitempty"` + WorkerVersion string `json:"worker_version,omitempty"` + LastSeen time.Time `json:"-"` + MissedBeats int `json:"-"` } // WorkerRegistry tracks live workers from NATS heartbeats. diff --git a/internal/db/billable_events.go b/internal/db/billable_events.go index a0bc014a..3d42672e 100644 --- a/internal/db/billable_events.go +++ b/internal/db/billable_events.go @@ -19,9 +19,10 @@ import ( // Event types written to billable_events.event_type. Mirror the schema // CHECK constraint in migration 030. const ( - BillableEventReservedUsage = "reserved_usage" - BillableEventOverageUsage = "overage_usage" - BillableEventDiskOverageUsage = "disk_overage_usage" + BillableEventReservedUsage = "reserved_usage" + BillableEventOverageUsage = "overage_usage" + BillableEventBurstUsage = "burst_usage" + BillableEventDiskOverageUsage = "disk_overage_usage" ) // Delivery states. Mirror the schema CHECK constraint in migration 030. @@ -33,22 +34,23 @@ const ( // BillableEvent is one outbox row. // -// `MemoryMB` is 0 for `reserved_usage` and `disk_overage_usage` (sentinel -// for "not a sandbox tier"), and the running sandbox tier for +// `MemoryMB` is 0 for `reserved_usage`, `burst_usage`, and +// `disk_overage_usage` (sentinel for "not a sandbox tier"), and the +// running sandbox tier for // `overage_usage` (one row per tier per bucket via the proportional split // rule — see ws-pricing/work/001 "Per-second integration walk"). type BillableEvent struct { - ID uuid.UUID `json:"id"` - OrgID uuid.UUID `json:"orgId"` - EventType string `json:"eventType"` - MemoryMB int `json:"memoryMB"` - GBSeconds float64 `json:"gbSeconds"` - BucketStart time.Time `json:"bucketStart"` - BucketEnd time.Time `json:"bucketEnd"` - DeliveryState string `json:"deliveryState"` - StripeEventID *string `json:"stripeEventId,omitempty"` - CreatedAt time.Time `json:"createdAt"` - DeliveredAt *time.Time `json:"deliveredAt,omitempty"` + ID uuid.UUID `json:"id"` + OrgID uuid.UUID `json:"orgId"` + EventType string `json:"eventType"` + MemoryMB int `json:"memoryMB"` + GBSeconds float64 `json:"gbSeconds"` + BucketStart time.Time `json:"bucketStart"` + BucketEnd time.Time `json:"bucketEnd"` + DeliveryState string `json:"deliveryState"` + StripeEventID *string `json:"stripeEventId,omitempty"` + CreatedAt time.Time `json:"createdAt"` + DeliveredAt *time.Time `json:"deliveredAt,omitempty"` } // UpsertBillableEvent inserts a new outbox row, or no-ops if a row with @@ -304,7 +306,18 @@ func (s *Store) GetReservedGBForBucket(ctx context.Context, orgID uuid.UUID, buc // metering. See the same NOT EXISTS in GetOrgUsage for rationale. func (s *Store) GetScaleEventsForBucket(ctx context.Context, orgID uuid.UUID, bucketStart, bucketEnd time.Time) ([]ScaleEvent, error) { rows, err := s.pool.Query(ctx, ` - SELECT id, sandbox_id, org_id, memory_mb, cpu_percent, disk_mb, started_at, ended_at + SELECT id, sandbox_id, org_id, memory_mb, cpu_percent, disk_mb, + EXISTS ( + SELECT 1 + FROM sandbox_sessions ss + WHERE ss.sandbox_id = se.sandbox_id + AND ( + COALESCE((ss.config->>'burst')::boolean, false) + OR COALESCE((ss.config->>'resumable')::boolean, false) + OR ss.config->>'sandboxFamily' IN ('spot', 'resumable', 'burst') + ) + ) AS burst, + started_at, ended_at FROM sandbox_scale_events se WHERE org_id = $1 AND started_at < $3 @@ -330,7 +343,7 @@ func (s *Store) GetScaleEventsForBucket(ctx context.Context, orgID uuid.UUID, bu for rows.Next() { var e ScaleEvent var orgUUID uuid.UUID - if err := rows.Scan(&e.ID, &e.SandboxID, &orgUUID, &e.MemoryMB, &e.CPUPct, &e.DiskMB, &e.StartedAt, &e.EndedAt); err != nil { + if err := rows.Scan(&e.ID, &e.SandboxID, &orgUUID, &e.MemoryMB, &e.CPUPct, &e.DiskMB, &e.Burst, &e.StartedAt, &e.EndedAt); err != nil { return nil, fmt.Errorf("scan scale event: %w", err) } e.OrgID = orgUUID.String() diff --git a/internal/db/migration_pgfixture_test.go b/internal/db/migration_pgfixture_test.go new file mode 100644 index 00000000..fed2cbf9 --- /dev/null +++ b/internal/db/migration_pgfixture_test.go @@ -0,0 +1,138 @@ +//go:build pgfixture + +package db + +import ( + "context" + "encoding/json" + "testing" +) + +func TestCompleteMigrationRecoversPostQMPError_pgfixture(t *testing.T) { + ctx := context.Background() + store := openPgStore(t) + orgID := seedOrgWithCap(t, store, 16) + sandboxID := freshSandboxID("recover-post-qmp") + + if _, err := store.CreateSandboxSession(ctx, sandboxID, orgID, nil, "default", "us-east-2", "source-worker", json.RawMessage(`{}`), json.RawMessage(`{}`), nil); err != nil { + t.Fatalf("create sandbox session: %v", err) + } + if err := store.SetMigrating(ctx, sandboxID, "first-target"); err != nil { + t.Fatalf("set migrating: %v", err) + } + if err := store.FailMigrationPostQMP(ctx, sandboxID, PostQMPMigrationFailureMessage); err != nil { + t.Fatalf("fail post-qmp: %v", err) + } + + if err := store.CompleteMigration(ctx, sandboxID, "second-target"); err != nil { + t.Fatalf("complete migration should recover post-QMP error: %v", err) + } + + sess, err := store.GetSandboxSession(ctx, sandboxID) + if err != nil { + t.Fatalf("get sandbox session: %v", err) + } + if sess.Status != "running" { + t.Fatalf("expected running, got %q", sess.Status) + } + if sess.WorkerID != "second-target" { + t.Fatalf("expected worker second-target, got %q", sess.WorkerID) + } + if sess.MigratingToWorker != "" { + t.Fatalf("expected migrating_to_worker cleared, got %q", sess.MigratingToWorker) + } + if sess.ErrorMsg != nil { + t.Fatalf("expected error_msg cleared, got %q", *sess.ErrorMsg) + } + if sess.StoppedAt != nil { + t.Fatalf("expected stopped_at cleared, got %s", sess.StoppedAt) + } +} + +func TestCompleteMigrationDoesNotRecoverUnrelatedError_pgfixture(t *testing.T) { + ctx := context.Background() + store := openPgStore(t) + orgID := seedOrgWithCap(t, store, 16) + sandboxID := freshSandboxID("unrelated-error") + + if _, err := store.CreateSandboxSession(ctx, sandboxID, orgID, nil, "default", "us-east-2", "source-worker", json.RawMessage(`{}`), json.RawMessage(`{}`), nil); err != nil { + t.Fatalf("create sandbox session: %v", err) + } + errMsg := "some unrelated failure" + if err := store.UpdateSandboxSessionStatus(ctx, sandboxID, "error", &errMsg); err != nil { + t.Fatalf("mark unrelated error: %v", err) + } + + if err := store.CompleteMigration(ctx, sandboxID, "second-target"); err == nil { + t.Fatalf("expected unrelated error row not to recover") + } + + sess, err := store.GetSandboxSession(ctx, sandboxID) + if err != nil { + t.Fatalf("get sandbox session: %v", err) + } + if sess.Status != "error" { + t.Fatalf("expected error status preserved, got %q", sess.Status) + } + if sess.WorkerID != "source-worker" { + t.Fatalf("expected worker source-worker preserved, got %q", sess.WorkerID) + } + if sess.ErrorMsg == nil || *sess.ErrorMsg != errMsg { + t.Fatalf("expected unrelated error message preserved, got %#v", sess.ErrorMsg) + } +} + +func TestReconcileWorkerSessionsSkipsResumable_pgfixture(t *testing.T) { + ctx := context.Background() + store := openPgStore(t) + orgID := seedOrgWithCap(t, store, 16) + workerID := "worker-reconcile-resumable" + + resumableID := freshSandboxID("reconcile-resumable") + hibernatableID := freshSandboxID("reconcile-hibernatable") + stoppableID := freshSandboxID("reconcile-stoppable") + + if _, err := store.CreateSandboxSession(ctx, resumableID, orgID, nil, "default", "us-east-2", workerID, json.RawMessage(`{"resumable":true}`), json.RawMessage(`{}`), nil); err != nil { + t.Fatalf("create resumable session: %v", err) + } + if _, err := store.CreateSandboxSession(ctx, hibernatableID, orgID, nil, "default", "us-east-2", workerID, json.RawMessage(`{}`), json.RawMessage(`{}`), nil); err != nil { + t.Fatalf("create hibernatable session: %v", err) + } + if _, _, err := store.CreateHibernation(ctx, hibernatableID, orgID, "s3://checkpoint", 123, "us-east-2", "default", json.RawMessage(`{}`)); err != nil { + t.Fatalf("create hibernation: %v", err) + } + if _, err := store.CreateSandboxSession(ctx, stoppableID, orgID, nil, "default", "us-east-2", workerID, json.RawMessage(`{}`), json.RawMessage(`{}`), nil); err != nil { + t.Fatalf("create stoppable session: %v", err) + } + + hibernated, stopped, err := store.ReconcileWorkerSessions(ctx, workerID) + if err != nil { + t.Fatalf("reconcile worker sessions: %v", err) + } + if !containsOrphan(hibernated, hibernatableID) { + t.Fatalf("expected hibernatable sandbox to be hibernated, got %#v", hibernated) + } + if !containsOrphan(stopped, stoppableID) { + t.Fatalf("expected stoppable sandbox to be stopped, got %#v", stopped) + } + if containsOrphan(hibernated, resumableID) || containsOrphan(stopped, resumableID) { + t.Fatalf("resumable sandbox should not be reconciled, hibernated=%#v stopped=%#v", hibernated, stopped) + } + + resumable, err := store.GetSandboxSession(ctx, resumableID) + if err != nil { + t.Fatalf("get resumable session: %v", err) + } + if resumable.Status != "running" { + t.Fatalf("expected resumable session to stay running, got %q", resumable.Status) + } +} + +func containsOrphan(rows []OrphanedSandbox, sandboxID string) bool { + for _, row := range rows { + if row.SandboxID == sandboxID { + return true + } + } + return false +} diff --git a/internal/db/migrations/030_billable_events.up.sql b/internal/db/migrations/030_billable_events.up.sql index 443da3eb..d74f27a9 100644 --- a/internal/db/migrations/030_billable_events.up.sql +++ b/internal/db/migrations/030_billable_events.up.sql @@ -13,6 +13,8 @@ -- 'overage_usage' — instant usage above the reserved floor, -- emitted per running sandbox tier with the -- tier as memory_mb +-- 'burst_usage' — burst sandbox compute, billed as GB-seconds +-- independent of reserved capacity; memory_mb = 0 -- 'disk_overage_usage' — disk above the 20 GB allowance, org-level -- per bucket; memory_mb = 0 CREATE TABLE IF NOT EXISTS billable_events ( @@ -28,7 +30,7 @@ CREATE TABLE IF NOT EXISTS billable_events ( created_at TIMESTAMPTZ NOT NULL DEFAULT now(), delivered_at TIMESTAMPTZ, UNIQUE (org_id, event_type, memory_mb, bucket_start), - CHECK (event_type IN ('reserved_usage', 'overage_usage', 'disk_overage_usage')), + CHECK (event_type IN ('reserved_usage', 'overage_usage', 'burst_usage', 'disk_overage_usage')), CHECK (delivery_state IN ('pending', 'sent', 'failed')), CHECK (gb_seconds > 0), CHECK (memory_mb >= 0), diff --git a/internal/db/migrations/045_burst_billable_events.down.sql b/internal/db/migrations/045_burst_billable_events.down.sql new file mode 100644 index 00000000..73c29f22 --- /dev/null +++ b/internal/db/migrations/045_burst_billable_events.down.sql @@ -0,0 +1,6 @@ +ALTER TABLE billable_events + DROP CONSTRAINT IF EXISTS billable_events_event_type_check; + +ALTER TABLE billable_events + ADD CONSTRAINT billable_events_event_type_check + CHECK (event_type IN ('reserved_usage', 'overage_usage', 'disk_overage_usage')); diff --git a/internal/db/migrations/045_burst_billable_events.up.sql b/internal/db/migrations/045_burst_billable_events.up.sql new file mode 100644 index 00000000..bcc49b37 --- /dev/null +++ b/internal/db/migrations/045_burst_billable_events.up.sql @@ -0,0 +1,6 @@ +ALTER TABLE billable_events + DROP CONSTRAINT IF EXISTS billable_events_event_type_check; + +ALTER TABLE billable_events + ADD CONSTRAINT billable_events_event_type_check + CHECK (event_type IN ('reserved_usage', 'overage_usage', 'burst_usage', 'disk_overage_usage')); diff --git a/internal/db/store.go b/internal/db/store.go index 8586c9c6..5ae24aeb 100644 --- a/internal/db/store.go +++ b/internal/db/store.go @@ -142,6 +142,7 @@ func (s *Store) Migrate(ctx context.Context) error { {42, "migrations/042_global_sync_outbox.up.sql"}, {43, "migrations/043_credit_halt.up.sql"}, {44, "migrations/044_drop_secret_store_fk.up.sql"}, + {45, "migrations/045_burst_billable_events.up.sql"}, } for _, m := range migrations { @@ -187,14 +188,14 @@ type Org struct { UpdatedAt time.Time `json:"updatedAt"` // Custom domain fields - CustomDomain *string `json:"customDomain,omitempty"` - CFHostnameID *string `json:"cfHostnameId,omitempty"` - DomainVerificationStatus string `json:"domainVerificationStatus"` - DomainSSLStatus string `json:"domainSslStatus"` - VerificationTxtName *string `json:"verificationTxtName,omitempty"` - VerificationTxtValue *string `json:"verificationTxtValue,omitempty"` - SSLTxtName *string `json:"sslTxtName,omitempty"` - SSLTxtValue *string `json:"sslTxtValue,omitempty"` + CustomDomain *string `json:"customDomain,omitempty"` + CFHostnameID *string `json:"cfHostnameId,omitempty"` + DomainVerificationStatus string `json:"domainVerificationStatus"` + DomainSSLStatus string `json:"domainSslStatus"` + VerificationTxtName *string `json:"verificationTxtName,omitempty"` + VerificationTxtValue *string `json:"verificationTxtValue,omitempty"` + SSLTxtName *string `json:"sslTxtName,omitempty"` + SSLTxtValue *string `json:"sslTxtValue,omitempty"` // WorkOS organization fields WorkOSOrgID *string `json:"workosOrgId,omitempty"` @@ -208,10 +209,10 @@ type Org struct { FreeCreditsRemainingCents int64 `json:"freeCreditsRemainingCents"` // Stripe billing fields - StripeCustomerID *string `json:"stripeCustomerId,omitempty"` - StripeSubscriptionID *string `json:"stripeSubscriptionId,omitempty"` - LastUsageReportedAt time.Time `json:"lastUsageReportedAt"` - PriceLocked bool `json:"priceLocked"` + StripeCustomerID *string `json:"stripeCustomerId,omitempty"` + StripeSubscriptionID *string `json:"stripeSubscriptionId,omitempty"` + LastUsageReportedAt time.Time `json:"lastUsageReportedAt"` + PriceLocked bool `json:"priceLocked"` // Per-org billing pipeline selector. 'legacy' = UsageReporter ships // to Stripe; 'unified' = the phase-3 sender ships from @@ -630,25 +631,27 @@ func (s *Store) DeleteAPIKeyForOrg(ctx context.Context, id uuid.UUID, orgID uuid // --- Sandbox Session operations --- type SandboxSession struct { - ID uuid.UUID `json:"id"` - SandboxID string `json:"sandboxId"` - OrgID uuid.UUID `json:"orgId"` - UserID *uuid.UUID `json:"userId,omitempty"` - Template string `json:"template"` - Region string `json:"region"` - WorkerID string `json:"workerId"` - Status string `json:"status"` - Config json.RawMessage `json:"config"` - Metadata json.RawMessage `json:"metadata,omitempty"` - StartedAt time.Time `json:"startedAt"` - StoppedAt *time.Time `json:"stoppedAt,omitempty"` - ErrorMsg *string `json:"errorMsg,omitempty"` - BasedOnCheckpointID *uuid.UUID `json:"basedOnCheckpointId,omitempty"` - LastPatchSequence int `json:"lastPatchSequence"` - MigratingToWorker string `json:"migratingToWorker,omitempty"` - PatchError *string `json:"patchError,omitempty"` - GoldenVersion *string `json:"goldenVersion,omitempty"` -} + ID uuid.UUID `json:"id"` + SandboxID string `json:"sandboxId"` + OrgID uuid.UUID `json:"orgId"` + UserID *uuid.UUID `json:"userId,omitempty"` + Template string `json:"template"` + Region string `json:"region"` + WorkerID string `json:"workerId"` + Status string `json:"status"` + Config json.RawMessage `json:"config"` + Metadata json.RawMessage `json:"metadata,omitempty"` + StartedAt time.Time `json:"startedAt"` + StoppedAt *time.Time `json:"stoppedAt,omitempty"` + ErrorMsg *string `json:"errorMsg,omitempty"` + BasedOnCheckpointID *uuid.UUID `json:"basedOnCheckpointId,omitempty"` + LastPatchSequence int `json:"lastPatchSequence"` + MigratingToWorker string `json:"migratingToWorker,omitempty"` + PatchError *string `json:"patchError,omitempty"` + GoldenVersion *string `json:"goldenVersion,omitempty"` +} + +const PostQMPMigrationFailureMessage = "migration failed after QMP transfer; source VM gone, target failed to complete" func (s *Store) CreateSandboxSession(ctx context.Context, sandboxID string, orgID uuid.UUID, userID *uuid.UUID, template, region, workerID string, config, metadata json.RawMessage, secretStoreID *uuid.UUID) (*SandboxSession, error) { return s.CreateSandboxSessionWithStatus(ctx, sandboxID, orgID, userID, template, region, workerID, config, metadata, "running", secretStoreID) @@ -741,10 +744,17 @@ func (s *Store) CompleteMigration(ctx context.Context, sandboxID, newWorkerID st // Use status IN ('migrating', 'running', 'stopped') — a race with the source // worker's cleanup may have already set it to 'stopped' or reverted to 'running'. // The migration DID succeed (QEMU is running on the target), so force the update. + // Also recover the specific post-QMP error state: an earlier target-complete + // timeout may have marked the row error even though the target VM later became + // usable and participated in a subsequent successful migration. tag, err := s.pool.Exec(ctx, `UPDATE sandbox_sessions SET status = 'running', worker_id = $1, migrating_to_worker = '', stopped_at = NULL, error_msg = NULL - WHERE sandbox_id = $2 AND status IN ('migrating', 'running', 'stopped')`, - newWorkerID, sandboxID) + WHERE sandbox_id = $2 + AND ( + status IN ('migrating', 'running', 'stopped') + OR (status = 'error' AND error_msg = $3) + )`, + newWorkerID, sandboxID, PostQMPMigrationFailureMessage) if err != nil { return err } @@ -826,15 +836,23 @@ func (s *Store) RecoverStaleMigrations(ctx context.Context, maxAge time.Duration return int(tag.RowsAffected()), nil } -// MarkOrphanedSandboxes marks running sandboxes on dead workers as error -// and returns the affected (sandbox_id, org_id, worker_id) tuples so the -// caller can publish `stopped` lifecycle events. liveWorkers is the set -// of worker IDs currently registered. Without the returned IDs the maintenance -// loop's PG sweep would never reach D1 sandboxes_index, which is exactly -// the post-cutover ghost-row bug. +// MarkOrphanedSandboxes marks non-resumable running sandboxes on dead workers +// as error and returns the affected (sandbox_id, org_id, worker_id) tuples so +// the caller can publish `stopped` lifecycle events. liveWorkers is the set of +// worker IDs currently registered. Resumable sandboxes are deliberately skipped +// here: the scaler's dead-worker recovery loop needs those rows to remain +// running long enough to recreate them on another worker. func (s *Store) MarkOrphanedSandboxes(ctx context.Context, liveWorkers map[string]bool) ([]OrphanedSandbox, error) { rows, err := s.pool.Query(ctx, - `SELECT DISTINCT worker_id FROM sandbox_sessions WHERE status = 'running'`) + `SELECT DISTINCT worker_id FROM sandbox_sessions + WHERE status = 'running' + AND NOT ( + COALESCE((config->>'burst')::boolean, false) + OR + COALESCE((config->>'resumable')::boolean, false) + OR config->>'sandboxFamily' = 'spot' + OR config->>'sandboxFamily' = 'resumable' + )`) if err != nil { return nil, err } @@ -856,6 +874,13 @@ func (s *Store) MarkOrphanedSandboxes(ctx context.Context, liveWorkers map[strin upd, err := s.pool.Query(ctx, `UPDATE sandbox_sessions SET status = 'error', error_msg = 'worker lost', stopped_at = now() WHERE worker_id = $1 AND status = 'running' + AND NOT ( + COALESCE((config->>'burst')::boolean, false) + OR + COALESCE((config->>'resumable')::boolean, false) + OR config->>'sandboxFamily' = 'spot' + OR config->>'sandboxFamily' = 'resumable' + ) RETURNING sandbox_id, org_id, worker_id`, workerID) if err != nil { continue @@ -979,6 +1004,7 @@ func (s *Store) ListSandboxSessions(ctx context.Context, orgID uuid.UUID, status type WorkerSandboxRef struct { SandboxID string OrgID uuid.UUID + WorkerID string } // ListSandboxesByWorkerStatus returns the (sandbox_id, org_id) pairs for sessions @@ -1008,6 +1034,33 @@ func (s *Store) ListSandboxesByWorkerStatus(ctx context.Context, workerID, statu return refs, rows.Err() } +// ListRunningSandboxesOnDeadWorkers returns running sessions whose worker_id is +// no longer present in the live worker set. It is read-only; callers decide +// whether to recover or close each row. +func (s *Store) ListRunningSandboxesOnDeadWorkers(ctx context.Context, liveWorkers map[string]bool) ([]WorkerSandboxRef, error) { + rows, err := s.pool.Query(ctx, + `SELECT sandbox_id, org_id, worker_id FROM sandbox_sessions + WHERE status = 'running' + ORDER BY started_at DESC + LIMIT 1000`) + if err != nil { + return nil, fmt.Errorf("list running sandboxes: %w", err) + } + defer rows.Close() + var refs []WorkerSandboxRef + for rows.Next() { + var r WorkerSandboxRef + if err := rows.Scan(&r.SandboxID, &r.OrgID, &r.WorkerID); err != nil { + return nil, err + } + if liveWorkers[r.WorkerID] { + continue + } + refs = append(refs, r) + } + return refs, rows.Err() +} + // ListSandboxIDsByWorkerStatus returns sandbox IDs (the public sb-... form, // not the row UUID) for sessions on a specific worker with a specific status. // Used by the reconcile-on-reconnect sweep: when a worker rejoins after being @@ -1381,9 +1434,24 @@ func (s *Store) UpdateSandboxSessionForWake(ctx context.Context, sandboxID, newW return err } +// UpdateSandboxSessionForRecreate marks a resumable sandbox running on a new +// worker after disk-only recreation. Unlike wake, this path does not require a +// hibernation row because RAM is intentionally not preserved. +func (s *Store) UpdateSandboxSessionForRecreate(ctx context.Context, sandboxID, newWorkerID string) error { + _, err := s.pool.Exec(ctx, + `UPDATE sandbox_sessions + SET status = 'running', worker_id = $1, stopped_at = NULL, error_msg = NULL + WHERE sandbox_id = $2`, + newWorkerID, sandboxID) + return err +} + // ReconcileWorkerSessions marks stale "running" sessions for a worker on startup. // Sessions with an active checkpoint are set to "hibernated" (recoverable via wake-on-request). // Sessions without a checkpoint are set to "stopped" (VM is gone, no recovery possible). +// Resumable sessions are skipped: their disk is the recovery source of truth, +// so the control plane's reverse-reconcile/dead-worker path should recreate +// them instead of converting them to hibernated/stopped on worker restart. // Returns the affected sandbox identities so the caller can emit lifecycle // events; without them D1 drifts from PG every time a worker restarts. func (s *Store) ReconcileWorkerSessions(ctx context.Context, workerID string) (hibernated, stopped []OrphanedSandbox, err error) { @@ -1397,6 +1465,13 @@ func (s *Store) ReconcileWorkerSessions(ctx context.Context, workerID string) (h hibernatedRows, err := tx.Query(ctx, `UPDATE sandbox_sessions SET status = 'hibernated' WHERE worker_id = $1 AND status = 'running' + AND NOT ( + COALESCE((config->>'burst')::boolean, false) + OR + COALESCE((config->>'resumable')::boolean, false) + OR config->>'sandboxFamily' = 'spot' + OR config->>'sandboxFamily' = 'resumable' + ) AND sandbox_id IN ( SELECT sandbox_id FROM sandbox_hibernations WHERE restored_at IS NULL AND expired_at IS NULL @@ -1426,6 +1501,13 @@ func (s *Store) ReconcileWorkerSessions(ctx context.Context, workerID string) (h `UPDATE sandbox_sessions SET status = 'stopped', stopped_at = now(), error_msg = 'worker restarted' WHERE worker_id = $1 AND status = 'running' + AND NOT ( + COALESCE((config->>'burst')::boolean, false) + OR + COALESCE((config->>'resumable')::boolean, false) + OR config->>'sandboxFamily' = 'spot' + OR config->>'sandboxFamily' = 'resumable' + ) RETURNING sandbox_id, org_id, worker_id`, workerID) if err != nil { return hibernated, nil, fmt.Errorf("failed to reconcile stopped sessions: %w", err) diff --git a/internal/db/usage.go b/internal/db/usage.go index 5f172ce6..ae08d729 100644 --- a/internal/db/usage.go +++ b/internal/db/usage.go @@ -15,6 +15,7 @@ type ScaleEvent struct { MemoryMB int `json:"memoryMB"` CPUPct int `json:"cpuPercent"` DiskMB int `json:"diskMB"` + Burst bool `json:"burst"` StartedAt time.Time `json:"startedAt"` EndedAt *time.Time `json:"endedAt,omitempty"` } diff --git a/internal/preemption/aws.go b/internal/preemption/aws.go new file mode 100644 index 00000000..965711d3 --- /dev/null +++ b/internal/preemption/aws.go @@ -0,0 +1,172 @@ +// AWS spot interruption monitor. +// +// AWS publishes interruption notices on IMDSv2 at +// GET /latest/meta-data/spot/instance-action +// returning HTTP 200 with a JSON body shaped: +// { "action": "stop|terminate|hibernate", "time": "2026-05-19T12:34:56Z" } +// +// Until interruption is scheduled, the endpoint returns 404. We poll every +// `pollInterval` (default 5s). AWS guarantees the notice lands at least +// ~2 minutes before the instance is reclaimed, so 5s polling gives ample +// drain budget. +// +// IMDSv2 requires a session token (PUT /latest/api/token). Tokens are +// scoped per-request and expire in 6h, but we re-fetch on every poll for +// simplicity — IMDS is link-local, the round-trip cost is negligible +// compared to the network operations the worker is otherwise doing. + +package preemption + +import ( + "context" + "encoding/json" + "io" + "log" + "net/http" + "strings" + "time" +) + +type awsMonitor struct { + pollInterval time.Duration + imdsEndpoint string +} + +func (m *awsMonitor) Name() string { return "aws-imds" } + +func (m *awsMonitor) Watch(ctx context.Context) <-chan Notice { + ch := make(chan Notice, 1) + + go func() { + defer close(ch) + + // Standalone client so the global default client's settings can't + // stall our 1-second IMDS round-trips with unrelated transport + // configuration. + client := &http.Client{Timeout: 2 * time.Second} + + ticker := time.NewTicker(m.pollInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + notice, found, err := m.probe(ctx, client) + if err != nil { + // Don't spam — IMDS unreachability is interesting once + // per minute, not once per 5s. + log.Printf("preemption: aws probe error (will retry): %v", err) + continue + } + if !found { + continue + } + // Buffered channel — non-blocking send. If the consumer is + // somehow not reading, drop the second-and-later notices + // (the first one is what counts). + select { + case ch <- notice: + default: + } + // Keep polling — the action / time can change between + // notice issuance and reclaim, and we want the consumer to + // see the freshest data. Drain is idempotent. + } + } + }() + + return ch +} + +// probe issues one IMDSv2 token-then-get round trip. Returns (Notice, true, +// nil) when interruption is imminent, (zero, false, nil) when healthy, and +// (zero, false, err) on transport / decode failures. +func (m *awsMonitor) probe(ctx context.Context, client *http.Client) (Notice, bool, error) { + token, err := m.fetchToken(ctx, client) + if err != nil { + return Notice{}, false, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + m.imdsEndpoint+"/latest/meta-data/spot/instance-action", nil) + if err != nil { + return Notice{}, false, err + } + req.Header.Set("X-aws-ec2-metadata-token", token) + + resp, err := client.Do(req) + if err != nil { + return Notice{}, false, err + } + defer resp.Body.Close() + + // 404 = healthy. AWS specifies this endpoint returns 404 until + // interruption is scheduled. + if resp.StatusCode == http.StatusNotFound { + return Notice{}, false, nil + } + if resp.StatusCode != http.StatusOK { + // Drain body to keep the connection reusable. + _, _ = io.Copy(io.Discard, resp.Body) + return Notice{}, false, nil + } + + var body struct { + Action string `json:"action"` + Time string `json:"time"` + } + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + return Notice{}, false, err + } + + eta, err := time.Parse(time.RFC3339, body.Time) + if err != nil { + // AWS guarantees RFC3339; if they ever change format we'd rather + // fire a Notice with the zero time than miss the signal — the + // drain budget still applies, and the caller can fallback to + // "drain now". + log.Printf("preemption: aws unexpected time format %q: %v", body.Time, err) + } + + return Notice{ + Action: Action(strings.ToLower(body.Action)), + ETA: eta, + Source: "aws-imds", + }, true, nil +} + +// fetchToken obtains an IMDSv2 session token. The token is single-use here; +// see package doc for why that's fine. +func (m *awsMonitor) fetchToken(ctx context.Context, client *http.Client) (string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPut, + m.imdsEndpoint+"/latest/api/token", nil) + if err != nil { + return "", err + } + req.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "60") + + resp, err := client.Do(req) + if err != nil { + return "", err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + _, _ = io.Copy(io.Discard, resp.Body) + return "", &imdsHTTPError{code: resp.StatusCode} + } + + b, err := io.ReadAll(resp.Body) + if err != nil { + return "", err + } + return strings.TrimSpace(string(b)), nil +} + +type imdsHTTPError struct{ code int } + +func (e *imdsHTTPError) Error() string { + return "imds token fetch returned HTTP " + http.StatusText(e.code) +} diff --git a/internal/preemption/azure.go b/internal/preemption/azure.go new file mode 100644 index 00000000..4212a7ec --- /dev/null +++ b/internal/preemption/azure.go @@ -0,0 +1,47 @@ +// Azure Scheduled Events preemption monitor. +// +// Azure publishes scheduled-events on IMDS at +// GET /metadata/scheduledevents?api-version=2020-07-01 +// (requires header `Metadata: true`). Response shape: +// { +// "DocumentIncarnation": 42, +// "Events": [ +// { "EventId": "...", "EventType": "Preempt|Reboot|Redeploy|Freeze|Terminate", +// "ResourceType": "VirtualMachine", +// "Resources": [""], +// "EventStatus": "Scheduled|Started", +// "NotBefore": "Mon, 19 May 2026 12:34:56 GMT" } +// ] +// } +// +// "Preempt" is the spot-equivalent ("This VM is being preempted; you have +// ~30s"). "Terminate" is unscheduled-but-imminent. +// +// This is a stub for the PoC — Azure cells aren't running the new +// preemption-aware code path yet. Wiring this up later is mostly populating +// the Watch loop with the same probe-then-emit pattern as the AWS monitor. + +package preemption + +import ( + "context" + "log" + "time" +) + +type azureMonitor struct { + pollInterval time.Duration + imdsEndpoint string +} + +func (m *azureMonitor) Name() string { return "azure-scheduled-events" } + +func (m *azureMonitor) Watch(ctx context.Context) <-chan Notice { + ch := make(chan Notice, 1) + go func() { + defer close(ch) + log.Printf("preemption: azure monitor stubbed — TODO wire to /metadata/scheduledevents (see internal/preemption/azure.go)") + <-ctx.Done() + }() + return ch +} diff --git a/internal/preemption/monitor.go b/internal/preemption/monitor.go new file mode 100644 index 00000000..a150d32f --- /dev/null +++ b/internal/preemption/monitor.go @@ -0,0 +1,96 @@ +// Package preemption watches the cloud-specific spot-interruption signal and +// publishes a Notice on a channel once interruption is imminent. The worker +// binary subscribes to the channel and kicks off graceful drain: +// +// 1. Flip its Redis-heartbeat state to "draining" so the CP stops placing. +// 2. Hibernate every live sandbox to S3 in parallel (existing primitive). +// 3. Delete the heartbeat key; exit cleanly. +// +// On AWS the signal is IMDSv2 /latest/meta-data/spot/instance-action returning +// 200 with a JSON body. On Azure it's /metadata/scheduledevents. The shape of +// the work the worker does in response is identical — only the detection is +// cloud-specific. This package abstracts that. +// +// LocalCloud is decided by the OPENSANDBOX_CLOUD env var (set by cloud-init +// from terraform). If unset, the Monitor returned by NewMonitor is a no-op +// that never fires — safe default for dev / Azure-without-AWS. + +package preemption + +import ( + "context" + "log" + "os" + "time" +) + +// Action enumerates what the cloud says it will do to the instance. +type Action string + +const ( + ActionTerminate Action = "terminate" + ActionStop Action = "stop" + ActionHibernate Action = "hibernate" // AWS-only; Azure preempts only with "preempt" +) + +// Notice carries the advance-warning details. ETA is when the cloud says it +// will act; the worker has roughly ETA - now to drain. AWS guarantees ~2 min +// for spot interruption; Azure scheduled events give ~30s for preempt and +// up to 15 min for reboot/freeze. +type Notice struct { + Action Action + ETA time.Time + // Source describes where the notice came from for logging — "aws-imds", + // "azure-scheduled-events", or "redis" for CP-fanout fallback. + Source string +} + +// Monitor watches for preemption and emits a Notice on its channel when one +// arrives. Implementations must: +// - Return a buffered channel (size >= 1) so a slow consumer doesn't lose +// the notice. +// - Survive transient errors (network blips polling IMDS) — log + retry. +// - Stop cleanly when the provided context is canceled. +type Monitor interface { + Name() string + Watch(ctx context.Context) <-chan Notice +} + +// NewMonitor returns the cloud-appropriate monitor or a no-op if no cloud +// is configured. Callers should always call Watch — even the no-op returns +// a channel they can select on, simplifying wire-up. +func NewMonitor() Monitor { + switch os.Getenv("OPENSANDBOX_CLOUD") { + case "aws": + return &awsMonitor{ + pollInterval: 5 * time.Second, + imdsEndpoint: "http://169.254.169.254", + } + case "azure": + return &azureMonitor{ + pollInterval: 5 * time.Second, + imdsEndpoint: "http://169.254.169.254", + } + default: + log.Printf("preemption: no cloud configured (OPENSANDBOX_CLOUD unset) — preemption notices disabled") + return &noopMonitor{} + } +} + +// noopMonitor never emits anything. Used for dev and any deployment where +// the cloud's preemption signal isn't worth wiring up (on-demand-only, +// bare-metal-colo, etc.). +type noopMonitor struct{} + +func (noopMonitor) Name() string { return "noop" } +func (noopMonitor) Watch(ctx context.Context) <-chan Notice { + ch := make(chan Notice, 1) + // Channel left open; the caller's select will simply never fire on + // this case. Closing on ctx.Done would also work but would make + // receivers reading via "n, ok := <-ch" misinterpret it as a notice. + go func() { + <-ctx.Done() + close(ch) + }() + return ch +} diff --git a/internal/qemu/golden_version.go b/internal/qemu/golden_version.go index 2fca2247..bb4e2047 100644 --- a/internal/qemu/golden_version.go +++ b/internal/qemu/golden_version.go @@ -5,6 +5,8 @@ import ( "fmt" "io" "os" + "path/filepath" + "strings" ) // computeGoldenVersion computes a version hash for a golden snapshot by hashing @@ -18,6 +20,10 @@ import ( // serve stale overlays against fresh bases, producing ext4 directory-block // checksum errors in guests. func computeGoldenVersion(baseImagePath string) (string, error) { + if v, err := readGoldenVersionSidecar(baseImagePath); err == nil && v != "" { + return v, nil + } + f, err := os.Open(baseImagePath) if err != nil { return "", fmt.Errorf("open base image: %w", err) @@ -31,6 +37,21 @@ func computeGoldenVersion(baseImagePath string) (string, error) { return fmt.Sprintf("%x", h.Sum(nil))[:16], nil } +func readGoldenVersionSidecar(baseImagePath string) (string, error) { + b, err := os.ReadFile(filepath.Join(filepath.Dir(baseImagePath), "golden-version")) + if err != nil { + return "", err + } + v := strings.TrimSpace(string(b)) + if v == "" { + return "", nil + } + if len(v) > 16 { + v = v[:16] + } + return v, nil +} + // ComputeGoldenVersion is the exported entry point used by cmd/worker's // "golden-version" subcommand so Packer invokes the same hash function // the runtime uses for archive-key lookups. diff --git a/internal/qemu/manager.go b/internal/qemu/manager.go index 7c63aa4a..f75bbe51 100644 --- a/internal/qemu/manager.go +++ b/internal/qemu/manager.go @@ -205,6 +205,64 @@ func quiesceAndCloseAgent(ctx context.Context, agent *AgentClient) error { return nil } +// resetAndCloseAgentTransportForSnapshot syncs the guest and resets the +// virtio-serial listener without freezing filesystems. Golden snapshots restore +// into live sandboxes immediately, so unlike hibernate we must not capture a +// frozen root filesystem. +func resetAndCloseAgentTransportForSnapshot(ctx context.Context, agent *AgentClient) error { + if agent == nil { + return nil + } + + prepareOnce := func() error { + rpcCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + _, err := agent.PrepareHibernate(rpcCtx, &pb.PrepareHibernateRequest{}) + return err + } + err := prepareOnce() + if err != nil && IsTransportError(err) { + log.Printf("qemu: snapshot PrepareHibernate transport error (%v), redialing", err) + if rdErr := agent.Redial(); rdErr == nil { + err = prepareOnce() + } else { + log.Printf("qemu: snapshot PrepareHibernate redial failed: %v (orig: %v)", rdErr, err) + } + } + if err != nil { + if st, ok := status.FromError(err); !ok || st.Code() != codes.Unimplemented { + log.Printf("qemu: snapshot PrepareHibernate RPC failed: %v (falling back to legacy path)", err) + } + execOnce := func() error { + execCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + _, e := agent.Exec(execCtx, &pb.ExecRequest{ + Command: "/bin/sh", + Args: []string{"-c", "sync; blockdev --flushbufs /dev/vda 2>/dev/null; blockdev --flushbufs /dev/vdb 2>/dev/null; sync; kill -USR1 1"}, + RunAsRoot: true, + }) + return e + } + fallbackErr := execOnce() + if fallbackErr != nil && IsTransportError(fallbackErr) { + log.Printf("qemu: snapshot prepare fallback Exec transport error (%v), redialing", fallbackErr) + if rdErr := agent.Redial(); rdErr == nil { + fallbackErr = execOnce() + } else { + log.Printf("qemu: snapshot prepare fallback redial failed: %v", rdErr) + } + } + if fallbackErr != nil { + return fmt.Errorf("%w: PrepareHibernate=%v, fallback Exec=%v", ErrAgentUnresponsive, err, fallbackErr) + } + time.Sleep(1 * time.Second) + } + + _ = agent.Close() + time.Sleep(200 * time.Millisecond) + return nil +} + // Compile-time check that Manager implements sandbox.Manager. var _ sandbox.Manager = (*Manager)(nil) @@ -290,6 +348,7 @@ type Config struct { DataDir string // base data directory (e.g., /data) KernelPath string // path to vmlinux kernel ImagesDir string // path to base rootfs images + GoldenDir string // path to worker-local golden snapshot cache QEMUBin string // path to qemu-system-x86_64 binary AgentBinaryPath string // path to osb-agent binary on host (for hot-upgrade) AgentVersion string // expected agent version (for hot-upgrade check) @@ -367,6 +426,9 @@ func NewManager(cfg Config) (*Manager, error) { if cfg.ImagesDir == "" { cfg.ImagesDir = filepath.Join(cfg.DataDir, "images") } + if cfg.GoldenDir == "" { + cfg.GoldenDir = filepath.Join(cfg.DataDir, "golden") + } if cfg.QEMUBin == "" { cfg.QEMUBin = "qemu-system-x86_64" } @@ -457,6 +519,13 @@ func (m *Manager) GoldenVersion() string { return m.goldenVersion } +func (m *Manager) goldenSnapshotDir() string { + if m.cfg.GoldenDir != "" { + return m.cfg.GoldenDir + } + return filepath.Join(m.cfg.DataDir, "golden") +} + // MemoryAllocatedBytes returns the sum of memory committed to currently-running // sandboxes, in bytes. Used by the worker's resource-stats tick to report // oversubscription independent of actual guest workload. @@ -704,7 +773,7 @@ func (m *Manager) PrepareGoldenSnapshot() error { log.Printf("qemu: blobstore golden fetch failed (will try local-only): %v", err) } - goldenDir := filepath.Join(m.cfg.DataDir, "golden") + goldenDir := m.goldenSnapshotDir() memFile := filepath.Join(goldenDir, "mem") rootfsFile := filepath.Join(goldenDir, "rootfs.qcow2") @@ -904,19 +973,19 @@ func (m *Manager) PrepareGoldenSnapshot() error { log.Printf("qemu: golden: /home/sandbox unmounted and synced") } - // Close agent connection before migration. Use a timeout because gRPC's - // graceful close over vsock can hang if vhost-vsock doesn't drain cleanly. - closeDone := make(chan struct{}) - go func() { - agentClient.Close() - close(closeDone) - }() - select { - case <-closeDone: - case <-time.After(2 * time.Second): - log.Printf("qemu: golden: agent close timed out, proceeding anyway") + // Reset the in-guest virtio-serial listener before snapshotting. A plain + // host-side Close can leave the golden image with a stale gRPC parser state; + // restored VMs then read HTTP/2 frames but never answer Ping. + prepCtx, prepCancel := context.WithTimeout(context.Background(), 15*time.Second) + if err := resetAndCloseAgentTransportForSnapshot(prepCtx, agentClient); err != nil { + prepCancel() + qmpClient.Close() + cmd.Process.Kill() + cmd.Wait() + return fmt.Errorf("golden prepare agent transport: %w", err) } - time.Sleep(500 * time.Millisecond) + prepCancel() + log.Printf("qemu: golden: agent transport reset") // QMP stop + migrate log.Printf("qemu: golden: sending QMP stop...") @@ -994,10 +1063,10 @@ func (m *Manager) PrepareGoldenSnapshot() error { // Returns the old and new golden version strings. func (m *Manager) RebuildGoldenSnapshot() (oldVersion, newVersion string, err error) { oldVersion = m.goldenVersion - goldenDir := filepath.Join(m.cfg.DataDir, "golden") + goldenDir := m.goldenSnapshotDir() // Build new golden in a staging directory - stagingDir := filepath.Join(m.cfg.DataDir, "golden-staging") + stagingDir := goldenDir + "-staging" os.RemoveAll(stagingDir) // clean up any prior failed attempt // Temporarily point goldenDir to staging so PrepareGoldenSnapshot builds there @@ -1005,7 +1074,7 @@ func (m *Manager) RebuildGoldenSnapshot() (oldVersion, newVersion string, err er m.goldenDir = "" // Rename current golden out of the way so PrepareGoldenSnapshot sees no existing snapshot - backupDir := filepath.Join(m.cfg.DataDir, "golden-old") + backupDir := goldenDir + "-old" os.RemoveAll(backupDir) if err := os.Rename(goldenDir, backupDir); err != nil && !os.IsNotExist(err) { m.goldenDir = oldGoldenDir @@ -1311,15 +1380,16 @@ func (m *Manager) createFromGolden(ctx context.Context, cfg types.SandboxConfig, mountCtx, mountCancel := context.WithTimeout(context.Background(), 10*time.Second) _, mountErr := agentClient.Exec(mountCtx, &pb.ExecRequest{ Command: "/bin/sh", - Args: []string{"-c", strings.Join([]string{ - "echo 3 > /proc/sys/vm/drop_caches", - "echo 3 > /proc/sys/vm/drop_caches", - "mount /dev/vdb /home/sandbox 2>/dev/null || true", - "resize2fs /dev/vdb 2>/dev/null || true", - "chown 1000:1000 /home/sandbox", - "mkdir -p /home/sandbox/.osb-apt-cache /var/cache/apt/archives", - "mountpoint -q /var/cache/apt/archives || mount --bind /home/sandbox/.osb-apt-cache /var/cache/apt/archives 2>/dev/null || true", - }, " && ")}, + Args: []string{"-c", ` +set -u +echo 3 > /proc/sys/vm/drop_caches +echo 3 > /proc/sys/vm/drop_caches +mount /dev/vdb /home/sandbox 2>/dev/null || true +(timeout 30s resize2fs /dev/vdb >/dev/null 2>&1 || true) & +chown 1000:1000 /home/sandbox +mkdir -p /home/sandbox/.osb-apt-cache /var/cache/apt/archives +mountpoint -q /var/cache/apt/archives || mount --bind /home/sandbox/.osb-apt-cache /var/cache/apt/archives 2>/dev/null || true +`}, RunAsRoot: true, }) mountCancel() @@ -1600,8 +1670,13 @@ func (m *Manager) Create(ctx context.Context, cfg types.SandboxConfig) (sb *type id = "sb-" + uuid.New().String()[:8] } + sandboxDir := filepath.Join(m.cfg.DataDir, "sandboxes", id) + rootfsPath := filepath.Join(sandboxDir, "rootfs.qcow2") + workspacePath := filepath.Join(sandboxDir, "workspace.qcow2") + reuseExistingDrives := cfg.IsResumable() && fileExists(rootfsPath) && fileExists(workspacePath) + // Fast path: restore from golden snapshot if available and using default template - if m.goldenDir != "" && template == "default" && cfg.TemplateRootfsKey == "" { + if !reuseExistingDrives && m.goldenDir != "" && template == "default" && cfg.TemplateRootfsKey == "" { sb, err := m.createFromGolden(ctx, cfg, id) if err != nil { log.Printf("qemu: golden restore failed for %s, falling back to cold boot: %v", id, err) @@ -1611,15 +1686,13 @@ func (m *Manager) Create(ctx context.Context, cfg types.SandboxConfig) (sb *type } } - sandboxDir := filepath.Join(m.cfg.DataDir, "sandboxes", id) if err := os.MkdirAll(sandboxDir, 0755); err != nil { return nil, fmt.Errorf("mkdir sandbox dir: %w", err) } - rootfsPath := filepath.Join(sandboxDir, "rootfs.qcow2") - workspacePath := filepath.Join(sandboxDir, "workspace.qcow2") - - if cfg.TemplateRootfsKey != "" { + if reuseExistingDrives { + log.Printf("qemu: create %s: reusing existing resumable drives", id) + } else if cfg.TemplateRootfsKey != "" { srcRootfs := strings.TrimPrefix(cfg.TemplateRootfsKey, "local://") srcWorkspace := strings.TrimPrefix(cfg.TemplateWorkspaceKey, "local://") log.Printf("qemu: create %s from snapshot template (rootfs=%s, workspace=%s)", id, srcRootfs, srcWorkspace) @@ -2101,9 +2174,17 @@ type HibernateAllResult struct { // HibernateAll hibernates all running VMs concurrently. func (m *Manager) HibernateAll(ctx context.Context, checkpointStore *storage.CheckpointStore) []HibernateAllResult { + return m.HibernateAllExcept(ctx, checkpointStore, nil) +} + +// HibernateAllExcept hibernates all running VMs except IDs in skip. +func (m *Manager) HibernateAllExcept(ctx context.Context, checkpointStore *storage.CheckpointStore, skip map[string]struct{}) []HibernateAllResult { m.mu.RLock() ids := make([]string, 0, len(m.vms)) for id := range m.vms { + if _, ok := skip[id]; ok { + continue + } ids = append(ids, id) } m.mu.RUnlock() diff --git a/internal/qemu/migration.go b/internal/qemu/migration.go index afb74313..3dfc597f 100644 --- a/internal/qemu/migration.go +++ b/internal/qemu/migration.go @@ -627,6 +627,11 @@ func (m *Manager) CompleteIncomingMigration(ctx context.Context, sandboxID strin // type so the worker gRPC layer can refer to it without an import cycle. func (m *Manager) PreCopyDrives(ctx context.Context, sandboxID string, checkpointStore *storage.CheckpointStore) (rootfsKey, workspaceKey, goldenVer string, baseCPU, baseMem, actualMem int, secrets sandbox.MigrationSecrets, err error) { + if checkpointStore == nil { + err = fmt.Errorf("checkpoint store not configured") + return + } + m.mu.RLock() vm, exists := m.vms[sandboxID] var pid int @@ -813,6 +818,10 @@ func (m *Manager) LiveMigrate(ctx context.Context, sandboxID, incomingAddr strin // uploadFile uploads a local file to S3. func (mc *MigrationCoordinator) uploadFile(ctx context.Context, localPath, s3Key string) (int64, error) { + if mc.checkpointStore == nil { + return 0, fmt.Errorf("checkpoint store not configured") + } + info, err := os.Stat(localPath) if err != nil { return 0, err diff --git a/internal/qemu/reset.go b/internal/qemu/reset.go index d4a0f8ec..10011b38 100644 --- a/internal/qemu/reset.go +++ b/internal/qemu/reset.go @@ -10,6 +10,7 @@ import ( "path/filepath" "time" + "github.com/opensandbox/opensandbox/pkg/types" pb "github.com/opensandbox/opensandbox/proto/agent" ) @@ -334,3 +335,234 @@ func (m *Manager) PowerCycleSandbox(ctx context.Context, sandboxID string) (host sandboxID, time.Since(t0).Milliseconds(), freshPort, netCfg.TAPName) return freshPort, nil } + +// StartExistingSandbox cold-boots an existing sandbox directory on this worker. +// It is the disk-only resumable recovery primitive: no savevm/loadvm and no RAM +// preservation. Only the workspace/data disk is durable; the rootfs is reset +// from this worker's clean golden image before boot. This avoids cold-booting a +// rootfs qcow2 that was last used by a VM restored from a golden memory snapshot +// or killed mid-flight, both of which can leave ext4 metadata unsafe to mount +// as a fresh boot device. +func (m *Manager) StartExistingSandbox(ctx context.Context, sandboxID string, cfg types.SandboxConfig) (*types.Sandbox, error) { + t0 := time.Now() + + m.mu.Lock() + if _, exists := m.vms[sandboxID]; exists { + m.mu.Unlock() + return nil, fmt.Errorf("sandbox %s is already running on this worker", sandboxID) + } + m.mu.Unlock() + + sandboxDir := filepath.Join(m.cfg.DataDir, "sandboxes", sandboxID) + workspacePath := detectDrivePath(sandboxDir, "workspace") + if !fileExists(workspacePath) { + return nil, fmt.Errorf("sandbox %s: existing workspace drive missing on this worker (dir=%s)", sandboxID, sandboxDir) + } + rootfsPath := filepath.Join(sandboxDir, "rootfs.qcow2") + goldenRootfs := filepath.Join(m.goldenDir, "rootfs.qcow2") + if !fileExists(goldenRootfs) { + return nil, fmt.Errorf("sandbox %s: golden rootfs missing on this worker (%s)", sandboxID, goldenRootfs) + } + tmpRootfs := filepath.Join(sandboxDir, fmt.Sprintf(".rootfs.%d.tmp", time.Now().UnixNano())) + if err := copyFileReflink(goldenRootfs, tmpRootfs); err != nil { + os.Remove(tmpRootfs) + return nil, fmt.Errorf("reset rootfs from golden: %w", err) + } + if err := os.Rename(tmpRootfs, rootfsPath); err != nil { + os.Remove(tmpRootfs) + return nil, fmt.Errorf("replace rootfs from golden: %w", err) + } + log.Printf("qemu: StartExistingSandbox %s: rootfs reset from golden %s (%dms)", + sandboxID, m.goldenVersion, time.Since(t0).Milliseconds()) + + var meta SandboxMeta + if data, err := os.ReadFile(filepath.Join(sandboxDir, "sandbox-meta.json")); err == nil { + _ = json.Unmarshal(data, &meta) + } + template := cfg.Template + if template == "" { + template = meta.Template + } + if template == "" { + template = "default" + } + cpus := cfg.CpuCount + if cpus <= 0 { + cpus = meta.CpuCount + } + if cpus <= 0 { + cpus = m.cfg.DefaultCPUs + } + memMB := cfg.MemoryMB + if memMB <= 0 { + memMB = meta.MemoryMB + } + if memMB <= 0 { + memMB = m.cfg.DefaultMemoryMB + } + guestPort := cfg.Port + if guestPort == 0 { + guestPort = meta.GuestPort + } + if guestPort == 0 { + guestPort = m.cfg.DefaultPort + } + + netCfg, err := m.subnets.Allocate() + if err != nil { + return nil, fmt.Errorf("allocate subnet: %w", err) + } + if err := CreateTAP(netCfg); err != nil { + m.subnets.Release(netCfg.TAPName) + return nil, fmt.Errorf("create TAP: %w", err) + } + hostPort, err := FindFreePort() + if err != nil { + DeleteTAP(netCfg.TAPName) + m.subnets.Release(netCfg.TAPName) + return nil, fmt.Errorf("find free port: %w", err) + } + netCfg.HostPort = hostPort + netCfg.GuestPort = guestPort + if err := AddDNAT(netCfg); err != nil { + DeleteTAP(netCfg.TAPName) + m.subnets.Release(netCfg.TAPName) + return nil, fmt.Errorf("add DNAT: %w", err) + } + if err := AddMetadataDNAT(netCfg.TAPName, netCfg.HostIP); err != nil { + log.Printf("qemu: StartExistingSandbox %s: metadata DNAT failed: %v", sandboxID, err) + } + + guestMAC := generateMAC(sandboxID) + guestCID := m.allocateCID() + bootArgs := fmt.Sprintf( + "console=ttyS0 reboot=k panic=1 root=/dev/vda rw ip=%s::%s:%s::eth0:off init=/sbin/init osb.gateway=%s", + netCfg.GuestIP, netCfg.HostIP, netCfg.Mask, netCfg.HostIP, + ) + qmpSockPath := filepath.Join(sandboxDir, "qmp.sock") + agentSockPath := filepath.Join(sandboxDir, "agent.sock") + os.Remove(qmpSockPath) + os.Remove(agentSockPath) + + logFile, err := os.Create(filepath.Join(sandboxDir, "qemu.log")) + if err != nil { + m.cleanupVM(netCfg, "") + return nil, fmt.Errorf("create log: %w", err) + } + args := m.buildQEMUArgs(cpus, memMB, rootfsPath, workspacePath, + netCfg.TAPName, guestMAC, agentSockPath, qmpSockPath, bootArgs) + cmd := exec.Command(m.cfg.QEMUBin, args...) + cmd.Stdout = logFile + cmd.Stderr = logFile + if err := cmd.Start(); err != nil { + logFile.Close() + m.cleanupVM(netCfg, "") + return nil, fmt.Errorf("start QEMU: %w", err) + } + logFile.Close() + + qmpClient, err := waitForQMP(qmpSockPath, 30*time.Second) + if err != nil { + cmd.Process.Kill() + cmd.Wait() + m.cleanupVM(netCfg, "") + return nil, fmt.Errorf("QMP connect: %w", err) + } + agentClient, err := m.waitForAgentSocket(ctx, agentSockPath, 60*time.Second) + if err != nil { + qmpClient.Close() + cmd.Process.Kill() + cmd.Wait() + m.cleanupVM(netCfg, "") + return nil, fmt.Errorf("agent connect: %w", err) + } + + if err := syncGuestClock(ctx, agentClient); err != nil { + log.Printf("qemu: StartExistingSandbox %s: clock sync failed: %v", sandboxID, err) + } + mountCtx, mountCancel := context.WithTimeout(ctx, 15*time.Second) + _, mountErr := agentClient.Exec(mountCtx, &pb.ExecRequest{ + Command: "/bin/sh", + Args: []string{"-c", "mount /dev/vdb /home/sandbox 2>/dev/null || true; resize2fs /dev/vdb 2>/dev/null || true; chown 1000:1000 /home/sandbox"}, + RunAsRoot: true, + }) + mountCancel() + if mountErr != nil { + log.Printf("qemu: StartExistingSandbox %s: mount /home/sandbox failed: %v", sandboxID, mountErr) + } + m.setupAptCacheBindMount(ctx, sandboxID, agentClient) + m.reinstallProxyCA(ctx, sandboxID, agentClient) + + envsToInject := m.sealSandboxEnvs(ctx, sandboxID, netCfg, agentClient, cfg) + if len(envsToInject) > 0 { + envCtx, envCancel := context.WithTimeout(ctx, 5*time.Second) + if err := agentClient.SetEnvs(envCtx, envsToInject); err != nil { + log.Printf("qemu: StartExistingSandbox %s: SetEnvs failed: %v", sandboxID, err) + } + envCancel() + } + + now := time.Now() + timeout := time.Duration(cfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 300 * time.Second + } + vm := &VMInstance{ + ID: sandboxID, + Template: template, + Status: types.SandboxStatusRunning, + StartedAt: now, + EndAt: now.Add(timeout), + CpuCount: cpus, + MemoryMB: memMB, + baseMemoryMB: memMB, + HostPort: hostPort, + GuestPort: guestPort, + pid: cmd.Process.Pid, + cmd: cmd, + network: netCfg, + sandboxDir: sandboxDir, + agent: agentClient, + qmpSockPath: qmpSockPath, + agentSockPath: agentSockPath, + qmp: qmpClient, + guestMAC: guestMAC, + guestCID: guestCID, + bootArgs: bootArgs, + goldenVersion: m.goldenVersion, + } + + m.mu.Lock() + m.vms[sandboxID] = vm + m.mu.Unlock() + + if m.onSandboxReady != nil { + m.onSandboxReady(sandboxID, netCfg.GuestIP, template, vm.StartedAt) + } + + sbMeta := SandboxMeta{ + SandboxID: sandboxID, + Template: template, + CpuCount: cpus, + MemoryMB: memMB, + GuestPort: guestPort, + } + if metaJSON, err := json.Marshal(sbMeta); err == nil { + if writeErr := os.WriteFile(filepath.Join(sandboxDir, "sandbox-meta.json"), metaJSON, 0644); writeErr != nil { + log.Printf("qemu: WARNING: failed to write sandbox-meta.json for %s: %v", sandboxDir, writeErr) + } + } + + log.Printf("qemu: StartExistingSandbox %s: complete (%dms, port=%d→%d, tap=%s)", + sandboxID, time.Since(t0).Milliseconds(), hostPort, guestPort, netCfg.TAPName) + return &types.Sandbox{ + ID: sandboxID, + Template: template, + Status: types.SandboxStatusRunning, + StartedAt: now, + EndAt: now.Add(timeout), + CpuCount: cpus, + MemoryMB: memMB, + HostPort: hostPort, + }, nil +} diff --git a/internal/worker/grpc_server.go b/internal/worker/grpc_server.go index 87bbfeaf..e5e92a69 100644 --- a/internal/worker/grpc_server.go +++ b/internal/worker/grpc_server.go @@ -9,6 +9,7 @@ import ( "os" "os/exec" "path/filepath" + "strconv" "strings" "time" @@ -88,6 +89,8 @@ type GRPCServer struct { // region is the worker's region label, used to tag operation metrics. // Set via SetRegion at startup. Empty = "unknown". region string + + createSem chan struct{} } // SetRegion stamps the worker's region onto operation metrics emitted from @@ -140,11 +143,42 @@ func NewGRPCServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, execMgr *san checkpointStore: checkpointStore, store: store, server: grpc.NewServer(serverOpts...), + createSem: make(chan struct{}, workerCreateConcurrency()), } pb.RegisterSandboxWorkerServer(s.server, s) return s } +func workerCreateConcurrency() int { + const defaultCreateConcurrency = 4 + raw := strings.TrimSpace(os.Getenv("OPENSANDBOX_WORKER_CREATE_CONCURRENCY")) + if raw == "" { + return defaultCreateConcurrency + } + n, err := strconv.Atoi(raw) + if err != nil || n <= 0 { + log.Printf("grpc: invalid OPENSANDBOX_WORKER_CREATE_CONCURRENCY=%q, using %d", raw, defaultCreateConcurrency) + return defaultCreateConcurrency + } + return n +} + +func (s *GRPCServer) acquireCreateSlot(ctx context.Context, sandboxID string) (func(), error) { + if s.createSem == nil { + return func() {}, nil + } + start := time.Now() + select { + case s.createSem <- struct{}{}: + if waited := time.Since(start); waited > 250*time.Millisecond { + log.Printf("grpc: CreateSandbox %s waited %s for worker create slot", sandboxID, waited.Round(time.Millisecond)) + } + return func() { <-s.createSem }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + // Start starts the gRPC server on the given address. func (s *GRPCServer) Start(addr string) error { lis, err := net.Listen("tcp", addr) @@ -204,6 +238,12 @@ func parseSecretAllowedHosts(m map[string]string) map[string][]string { } func (s *GRPCServer) CreateSandbox(ctx context.Context, req *pb.CreateSandboxRequest) (*pb.CreateSandboxResponse, error) { + releaseCreateSlot, err := s.acquireCreateSlot(ctx, req.SandboxId) + if err != nil { + return nil, status.Errorf(codes.Unavailable, "worker create queue cancelled: %v", err) + } + defer releaseCreateSlot() + cfg := types.SandboxConfig{ Template: req.Template, Timeout: int(req.Timeout), @@ -220,6 +260,11 @@ func (s *GRPCServer) CreateSandbox(ctx context.Context, req *pb.CreateSandboxReq SecretEnvs: req.SecretEnvs, DiskMB: int(req.DiskMb), } + if cfg.Envs["OPENSANDBOX_BURST"] == "true" || cfg.Envs["OPENSANDBOX_RESUMABLE"] == "true" { + cfg.Burst = true + cfg.Resumable = true + cfg.SandboxFamily = types.SandboxFamilySpot + } // Warm fork: if checkpoint_id is set, fork from the local checkpoint cache. // ForkFromCheckpoint uses the local cache directly — no S3 needed. diff --git a/internal/worker/handlers.go b/internal/worker/handlers.go index 9b5deee8..57067186 100644 --- a/internal/worker/handlers.go +++ b/internal/worker/handlers.go @@ -21,6 +21,92 @@ import ( // jsonMarshal is a helper to marshal JSON for agent session stdin commands. var jsonMarshal = json.Marshal +type restartNoticeRequest struct { + NoticeSeconds int `json:"noticeSeconds"` + ETA string `json:"eta,omitempty"` +} + +type recreateSandboxRequest struct { + SandboxID string `json:"sandboxId"` + Config types.SandboxConfig `json:"config"` +} + +type existingSandboxStarter interface { + StartExistingSandbox(ctx context.Context, sandboxID string, cfg types.SandboxConfig) (*types.Sandbox, error) +} + +func (s *HTTPServer) adminRestartNotice(c echo.Context) error { + var req restartNoticeRequest + if c.Request().Body != nil { + _ = c.Bind(&req) + } + + notice := 25 * time.Second + if req.NoticeSeconds > 0 { + if req.NoticeSeconds > 120 { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "noticeSeconds must be <= 120"}) + } + notice = time.Duration(req.NoticeSeconds) * time.Second + } + + var eta time.Time + if req.ETA != "" { + parsed, err := time.Parse(time.RFC3339, req.ETA) + if err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "eta must be RFC3339"}) + } + eta = parsed + } + + notified := NotifyResumableSandboxesBeforeRestart(c.Request().Context(), s.manager, s.store, s.sandboxDBs, notice, eta) + return c.JSON(http.StatusAccepted, map[string]any{ + "notified": notified, + "noticeSeconds": int(notice.Seconds()), + }) +} + +func (s *HTTPServer) adminRecreateSandbox(c echo.Context) error { + var req recreateSandboxRequest + if err := c.Bind(&req); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "invalid request body: " + err.Error()}) + } + if req.SandboxID == "" { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "sandboxId required"}) + } + starter, ok := s.manager.(existingSandboxStarter) + if !ok { + return c.JSON(http.StatusNotImplemented, map[string]string{"error": "backend does not support recreating existing sandbox disks"}) + } + sb, err := starter.StartExistingSandbox(c.Request().Context(), req.SandboxID, req.Config) + if err != nil { + return c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()}) + } + if s.router != nil { + timeout := req.Config.Timeout + if timeout < 0 { + timeout = 0 + } + s.router.Register(req.SandboxID, time.Duration(timeout)*time.Second) + if timeout == 0 { + // Register applies the worker default timeout for zero values. For + // disk-only resumable recreates, zero means persistent: keep the + // router entry for request routing, but remove the idle timer. + s.router.SetTimeout(req.SandboxID, 0) + } + } + if s.sandboxDBs != nil { + if sdb, dbErr := s.sandboxDBs.Get(req.SandboxID); dbErr == nil { + _ = sdb.LogEvent("resumed", map[string]string{ + "sandbox_id": req.SandboxID, + "restart_reason": "resumable_recreate", + "preserves_disk": "true", + "preserves_memory": "false", + }) + } + } + return c.JSON(http.StatusOK, sb) +} + func (s *HTTPServer) setTimeout(c echo.Context) error { if s.router == nil { return c.JSON(http.StatusServiceUnavailable, map[string]string{ diff --git a/internal/worker/http_server.go b/internal/worker/http_server.go index 14eae37f..4063dae2 100644 --- a/internal/worker/http_server.go +++ b/internal/worker/http_server.go @@ -6,6 +6,7 @@ import ( "github.com/labstack/echo/v4" "github.com/labstack/echo/v4/middleware" "github.com/opensandbox/opensandbox/internal/auth" + "github.com/opensandbox/opensandbox/internal/db" "github.com/opensandbox/opensandbox/internal/observability" "github.com/opensandbox/opensandbox/internal/obslog" "github.com/opensandbox/opensandbox/internal/proxy" @@ -23,10 +24,11 @@ type HTTPServer struct { sandboxDBs *sandbox.SandboxDBManager router *sandbox.SandboxRouter sandboxDomain string + store *db.Store } // NewHTTPServer creates a new worker HTTP server for direct SDK access. -func NewHTTPServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, execMgr *sandbox.ExecSessionManager, jwtIssuer *auth.JWTIssuer, sandboxDBs *sandbox.SandboxDBManager, sbProxy *proxy.SandboxProxy, sbRouter *sandbox.SandboxRouter, sandboxDomain string) *HTTPServer { +func NewHTTPServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, execMgr *sandbox.ExecSessionManager, jwtIssuer *auth.JWTIssuer, sandboxDBs *sandbox.SandboxDBManager, sbProxy *proxy.SandboxProxy, sbRouter *sandbox.SandboxRouter, sandboxDomain string, store *db.Store) *HTTPServer { e := echo.New() e.HideBanner = true e.HidePort = true @@ -40,6 +42,7 @@ func NewHTTPServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, execMgr *san sandboxDBs: sandboxDBs, router: sbRouter, sandboxDomain: sandboxDomain, + store: store, } // Global middleware. Sentry goes first so it can observe panics and @@ -63,6 +66,8 @@ func NewHTTPServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, execMgr *san e.GET("/health", func(c echo.Context) error { return c.JSON(http.StatusOK, map[string]string{"status": "ok", "role": "worker"}) }) + e.POST("/admin/resumable/restart-notice", s.adminRestartNotice) + e.POST("/admin/resumable/recreate", s.adminRecreateSandbox) // All sandbox routes require JWT auth api := e.Group("") diff --git a/internal/worker/redis_heartbeat.go b/internal/worker/redis_heartbeat.go index 1b53e5bd..0f9854a3 100644 --- a/internal/worker/redis_heartbeat.go +++ b/internal/worker/redis_heartbeat.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log" + "sync" "time" "github.com/redis/go-redis/v9" @@ -12,15 +13,16 @@ import ( // redisHeartbeatPayload is the JSON structure published to Redis. type redisHeartbeatPayload struct { - WorkerID string `json:"worker_id"` - MachineID string `json:"machine_id,omitempty"` // EC2 instance ID (e.g. i-099088f8ac4a34ef3) - Region string `json:"region"` - GRPCAddr string `json:"grpc_addr"` - HTTPAddr string `json:"http_addr"` - Capacity int `json:"capacity"` - Current int `json:"current"` - CPUPct float64 `json:"cpu_pct"` - MemPct float64 `json:"mem_pct"` + WorkerID string `json:"worker_id"` + MachineID string `json:"machine_id,omitempty"` // EC2 instance ID (e.g. i-099088f8ac4a34ef3) + Pool string `json:"pool,omitempty"` + Region string `json:"region"` + GRPCAddr string `json:"grpc_addr"` + HTTPAddr string `json:"http_addr"` + Capacity int `json:"capacity"` + Current int `json:"current"` + CPUPct float64 `json:"cpu_pct"` + MemPct float64 `json:"mem_pct"` DiskPct float64 `json:"disk_pct"` TotalMemoryMB int `json:"total_memory_mb,omitempty"` CommittedMemoryMB int `json:"committed_memory_mb,omitempty"` @@ -51,20 +53,22 @@ type SandboxStatsWire struct { // 1. SETs worker:{id} with a 30s TTL (auto-expires if worker dies) // 2. PUBLISHes to workers:heartbeat for real-time server notification type RedisHeartbeat struct { - rdb *redis.Client - workerID string - machineID string - region string - grpcAddr string - httpAddr string - getStats func() (capacity, current int, cpuPct, memPct, diskPct float64) - getMemoryInfo func() (totalMB, committedMB int) // optional: committed memory for dynamic capacity - getSandboxStats func() map[string]SandboxStatsWire // optional: per-sandbox stats for autoscaler - onReconnect func() // called when heartbeat succeeds after a previous failure - goldenVersion string - workerVersion string - wasDown bool // true if the last publish failed (used to detect reconnect) - stop chan struct{} + rdb *redis.Client + workerID string + machineID string + pool string + region string + grpcAddr string + httpAddr string + getStats func() (capacity, current int, cpuPct, memPct, diskPct float64) + getMemoryInfo func() (totalMB, committedMB int) // optional: committed memory for dynamic capacity + getSandboxStats func() map[string]SandboxStatsWire // optional: per-sandbox stats for autoscaler + onReconnect func() // called when heartbeat succeeds after a previous failure + goldenVersion string + workerVersion string + wasDown bool // true if the last publish failed (used to detect reconnect) + stop chan struct{} + stopOnce sync.Once // guards close(stop) + rdb.Del — Stop() may be called from preemption handler and defer } // NewRedisHeartbeat creates a new heartbeat publisher. @@ -99,6 +103,12 @@ func NewRedisHeartbeat(redisURL, workerID, region, grpcAddr, httpAddr string) (* }, nil } +// SetPool sets the placement pool advertised by this worker. Empty defaults +// to ondemand on the control plane for backward compatibility. +func (h *RedisHeartbeat) SetPool(pool string) { + h.pool = pool +} + // SetMachineID sets the EC2 instance ID for the heartbeat (used by scaler for drain/terminate). func (h *RedisHeartbeat) SetMachineID(id string) { h.machineID = id @@ -159,15 +169,16 @@ func (h *RedisHeartbeat) publish() { capacity, current, cpuPct, memPct, diskPct := h.getStats() payload := redisHeartbeatPayload{ - WorkerID: h.workerID, - MachineID: h.machineID, - Region: h.region, - GRPCAddr: h.grpcAddr, - HTTPAddr: h.httpAddr, - Capacity: capacity, - Current: current, - CPUPct: cpuPct, - MemPct: memPct, + WorkerID: h.workerID, + MachineID: h.machineID, + Pool: h.pool, + Region: h.region, + GRPCAddr: h.grpcAddr, + HTTPAddr: h.httpAddr, + Capacity: capacity, + Current: current, + CPUPct: cpuPct, + MemPct: memPct, DiskPct: diskPct, GoldenVersion: h.goldenVersion, WorkerVersion: h.workerVersion, @@ -218,14 +229,18 @@ func (h *RedisHeartbeat) publish() { } // Stop stops the heartbeat publisher and closes the Redis connection. +// Idempotent — safe to call from both the preemption-handler goroutine and +// the normal `defer hb.Stop()` shutdown path. func (h *RedisHeartbeat) Stop() { - close(h.stop) + h.stopOnce.Do(func() { + close(h.stop) - // Remove the key so the server knows we're gone immediately - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - h.rdb.Del(ctx, "worker:"+h.workerID) + // Remove the key so the server knows we're gone immediately + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + h.rdb.Del(ctx, "worker:"+h.workerID) - h.rdb.Close() - log.Println("redis_heartbeat: stopped") + h.rdb.Close() + log.Println("redis_heartbeat: stopped") + }) } diff --git a/internal/worker/resumable_notice.go b/internal/worker/resumable_notice.go new file mode 100644 index 00000000..05d98943 --- /dev/null +++ b/internal/worker/resumable_notice.go @@ -0,0 +1,164 @@ +package worker + +import ( + "context" + "encoding/json" + "fmt" + "log" + "sync" + "time" + + "github.com/opensandbox/opensandbox/internal/db" + "github.com/opensandbox/opensandbox/internal/sandbox" + "github.com/opensandbox/opensandbox/pkg/types" +) + +// NotifyResumableSandboxesBeforeRestart runs the same in-guest notice hook used +// by cloud preemption handling. It gives resumable sandboxes a bounded window +// to flush state to disk before the host restarts. +func NotifyResumableSandboxesBeforeRestart(ctx context.Context, mgr sandbox.Manager, store *db.Store, sandboxDBs *sandbox.SandboxDBManager, noticeWindow time.Duration, eta time.Time) int { + if mgr == nil { + return 0 + } + wait := noticeWindow + if !eta.IsZero() { + until := time.Until(eta) + switch { + case until <= 0: + wait = 0 + case until <= 5*time.Second: + wait = until + case until-5*time.Second < wait: + wait = until - 5*time.Second + } + } + if wait < 0 { + wait = 0 + } + + listCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + sandboxes, err := mgr.List(listCtx) + cancel() + if err != nil { + log.Printf("opensandbox-worker: resumable restart notice: failed to list sandboxes: %v", err) + sleepContext(ctx, wait) + return 0 + } + + started := time.Now() + noticeSeconds := int(wait.Seconds()) + if noticeSeconds < 0 { + noticeSeconds = 0 + } + hookTimeout := noticeSeconds + if hookTimeout <= 0 { + hookTimeout = 1 + } + if hookTimeout > 20 { + hookTimeout = 20 + } + + const restartNoticeScript = `notice="${OPENSANDBOX_RESUME_NOTICE_SECONDS:-25}" +export OPENSANDBOX_RESTART_NOTICE_SECONDS="$notice" +for hook in /etc/opencomputer/on-restart-notice /home/sandbox/.opencomputer/on-restart-notice; do + if [ -x "$hook" ]; then + "$hook" "$notice" + fi +done +sync` + + var wg sync.WaitGroup + notified := 0 + for _, sb := range sandboxes { + if sb.Status != "" && sb.Status != "running" { + continue + } + if !isResumableSandboxSession(ctx, store, sb.ID) { + continue + } + + notified++ + if sandboxDBs != nil { + if sdb, dbErr := sandboxDBs.Get(sb.ID); dbErr == nil { + _ = sdb.LogEvent("restart_notice", map[string]string{ + "sandbox_id": sb.ID, + "notice_seconds": fmt.Sprintf("%d", noticeSeconds), + "restart_reason": "worker_preemption", + "preserves_disk": "true", + "preserves_memory": "false", + }) + } + } + + sandboxID := sb.ID + wg.Add(1) + go func() { + defer wg.Done() + execCtx, execCancel := context.WithTimeout(ctx, time.Duration(hookTimeout)*time.Second) + defer execCancel() + _, err := mgr.Exec(execCtx, sandboxID, types.ProcessConfig{ + Command: "/bin/sh", + Args: []string{"-lc", restartNoticeScript}, + Env: map[string]string{ + "OPENSANDBOX_BURST": "true", + "OPENSANDBOX_RESUMABLE": "true", + "OPENSANDBOX_RESUME_NOTICE_SECONDS": fmt.Sprintf("%d", noticeSeconds), + }, + Timeout: hookTimeout, + }) + if err != nil { + log.Printf("opensandbox-worker: resumable restart notice: hook failed for %s: %v", sandboxID, err) + } + }() + } + + if notified == 0 { + log.Printf("opensandbox-worker: resumable restart notice: no resumable sandboxes found") + sleepContext(ctx, wait) + return 0 + } + log.Printf("opensandbox-worker: resumable restart notice: notifying %d sandboxes with %ds notice", notified, noticeSeconds) + wg.Wait() + + remaining := wait - time.Since(started) + if remaining > 0 { + sleepContext(ctx, remaining) + } + log.Printf("opensandbox-worker: resumable restart notice: completed for %d sandboxes", notified) + return notified +} + +func isResumableSandboxSession(ctx context.Context, store *db.Store, sandboxID string) bool { + if store == nil || sandboxID == "" { + return false + } + sessionCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + session, err := store.GetSandboxSession(sessionCtx, sandboxID) + if err != nil || session == nil { + if err != nil { + log.Printf("opensandbox-worker: resumable restart notice: failed to load sandbox session %s: %v", sandboxID, err) + } + return false + } + var cfg types.SandboxConfig + if len(session.Config) > 0 { + if err := json.Unmarshal(session.Config, &cfg); err != nil { + log.Printf("opensandbox-worker: resumable restart notice: failed to parse config for %s: %v", sandboxID, err) + return false + } + } + return cfg.IsResumable() +} + +func sleepContext(ctx context.Context, d time.Duration) { + if d <= 0 { + return + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + case <-timer.C: + } +} diff --git a/pkg/types/sandbox.go b/pkg/types/sandbox.go index 9919dd47..fcfb2908 100644 --- a/pkg/types/sandbox.go +++ b/pkg/types/sandbox.go @@ -16,42 +16,61 @@ const ( SandboxStatusHibernated SandboxStatus = "hibernated" ) +const ( + SandboxFamilyDefault = "" + SandboxFamilySpot = "spot" +) + // Sandbox represents a running sandbox instance. type Sandbox struct { - ID string `json:"sandboxID"` - Template string `json:"templateID,omitempty"` - Alias string `json:"alias,omitempty"` - ClientID string `json:"clientID,omitempty"` - Status SandboxStatus `json:"status"` - StartedAt time.Time `json:"startedAt"` - EndAt time.Time `json:"endAt"` - Metadata map[string]string `json:"metadata,omitempty"` - CpuCount int `json:"cpuCount"` - MemoryMB int `json:"memoryMB"` - MachineID string `json:"machineID,omitempty"` + ID string `json:"sandboxID"` + Template string `json:"templateID,omitempty"` + Alias string `json:"alias,omitempty"` + ClientID string `json:"clientID,omitempty"` + Status SandboxStatus `json:"status"` + StartedAt time.Time `json:"startedAt"` + EndAt time.Time `json:"endAt"` + Metadata map[string]string `json:"metadata,omitempty"` + CpuCount int `json:"cpuCount"` + MemoryMB int `json:"memoryMB"` + SandboxFamily string `json:"sandboxFamily,omitempty"` + Burst bool `json:"burst,omitempty"` + MachineID string `json:"machineID,omitempty"` // ConnectURL and Token are currently unused by SDKs. All data-plane traffic // flows through the control plane's SandboxAPIProxy, which proxies to workers // over the internal VPC network. Direct worker access support coming in a future release. - ConnectURL string `json:"connectURL,omitempty"` - Token string `json:"token,omitempty"` - HostPort int `json:"hostPort,omitempty"` // Mapped host port for the sandbox's container port + ConnectURL string `json:"connectURL,omitempty"` + Token string `json:"token,omitempty"` + HostPort int `json:"hostPort,omitempty"` // Mapped host port for the sandbox's container port } // SandboxConfig is the request body for creating a sandbox. type SandboxConfig struct { - Template string `json:"templateID,omitempty"` - Alias string `json:"alias,omitempty"` - Metadata map[string]string `json:"metadata,omitempty"` - Timeout int `json:"timeout,omitempty"` // seconds, default 300 - CpuCount int `json:"cpuCount,omitempty"` // default 1 - MemoryMB int `json:"memoryMB,omitempty"` // default 256 - DiskMB int `json:"diskMB,omitempty"` // workspace disk in MB (default 20480) - Envs map[string]string `json:"envs,omitempty"` - Port int `json:"port,omitempty"` // container port to expose via subdomain (default 80) + Template string `json:"templateID,omitempty"` + Alias string `json:"alias,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` + Timeout int `json:"timeout,omitempty"` // seconds, default 300 + CpuCount int `json:"cpuCount,omitempty"` // default 1 + MemoryMB int `json:"memoryMB,omitempty"` // default 256 + DiskMB int `json:"diskMB,omitempty"` // workspace disk in MB (default 20480) + // SandboxFamily selects an internal placement family. Empty is the default + // on-demand family. "spot" routes through burst spare-capacity workers. + SandboxFamily string `json:"sandboxFamily,omitempty"` + // Burst selects the alpha lower-cost Burst Sandbox tier. Burst Sandboxes + // preserve disk across infrastructure restarts but do not guarantee running + // process or memory survival. Internally this maps to the spot placement + // family while that tier is backed by spare cloud capacity. + Burst bool `json:"burst,omitempty"` + // Resumable is the legacy public field name for Burst Sandboxes. Keep + // accepting it for existing callers and rows, but prefer Burst for new API + // requests/responses. + Resumable bool `json:"resumable,omitempty"` + Envs map[string]string `json:"envs,omitempty"` + Port int `json:"port,omitempty"` // container port to expose via subdomain (default 80) // NetworkEnabled is a pointer so we can distinguish "unset" from // "explicitly false". Unset defaults to true (see IsNetworkEnabled). - NetworkEnabled *bool `json:"networkEnabled,omitempty"` - ImageRef string `json:"imageRef,omitempty"` // resolved ECR URI for custom templates + NetworkEnabled *bool `json:"networkEnabled,omitempty"` + ImageRef string `json:"imageRef,omitempty"` // resolved ECR URI for custom templates // Sandbox snapshot template: S3 keys for rootfs and workspace drives. // When set, the sandbox boots from these drives instead of the standard base image. TemplateRootfsKey string `json:"templateRootfsKey,omitempty"` @@ -188,6 +207,44 @@ func ValidateResourceTier(cfg *SandboxConfig) error { return fmt.Errorf("cpuCount %d and memoryMB %d do not match an allowed tier; valid combinations: 1/1024, 1/4096, 2/8192, 4/16384, 8/32768, 16/65536", cfg.CpuCount, cfg.MemoryMB) } +// ApplySandboxFamilyDefaultsAndValidate normalizes alpha sandbox-family options +// before regular resource-tier validation. +func ApplySandboxFamilyDefaultsAndValidate(cfg *SandboxConfig) error { + if cfg.Burst || cfg.Resumable { + if cfg.SandboxFamily == "" || cfg.SandboxFamily == "default" { + cfg.SandboxFamily = SandboxFamilySpot + } + } + switch cfg.SandboxFamily { + case SandboxFamilyDefault: + return nil + case "default": + cfg.SandboxFamily = SandboxFamilyDefault + return nil + case "burst": + cfg.SandboxFamily = SandboxFamilySpot + cfg.Burst = true + return ApplySandboxFamilyDefaultsAndValidate(cfg) + case "resumable": + cfg.SandboxFamily = SandboxFamilySpot + cfg.Burst = true + return ApplySandboxFamilyDefaultsAndValidate(cfg) + case SandboxFamilySpot: + cfg.Burst = true + return nil + default: + return fmt.Errorf("unsupported sandboxFamily %q", cfg.SandboxFamily) + } +} + +func (c SandboxConfig) IsSpotFamily() bool { + return c.SandboxFamily == SandboxFamilySpot +} + +func (c SandboxConfig) IsResumable() bool { + return c.Burst || c.Resumable || c.SandboxFamily == SandboxFamilySpot +} + // SandboxListResponse is the response for listing sandboxes. type SandboxListResponse struct { Sandboxes []Sandbox `json:"sandboxes"` diff --git a/pkg/types/sandbox_family_test.go b/pkg/types/sandbox_family_test.go new file mode 100644 index 00000000..ff87870d --- /dev/null +++ b/pkg/types/sandbox_family_test.go @@ -0,0 +1,87 @@ +package types + +import ( + "strings" + "testing" +) + +func TestApplySandboxFamilyDefaultsAndValidateSpotMarksBurst(t *testing.T) { + cfg := SandboxConfig{SandboxFamily: SandboxFamilySpot} + + if err := ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !cfg.Burst { + t.Fatalf("expected internal spot family to mark sandbox burst") + } +} + +func TestApplySandboxFamilyDefaultsAndValidateBurstFlagMapsToSpot(t *testing.T) { + cfg := SandboxConfig{Burst: true} + + if err := ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.SandboxFamily != SandboxFamilySpot { + t.Fatalf("expected burst to map to internal spot family, got %q", cfg.SandboxFamily) + } + if cfg.CpuCount != 0 || cfg.MemoryMB != 0 { + t.Fatalf("expected burst not to force resources, got cpu=%d memory=%d", cfg.CpuCount, cfg.MemoryMB) + } +} + +func TestApplySandboxFamilyDefaultsAndValidateLegacyResumableFlagMapsToSpot(t *testing.T) { + cfg := SandboxConfig{Resumable: true} + + if err := ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.SandboxFamily != SandboxFamilySpot { + t.Fatalf("expected resumable to map to internal spot family, got %q", cfg.SandboxFamily) + } + if cfg.CpuCount != 0 || cfg.MemoryMB != 0 { + t.Fatalf("expected resumable not to force resources, got cpu=%d memory=%d", cfg.CpuCount, cfg.MemoryMB) + } +} + +func TestApplySandboxFamilyDefaultsAndValidateBurstFamilyAlias(t *testing.T) { + cfg := SandboxConfig{SandboxFamily: "burst"} + + if err := ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.SandboxFamily != SandboxFamilySpot || !cfg.Burst { + t.Fatalf("expected burst alias to map to internal spot family, got family=%q burst=%v", cfg.SandboxFamily, cfg.Burst) + } +} + +func TestApplySandboxFamilyDefaultsAndValidateLegacyResumableFamilyAlias(t *testing.T) { + cfg := SandboxConfig{SandboxFamily: "resumable"} + + if err := ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.SandboxFamily != SandboxFamilySpot || !cfg.Burst { + t.Fatalf("expected legacy resumable alias to map to internal spot family, got family=%q burst=%v", cfg.SandboxFamily, cfg.Burst) + } +} + +func TestApplySandboxFamilyDefaultsAndValidateBurstAllowsLargerTier(t *testing.T) { + cfg := SandboxConfig{SandboxFamily: SandboxFamilySpot, CpuCount: 1, MemoryMB: 4096} + + if err := ApplySandboxFamilyDefaultsAndValidate(&cfg); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if err := ValidateResourceTier(&cfg); err != nil { + t.Fatalf("expected larger burst tier to validate normally: %v", err) + } +} + +func TestApplySandboxFamilyDefaultsAndValidateRejectsUnknownFamily(t *testing.T) { + cfg := SandboxConfig{SandboxFamily: "gpu"} + + err := ApplySandboxFamilyDefaultsAndValidate(&cfg) + if err == nil || !strings.Contains(err.Error(), "unsupported sandboxFamily") { + t.Fatalf("expected unsupported family rejection, got %v", err) + } +} diff --git a/scripts/integration-tests/05-burst-edge-routing-smoke.ts b/scripts/integration-tests/05-burst-edge-routing-smoke.ts new file mode 100644 index 00000000..6bd25b76 --- /dev/null +++ b/scripts/integration-tests/05-burst-edge-routing-smoke.ts @@ -0,0 +1,196 @@ +/** + * Smoke test for Burst sandbox edge routing. + * + * Creates a Burst sandbox through the main public API, verifies basic exec, + * then proves the sandbox landed in the expected cell by either: + * - listing sandboxes through the direct cell/LB API, or + * - reading sandbox_sessions from Postgres. + * + * Usage: + * OPENCOMPUTER_API_KEY=osb_... \ + * OPENCOMPUTER_VERIFY_MODE=postgres \ + * OPENCOMPUTER_DATABASE_URL=postgres://... \ + * npx tsx scripts/integration-tests/05-burst-edge-routing-smoke.ts + * + * Or verify via the direct cell API: + * OPENCOMPUTER_API_KEY=osb_... \ + * OPENCOMPUTER_VERIFY_MODE=direct-list \ + * OPENCOMPUTER_DIRECT_API_URL=https://oc-alb-aws-us-east-2-burst-prod-972530125.us-east-2.elb.amazonaws.com \ + * OPENCOMPUTER_DIRECT_API_KEY=osb_... \ + * OPENCOMPUTER_DIRECT_TLS_INSECURE=1 \ + * npx tsx scripts/integration-tests/05-burst-edge-routing-smoke.ts + */ + +import { Sandbox } from "../../sdks/typescript/src"; +import { execFileSync } from "node:child_process"; + +const MAIN_API_URL = process.env.OPENCOMPUTER_API_URL || "https://app.opencomputer.dev"; +const MAIN_API_KEY = process.env.OPENCOMPUTER_API_KEY || ""; +const VERIFY_MODE = process.env.OPENCOMPUTER_VERIFY_MODE || "direct-list"; +const DIRECT_API_URL = + process.env.OPENCOMPUTER_DIRECT_API_URL || + "https://oc-alb-aws-us-east-2-burst-prod-972530125.us-east-2.elb.amazonaws.com"; +const DIRECT_API_KEY = process.env.OPENCOMPUTER_DIRECT_API_KEY || MAIN_API_KEY; +const DATABASE_URL = process.env.OPENCOMPUTER_DATABASE_URL || ""; + +if (process.env.OPENCOMPUTER_DIRECT_TLS_INSECURE === "1") { + process.env.NODE_TLS_REJECT_UNAUTHORIZED = "0"; +} + +function apiBase(url: string): string { + const base = url.replace(/\/+$/, ""); + return base.endsWith("/api") ? base : `${base}/api`; +} + +function requireEnv(name: string, value: string): void { + if (!value) { + throw new Error(`${name} is required`); + } +} + +interface SessionRow { + sandboxID: string; + status: string; + workerID: string; + region: string; + template: string; + config: Record; +} + +async function listDirectSandboxes(): Promise>> { + const resp = await fetch(`${apiBase(DIRECT_API_URL)}/sandboxes`, { + headers: DIRECT_API_KEY ? { "X-API-Key": DIRECT_API_KEY } : {}, + }); + if (!resp.ok) { + const text = await resp.text(); + throw new Error(`direct sandbox list failed: ${resp.status} ${text}`); + } + const data = await resp.json(); + if (!Array.isArray(data)) { + throw new Error(`direct sandbox list returned non-array response: ${JSON.stringify(data)}`); + } + return data; +} + +async function waitForDirectList(sandboxID: string, timeoutMs = 30_000): Promise> { + const deadline = Date.now() + timeoutMs; + let lastCount = 0; + + while (Date.now() < deadline) { + const sandboxes = await listDirectSandboxes(); + lastCount = sandboxes.length; + const found = sandboxes.find((sb) => sb.sandboxID === sandboxID || sb.id === sandboxID); + if (found) { + return found; + } + await new Promise((resolve) => setTimeout(resolve, 1_000)); + } + + throw new Error(`sandbox ${sandboxID} was not visible from direct cell list after ${timeoutMs}ms (last count=${lastCount})`); +} + +function queryPostgresSession(sandboxID: string): SessionRow { + if (!/^[a-zA-Z0-9_-]+$/.test(sandboxID)) { + throw new Error(`refusing unsafe sandbox id for SQL lookup: ${sandboxID}`); + } + + const query = ` + SELECT json_build_object( + 'sandboxID', sandbox_id, + 'status', status, + 'workerID', worker_id, + 'region', region, + 'template', template, + 'config', config + )::text + FROM sandbox_sessions + WHERE sandbox_id = '${sandboxID}' + ORDER BY started_at DESC + LIMIT 1; + `; + + const out = execFileSync("psql", [DATABASE_URL, "-XAt", "-c", query], { + encoding: "utf8", + stdio: ["ignore", "pipe", "pipe"], + }).trim(); + + if (!out) { + throw new Error(`sandbox ${sandboxID} was not found in sandbox_sessions`); + } + return JSON.parse(out) as SessionRow; +} + +async function main() { + requireEnv("OPENCOMPUTER_API_KEY", MAIN_API_KEY); + + console.log(`main API: ${MAIN_API_URL}`); + console.log(`verify mode: ${VERIFY_MODE}`); + + if (VERIFY_MODE === "direct-list") { + requireEnv("OPENCOMPUTER_DIRECT_API_KEY", DIRECT_API_KEY); + console.log(`direct API: ${DIRECT_API_URL}`); + const before = await listDirectSandboxes(); + console.log(`direct list before create: ${before.length} running sandbox(es)`); + } else if (VERIFY_MODE === "postgres") { + requireEnv("OPENCOMPUTER_DATABASE_URL", DATABASE_URL); + } else { + throw new Error(`unsupported OPENCOMPUTER_VERIFY_MODE=${VERIFY_MODE}`); + } + + const sandbox = await Sandbox.create({ + apiUrl: MAIN_API_URL, + apiKey: MAIN_API_KEY, + template: "base", + burst: true, + timeout: 0, + }); + + console.log(`created burst sandbox: ${sandbox.sandboxId} (${sandbox.status})`); + + try { + const result = await sandbox.exec.run("echo hello world", { timeout: 30 }); + const stdout = result.stdout.trim(); + if (result.exitCode !== 0 || stdout !== "hello world") { + throw new Error(`exec failed: exit=${result.exitCode}, stdout=${JSON.stringify(stdout)}, stderr=${JSON.stringify(result.stderr)}`); + } + console.log("exec check: ok"); + + if (VERIFY_MODE === "direct-list") { + const direct = await waitForDirectList(sandbox.sandboxId); + console.log(`direct list check: ok (${sandbox.sandboxId})`); + if (direct.workerID || direct.workerId || direct.worker_id) { + console.log(`worker: ${String(direct.workerID || direct.workerId || direct.worker_id)}`); + } + if (direct.region) { + console.log(`region: ${String(direct.region)}`); + } + } else { + const row = queryPostgresSession(sandbox.sandboxId); + if (row.status !== "running") { + throw new Error(`expected running session, got ${row.status}`); + } + if (!row.workerID) { + throw new Error("session has no workerID"); + } + if (row.config?.burst !== true && row.config?.resumable !== true && row.config?.sandboxFamily !== "spot") { + throw new Error(`session config does not look like Burst: ${JSON.stringify(row.config)}`); + } + console.log(`postgres check: ok (${row.sandboxID})`); + console.log(`worker: ${row.workerID}`); + console.log(`region: ${row.region}`); + console.log(`template: ${row.template}`); + } + } finally { + try { + await sandbox.kill(); + console.log(`cleaned up: ${sandbox.sandboxId}`); + } catch (err) { + console.error(`cleanup failed for ${sandbox.sandboxId}:`, err); + } + } +} + +main().catch((err) => { + console.error(err); + process.exit(1); +}); diff --git a/scripts/test-spot-drain-health.sh b/scripts/test-spot-drain-health.sh new file mode 100755 index 00000000..e5f54de8 --- /dev/null +++ b/scripts/test-spot-drain-health.sh @@ -0,0 +1,232 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Live-migration drain health test for spot/preemption drills. +# +# Required env: +# OPENSANDBOX_API_URL e.g. https://spot-poc.opencomputer.app or http://127.0.0.1:8080 +# OPENSANDBOX_API_KEY +# SOURCE_WORKER_ID worker to evacuate +# TARGET_WORKER_ID spare worker expected to receive the sandboxes +# +# Optional env: +# COUNT number of sandboxes to create (default: 40) +# TEMPLATE sandbox template (default: default) +# MEMORY_MB sandbox memory (default: 1024) +# CPU_COUNT sandbox vCPU count (default: 1) +# DISK_MB sandbox disk MB (default: 20480) +# POLL_SECONDS drain poll interval (default: 2) +# DRAIN_TIMEOUT_SECONDS drain timeout (default: 600) +# KEEP_SANDBOXES set to 1 to leave sandboxes running on exit (default: 0) +# FORCE_SOURCE_PLACEMENT set to 0 to skip temporarily draining target during create (default: 1) + +API_URL="${OPENSANDBOX_API_URL:?OPENSANDBOX_API_URL is required}" +API_KEY="${OPENSANDBOX_API_KEY:?OPENSANDBOX_API_KEY is required}" +SOURCE_WORKER_ID="${SOURCE_WORKER_ID:?SOURCE_WORKER_ID is required}" +TARGET_WORKER_ID="${TARGET_WORKER_ID:?TARGET_WORKER_ID is required}" + +COUNT="${COUNT:-40}" +TEMPLATE="${TEMPLATE:-default}" +MEMORY_MB="${MEMORY_MB:-1024}" +CPU_COUNT="${CPU_COUNT:-1}" +DISK_MB="${DISK_MB:-20480}" +POLL_SECONDS="${POLL_SECONDS:-2}" +DRAIN_TIMEOUT_SECONDS="${DRAIN_TIMEOUT_SECONDS:-600}" +KEEP_SANDBOXES="${KEEP_SANDBOXES:-0}" +FORCE_SOURCE_PLACEMENT="${FORCE_SOURCE_PLACEMENT:-1}" + +AUTH_HEADER="X-API-Key: ${API_KEY}" +CONTENT_HEADER="Content-Type: application/json" +SANDBOX_FILE="$(mktemp -t spot-drain-sandboxes.XXXXXX)" +FAIL_FILE="$(mktemp -t spot-drain-failures.XXXXXX)" + +cleanup() { + api POST "/admin/workers/${SOURCE_WORKER_ID}/drain?drain=false" >/dev/null 2>&1 || true + api POST "/admin/workers/${TARGET_WORKER_ID}/drain?drain=false" >/dev/null 2>&1 || true + + if [[ "$KEEP_SANDBOXES" == "1" ]]; then + echo "Leaving sandboxes running. IDs: $SANDBOX_FILE" + return + fi + + if [[ -s "$SANDBOX_FILE" ]]; then + echo "Cleaning up $(wc -l < "$SANDBOX_FILE" | tr -d ' ') sandboxes..." + while read -r sandbox_id; do + [[ -n "$sandbox_id" ]] || continue + curl -fsS -X DELETE -H "$AUTH_HEADER" "${API_URL}/api/sandboxes/${sandbox_id}" >/dev/null || true + done < "$SANDBOX_FILE" + fi +} +trap cleanup EXIT + +api() { + local method="$1" + local path="$2" + local body="${3:-}" + + if [[ -n "$body" ]]; then + curl -fsS -X "$method" -H "$AUTH_HEADER" -H "$CONTENT_HEADER" -d "$body" "${API_URL}${path}" + else + curl -fsS -X "$method" -H "$AUTH_HEADER" "${API_URL}${path}" + fi +} + +json_value() { + local expr="$1" + python3 -c " +import json, sys +try: + data = json.load(sys.stdin) + value = ${expr} + print('' if value is None else value) +except Exception: + print('') +" +} + +tracked_counts() { + local source_count=0 target_count=0 other_count=0 non_running=0 total=0 + local sandbox_id meta status worker + + while read -r sandbox_id; do + [[ -n "$sandbox_id" ]] || continue + total=$((total + 1)) + meta=$(api GET "/api/sandboxes/${sandbox_id}" || echo '{}') + status=$(printf "%s" "$meta" | json_value 'data.get("status")') + worker=$(printf "%s" "$meta" | json_value 'data.get("workerID")') + + if [[ "$status" != "running" ]]; then + non_running=$((non_running + 1)) + elif [[ "$worker" == "$SOURCE_WORKER_ID" ]]; then + source_count=$((source_count + 1)) + elif [[ "$worker" == "$TARGET_WORKER_ID" ]]; then + target_count=$((target_count + 1)) + else + other_count=$((other_count + 1)) + fi + done < "$SANDBOX_FILE" + + echo "source=$source_count target=$target_count other=$other_count non_running=$non_running total=$total" +} + +create_sandbox() { + local body response sandbox_id + body=$(printf '{"template":"%s","memoryMB":%s,"cpuCount":%s,"diskMB":%s}' \ + "$TEMPLATE" "$MEMORY_MB" "$CPU_COUNT" "$DISK_MB") + response=$(api POST /api/sandboxes "$body") + sandbox_id=$(printf "%s" "$response" | json_value 'data.get("sandboxID") or data.get("id")') + if [[ -z "$sandbox_id" ]]; then + echo "create failed: $response" >&2 + return 1 + fi + echo "$sandbox_id" +} + +exec_ok() { + local sandbox_id="$1" + local response exit_code + response=$(api POST "/api/sandboxes/${sandbox_id}/exec/run" \ + '{"cmd":"/bin/true","args":[],"timeout":5}' || true) + exit_code=$(printf "%s" "$response" | json_value 'data.get("exitCode")') + [[ "$exit_code" == "0" ]] +} + +pty_ok() { + local sandbox_id="$1" + local response session_id + response=$(api POST "/api/sandboxes/${sandbox_id}/pty" \ + '{"cols":80,"rows":24,"shell":"/bin/bash"}' || true) + session_id=$(printf "%s" "$response" | json_value 'data.get("sessionID")') + if [[ -z "$session_id" ]]; then + return 1 + fi + api DELETE "/api/sandboxes/${sandbox_id}/pty/${session_id}" >/dev/null || true +} + +worker_of() { + local sandbox_id="$1" + api GET "/api/sandboxes/${sandbox_id}" | json_value 'data.get("workerID")' +} + +health_pass() { + local label="$1" + local failures=0 total=0 worker exec_status pty_status + : > "$FAIL_FILE" + + echo "Health pass: $label" + while read -r sandbox_id; do + [[ -n "$sandbox_id" ]] || continue + total=$((total + 1)) + worker=$(worker_of "$sandbox_id") + + exec_status=ok + if ! exec_ok "$sandbox_id"; then + exec_status=fail + fi + + pty_status=ok + if ! pty_ok "$sandbox_id"; then + pty_status=fail + fi + + printf "%s worker=%s exec=%s pty=%s\n" "$sandbox_id" "$worker" "$exec_status" "$pty_status" + if [[ "$exec_status" != "ok" || "$pty_status" != "ok" ]]; then + failures=$((failures + 1)) + printf "%s worker=%s exec=%s pty=%s\n" "$sandbox_id" "$worker" "$exec_status" "$pty_status" >> "$FAIL_FILE" + fi + done < "$SANDBOX_FILE" + + echo "Health summary: label=$label total=$total failures=$failures" + if [[ "$failures" -ne 0 ]]; then + cat "$FAIL_FILE" >&2 + return 1 + fi +} + +echo "Clearing drain markers on source and target workers..." +api POST "/admin/workers/${SOURCE_WORKER_ID}/drain?drain=false" >/dev/null || true +api POST "/admin/workers/${TARGET_WORKER_ID}/drain?drain=false" >/dev/null || true + +if [[ "$FORCE_SOURCE_PLACEMENT" == "1" ]]; then + echo "Temporarily draining target worker during create: $TARGET_WORKER_ID" + api POST "/admin/workers/${TARGET_WORKER_ID}/drain" >/dev/null +fi + +echo "Creating $COUNT sandboxes..." +for i in $(seq 1 "$COUNT"); do + sandbox_id=$(create_sandbox) + echo "$sandbox_id" | tee -a "$SANDBOX_FILE" + sleep 0.3 +done + +if [[ "$FORCE_SOURCE_PLACEMENT" == "1" ]]; then + echo "Clearing target worker drain marker before evacuation: $TARGET_WORKER_ID" + api POST "/admin/workers/${TARGET_WORKER_ID}/drain?drain=false" >/dev/null +fi + +echo "Initial tracked counts: $(tracked_counts)" +health_pass before-drain + +echo "Evacuating source worker: $SOURCE_WORKER_ID" +api POST "/admin/workers/${SOURCE_WORKER_ID}/evacuate" >/dev/null + +start_epoch=$(date +%s) +while true; do + counts=$(tracked_counts) + now=$(date -u "+%Y-%m-%d %H:%M:%S UTC") + elapsed=$(( $(date +%s) - start_epoch )) + echo "$now elapsed=${elapsed}s $counts" + + if [[ "$counts" == source=0* ]]; then + break + fi + if [[ "$elapsed" -gt "$DRAIN_TIMEOUT_SECONDS" ]]; then + echo "Drain timed out after ${elapsed}s" >&2 + exit 1 + fi + sleep "$POLL_SECONDS" +done + +echo "Final tracked counts: $(tracked_counts)" +health_pass after-drain +echo "PASS: all $COUNT sandboxes retained exec and PTY after drain" diff --git a/sdks/python/opencomputer/__init__.py b/sdks/python/opencomputer/__init__.py index ae6b3471..4e98c0c9 100644 --- a/sdks/python/opencomputer/__init__.py +++ b/sdks/python/opencomputer/__init__.py @@ -1,6 +1,6 @@ """OpenComputer Python SDK - cloud sandbox platform.""" -from opencomputer.sandbox import Sandbox, ScalingLockedError, PlanLimitError +from opencomputer.sandbox import Sandbox, ScalingLockedError, PlanLimitError, SandboxFamilyLimitError from opencomputer.agent import Agent, AgentEvent, AgentSession, AgentSessionInfo from opencomputer.filesystem import Filesystem from opencomputer.exec import Exec, ProcessResult, ExecSession, ExecSessionInfo @@ -29,6 +29,7 @@ "Sandbox", "ScalingLockedError", "PlanLimitError", + "SandboxFamilyLimitError", "Agent", "AgentEvent", "AgentSession", diff --git a/sdks/python/opencomputer/sandbox.py b/sdks/python/opencomputer/sandbox.py index 5fe14e11..f4b1485e 100644 --- a/sdks/python/opencomputer/sandbox.py +++ b/sdks/python/opencomputer/sandbox.py @@ -33,6 +33,14 @@ class PlanLimitError(Exception): """ +class SandboxFamilyLimitError(Exception): + """Raised when a resource-changing call is blocked by a sandbox family. + Kept for compatibility with older API responses. + """ + + code = "sandbox_family_scale_disabled" + + def _raise_scaling_error(resp: httpx.Response, action: str) -> None: """Inspect a non-OK scaling response and raise the most specific error. Falls back to ``raise_for_status`` so callers still see HTTP details for @@ -43,6 +51,8 @@ def _raise_scaling_error(resp: httpx.Response, action: str) -> None: body = {} if resp.status_code == 403 and isinstance(body, dict) and body.get("code") == "scaling_locked": raise ScalingLockedError(body.get("error", "scaling is locked on this sandbox")) + if resp.status_code == 403 and isinstance(body, dict) and body.get("code") == "sandbox_family_scale_disabled": + raise SandboxFamilyLimitError(body.get("error", "sandbox family does not allow scaling")) if resp.status_code == 402: msg = body.get("error", "plan limit exceeded") if isinstance(body, dict) else "plan limit exceeded" raise PlanLimitError(msg) @@ -75,6 +85,9 @@ async def create( api_url: str | None = None, envs: dict[str, str] | None = None, metadata: dict[str, str] | None = None, + burst: bool | None = None, + resumable: bool | None = None, + sandbox_family: str | None = None, disk_mb: int | None = None, memory_mb: int | None = None, secret_store: str | None = None, @@ -91,6 +104,11 @@ async def create( api_url: API URL (or OPENCOMPUTER_API_URL env var). envs: Environment variables to inject. Overrides store secrets. metadata: Custom metadata key-value pairs. + burst: Create a Burst Sandbox. Disk is preserved across + infrastructure restarts; processes may restart. + resumable: Deprecated alias for ``burst``. + sandbox_family: Internal/legacy placement family. Prefer + ``burst=True`` for public API usage. disk_mb: Workspace disk size in MB (default 20480 = 20GB). Any additional GB above 20GB is metered at a per-second rate comparable to EBS gp3. Closed beta: requests above 20GB @@ -134,6 +152,12 @@ async def create( body["envs"] = envs if metadata: body["metadata"] = metadata + if burst is not None: + body["burst"] = burst + elif resumable is not None: + body["resumable"] = resumable + if sandbox_family: + body["sandboxFamily"] = sandbox_family if disk_mb is not None: body["diskMB"] = disk_mb if memory_mb is not None: diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index d2498465..f77cd463 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "opencomputer-sdk" -version = "0.6.0" +version = "0.6.1" description = "Python SDK for OpenComputer - cloud sandbox platform" readme = "README.md" requires-python = ">=3.10" diff --git a/sdks/typescript/package-lock.json b/sdks/typescript/package-lock.json index 6cf6c844..e54228f2 100644 --- a/sdks/typescript/package-lock.json +++ b/sdks/typescript/package-lock.json @@ -1,12 +1,12 @@ { "name": "@opencomputer/sdk", - "version": "0.6.0", + "version": "0.6.1", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@opencomputer/sdk", - "version": "0.6.0", + "version": "0.6.1", "license": "MIT", "devDependencies": { "@types/node": "^25.3.0", diff --git a/sdks/typescript/package.json b/sdks/typescript/package.json index 9f55155b..2caa0f8f 100644 --- a/sdks/typescript/package.json +++ b/sdks/typescript/package.json @@ -1,6 +1,6 @@ { "name": "@opencomputer/sdk", - "version": "0.6.0", + "version": "0.6.1", "description": "TypeScript SDK for OpenComputer - cloud sandbox platform", "type": "module", "main": "dist/index.js", diff --git a/sdks/typescript/src/index.ts b/sdks/typescript/src/index.ts index 1bd215b0..e3512d49 100644 --- a/sdks/typescript/src/index.ts +++ b/sdks/typescript/src/index.ts @@ -2,6 +2,7 @@ export { Sandbox, ScalingLockedError, PlanLimitError, + SandboxFamilyLimitError, type SandboxOpts, type CheckpointInfo, type PatchInfo, diff --git a/sdks/typescript/src/sandbox.ts b/sdks/typescript/src/sandbox.ts index b1f9cb1a..c6918e55 100644 --- a/sdks/typescript/src/sandbox.ts +++ b/sdks/typescript/src/sandbox.ts @@ -13,6 +13,12 @@ function resolveApiUrl(url: string): string { export interface SandboxOpts { template?: string; + /** Create a Burst Sandbox. Disk is preserved across infrastructure restarts; processes may restart. */ + burst?: boolean; + /** @deprecated Use `burst` instead. */ + resumable?: boolean; + /** Internal/legacy placement family. Prefer `burst: true` for public API usage. */ + sandboxFamily?: "spot"; /** * Idle timeout in seconds after which the sandbox auto-hibernates. * Default: `0` (persistent — never auto-hibernate). @@ -46,6 +52,9 @@ interface SandboxData { sandboxID: string; status: string; templateID?: string; + sandboxFamily?: string; + burst?: boolean; + resumable?: boolean; connectURL?: string; token?: string; sandboxDomain?: string; @@ -117,6 +126,18 @@ export class PlanLimitError extends Error { } } +/** + * Thrown when a resource-changing call is blocked by the sandbox's family. + * Kept for compatibility with older API responses. + */ +export class SandboxFamilyLimitError extends Error { + readonly code = "sandbox_family_scale_disabled"; + constructor(message?: string) { + super(message ?? "sandbox family does not allow scaling"); + this.name = "SandboxFamilyLimitError"; + } +} + /** * Inspect a non-OK response from a scaling endpoint and throw the most * specific error type. Falls back to a generic Error when the response @@ -133,6 +154,9 @@ async function throwScalingError(resp: Response, action: string): Promise if (resp.status === 403 && body.code === "scaling_locked") { throw new ScalingLockedError(body.error); } + if (resp.status === 403 && body.code === "sandbox_family_scale_disabled") { + throw new SandboxFamilyLimitError(body.error); + } if (resp.status === 402) { throw new PlanLimitError(body.error); } @@ -245,6 +269,9 @@ export class Sandbox { }; if (opts.envs) body.envs = opts.envs; if (opts.metadata) body.metadata = opts.metadata; + if (opts.burst != null) body.burst = opts.burst; + if (opts.resumable != null) body.resumable = opts.resumable; + if (opts.sandboxFamily) body.sandboxFamily = opts.sandboxFamily; if (opts.cpuCount != null) body.cpuCount = opts.cpuCount; if (opts.memoryMB != null) body.memoryMB = opts.memoryMB; if (opts.diskMB != null) body.diskMB = opts.diskMB;