reconciler: Add WaitUntilReconciled method

Add `Reconciler.WaitUntilReconciled` method to allow waiting for the
reconciler to catch up processing to a given revision.

Example usage:
```
wtxn := db.WriteTxn(table)
table.Insert(wtxn, &Obj{ID: 1, Status: reconciler.StatusPending()})
table.Insert(wtxn, &Obj{ID: 2, Status: reconciler.StatusPending()})
table.Insert(wtxn, &Obj{ID: 3, Status: reconciler.StatusPending()})
revToWaitFor := table.Revision(wtxn)
wtxn.Commit()

// Block until reconciler has catched up to [revToWaitFor] or [ctx]
// is cancelled.
myReconciler.WaitUntilReconciled(ctx, revToWaitFor)
```

Signed-off-by: Jussi Maki <jussi.maki@isovalent.com>

Author: joamaki

reconciler/builder.go

@@ -78,6 +78,7 @@ func Register[Obj comparable](
 		retries:              newRetries(cfg.RetryBackoffMinDuration, cfg.RetryBackoffMaxDuration, objectToKey),
 		externalPruneTrigger: make(chan struct{}, 1),
 		primaryIndexer:       idx,
+		progress:             newProgressTracker(),
 	}
 
 	params.JobGroup.Add(job.OneShot("reconcile", r.reconcileLoop))

reconciler/incremental.go

@@ -45,13 +45,13 @@ type opResult struct {
 	id       uint64 // the "pending" identifier
 }
 
-func (incr *incremental[Obj]) run(ctx context.Context, txn statedb.ReadTxn, changes iter.Seq2[statedb.Change[Obj], statedb.Revision]) []error {
+func (incr *incremental[Obj]) run(ctx context.Context, txn statedb.ReadTxn, changes iter.Seq2[statedb.Change[Obj], statedb.Revision]) (errs []error, lastRev statedb.Revision, changed bool) {
 	// Reconcile new and changed objects using either Operations
 	// or BatchOperations.
 	if incr.config.BatchOperations != nil {
-		incr.batch(ctx, txn, changes)
+		lastRev, changed = incr.batch(ctx, txn, changes)
 	} else {
-		incr.single(ctx, txn, changes)
+		lastRev, changed = incr.single(ctx, txn, changes)
 	}
 
 	// Process objects that need to be retried that were not cleared.
@@ -63,19 +63,22 @@ func (incr *incremental[Obj]) run(ctx context.Context, txn statedb.ReadTxn, chan
 	// Since all failures are retried, we can return the errors from the retry
 	// queue which includes both errors occurred in this round and the old
 	// errors.
-	errs := incr.retries.errors()
+	errs = incr.retries.errors()
 	incr.metrics.ReconciliationErrors(incr.moduleID, newErrors, len(errs))
 
 	// Prepare for next round
 	incr.numReconciled = 0
 	clear(incr.results)
 
-	return errs
+	return errs, lastRev, changed
 }
 
-func (incr *incremental[Obj]) single(ctx context.Context, txn statedb.ReadTxn, changes iter.Seq2[statedb.Change[Obj], statedb.Revision]) {
+func (incr *incremental[Obj]) single(ctx context.Context, txn statedb.ReadTxn, changes iter.Seq2[statedb.Change[Obj], statedb.Revision]) (lastRev statedb.Revision, changed bool) {
 	// Iterate in revision order through new and changed objects.
 	for change, rev := range changes {
+		lastRev = rev
+		changed = true
+
 		obj := change.Object
 
 		status := incr.config.GetObjectStatus(obj)
@@ -95,14 +98,19 @@ func (incr *incremental[Obj]) single(ctx context.Context, txn statedb.ReadTxn, c
 			break
 		}
 	}
+
+	return
 }
 
-func (incr *incremental[Obj]) batch(ctx context.Context, txn statedb.ReadTxn, changes iter.Seq2[statedb.Change[Obj], statedb.Revision]) {
+func (incr *incremental[Obj]) batch(ctx context.Context, txn statedb.ReadTxn, changes iter.Seq2[statedb.Change[Obj], statedb.Revision]) (lastRev statedb.Revision, changed bool) {
 	ops := incr.config.BatchOperations
 	updateBatch := []BatchEntry[Obj]{}
 	deleteBatch := []BatchEntry[Obj]{}
 
 	for change, rev := range changes {
+		lastRev = rev
+		changed = true
+
 		obj := change.Object
 
 		status := incr.config.GetObjectStatus(obj)
@@ -167,6 +175,8 @@ func (incr *incremental[Obj]) batch(ctx context.Context, txn statedb.ReadTxn, ch
 			incr.results[entry.Object] = opResult{rev: entry.Revision, id: status.ID, err: entry.Result, original: entry.original}
 		}
 	}
+
+	return
 }
 
 func (incr *incremental[Obj]) processRetries(ctx context.Context, txn statedb.ReadTxn) {

reconciler/progress.go

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Cilium
+
+package reconciler
+
+import (
+	"context"
+	"sync"
+
+	"github.com/cilium/statedb"
+)
+
+// progressTracker tracks the highest revision observed as reconciled and
+// allows callers to wait until a target revision is reached.
+type progressTracker struct {
+	mu       sync.Mutex
+	revision statedb.Revision
+	watch    chan struct{}
+}
+
+func newProgressTracker() *progressTracker {
+	return &progressTracker{
+		watch: make(chan struct{}),
+	}
+}
+
+func (p *progressTracker) update(rev statedb.Revision) {
+	p.mu.Lock()
+	if rev > p.revision {
+		p.revision = rev
+		close(p.watch)
+		p.watch = make(chan struct{})
+	}
+	p.mu.Unlock()
+}
+
+func (p *progressTracker) wait(ctx context.Context, rev statedb.Revision) (statedb.Revision, error) {
+	for {
+		p.mu.Lock()
+		current := p.revision
+		watch := p.watch
+		p.mu.Unlock()
+
+		if current >= rev {
+			return current, nil
+		}
+		select {
+		case <-ctx.Done():
+			return current, ctx.Err()
+		case <-watch:
+		}
+	}
+}

reconciler/progress_test.go

@@ -0,0 +1,205 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Cilium
+
+package reconciler_test
+
+import (
+	"context"
+	"fmt"
+	"iter"
+	"log/slog"
+	"sync"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/cilium/hive"
+	"github.com/cilium/hive/cell"
+	"github.com/cilium/hive/hivetest"
+	"github.com/cilium/hive/job"
+	"github.com/cilium/statedb"
+	"github.com/cilium/statedb/index"
+	"github.com/cilium/statedb/reconciler"
+	"github.com/stretchr/testify/require"
+)
+
+type waitObject struct {
+	ID     uint64
+	Status reconciler.Status
+}
+
+// TableHeader implements statedb.TableWritable.
+func (w waitObject) TableHeader() []string {
+	return []string{"ID", "Status"}
+}
+
+// TableRow implements statedb.TableWritable.
+func (w waitObject) TableRow() []string {
+	return []string{
+		fmt.Sprintf("%d", w.ID),
+		w.Status.String(),
+	}
+}
+
+var _ statedb.TableWritable = waitObject{}
+
+func (w *waitObject) Clone() *waitObject {
+	w2 := *w
+	return &w2
+}
+
+func (w *waitObject) GetStatus() reconciler.Status {
+	return w.Status
+}
+
+func (w *waitObject) SetStatus(status reconciler.Status) *waitObject {
+	w.Status = status
+	return w
+}
+
+var waitObjectIDIndex = statedb.Index[*waitObject, uint64]{
+	Name: "id",
+	FromObject: func(obj *waitObject) index.KeySet {
+		return index.NewKeySet(index.Uint64(obj.ID))
+	},
+	FromKey: index.Uint64,
+	Unique:  true,
+}
+
+type waitOps struct {
+	started     chan struct{}
+	unblock     chan struct{}
+	startedOnce sync.Once
+}
+
+func newWaitOps() *waitOps {
+	return &waitOps{
+		started: make(chan struct{}),
+		unblock: make(chan struct{}),
+	}
+}
+
+// Delete implements reconciler.Operations.
+func (*waitOps) Delete(context.Context, statedb.ReadTxn, statedb.Revision, *waitObject) error {
+	return nil
+}
+
+// Prune implements reconciler.Operations.
+func (*waitOps) Prune(context.Context, statedb.ReadTxn, iter.Seq2[*waitObject, statedb.Revision]) error {
+	return nil
+}
+
+// Update implements reconciler.Operations.
+func (w *waitOps) Update(ctx context.Context, txn statedb.ReadTxn, rev statedb.Revision, obj *waitObject) error {
+	w.startedOnce.Do(func() {
+		close(w.started)
+	})
+	select {
+	case <-w.unblock:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+var _ reconciler.Operations[*waitObject] = &waitOps{}
+
+func TestWaitUntilReconciled(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		var (
+			table statedb.RWTable[*waitObject]
+			db    *statedb.DB
+			r     reconciler.Reconciler[*waitObject]
+		)
+		ops := newWaitOps()
+
+		hive := hive.New(
+			statedb.Cell,
+			job.Cell,
+			cell.Provide(
+				cell.NewSimpleHealth,
+				reconciler.NewExpVarMetrics,
+				func(r job.Registry, h cell.Health, lc cell.Lifecycle) job.Group {
+					return r.NewGroup(h, lc)
+				},
+			),
+			cell.Invoke(func(db_ *statedb.DB) (err error) {
+				db = db_
+				table, err = statedb.NewTable(db, "wait-objects", waitObjectIDIndex)
+				return err
+			}),
+			cell.Module("test", "test",
+				cell.Invoke(func(params reconciler.Params) error {
+					var err error
+					r, err = reconciler.Register(
+						params,
+						table,
+						(*waitObject).Clone,
+						(*waitObject).SetStatus,
+						(*waitObject).GetStatus,
+						ops,
+						nil,
+						reconciler.WithoutPruning(),
+					)
+					return err
+				}),
+			),
+		)
+
+		log := hivetest.Logger(t, hivetest.LogLevel(slog.LevelError))
+		require.NoError(t, hive.Start(log, context.TODO()), "Start")
+		defer func() {
+			require.NoError(t, hive.Stop(log, context.TODO()), "Stop")
+		}()
+
+		waitCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		// Won't block if we query with 0 revision.
+		_, err := r.WaitUntilReconciled(waitCtx, 0)
+		require.NoError(t, err)
+
+		// Insert an object and wait for it to be reconciled.
+		wtxn := db.WriteTxn(table)
+		table.Insert(wtxn, &waitObject{
+			ID:     1,
+			Status: reconciler.StatusPending(),
+		})
+		revision := table.Revision(wtxn)
+		wtxn.Commit()
+
+		type waitResult struct {
+			rev statedb.Revision
+			err error
+		}
+		done := make(chan waitResult, 1)
+		go func() {
+			rev, err := r.WaitUntilReconciled(waitCtx, revision)
+			done <- waitResult{rev: rev, err: err}
+		}()
+
+		synctest.Wait()
+		select {
+		case <-ops.started:
+		default:
+			t.Fatal("expected update to start")
+		}
+
+		select {
+		case result := <-done:
+			t.Fatalf("WaitUntilReconciled returned early: %v", result.err)
+		default:
+		}
+
+		close(ops.unblock)
+		synctest.Wait()
+
+		select {
+		case result := <-done:
+			require.NoError(t, result.err)
+			require.Equal(t, revision, result.rev)
+		default:
+			t.Fatal("expected WaitUntilReconciled to complete")
+		}
+	})
+}

reconciler/reconciler.go

@@ -19,6 +19,7 @@ type reconciler[Obj comparable] struct {
 	retries              *retries
 	externalPruneTrigger chan struct{}
 	primaryIndexer       statedb.Indexer[Obj]
+	progress             *progressTracker
 }
 
 func (r *reconciler[Obj]) Prune() {
@@ -28,6 +29,10 @@ func (r *reconciler[Obj]) Prune() {
 	}
 }
 
+func (r *reconciler[Obj]) WaitUntilReconciled(ctx context.Context, untilRevision statedb.Revision) (statedb.Revision, error) {
+	return r.progress.wait(ctx, untilRevision)
+}
+
 func (r *reconciler[Obj]) reconcileLoop(ctx context.Context, health cell.Health) error {
 	var pruneTickerChan <-chan time.Time
 	if r.config.PruneInterval > 0 {
@@ -101,7 +106,10 @@ func (r *reconciler[Obj]) reconcileLoop(ctx context.Context, health cell.Health)
 
 		// Perform incremental reconciliation and retries of previously failed
 		// objects.
-		errs := incremental.run(ctx, txn, changes)
+		errs, lastRevision, changed := incremental.run(ctx, txn, changes)
+		if changed {
+			r.progress.update(lastRevision)
+		}
 
 		if tableInitialized && (prune || externalPrune) {
 			if err := r.prune(ctx, txn); err != nil {

reconciler/types.go

@@ -34,6 +34,12 @@ type Reconciler[Obj any] interface {
 	// that something has gone wrong in the reconciliation target and full
 	// reconciliation is needed to recover.
 	Prune()
+
+	// WaitUntilReconciled blocks until the reconciler has processed all
+	// table changes up to untilRevision. Returns the latest processed
+	// revision and ctx.Err() if the context is cancelled.
+	// Note: errors from Update/Delete are treated as reconciled.
+	WaitUntilReconciled(ctx context.Context, untilRevision statedb.Revision) (statedb.Revision, error)
 }
 
 // Params are the reconciler dependencies that are independent of the