perf: improve reconcile throughput under registry latency and large namespaces

cocoonset/agents.go

@@ -32,40 +32,77 @@ const subAgentCreateConcurrency = 8
 // (non-zero when a sub-agent is in rebuild backoff and the caller should
 // re-reconcile when backoff elapses). Missing slots are created concurrently
 // so batch scale-ups do not serialize N apiserver round trips.
-func (r *Reconciler) ensureSubAgents(ctx context.Context, cs *cocoonv1.CocoonSet, classified classifiedPods, mainVMName, mainNodeName string) (bool, time.Duration, error) {
+func (r *Reconciler) ensureSubAgents(ctx context.Context, cs *cocoonv1.CocoonSet, classified classifiedPods, mainVMName, mainNodeName string, intent restoreIntent) (bool, time.Duration, error) {
 	logger := log.WithFunc("cocoonset.Reconciler.ensureSubAgents")
 	changed := false
 	var requeueAfter time.Duration
 
-	restorable, err := r.podsRestorableByCR(ctx, cs.Namespace)
+	var missing []int32
+	for slot := int32(1); slot <= cs.Spec.Agent.Replicas; slot++ {
+		pod, exists := classified.sub[slot]
+		if !exists {
+			missing = append(missing, slot)
+			continue
+		}
+		deleted, wait, err := r.triageSubAgent(ctx, logger, pod, cs, slot)
+		if err != nil {
+			return changed, requeueAfter, err
+		}
+		if deleted {
+			changed = true
+		}
+		if wait > 0 && (requeueAfter == 0 || wait < requeueAfter) {
+			requeueAfter = wait
+		}
+	}
+
+	created, err := r.createSubAgents(ctx, logger, cs, missing, mainVMName, mainNodeName, intent)
+	if created {
+		changed = true
+	}
 	if err != nil {
 		return changed, requeueAfter, err
 	}
 
+	for _, slot := range slices.Sorted(maps.Keys(classified.sub)) {
+		if slot <= cs.Spec.Agent.Replicas {
+			continue
+		}
+		pod := classified.sub[slot]
+		if err := r.Delete(ctx, pod); err != nil {
+			if apierrors.IsNotFound(err) {
+				continue
+			}
+			return changed, requeueAfter, fmt.Errorf("delete extra sub-agent slot %d: %w", slot, err)
+		}
+		logger.Infof(ctx, "deleted extra sub-agent %s/%s", pod.Namespace, pod.Name)
+		changed = true
+	}
+	return changed, requeueAfter, nil
+}
+
+// createSubAgents builds the missing slots concurrently so a batch scale-up does
+// not serialize N apiserver round trips.
+func (r *Reconciler) createSubAgents(ctx context.Context, logger *log.Fields, cs *cocoonv1.CocoonSet, missing []int32, mainVMName, mainNodeName string, intent restoreIntent) (bool, error) {
+	if len(missing) == 0 {
+		return false, nil
+	}
+	// Resolved before the fan-out so the goroutines only ever read it.
+	restorable, err := intent()
+	if err != nil {
+		return false, err
+	}
 	g, gctx := errgroup.WithContext(ctx)
 	g.SetLimit(subAgentCreateConcurrency)
 	var created atomic.Bool
-	for slot := int32(1); slot <= cs.Spec.Agent.Replicas; slot++ {
-		if pod, exists := classified.sub[slot]; exists {
-			deleted, wait, err := r.triageSubAgent(ctx, logger, pod, cs, slot)
-			if err != nil {
-				return changed, requeueAfter, err
-			}
-			if deleted {
-				changed = true
-			}
-			if wait > 0 && (requeueAfter == 0 || wait < requeueAfter) {
-				requeueAfter = wait
-			}
-			continue
-		}
+	for _, slot := range missing {
 		g.Go(func() error {
 			subPod, err := buildAgentPod(cs, slot, mainVMName, mainNodeName, r.Scheme)
 			if err != nil {
 				return fmt.Errorf("build sub-agent slot %d: %w", slot, err)
 			}
-			_, intent := restorable[subPod.Name]
-			if err := r.markRestoreIfHibernated(gctx, subPod, intent); err != nil {
+			_, wantRestore := restorable[subPod.Name]
+			if err := r.markRestoreIfHibernated(gctx, subPod, wantRestore); err != nil {
 				return fmt.Errorf("mark restore sub-agent slot %d: %w", slot, err)
 			}
 			if err := r.Create(gctx, subPod); err != nil {
@@ -79,27 +116,8 @@ func (r *Reconciler) ensureSubAgents(ctx context.Context, cs *cocoonv1.CocoonSet
 			return nil
 		})
 	}
-	if err := g.Wait(); err != nil {
-		return changed || created.Load(), requeueAfter, err
-	}
-	if created.Load() {
-		changed = true
-	}
-	for _, slot := range slices.Sorted(maps.Keys(classified.sub)) {
-		if slot <= cs.Spec.Agent.Replicas {
-			continue
-		}
-		pod := classified.sub[slot]
-		if err := r.Delete(ctx, pod); err != nil {
-			if apierrors.IsNotFound(err) {
-				continue
-			}
-			return changed, requeueAfter, fmt.Errorf("delete extra sub-agent slot %d: %w", slot, err)
-		}
-		logger.Infof(ctx, "deleted extra sub-agent %s/%s", pod.Namespace, pod.Name)
-		changed = true
-	}
-	return changed, requeueAfter, nil
+	waitErr := g.Wait()
+	return created.Load(), waitErr
 }
 
 // triageSubAgent deletes pod when it is terminal or has drifted from spec.

cocoonset/pods_test.go

@@ -559,7 +559,7 @@ func TestPodSpecMatchesAgentIgnoresAffinity(t *testing.T) {
 	}
 }
 
-func testScheme(t *testing.T) *runtime.Scheme {
+func testScheme(t testing.TB) *runtime.Scheme {
 	t.Helper()
 	scheme := runtime.NewScheme()
 	utilruntime.Must(clientgoscheme.AddToScheme(scheme))

cocoonset/reconcile_bench_test.go

@@ -0,0 +1,112 @@
+package cocoonset
+
+import (
+	"context"
+	"fmt"
+	"slices"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	ctrlfake "sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	cocoonv1 "github.com/cocoonstack/cocoon-common/apis/v1"
+)
+
+const (
+	benchCocoonSets    = 24
+	benchRegistryDelay = 5 * time.Millisecond
+)
+
+// BenchmarkReconcileThroughput drives independent CocoonSets that each block on
+// one fixed-latency registry probe; concurrency 1 is the pre-change behavior.
+func BenchmarkReconcileThroughput(b *testing.B) {
+	for _, concurrency := range []int{1, 4, 8} {
+		b.Run(fmt.Sprintf("concurrency=%d", concurrency), func(b *testing.B) {
+			var wall time.Duration
+			var latencies []time.Duration
+			for b.Loop() {
+				b.StopTimer()
+				r, sets := newBenchFixture(b)
+				b.StartTimer()
+
+				start := time.Now()
+				latencies = append(latencies, runReconciles(b, r, sets, concurrency)...)
+				wall += time.Since(start)
+			}
+			slices.Sort(latencies)
+			b.ReportMetric(float64(wall.Milliseconds())/float64(b.N), "wall-ms/op")
+			b.ReportMetric(float64(percentile(latencies, 95).Milliseconds()), "p95-ms")
+		})
+	}
+}
+
+// runReconciles drains sets through a worker pool, the shape controller-runtime
+// gives MaxConcurrentReconciles.
+func runReconciles(b *testing.B, r *Reconciler, sets []*cocoonv1.CocoonSet, concurrency int) []time.Duration {
+	b.Helper()
+	queue := make(chan *cocoonv1.CocoonSet, len(sets))
+	for _, cs := range sets {
+		queue <- cs
+	}
+	close(queue)
+
+	latencies := make([]time.Duration, len(sets))
+	var idx atomic.Int32
+	var wg sync.WaitGroup
+	for range concurrency {
+		wg.Go(func() {
+			for cs := range queue {
+				start := time.Now()
+				if _, err := r.Reconcile(context.Background(), reqFor(cs)); err != nil {
+					b.Errorf("Reconcile %s: %v", cs.Name, err)
+				}
+				latencies[idx.Add(1)-1] = time.Since(start)
+			}
+		})
+	}
+	wg.Wait()
+	return latencies
+}
+
+// newBenchFixture builds independent CocoonSets each missing its main agent and
+// carrying a hibernated CR, so every reconcile pays exactly one registry probe.
+func newBenchFixture(b *testing.B) (*Reconciler, []*cocoonv1.CocoonSet) {
+	b.Helper()
+	scheme := testScheme(b)
+	objs := make([]client.Object, 0, benchCocoonSets*2)
+	sets := make([]*cocoonv1.CocoonSet, 0, benchCocoonSets)
+	for i := range benchCocoonSets {
+		cs := &cocoonv1.CocoonSet{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:       fmt.Sprintf("bench-%d", i),
+				Namespace:  "ns",
+				UID:        types.UID(fmt.Sprintf("uid-%d", i)),
+				Finalizers: []string{finalizerName},
+			},
+			Spec: cocoonv1.CocoonSetSpec{
+				Agent: cocoonv1.AgentSpec{Image: "ghcr.io/cocoonstack/cocoon/ubuntu:24.04"},
+			},
+		}
+		sets = append(sets, cs)
+		objs = append(objs, cs, hibernatedFor(cs))
+	}
+	cli := ctrlfake.NewClientBuilder().
+		WithScheme(scheme).
+		WithObjects(objs...).
+		WithStatusSubresource(&cocoonv1.CocoonSet{}).
+		Build()
+	return &Reconciler{Client: cli, Scheme: scheme, Registry: &fakeRegistry{delay: benchRegistryDelay}}, sets
+}
+
+func percentile(sorted []time.Duration, p int) time.Duration {
+	if len(sorted) == 0 {
+		return 0
+	}
+	i := min(len(sorted)*p/100, len(sorted)-1)
+	return sorted[i]
+}

cocoonset/reconciler.go

@@ -13,6 +13,7 @@ import (
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/builder"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
 
@@ -36,15 +37,22 @@ type Reconciler struct {
 	Scheme   *runtime.Scheme
 	Registry snapshot.Registry
 	Recorder record.EventRecorder
+	// Concurrency caps in-flight reconciles; at 1 one slow registry probe
+	// stalls every other CocoonSet.
+	Concurrency int
 }
 
 // SetupWithManager registers the reconciler. `For` uses GenerationChangedPredicate
 // to avoid status-update loops; Owns filters pod events to creation, deletion,
 // and readiness transitions to prevent reconcile storms from VK status churn.
 func (r *Reconciler) SetupWithManager(_ context.Context, mgr ctrl.Manager) error {
+	if r.Concurrency < 1 {
+		return fmt.Errorf("cocoonset concurrency must be at least 1, got %d", r.Concurrency)
+	}
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&cocoonv1.CocoonSet{}, builder.WithPredicates(predicate.GenerationChangedPredicate{})).
 		Owns(&corev1.Pod{}, builder.WithPredicates(podRelevantChange{})).
+		WithOptions(controller.Options{MaxConcurrentReconciles: r.Concurrency}).
 		Complete(r)
 }
 
@@ -122,8 +130,9 @@ func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Resu
 		}
 		return ctrl.Result{Requeue: true}, nil
 	}
+	intent := r.newRestoreIntent(ctx, cs.Namespace)
 	if classified.main == nil {
-		return r.createMainAgent(ctx, &cs)
+		return r.createMainAgent(ctx, &cs, intent)
 	}
 
 	// Sub-agents fork from main and need it live before creation.
@@ -135,11 +144,11 @@ func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Resu
 	mainVMName := meta.ParseVMSpec(classified.main).VMName
 	mainNodeName := classified.main.Spec.NodeName
 
-	subChanged, subRequeue, err := r.ensureSubAgents(ctx, &cs, classified, mainVMName, mainNodeName)
+	subChanged, subRequeue, err := r.ensureSubAgents(ctx, &cs, classified, mainVMName, mainNodeName, intent)
 	if err != nil {
 		return ctrl.Result{}, err
 	}
-	tbChanged, err := r.ensureToolboxes(ctx, &cs, classified)
+	tbChanged, err := r.ensureToolboxes(ctx, &cs, classified, intent)
 	if err != nil {
 		return ctrl.Result{}, err
 	}
@@ -157,18 +166,18 @@ func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Resu
 // restore-from-hibernate when the agent is hibernated so a cross-node recreate
 // restores from the :hibernate snapshot instead of booting fresh. It always
 // requeues so sub-agents fork off the now-created main.
-func (r *Reconciler) createMainAgent(ctx context.Context, cs *cocoonv1.CocoonSet) (ctrl.Result, error) {
+func (r *Reconciler) createMainAgent(ctx context.Context, cs *cocoonv1.CocoonSet, intent restoreIntent) (ctrl.Result, error) {
 	logger := log.WithFunc("cocoonset.Reconciler.createMainAgent")
 	mainPod, err := buildAgentPod(cs, 0, "", "", r.Scheme)
 	if err != nil {
 		return ctrl.Result{}, fmt.Errorf("build main agent: %w", err)
 	}
-	restorable, err := r.podsRestorableByCR(ctx, cs.Namespace)
+	restorable, err := intent()
 	if err != nil {
 		return ctrl.Result{}, err
 	}
-	_, intent := restorable[mainPod.Name]
-	if err := r.markRestoreIfHibernated(ctx, mainPod, intent); err != nil {
+	_, wantRestore := restorable[mainPod.Name]
+	if err := r.markRestoreIfHibernated(ctx, mainPod, wantRestore); err != nil {
 		return ctrl.Result{}, err
 	}
 	if err := r.Create(ctx, mainPod); err != nil {