Add iteration support to etcd (#52199)

Introduces a new
`Items(context.Context, backend.IterateParams) iter.Seq2[backend.Item, error]`
to the etcd backend. The core logic for this function was moved from
GetRange, and GetRange now defers to calling Items and collecting a slice
of items.

The motivation for this change is to attempt to reduce the
complexity of pagination for consumers of GetRange. While
backend.IterateRange and backend.StreamRange do already exists to
serve the same purpose, they too have suffered from pagination bugs
in the past. Additionally, this aims to unify iteration to using
the iter package instead of the various homegrown iteration
mechanisms that exist throughout the code base.

There is also one distinct difference to the iteration api: it
allows retrieving items within the specified range in ascending
or descending key order. While there may not be many use cases for
this today, exposing this in the api from it's inception is easier
than adding functional options, or ItemsAscending/ItemsDescending
later.

Author: rosstimothy

lib/backend/backend.go

@@ -23,6 +23,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"iter"
 	"sort"
 	"time"
 
@@ -113,6 +114,34 @@ type Backend interface {
 	CloseWatchers()
 }
 
+// IterateParams are parameters that are provided to
+// [BackendWithItems.Items] to alter the iteration behavior.
+type IterateParams struct {
+	// StartKey is the minimum key in the range yielded by the iteration. This key
+	// will be included in the results if it exists.
+	StartKey Key
+	// EndKey is the maximum key in the range yielded by the iteration. This key
+	// will be included in the results if it exists.
+	EndKey Key
+	// Descending makes the iteration yield items from the biggest to the smallest
+	// key (i.e. from EndKey to StartKey). If unset, the iteration will proceed in the
+	// usual ascending order (i.e. from StartKey to EndKey).
+	Descending bool
+	// Limit is an optional maximum number of items to retrieve during iteration.
+	Limit int
+}
+
+// BackendWithItems is a temporary interface that will be added to [backend.Backend]
+// once all concrete backend implementations satisfy the new interface.
+// TODO(tross): REMEMBER TO DELETE THIS
+type BackendWithItems interface {
+	Backend
+
+	// Items produces an iterator of backend items in the range, and order
+	// described in the provided [IterateParams].
+	Items(ctx context.Context, params IterateParams) iter.Seq2[Item, error]
+}
+
 // New initializes a new [Backend] implementation based on the service config.
 func New(ctx context.Context, backend string, params Params) (Backend, error) {
 	registryMu.RLock()

lib/backend/etcdbk/etcd.go

@@ -25,9 +25,9 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 	"errors"
+	"iter"
 	"log/slog"
 	"os"
-	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -649,6 +649,84 @@ func (b *EtcdBackend) NewWatcher(ctx context.Context, watch backend.Watch) (back
 	return b.buf.NewWatcher(ctx, watch)
 }
 
+func (b *EtcdBackend) Items(ctx context.Context, params backend.IterateParams) iter.Seq2[backend.Item, error] {
+	if params.StartKey.IsZero() {
+		err := trace.BadParameter("missing parameter startKey")
+		return func(yield func(backend.Item, error) bool) { yield(backend.Item{}, err) }
+	}
+	if params.EndKey.IsZero() {
+		err := trace.BadParameter("missing parameter endKey")
+		return func(yield func(backend.Item, error) bool) { yield(backend.Item{}, err) }
+	}
+
+	sort := clientv3.SortAscend
+	if params.Descending {
+		sort = clientv3.SortDescend
+	}
+
+	const defaultPageSize = 1000
+	return func(yield func(backend.Item, error) bool) {
+		inclusiveStartKey := b.prependPrefix(params.StartKey)
+		// etcd's range query includes the start point and excludes the end point,
+		// but Backend.GetRange is supposed to be inclusive at both ends, so we
+		// query until the very next key in lexicographic order (i.e., the same key
+		// followed by a 0 byte)
+		endKey := b.prependPrefix(params.EndKey) + "\x00"
+		count := 0
+
+		pageSize := defaultPageSize
+		for {
+			if params.Limit > backend.NoLimit {
+				pageSize = min(params.Limit-count, defaultPageSize)
+			}
+
+			start := b.clock.Now()
+			re, err := b.clients.Next().Get(ctx, inclusiveStartKey,
+				clientv3.WithRange(endKey),
+				clientv3.WithSort(clientv3.SortByKey, sort),
+				clientv3.WithLimit(int64(pageSize)),
+			)
+			batchReadLatencies.Observe(time.Since(start).Seconds())
+			batchReadRequests.Inc()
+			if err := convertErr(err); err != nil {
+				yield(backend.Item{}, trace.Wrap(err))
+				return
+			}
+
+			if len(re.Kvs) == 0 {
+				return
+			}
+
+			for _, kv := range re.Kvs {
+				value, err := unmarshal(kv.Value)
+				if err != nil {
+					yield(backend.Item{}, trace.Wrap(err))
+					return
+				}
+
+				if !yield(backend.Item{
+					Key:      b.trimPrefix(kv.Key),
+					Value:    value,
+					Revision: toBackendRevision(kv.ModRevision),
+				}, nil) {
+					return
+				}
+
+				count++
+				if params.Limit != backend.NoLimit && count >= params.Limit {
+					return
+				}
+			}
+
+			if params.Descending {
+				endKey = string(re.Kvs[len(re.Kvs)-1].Key)
+			} else {
+				inclusiveStartKey = string(re.Kvs[len(re.Kvs)-1].Key) + "\x00"
+			}
+		}
+	}
+}
+
 // GetRange returns query range
 func (b *EtcdBackend) GetRange(ctx context.Context, startKey, endKey backend.Key, limit int) (*backend.GetResult, error) {
 	if startKey.IsZero() {
@@ -657,35 +735,18 @@ func (b *EtcdBackend) GetRange(ctx context.Context, startKey, endKey backend.Key
 	if endKey.IsZero() {
 		return nil, trace.BadParameter("missing parameter endKey")
 	}
-	// etcd's range query includes the start point and excludes the end point,
-	// but Backend.GetRange is supposed to be inclusive at both ends, so we
-	// query until the very next key in lexicographic order (i.e., the same key
-	// followed by a 0 byte)
-	opts := []clientv3.OpOption{clientv3.WithRange(b.prependPrefix(endKey) + "\x00")}
-	if limit > 0 {
-		opts = append(opts, clientv3.WithLimit(int64(limit)))
-	}
-	start := b.clock.Now()
-	re, err := b.clients.Next().Get(ctx, b.prependPrefix(startKey), opts...)
-	batchReadLatencies.Observe(time.Since(start).Seconds())
-	batchReadRequests.Inc()
-	if err := convertErr(err); err != nil {
-		return nil, trace.Wrap(err)
-	}
-	items := make([]backend.Item, 0, len(re.Kvs))
-	for _, kv := range re.Kvs {
-		value, err := unmarshal(kv.Value)
+
+	var result backend.GetResult
+	for item, err := range b.Items(ctx, backend.IterateParams{StartKey: startKey, EndKey: endKey, Limit: limit}) {
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
-		items = append(items, backend.Item{
-			Key:      b.trimPrefix(kv.Key),
-			Value:    value,
-			Revision: toBackendRevision(kv.ModRevision),
-		})
+		result.Items = append(result.Items, item)
+		if limit != backend.NoLimit && len(result.Items) > limit {
+			return nil, trace.BadParameter("item iterator produced more items than requested (this is a bug). limit=%d, recevied=%d", limit, len(result.Items))
+		}
 	}
-	sort.Sort(backend.Items(items))
-	return &backend.GetResult{Items: items}, nil
+	return &result, nil
 }
 
 func toBackendRevision(rev int64) string {