Add ability to run Reconciler steps in Parallel (#66243)

lib/services/reconciler.go

@@ -21,11 +21,13 @@ package services
 import (
 	"context"
 	"log/slog"
+	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/gravitational/trace"
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/types"
@@ -75,6 +77,13 @@ type GenericReconcilerConfig[K comparable, T any] struct {
 	// disallowed to enforce segregation between of resources from different
 	// sources.
 	AllowOriginChanges bool
+	// Concurrency sets the number of goroutines used to process resources
+	// during reconciliation. When set to 0 or 1, resources are processed
+	// sequentially. When set to a value greater than 1, resources are
+	// processed concurrently using up to that many goroutines.
+	// The OnCreate, OnUpdate, OnDelete, Matcher, and CompareResources
+	// callbacks must be safe for concurrent use when Concurrency > 1.
+	Concurrency int
 }
 
 // CheckAndSetDefaults validates the reconciler configuration and sets defaults.
@@ -103,6 +112,9 @@ func (c *GenericReconcilerConfig[K, T]) CheckAndSetDefaults() error {
 	if c.Logger == nil {
 		c.Logger = slog.With(teleport.ComponentKey, "reconciler")
 	}
+	if c.Concurrency < 1 {
+		c.Concurrency = 1
+	}
 	if c.Metrics == nil {
 		var err error
 		// If we are not given metrics, we create our own so we don't
@@ -290,23 +302,43 @@ func (r *GenericReconciler[K, T]) Reconcile(ctx context.Context) error {
 	r.logger.DebugContext(ctx, "Reconciling current resources with new resources",
 		"current_resource_count", len(currentResources), "new_resource_count", len(newResources))
 
-	var errs []error
-
 	start := time.Now()
+
+	var g errgroup.Group
+	g.SetLimit(r.cfg.Concurrency)
+
+	var (
+		mu   sync.Mutex
+		errs []error
+	)
+
 	// Process already registered resources to see if any of them were removed.
 	for key, current := range currentResources {
-		if err := r.processRegisteredResource(ctx, newResources, key, current); err != nil {
-			errs = append(errs, trace.Wrap(err))
-		}
+		g.Go(func() error {
+			if err := r.processRegisteredResource(ctx, newResources, key, current); err != nil {
+				mu.Lock()
+				errs = append(errs, trace.Wrap(err))
+				mu.Unlock()
+			}
+			return nil
+		})
 	}
 
 	// Add new resources if there are any or refresh those that were updated.
 	for key, newResource := range newResources {
-		if err := r.processNewResource(ctx, currentResources, key, newResource); err != nil {
-			errs = append(errs, trace.Wrap(err))
-		}
+		g.Go(func() error {
+			if err := r.processNewResource(ctx, currentResources, key, newResource); err != nil {
+				mu.Lock()
+				errs = append(errs, trace.Wrap(err))
+				mu.Unlock()
+			}
+			return nil
+		})
 	}
 
+	// Error are collected separately.
+	_ = g.Wait()
+
 	if r.stats.hasChanges() {
 		r.logger.InfoContext(ctx, "Reconciliation completed",
 			"kind", r.resourceKind(currentResources, newResources),

lib/services/reconciler_test.go

@@ -20,7 +20,9 @@ package services
 
 import (
 	"context"
+	"fmt"
 	"maps"
+	"sync"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -384,6 +386,90 @@ func TestGenericReconciler(t *testing.T) {
 	require.ElementsMatch(t, expectedDeleteCalls, onDeleteCalls)
 }
 
+// TestGenericReconcilerConcurrent verifies that the concurrent reconciliation
+// when the Parallel option is set.
+func TestGenericReconcilerConcurrent(t *testing.T) {
+	t.Parallel()
+
+	const n = 100
+	labels := map[string]string{"env": "prod"}
+
+	currentResources := make(map[int]testResource, n)
+	for i := 0; i < n; i++ {
+		currentResources[i] = makeDynamicResource(fmt.Sprintf("res%d", i), maps.Clone(labels))
+	}
+
+	newResources := make(map[int]testResource, 2*n)
+	for i := 0; i < n/2; i++ {
+		newResources[i] = makeDynamicResource(fmt.Sprintf("res%d", i), map[string]string{"env": "stage"})
+	}
+	for i := n; i < 2*n; i++ {
+		newResources[i] = makeDynamicResource(fmt.Sprintf("res%d", i), maps.Clone(labels))
+	}
+
+	var (
+		mu            sync.Mutex
+		onCreateCalls []testResource
+		onUpdateCalls []updateCall
+		onDeleteCalls []testResource
+	)
+
+	r, err := NewGenericReconciler(GenericReconcilerConfig[int, testResource]{
+		Matcher: func(tr testResource) bool { return true },
+		CompareResources: func(tr1, tr2 testResource) int {
+			return EqualFromBool(cmp.Equal(tr1, tr2, cmpopts.IgnoreUnexported(headerv1.Metadata{})))
+		},
+		GetCurrentResources: func() map[int]testResource { return currentResources },
+		GetNewResources:     func() map[int]testResource { return newResources },
+		OnCreate: func(ctx context.Context, tr testResource) error {
+			mu.Lock()
+			defer mu.Unlock()
+			onCreateCalls = append(onCreateCalls, tr)
+			return nil
+		},
+		OnUpdate: func(ctx context.Context, tr, old testResource) error {
+			mu.Lock()
+			defer mu.Unlock()
+			onUpdateCalls = append(onUpdateCalls, updateCall{new: tr, old: old})
+			return nil
+		},
+		OnDelete: func(ctx context.Context, tr testResource) error {
+			mu.Lock()
+			defer mu.Unlock()
+			onDeleteCalls = append(onDeleteCalls, tr)
+			return nil
+		},
+		Concurrency: 10,
+	})
+	require.NoError(t, err)
+	require.NoError(t, r.Reconcile(context.Background()))
+
+	// 100 new resources (IDs 100–199) should be created.
+	var expectedCreates []testResource
+	for i := n; i < 2*n; i++ {
+		expectedCreates = append(expectedCreates, makeDynamicResource(fmt.Sprintf("res%d", i), maps.Clone(labels)))
+	}
+	require.ElementsMatch(t, expectedCreates, onCreateCalls)
+
+	// 50 resources (IDs 0–49) should be updated.
+	var expectedUpdates []updateCall
+	for i := 0; i < n/2; i++ {
+		name := fmt.Sprintf("res%d", i)
+		expectedUpdates = append(expectedUpdates, updateCall{
+			new: makeDynamicResource(name, map[string]string{"env": "stage"}),
+			old: makeDynamicResource(name, maps.Clone(labels)),
+		})
+	}
+	require.ElementsMatch(t, expectedUpdates, onUpdateCalls)
+
+	// 50 resources (IDs 50–99) absent from new set should be deleted.
+	var expectedDeletes []testResource
+	for i := n / 2; i < n; i++ {
+		expectedDeletes = append(expectedDeletes, makeDynamicResource(fmt.Sprintf("res%d", i), maps.Clone(labels)))
+	}
+	require.ElementsMatch(t, expectedDeletes, onDeleteCalls)
+}
+
 func makeStaticResource(name string, labels map[string]string) testResource {
 	return makeResource(name, labels, map[string]string{
 		types.OriginLabel: types.OriginConfigFile,