open-lambda · AmiBuch · Feb 19, 2026 · Feb 19, 2026 · Mar 2, 2026 · Mar 2, 2026
diff --git a/go/worker/sandbox/mock.go b/go/worker/sandbox/mock.go
@@ -0,0 +1,122 @@
+package sandbox
+
+import (
+	"fmt"
+	"net/http"
+	"sync"
+	"sync/atomic"
+)
+
+// MockSandbox is a test double for Sandbox.
+// Exported fields control error injection; state fields track lifecycle.
+type MockSandbox struct {
+	mu        sync.Mutex
+	id        string
+	paused    bool
+	destroyed bool
+
+	// Set these before calling Get/Put to inject errors.
+	PauseErr   error
+	UnpauseErr error
+}
+
+var mockIDCounter int64
+
+// NewMockSandbox creates a MockSandbox with the given ID.
+func NewMockSandbox(id string) *MockSandbox {
+	return &MockSandbox{id: id, paused: true}
+}
+
+func (m *MockSandbox) ID() string { return m.id }
+
+func (m *MockSandbox) Destroy(reason string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.destroyed = true
+}
+
+func (m *MockSandbox) DestroyIfPaused(reason string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.paused {
+		m.destroyed = true
+	}
+}
+
+func (m *MockSandbox) Pause() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.PauseErr != nil {
+		return m.PauseErr
+	}
+	m.paused = true
+	return nil
+}
+
+func (m *MockSandbox) Unpause() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.UnpauseErr != nil {
+		return m.UnpauseErr
+	}
+	m.paused = false
+	return nil
+}
+
+func (m *MockSandbox) Client() *http.Client   { return nil }
+func (m *MockSandbox) Meta() *SandboxMeta      { return nil }
+func (m *MockSandbox) GetRuntimeLog() string   { return "" }
+func (m *MockSandbox) GetProxyLog() string     { return "" }
+func (m *MockSandbox) DebugString() string     { return fmt.Sprintf("mock:%s", m.id) }
+func (m *MockSandbox) fork(dst Sandbox) error  { return nil }
+func (m *MockSandbox) childExit(child Sandbox) {}
+
+// IsDestroyed returns whether Destroy has been called.
+func (m *MockSandbox) IsDestroyed() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.destroyed
+}
+
+// IsPaused returns the current pause state.
+func (m *MockSandbox) IsPaused() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.paused
+}
+
+// MockSandboxPool is a test double for SandboxPool.
+// It creates MockSandbox instances with auto-incremented IDs.
+type MockSandboxPool struct {
+	mu      sync.Mutex
+	Created []*MockSandbox
+
+	// Set before calling Get to make pool.Create fail.
+	CreateErr error
+}
+
+func (p *MockSandboxPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta) (Sandbox, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.CreateErr != nil {
+		return nil, p.CreateErr
+	}
+	id := fmt.Sprintf("mock-%d", atomic.AddInt64(&mockIDCounter, 1))
+	sb := NewMockSandbox(id)
+	sb.paused = false // Pool.Create returns unpaused sandboxes
+	p.Created = append(p.Created, sb)
+	return sb, nil
+}
+
+func (p *MockSandboxPool) Cleanup()                            {}
+func (p *MockSandboxPool) AddListener(handler SandboxEventFunc) {}
+func (p *MockSandboxPool) DebugString() string                  { return "mock-pool" }
+
+// CreatedSandboxes returns a snapshot of all sandboxes created by this pool.
+func (p *MockSandboxPool) CreatedSandboxes() []*MockSandbox {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	out := make([]*MockSandbox, len(p.Created))
+	copy(out, p.Created)
+	return out
+}
diff --git a/go/worker/sandboxset/api.go b/go/worker/sandboxset/api.go
@@ -0,0 +1,126 @@
+// Package sandboxset provides a thread-safe pool of sandboxes for a single
+// Lambda function.
+//
+// A SandboxSet replaces per-instance goroutines with a simple pool.
+// Callers just ask for a sandbox and don't worry about whether it is
+// freshly created or recycled from a previous request.
+//
+// Sandbox lifecycle inside a SandboxSet:
+//
+//	[created]
+//	    |
+//	    v
+//	[paused]  <---+
+//	    |         |
+//	    v         |
+//	[in-use]  ----+  (Put)
+//	    |
+//	    v
+//	[destroyed]     (Destroy / Close / error)
+//
+// Usage:
+//
+//	set, err := sandboxset.New(&sandboxset.Config{
+//	    Pool:        myPool,
+//	    CodeDir:     "/path/to/lambda",
+//	    ScratchDirs: myScratchDirs,
+//	})
+//
+//	sb, err := set.Get()
+//	// ... handle request ...
+//	set.Put(sb)
+package sandboxset
+
+import (
+	"github.com/open-lambda/open-lambda/go/common"
+	"github.com/open-lambda/open-lambda/go/worker/sandbox"
+)
+
+/*
+A SandboxSet manages a pool of sandboxes for one Lambda function.
+All methods are safe to call from multiple goroutines.
+
+The design mirrors the C process API: Get (create), Put (exit),
+Destroy (kill), Close (cleanup). There are no warm-up, shrink, or
+stats methods yet — those can be added in later PRs without
+changing the core interface.
+*/
+type SandboxSet interface {
+	// Return an unpaused sandbox ready to handle a request.
+	//
+	// If the pool has an idle sandbox, it is unpaused and returned.
+	// If Unpause fails (e.g., the SOCK container died while paused),
+	// that sandbox is destroyed and Get tries the next idle one or
+	// creates a fresh sandbox.
+	//
+	// A fresh scratch directory is created for each new sandbox
+	// via Config.ScratchDirs. Reused sandboxes keep their
+	// existing scratch directory from when they were first created.
+	Get() (sandbox.Sandbox, error)
+
+	// Return a sandbox to the pool after a successful request.
+	//
+	// The sandbox is paused and becomes available for the next Get.
+	// If Pause fails (e.g., the container died during the request),
+	// the sandbox is destroyed automatically — a bad sandbox never
+	// re-enters the pool.
+	//
+	// Passing a sandbox that is not in the pool returns an error
+	// but is otherwise harmless.
+	Put(sb sandbox.Sandbox) error
+
+	// Permanently remove a sandbox from the pool and destroy it.
+	//
+	// Use this when a request produced an unrecoverable error and
+	// the sandbox should not be reused. "reason" is a
+	// human-readable explanation that shows up in later error
+	// messages (same convention as sandbox.Sandbox.Destroy).
+	//
+	// If the sandbox is not in the pool it is still destroyed —
+	// resources are always freed. The returned error is
+	// informational only.
+	Destroy(sb sandbox.Sandbox, reason string) error
+
+	// Destroy all sandboxes in the pool and mark the set as closed.
+	//
+	// Callers who still hold sandbox references from a previous Get
+	// will find them already dead, which is safe: per the Sandbox
+	// contract, methods on a destroyed sandbox are harmless no-ops
+	// that return errors.
+	//
+	// Calling Close a second time returns an error.
+	Close() error
+}
+
+// Config holds the parameters needed to create a SandboxSet.
+type Config struct {
+	// Pool creates and destroys the underlying sandboxes.
+	Pool sandbox.SandboxPool
+
+	// Parent sandbox to fork from (may be nil). When nil, new
+	// sandboxes are created from scratch. Not all SandboxPool
+	// implementations support forking.
+	Parent sandbox.Sandbox
+
+	// IsLeaf marks sandboxes as non-forkable, meaning they will
+	// not be used as parents for future forks.
+	IsLeaf bool
+
+	// CodeDir is the directory containing the Lambda handler code.
+	CodeDir string
+
+	// Meta holds runtime configuration (memory limits, packages,
+	// imports, etc.). Nil means the pool fills in defaults.
+	Meta *sandbox.SandboxMeta
+
+	// ScratchDirs creates a unique writable directory for each
+	// new sandbox. The set calls ScratchDirs.Make internally
+	// so that Get can remain argument-free.
+	ScratchDirs *common.DirMaker
+}
+
+// New creates a SandboxSet from cfg. Returns an error if any of
+// Pool, CodeDir, or ScratchDirs are missing.
+func New(cfg *Config) (SandboxSet, error) {
+	return newSandboxSet(cfg)
+}
diff --git a/go/worker/sandboxset/close.go b/go/worker/sandboxset/close.go
@@ -0,0 +1,23 @@
+package sandboxset
+
+import "fmt"
+
+// Close implements SandboxSet.
+//
+// All sandboxes are snapshot under the lock, then destroyed outside it.
+func (s *sandboxSetImpl) Close() error {
+	s.mu.Lock()
+	if s.closed {
+		s.mu.Unlock()
+		return fmt.Errorf("sandboxset: already closed")
+	}
+	s.closed = true
+	pool := s.pool
+	s.pool = nil
+	s.mu.Unlock()
+
+	for _, w := range pool {
+		w.sb.Destroy("sandboxset closed")
+	}
+	return nil
+}
diff --git a/go/worker/sandboxset/destroy.go b/go/worker/sandboxset/destroy.go
@@ -0,0 +1,34 @@
+package sandboxset
+
+import (
+	"fmt"
+
+	"github.com/open-lambda/open-lambda/go/worker/sandbox"
+)
+
+// Destroy implements SandboxSet.
+//
+// The wrapper is spliced out of the pool under a short lock using O(1)
+// swap-with-tail. Destroy is called outside the lock to keep critical
+// sections short. The sandbox is always destroyed even if it was not
+// found in the pool.
+func (s *sandboxSetImpl) Destroy(sb sandbox.Sandbox, reason string) error {
+	s.mu.Lock()
+	found := false
+	for i, w := range s.pool {
+		if w.sb.ID() == sb.ID() {
+			s.pool[i] = s.pool[len(s.pool)-1]
+			s.pool = s.pool[:len(s.pool)-1]
+			found = true
+			break
+		}
+	}
+	s.mu.Unlock()
+
+	sb.Destroy(reason)
+
+	if !found {
+		return fmt.Errorf("sandboxset: sandbox %s not found in pool (still destroyed)", sb.ID())
+	}
+	return nil
+}
diff --git a/go/worker/sandboxset/get.go b/go/worker/sandboxset/get.go
@@ -0,0 +1,59 @@
+package sandboxset
+
+import (
+	"fmt"
+
+	"github.com/open-lambda/open-lambda/go/worker/sandbox"
+)
+
+// Get implements SandboxSet.
+//
+// Fast path: an idle sandbox is claimed under a short lock, then Unpause
+// runs outside the lock so the pool is not stalled during I/O.
+//
+// Slow path: no idle sandbox exists, so a new one is created without
+// holding the lock.
+func (s *sandboxSetImpl) Get() (sandbox.Sandbox, error) {
+	s.mu.Lock()
+	if s.closed {
+		s.mu.Unlock()
+		return nil, fmt.Errorf("sandboxset: closed")
+	}
+
+	// Fast path: claim an idle sandbox.
+	var claimed sandbox.Sandbox
+	for _, w := range s.pool {
+		if !w.inUse {
+			w.inUse = true
+			claimed = w.sb
+			break
+		}
+	}
+	s.mu.Unlock()
+
+	if claimed != nil {
+		// Unpause outside the lock (split-lock pattern).
+		if err := claimed.Unpause(); err != nil {
+			_ = s.Destroy(claimed, fmt.Sprintf("unpause: %v", err))
+			return s.Get()
+		}
+		return claimed, nil
+	}
+
+	// Slow path: create a new sandbox without holding the lock.
+	scratchDir := s.cfg.ScratchDirs.Make("sb")
+	sb, err := s.cfg.Pool.Create(
+		s.cfg.Parent, s.cfg.IsLeaf,
+		s.cfg.CodeDir, scratchDir,
+		s.cfg.Meta,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("sandboxset: create: %w", err)
+	}
+
+	s.mu.Lock()
+	s.pool = append(s.pool, &sandboxWrapper{sb: sb, inUse: true})
+	s.mu.Unlock()
+
+	return sb, nil
+}
diff --git a/go/worker/sandboxset/put.go b/go/worker/sandboxset/put.go
@@ -0,0 +1,30 @@
+package sandboxset
+
+import (
+	"fmt"
+
+	"github.com/open-lambda/open-lambda/go/worker/sandbox"
+)
+
+// Put implements SandboxSet.
+//
+// The sandbox is paused and its wrapper is flipped back to idle.
+// If Pause fails, the sandbox is destroyed rather than silently
+// recycled — a bad sandbox should never re-enter the pool.
+func (s *sandboxSetImpl) Put(sb sandbox.Sandbox) error {
+	if err := sb.Pause(); err != nil {
+		_ = s.Destroy(sb, fmt.Sprintf("pause failed: %v", err))
+		return fmt.Errorf("sandboxset: sandbox %s destroyed because Pause failed: %w", sb.ID(), err)
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for _, w := range s.pool {
+		if w.sb.ID() == sb.ID() {
+			w.inUse = false
+			return nil
+		}
+	}
+	return fmt.Errorf("sandboxset: sandbox %s not found in pool", sb.ID())
+}