Merge pull request #26 from diggerhq/feat/checkpointing

Add manual checkpoint system for sandboxes

Author: motatoes

cmd/server/main.go

@@ -157,7 +157,7 @@ func main() {
 			WorkerID:        workerID,
 			OnHibernate: func(sandboxID string, result *sandbox.HibernateResult) {
 				log.Printf("opensandbox: sandbox %s auto-hibernated (key=%s, size=%d bytes)",
-					sandboxID, result.CheckpointKey, result.SizeBytes)
+					sandboxID, result.HibernationKey, result.SizeBytes)
 				if opts.Store != nil {
 					_ = opts.Store.UpdateSandboxSessionStatus(context.Background(), sandboxID, "hibernated", nil)
 				}

cmd/worker/main.go

@@ -52,6 +52,17 @@ func main() {
 	defer fcMgr.Close()
 	log.Println("opensandbox-worker: Firecracker VM manager initialized")
 
+	// Clean up orphaned Firecracker processes + TAP devices BEFORE starting golden snapshot.
+	// Must run first to avoid killing the golden snapshot VM (race condition).
+	fcMgr.CleanupOrphanedProcesses()
+
+	// Prepare golden snapshot for fast default VM creation (~500ms vs ~2s cold boot)
+	go func() {
+		if err := fcMgr.PrepareGoldenSnapshot(); err != nil {
+			log.Printf("opensandbox-worker: golden snapshot preparation failed: %v (cold boot fallback active)", err)
+		}
+	}()
+
 	// The Firecracker manager implements sandbox.Manager
 	var mgr sandbox.Manager = fcMgr
 
@@ -149,9 +160,6 @@ func main() {
 			defer store.Close()
 			log.Println("opensandbox-worker: PostgreSQL store connected (auto-wake enabled)")
 
-			// Kill orphaned Firecracker processes + TAP devices from previous run
-			fcMgr.CleanupOrphanedProcesses()
-
 			// Local NVMe recovery: scan for sandbox data left from a previous run
 			recoveries := fcMgr.RecoverLocalSandboxes()
 			if len(recoveries) > 0 {
@@ -163,14 +171,14 @@ func main() {
 						continue
 					}
 					if r.HasSnapshot {
-						// Full snapshot on NVMe — create checkpoint record so doWake finds local files
-						_, _ = store.CreateCheckpoint(ctx, r.SandboxID, session.OrgID,
+						// Full snapshot on NVMe — create hibernation record so doWake finds local files
+						_, _ = store.CreateHibernation(ctx, r.SandboxID, session.OrgID,
 							"local://"+r.SandboxID, 0, session.Region, session.Template, session.Config)
 						_ = store.UpdateSandboxSessionStatus(ctx, r.SandboxID, "hibernated", nil)
 						snapshotCount++
 					} else {
-						// Workspace only — create local sentinel checkpoint for cold boot
-						_, _ = store.CreateCheckpoint(ctx, r.SandboxID, session.OrgID,
+						// Workspace only — create local sentinel hibernation for cold boot
+						_, _ = store.CreateHibernation(ctx, r.SandboxID, session.OrgID,
 							"local://"+r.SandboxID, 0, session.Region, session.Template, session.Config)
 						_ = store.UpdateSandboxSessionStatus(ctx, r.SandboxID, "hibernated", nil)
 						workspaceCount++
@@ -199,13 +207,13 @@ func main() {
 		WorkerID:        cfg.WorkerID,
 		OnHibernate: func(sandboxID string, result *sandbox.HibernateResult) {
 			log.Printf("opensandbox-worker: sandbox %s auto-hibernated (key=%s, size=%d bytes)",
-				sandboxID, result.CheckpointKey, result.SizeBytes)
+				sandboxID, result.HibernationKey, result.SizeBytes)
 			if store != nil {
-				// Create checkpoint record so wake-on-request can find it
+				// Create hibernation record so wake-on-request can find it
 				session, err := store.GetSandboxSession(context.Background(), sandboxID)
 				if err == nil {
-					_, _ = store.CreateCheckpoint(context.Background(), sandboxID, session.OrgID,
-						result.CheckpointKey, result.SizeBytes, session.Region, session.Template, session.Config)
+					_, _ = store.CreateHibernation(context.Background(), sandboxID, session.OrgID,
+						result.HibernationKey, result.SizeBytes, session.Region, session.Template, session.Config)
 				}
 				_ = store.UpdateSandboxSessionStatus(context.Background(), sandboxID, "hibernated", nil)
 			}
@@ -346,12 +354,12 @@ func main() {
 					}
 					continue
 				}
-				log.Printf("opensandbox-worker: hibernated %s (key=%s)", r.SandboxID, r.CheckpointKey)
+				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 {
-						_, _ = store.CreateCheckpoint(context.Background(), r.SandboxID, session.OrgID,
-							r.CheckpointKey, 0, session.Region, session.Template, session.Config)
+						_, _ = store.CreateHibernation(context.Background(), r.SandboxID, session.OrgID,
+							r.HibernationKey, 0, session.Region, session.Template, session.Config)
 						_ = store.UpdateSandboxSessionStatus(context.Background(), r.SandboxID, "hibernated", nil)
 					}
 				}

internal/agent/stats.go

@@ -3,10 +3,10 @@ package agent
 import (
 	"bufio"
 	"context"
-	"fmt"
 	"os"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 
 	pb "github.com/opensandbox/opensandbox/proto/agent"
@@ -182,14 +182,10 @@ func (s *Server) SyncFS(ctx context.Context, req *pb.SyncFSRequest) (*pb.SyncFSR
 	return &pb.SyncFSResponse{}, nil
 }
 
-// syncFS calls sync(2) to flush all filesystems.
+// syncFS calls sync(2) to flush all filesystem buffers.
+// This syncs ALL mounted filesystems (rootfs + workspace), ensuring dirty pages
+// are written to their backing ext4 images before snapshot/checkpoint.
 func syncFS() error {
-	f, err := os.Open("/")
-	if err != nil {
-		return fmt.Errorf("open /: %w", err)
-	}
-	defer f.Close()
-	// sync(2) — we use SyncFileRange isn't available on all filesystems,
-	// just use Sync on the root fd which triggers global sync
-	return f.Sync()
+	syscall.Sync()
+	return nil
 }

internal/api/commands.go

@@ -77,6 +77,24 @@ func (s *Server) runCommand(c echo.Context) error {
 }
 
 func (s *Server) runCommandRemote(c echo.Context, sandboxID string) error {
+	// Wait for sandbox if it's being created asynchronously
+	if v, ok := s.pendingCreates.Load(sandboxID); ok {
+		pending := v.(*pendingCreate)
+		select {
+		case <-pending.ready:
+			if pending.err != nil {
+				return c.JSON(http.StatusInternalServerError, map[string]string{
+					"error": "sandbox creation failed: " + pending.err.Error(),
+				})
+			}
+			s.pendingCreates.Delete(sandboxID)
+		case <-c.Request().Context().Done():
+			return c.JSON(http.StatusGatewayTimeout, map[string]string{
+				"error": "timed out waiting for sandbox creation",
+			})
+		}
+	}
+
 	if s.store == nil {
 		return c.JSON(http.StatusServiceUnavailable, map[string]string{
 			"error": "database not configured",

internal/api/dashboard.go

@@ -708,14 +708,14 @@ func (s *Server) dashboardGetSession(c echo.Context) error {
 		}
 	}
 
-	// If hibernated, include checkpoint info
+	// If hibernated, include hibernation info
 	if session.Status == "hibernated" {
-		checkpoint, err := s.store.GetActiveCheckpoint(c.Request().Context(), sandboxID)
+		hibernation, err := s.store.GetActiveHibernation(c.Request().Context(), sandboxID)
 		if err == nil {
-			resp["checkpoint"] = map[string]interface{}{
-				"checkpointKey": checkpoint.CheckpointKey,
-				"sizeBytes":     checkpoint.SizeBytes,
-				"hibernatedAt":  checkpoint.HibernatedAt,
+			resp["hibernation"] = map[string]interface{}{
+				"hibernationKey": hibernation.HibernationKey,
+				"sizeBytes":      hibernation.SizeBytes,
+				"hibernatedAt":   hibernation.HibernatedAt,
 			}
 		}
 	}

internal/api/router.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"github.com/labstack/echo/v4"
 	"github.com/labstack/echo/v4/middleware"
@@ -43,6 +44,13 @@ type Server struct {
 	sandboxDomain   string                            // base domain for sandbox subdomains
 	ecrConfig       *ecr.Config                       // nil if ECR not configured
 	cfClient        *cloudflare.Client                // nil if Cloudflare not configured
+	pendingCreates  sync.Map                          // map[sandboxID]*pendingCreate — async sandbox creation tracking
+}
+
+// pendingCreate tracks an async sandbox creation.
+type pendingCreate struct {
+	ready chan struct{} // closed when creation completes
+	err   error        // set before closing ready
 }
 
 // ServerOpts holds optional dependencies for the API server.
@@ -127,6 +135,13 @@ func NewServer(mgr sandbox.Manager, ptyMgr *sandbox.PTYManager, apiKey string, o
 	api.POST("/sandboxes/:id/hibernate", s.hibernateSandbox)
 	api.POST("/sandboxes/:id/wake", s.wakeSandbox)
 
+	// Checkpoints
+	api.POST("/sandboxes/:id/checkpoints", s.createCheckpoint)
+	api.GET("/sandboxes/:id/checkpoints", s.listCheckpoints)
+	api.POST("/sandboxes/:id/checkpoints/:checkpointId/restore", s.restoreCheckpoint)
+	api.POST("/sandboxes/from-checkpoint/:checkpointId", s.createFromCheckpoint)
+	api.DELETE("/sandboxes/:id/checkpoints/:checkpointId", s.deleteCheckpoint)
+
 	// Preview URLs (on-demand port-based)
 	api.POST("/sandboxes/:id/preview", s.createPreviewURL)
 	api.GET("/sandboxes/:id/preview", s.listPreviewURLs)