optimization: added a sandbox queue to retrieve warm sandboxes (#536)

nit

added mutex

addressing feedback

updated comments

linter

linter

comments

nit

fix tests

addressed comments

fix

update

nit

nit

lint

memory leak

lint

updated exclusivity test

addressed copilot comments

Author: vicentefb

cmd/agent-sandbox-controller/main.go

@@ -29,15 +29,15 @@ import (
 
 	"github.com/felixge/fgprof"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
-	ctrl "sigs.k8s.io/controller-runtime"
-	"sigs.k8s.io/controller-runtime/pkg/healthz"
-	"sigs.k8s.io/controller-runtime/pkg/log/zap"
-	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
-
 	"sigs.k8s.io/agent-sandbox/controllers"
 	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
 	extensionscontrollers "sigs.k8s.io/agent-sandbox/extensions/controllers"
+	"sigs.k8s.io/agent-sandbox/extensions/controllers/queue"
 	asmetrics "sigs.k8s.io/agent-sandbox/internal/metrics"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/healthz"
+	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 	//+kubebuilder:scaffold:imports
 )
 
@@ -227,11 +227,13 @@ func main() {
 	}
 
 	if extensions {
+		warmSandboxQueue := queue.NewSimpleSandboxQueue()
 		if err = (&extensionscontrollers.SandboxClaimReconciler{
-			Client:   mgr.GetClient(),
-			Scheme:   mgr.GetScheme(),
-			Recorder: mgr.GetEventRecorder("sandboxclaim-controller"),
-			Tracer:   instrumenter,
+			Client:           mgr.GetClient(),
+			Scheme:           mgr.GetScheme(),
+			WarmSandboxQueue: warmSandboxQueue,
+			Recorder:         mgr.GetEventRecorder("sandboxclaim-controller"),
+			Tracer:           instrumenter,
 		}).SetupWithManager(mgr, sandboxClaimConcurrentWorkers); err != nil {
 			setupLog.Error(err, "unable to create controller", "controller", "SandboxClaim")
 			os.Exit(1)

extensions/controllers/queue/simple_sandbox_queue.go

@@ -0,0 +1,141 @@
+// Copyright 2026 The Kubernetes Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package queue
+
+import (
+	"sync"
+
+	"k8s.io/apimachinery/pkg/types"
+)
+
+// SandboxKey uniquely identifies a sandbox in the queue.
+type SandboxKey types.NamespacedName
+
+// SandboxQueue defines the interface for managing a thread-safe,
+// highly concurrent queue of adoptable warm pool sandboxes.
+type SandboxQueue interface {
+	Add(templateHash string, item SandboxKey)
+	Get(templateHash string) (SandboxKey, bool)
+	RemoveQueue(templateHash string)
+	RemoveItem(templateHash string, item SandboxKey)
+}
+
+// SimpleSandboxQueue implements SandboxQueue using simple synchronized slices.
+type SimpleSandboxQueue struct {
+	// queues is a thread-safe dictionary from template hash to a synchronizedQueue
+	queues sync.Map
+}
+
+// NewSimpleSandboxQueue initializes a new SimpleSandboxQueue.
+func NewSimpleSandboxQueue() *SimpleSandboxQueue {
+	return &SimpleSandboxQueue{}
+}
+
+// Add pushes an item to the specific template's queue.
+func (s *SimpleSandboxQueue) Add(templateHash string, item SandboxKey) {
+	q, _ := s.queues.LoadOrStore(templateHash, newSynchronizedQueue())
+	q.(*synchronizedQueue).Push(item)
+}
+
+// Get pops an item from the specific template's queue.
+func (s *SimpleSandboxQueue) Get(templateHash string) (SandboxKey, bool) {
+	q, ok := s.queues.Load(templateHash)
+	if !ok {
+		return SandboxKey{}, false
+	}
+	return q.(*synchronizedQueue).Pop()
+}
+
+// RemoveItem deletes a specific sandbox from a template's queue.
+func (s *SimpleSandboxQueue) RemoveItem(templateHash string, item SandboxKey) {
+	if q, ok := s.queues.Load(templateHash); ok {
+		q.(*synchronizedQueue).Remove(item)
+	}
+}
+
+// Remove scans the slice and deletes the item to prevent Ghost Pods.
+func (q *synchronizedQueue) Remove(key SandboxKey) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if _, exists := q.set[key]; !exists {
+		return
+	}
+
+	delete(q.set, key)
+
+	for i, k := range q.items {
+		if k == key {
+			// Delete from slice
+			q.items = append(q.items[:i], q.items[i+1:]...)
+			break
+		}
+	}
+}
+
+// TODO(vicentefb): Implement queue cleanup mechanism.
+// We should remove the queue from the sync.Map when the corresponding
+// SandboxWarmPool for a given template is deleted to prevent memory leaks.
+type synchronizedQueue struct {
+	mu    sync.Mutex
+	items []SandboxKey
+	set   map[SandboxKey]struct{} // Used for O(1) deduplication
+}
+
+func newSynchronizedQueue() *synchronizedQueue {
+	return &synchronizedQueue{
+		items: make([]SandboxKey, 0),
+		set:   make(map[SandboxKey]struct{}),
+	}
+}
+
+// Push adds an item to the queue if it isn't already present.
+func (q *synchronizedQueue) Push(key SandboxKey) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if _, exists := q.set[key]; !exists {
+		q.set[key] = struct{}{}
+		q.items = append(q.items, key)
+	}
+}
+
+// Pop removes and returns the first item from the queue.
+func (q *synchronizedQueue) Pop() (SandboxKey, bool) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if len(q.items) == 0 {
+		return SandboxKey{}, false
+	}
+
+	// Grab the first item
+	item := q.items[0]
+
+	// This removes the pointer references so the Garbage Collector
+	// can free the strings in memory!
+	q.items[0] = SandboxKey{}
+
+	// Remove it from slice and set
+	q.items = q.items[1:]
+	delete(q.set, item)
+
+	return item, true
+}
+
+// RemoveQueue completely deletes a template's queue from the sync.Map
+// to prevent memory leaks when SandboxTemplates or WarmPools are deleted.
+func (s *SimpleSandboxQueue) RemoveQueue(templateHash string) {
+	s.queues.Delete(templateHash)
+}

extensions/controllers/queue/simple_sandbox_queue_test.go

@@ -0,0 +1,118 @@
+// Copyright 2026 The Kubernetes Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package queue
+
+import (
+	"testing"
+)
+
+func TestSimpleSandboxQueue_BasicOperations(t *testing.T) {
+	q := NewSimpleSandboxQueue()
+	hash := "template-hash-1"
+
+	key1 := SandboxKey{Namespace: "default", Name: "sb-1"}
+	key2 := SandboxKey{Namespace: "default", Name: "sb-2"}
+
+	// Test Add
+	q.Add(hash, key1)
+	q.Add(hash, key2)
+
+	// Test Get (Should be FIFO)
+	got1, ok1 := q.Get(hash)
+	if !ok1 || got1 != key1 {
+		t.Errorf("Expected %v, got %v (ok: %v)", key1, got1, ok1)
+	}
+
+	got2, ok2 := q.Get(hash)
+	if !ok2 || got2 != key2 {
+		t.Errorf("Expected %v, got %v (ok: %v)", key2, got2, ok2)
+	}
+
+	// Queue should now be empty
+	_, ok3 := q.Get(hash)
+	if ok3 {
+		t.Errorf("Expected queue to be empty, but got an item")
+	}
+}
+
+func TestSimpleSandboxQueue_RemoveItem_GhostPodFix(t *testing.T) {
+	q := NewSimpleSandboxQueue()
+	hash := "template-hash-1"
+
+	key1 := SandboxKey{Namespace: "default", Name: "sb-1"}
+	key2 := SandboxKey{Namespace: "default", Name: "sb-2"}
+	key3 := SandboxKey{Namespace: "default", Name: "sb-3"}
+
+	q.Add(hash, key1)
+	q.Add(hash, key2)
+	q.Add(hash, key3)
+
+	// Simulate the Kubelet deleting the middle pod (Ghost Pod scenario)
+	q.RemoveItem(hash, key2)
+
+	// First pop should still be key1
+	got1, _ := q.Get(hash)
+	if got1 != key1 {
+		t.Errorf("Expected %v, got %v", key1, got1)
+	}
+
+	// Second pop should be key3! (key2 was successfully removed)
+	got3, _ := q.Get(hash)
+	if got3 != key3 {
+		t.Errorf("Expected %v to skip deleted item and return %v, but got %v", hash, key3, got3)
+	}
+
+	// Queue should now be empty
+	_, ok := q.Get(hash)
+	if ok {
+		t.Errorf("Expected queue to be empty after Ghost Pod removal")
+	}
+}
+
+func TestSynchronizedQueue_Deduplication(t *testing.T) {
+	q := newSynchronizedQueue()
+	key := SandboxKey{Namespace: "default", Name: "duplicate-sb"}
+
+	// Push the exact same pod 3 times
+	q.Push(key)
+	q.Push(key)
+	q.Push(key)
+
+	// Verify it only stored it once
+	if len(q.items) != 1 {
+		t.Errorf("Expected length 1 due to O(1) deduplication, got %d", len(q.items))
+	}
+
+	// Verify the set also only has 1 item
+	if len(q.set) != 1 {
+		t.Errorf("Expected set length 1, got %d", len(q.set))
+	}
+}
+
+func TestSimpleSandboxQueue_RemoveQueue_MemoryLeakFix(t *testing.T) {
+	q := NewSimpleSandboxQueue()
+	hash := "template-hash-to-delete"
+	key1 := SandboxKey{Namespace: "default", Name: "sb-1"}
+
+	q.Add(hash, key1)
+
+	// Simulate SandboxTemplate deletion
+	q.RemoveQueue(hash)
+
+	// Verify the entire queue was wiped from the sync.Map
+	_, ok := q.Get(hash)
+	if ok {
+		t.Errorf("Expected queue to be completely removed, but it still existed")
+	}
+}