Cleanup enabling resource size estimate

staging/src/k8s.io/apiserver/pkg/storage/cacher/cacher.go

@@ -441,6 +441,13 @@ func NewCacherFromConfig(config Config) (*Cacher, error) {
 	cacher.watchCache = watchCache
 	cacher.reflector = reflector
 
+	if utilfeature.DefaultFeatureGate.Enabled(features.SizeBasedListCostEstimate) {
+		err := config.Storage.EnableResourceSizeEstimation(cacher.getKeys)
+		if err != nil {
+			return nil, fmt.Errorf("failed to enable resource size estimation: %w", err)
+		}
+	}
+
 	if utilfeature.DefaultFeatureGate.Enabled(features.ListFromCacheSnapshot) {
 		cacher.compactor = newCompactor(config.Storage, watchCache, config.Clock)
 		go cacher.compactor.Run(stopCh)
@@ -461,7 +468,6 @@ func NewCacherFromConfig(config Config) (*Cacher, error) {
 			}, time.Second, stopCh,
 		)
 	}()
-	config.Storage.SetKeysFunc(cacher.getKeys)
 	return cacher, nil
 }
 

staging/src/k8s.io/apiserver/pkg/storage/cacher/cacher_whitebox_test.go

@@ -196,7 +196,8 @@ func (d *dummyStorage) GuaranteedUpdate(_ context.Context, _ string, _ runtime.O
 func (d *dummyStorage) Stats(_ context.Context) (storage.Stats, error) {
 	return storage.Stats{}, fmt.Errorf("unimplemented")
 }
-func (d *dummyStorage) SetKeysFunc(storage.KeysFunc) {
+func (d *dummyStorage) EnableResourceSizeEstimation(storage.KeysFunc) error {
+	return nil
 }
 func (d *dummyStorage) ReadinessCheck() error {
 	return nil
@@ -3154,6 +3155,7 @@ func TestGetBookmarkAfterResourceVersionLockedFunc(t *testing.T) {
 }
 
 func TestWatchStreamSeparation(t *testing.T) {
+	featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.SizeBasedListCostEstimate, false)
 	server, etcdStorage := newEtcdTestStorage(t, etcd3testing.PathPrefix())
 	t.Cleanup(func() {
 		server.Terminate(t)

staging/src/k8s.io/apiserver/pkg/storage/cacher/delegator.go

@@ -99,8 +99,8 @@ func (c *CacheDelegator) GetCurrentResourceVersion(ctx context.Context) (uint64,
 	return c.storage.GetCurrentResourceVersion(ctx)
 }
 
-func (c *CacheDelegator) SetKeysFunc(keys storage.KeysFunc) {
-	c.storage.SetKeysFunc(keys)
+func (c *CacheDelegator) EnableResourceSizeEstimation(keys storage.KeysFunc) error {
+	return c.storage.EnableResourceSizeEstimation(keys)
 }
 
 func (c *CacheDelegator) Delete(ctx context.Context, key string, out runtime.Object, preconditions *storage.Preconditions, validateDeletion storage.ValidateObjectFunc, cachedExistingObject runtime.Object, opts storage.DeleteOptions) error {

staging/src/k8s.io/apiserver/pkg/storage/etcd3/stats.go

@@ -18,7 +18,6 @@ package etcd3
 
 import (
 	"context"
-	"errors"
 	"strings"
 	"sync"
 
@@ -34,11 +33,11 @@ import (
 
 const sizerRefreshInterval = time.Minute
 
-func newStatsCache(prefix string, getKeys storage.KeysFunc) *statsCache {
+func newResourceSizeEstimator(prefix string, getKeys storage.KeysFunc) *resourceSizeEstimator {
 	if prefix[len(prefix)-1] != '/' {
 		prefix += "/"
 	}
-	sc := &statsCache{
+	sc := &resourceSizeEstimator{
 		prefix:  prefix,
 		getKeys: getKeys,
 		stop:    make(chan struct{}),
@@ -52,23 +51,21 @@ func newStatsCache(prefix string, getKeys storage.KeysFunc) *statsCache {
 	return sc
 }
 
-// statsCache efficiently estimates the average object size
+// resourceSizeEstimator efficiently estimates the average object size
 // based on the last observed state of individual keys.
-// By plugging statsCache into GetList and Watch functions,
+// By plugging resourceSizeEstimator into GetList and Watch functions,
 // a fairly accurate estimate of object sizes can be maintained
 // without additional requests to the underlying storage.
 // To handle potential out-of-order or incomplete data,
 // it uses a per-key revision to identify the newer state.
 // This approach may leak keys if delete events are not observed,
 // thus we run a background goroutine to periodically cleanup keys if needed.
-type statsCache struct {
+type resourceSizeEstimator struct {
 	prefix         string
 	stop           chan struct{}
 	wg             sync.WaitGroup
 	lastKeyCleanup atomic.Pointer[time.Time]
-
-	getKeysLock sync.Mutex
-	getKeys     storage.KeysFunc
+	getKeys        storage.KeysFunc
 
 	keysLock sync.Mutex
 	keys     map[string]sizeRevision
@@ -79,10 +76,8 @@ type sizeRevision struct {
 	revision  int64
 }
 
-var errStatsDisabled = errors.New("key size stats disabled")
-
-func (sc *statsCache) Stats(ctx context.Context) (storage.Stats, error) {
-	keys, err := sc.GetKeys(ctx)
+func (sc *resourceSizeEstimator) Stats(ctx context.Context) (storage.Stats, error) {
+	keys, err := sc.getKeys(ctx)
 	if err != nil {
 		return storage.Stats{}, err
 	}
@@ -98,30 +93,12 @@ func (sc *statsCache) Stats(ctx context.Context) (storage.Stats, error) {
 	return stats, nil
 }
 
-func (sc *statsCache) GetKeys(ctx context.Context) ([]string, error) {
-	sc.getKeysLock.Lock()
-	getKeys := sc.getKeys
-	sc.getKeysLock.Unlock()
-
-	if getKeys == nil {
-		return nil, errStatsDisabled
-	}
-	// Don't execute getKeys under lock.
-	return getKeys(ctx)
-}
-
-func (sc *statsCache) SetKeysFunc(keys storage.KeysFunc) {
-	sc.getKeysLock.Lock()
-	defer sc.getKeysLock.Unlock()
-	sc.getKeys = keys
-}
-
-func (sc *statsCache) Close() {
+func (sc *resourceSizeEstimator) Close() {
 	close(sc.stop)
 	sc.wg.Wait()
 }
 
-func (sc *statsCache) run() {
+func (sc *resourceSizeEstimator) run() {
 	jitter := 0.5 // Period between [interval, interval * (1.0 + jitter)]
 	sliding := true
 	// wait.JitterUntilWithContext starts work immediately, so wait first.
@@ -132,16 +109,14 @@ func (sc *statsCache) run() {
 	wait.JitterUntilWithContext(wait.ContextForChannel(sc.stop), sc.cleanKeysIfNeeded, sizerRefreshInterval, jitter, sliding)
 }
 
-func (sc *statsCache) cleanKeysIfNeeded(ctx context.Context) {
+func (sc *resourceSizeEstimator) cleanKeysIfNeeded(ctx context.Context) {
 	lastKeyCleanup := sc.lastKeyCleanup.Load()
 	if lastKeyCleanup != nil && time.Since(*lastKeyCleanup) < sizerRefreshInterval {
 		return
 	}
-	keys, err := sc.GetKeys(ctx)
+	keys, err := sc.getKeys(ctx)
 	if err != nil {
-		if !errors.Is(err, errStatsDisabled) {
-			klog.InfoS("Error getting keys", "err", err)
-		}
+		klog.InfoS("Error getting keys", "err", err)
 		return
 	}
 	sc.keysLock.Lock()
@@ -149,7 +124,7 @@ func (sc *statsCache) cleanKeysIfNeeded(ctx context.Context) {
 	sc.cleanKeys(keys)
 }
 
-func (sc *statsCache) cleanKeys(keepKeys []string) {
+func (sc *resourceSizeEstimator) cleanKeys(keepKeys []string) {
 	newKeys := make(map[string]sizeRevision, len(keepKeys))
 	for _, key := range keepKeys {
 		// Handle cacher keys not having prefix.
@@ -171,14 +146,14 @@ func (sc *statsCache) cleanKeys(keepKeys []string) {
 	sc.lastKeyCleanup.Store(&now)
 }
 
-func (sc *statsCache) keySizes() (totalSize int64) {
+func (sc *resourceSizeEstimator) keySizes() (totalSize int64) {
 	for _, sizeRevision := range sc.keys {
 		totalSize += sizeRevision.sizeBytes
 	}
 	return totalSize
 }
 
-func (sc *statsCache) Update(kvs []*mvccpb.KeyValue) {
+func (sc *resourceSizeEstimator) Update(kvs []*mvccpb.KeyValue) {
 	sc.keysLock.Lock()
 	defer sc.keysLock.Unlock()
 	for _, kv := range kvs {
@@ -186,14 +161,14 @@ func (sc *statsCache) Update(kvs []*mvccpb.KeyValue) {
 	}
 }
 
-func (sc *statsCache) UpdateKey(kv *mvccpb.KeyValue) {
+func (sc *resourceSizeEstimator) UpdateKey(kv *mvccpb.KeyValue) {
 	sc.keysLock.Lock()
 	defer sc.keysLock.Unlock()
 
 	sc.updateKey(kv)
 }
 
-func (sc *statsCache) updateKey(kv *mvccpb.KeyValue) {
+func (sc *resourceSizeEstimator) updateKey(kv *mvccpb.KeyValue) {
 	key := string(kv.Key)
 	keySizeRevision := sc.keys[key]
 	if keySizeRevision.revision >= kv.ModRevision {
@@ -206,7 +181,7 @@ func (sc *statsCache) updateKey(kv *mvccpb.KeyValue) {
 	}
 }
 
-func (sc *statsCache) DeleteKey(kv *mvccpb.KeyValue) {
+func (sc *resourceSizeEstimator) DeleteKey(kv *mvccpb.KeyValue) {
 	sc.keysLock.Lock()
 	defer sc.keysLock.Unlock()
 

staging/src/k8s.io/apiserver/pkg/storage/etcd3/stats_test.go

@@ -27,7 +27,7 @@ import (
 
 func TestStatsCache(t *testing.T) {
 	ctx := t.Context()
-	store := newStatsCache("/prefix", func(ctx context.Context) ([]string, error) { return []string{}, nil })
+	store := newResourceSizeEstimator("/prefix", func(ctx context.Context) ([]string, error) { return []string{}, nil })
 	defer store.Close()
 
 	stats, err := store.Stats(ctx)

staging/src/k8s.io/apiserver/pkg/storage/etcd3/store.go

@@ -25,6 +25,7 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"go.etcd.io/etcd/api/v3/mvccpb"
@@ -90,8 +91,10 @@ type store struct {
 
 	resourcePrefix string
 	newListFunc    func() runtime.Object
-	stats          *statsCache
 	compactor      Compactor
+
+	collectorMux          sync.RWMutex
+	resourceSizeEstimator *resourceSizeEstimator
 }
 
 var _ storage.Interface = (*store)(nil)
@@ -189,13 +192,8 @@ func New(c *kubernetes.Client, compactor Compactor, codec runtime.Codec, newFunc
 		newListFunc:    newListFunc,
 		compactor:      compactor,
 	}
-	// Collecting stats requires properly set resourcePrefix to call getKeys.
-	if utilfeature.DefaultFeatureGate.Enabled(features.SizeBasedListCostEstimate) {
-		stats := newStatsCache(pathPrefix, nil)
-		s.stats = stats
-		w.stats = stats
-	}
 
+	w.getResourceSizeEstimator = s.getResourceSizeEstimator
 	w.getCurrentStorageRV = func(ctx context.Context) (uint64, error) {
 		return s.GetCurrentResourceVersion(ctx)
 	}
@@ -218,11 +216,18 @@ func (s *store) Versioner() storage.Versioner {
 }
 
 func (s *store) Close() {
-	if s.stats != nil {
-		s.stats.Close()
+	stats := s.getResourceSizeEstimator()
+	if stats != nil {
+		stats.Close()
 	}
 }
 
+func (s *store) getResourceSizeEstimator() *resourceSizeEstimator {
+	s.collectorMux.RLock()
+	defer s.collectorMux.RUnlock()
+	return s.resourceSizeEstimator
+}
+
 // Get implements storage.Interface.Get.
 func (s *store) Get(ctx context.Context, key string, opts storage.GetOptions, out runtime.Object) error {
 	preparedKey, err := s.prepareKey(key)
@@ -633,13 +638,12 @@ func getNewItemFunc(listObj runtime.Object, v reflect.Value) func() runtime.Obje
 	}
 }
 
-func (s *store) Stats(ctx context.Context) (stats storage.Stats, err error) {
-	if s.stats != nil {
-		stats, err := s.stats.Stats(ctx)
-		if !errors.Is(err, errStatsDisabled) {
-			return stats, err
-		}
+func (s *store) Stats(ctx context.Context) (storage.Stats, error) {
+	if collector := s.getResourceSizeEstimator(); collector != nil {
+		return collector.Stats(ctx)
 	}
+	// returning stats without resource size
+
 	startTime := time.Now()
 	prefix, err := s.prepareKey(s.resourcePrefix)
 	if err != nil {
@@ -658,10 +662,17 @@ func (s *store) Stats(ctx context.Context) (stats storage.Stats, err error) {
 	}, nil
 }
 
-func (s *store) SetKeysFunc(keys storage.KeysFunc) {
-	if s.stats != nil {
-		s.stats.SetKeysFunc(keys)
+func (s *store) EnableResourceSizeEstimation(getKeys storage.KeysFunc) error {
+	if getKeys == nil {
+		return errors.New("KeysFunc cannot be nil")
 	}
+	s.collectorMux.Lock()
+	defer s.collectorMux.Unlock()
+	if s.resourceSizeEstimator != nil {
+		return errors.New("resourceSizeEstimator already enabled")
+	}
+	s.resourceSizeEstimator = newResourceSizeEstimator(s.pathPrefix, getKeys)
+	return nil
 }
 
 func (s *store) getKeys(ctx context.Context) ([]string, error) {
@@ -780,6 +791,8 @@ func (s *store) GetList(ctx context.Context, key string, opts storage.ListOption
 		metricsOp = "list"
 	}
 
+	stats := s.getResourceSizeEstimator()
+
 	aggregator := s.listErrAggrFactory()
 	for {
 		startTime := time.Now()
@@ -821,8 +834,8 @@ func (s *store) GetList(ctx context.Context, key string, opts storage.ListOption
 		} else {
 			growSlice(v, 2048, len(getResp.Kvs))
 		}
-		if s.stats != nil {
-			s.stats.Update(getResp.Kvs)
+		if stats != nil {
+			stats.Update(getResp.Kvs)
 		}
 
 		// take items from the response until the bucket is full, filtering as we go

staging/src/k8s.io/apiserver/pkg/storage/etcd3/store_test.go

@@ -381,10 +381,13 @@ func TestListResourceVersionMatch(t *testing.T) {
 func TestStats(t *testing.T) {
 	for _, sizeBasedListCostEstimate := range []bool{true, false} {
 		t.Run(fmt.Sprintf("SizeBasedListCostEstimate=%v", sizeBasedListCostEstimate), func(t *testing.T) {
-			featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.SizeBasedListCostEstimate, sizeBasedListCostEstimate)
-			// Match transformer with cacher tests.
 			ctx, store, _ := testSetup(t)
-			store.SetKeysFunc(store.getKeys)
+			if sizeBasedListCostEstimate {
+				err := store.EnableResourceSizeEstimation(store.getKeys)
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
 			storagetesting.RunTestStats(ctx, t, store, store.codec, store.transformer, sizeBasedListCostEstimate)
 		})
 	}
@@ -991,8 +994,8 @@ func BenchmarkStatsCacheCleanKeys(b *testing.B) {
 	if err != nil {
 		b.Fatal(err)
 	}
-	if len(store.stats.keys) < namespaceCount*podPerNamespaceCount {
-		b.Fatalf("Unexpected number of keys in stats, want: %d, got: %d", namespaceCount*podPerNamespaceCount, len(store.stats.keys))
+	if len(store.resourceSizeEstimator.keys) < namespaceCount*podPerNamespaceCount {
+		b.Fatalf("Unexpected number of keys in stats, want: %d, got: %d", namespaceCount*podPerNamespaceCount, len(store.resourceSizeEstimator.keys))
 	}
 	// Get keys to measure only cleanupKeys time
 	keys, err := store.getKeys(ctx)
@@ -1001,10 +1004,10 @@ func BenchmarkStatsCacheCleanKeys(b *testing.B) {
 	}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		store.stats.cleanKeys(keys)
+		store.resourceSizeEstimator.cleanKeys(keys)
 	}
-	if len(store.stats.keys) < namespaceCount*podPerNamespaceCount {
-		b.Fatalf("Unexpected number of keys in stats, want: %d, got: %d", namespaceCount*podPerNamespaceCount, len(store.stats.keys))
+	if len(store.resourceSizeEstimator.keys) < namespaceCount*podPerNamespaceCount {
+		b.Fatalf("Unexpected number of keys in stats, want: %d, got: %d", namespaceCount*podPerNamespaceCount, len(store.resourceSizeEstimator.keys))
 	}
 }
 
@@ -1045,40 +1048,27 @@ func TestPrefixStats(t *testing.T) {
 	tcs := []struct {
 		name        string
 		estimate    bool
-		setKeys     bool
 		expectStats storage.Stats
 	}{
 		{
-			name:        "SizeBasedListCostEstimate=false,SetKeys=false",
-			setKeys:     false,
+			name:        "Estimate=false",
 			estimate:    false,
 			expectStats: storage.Stats{ObjectCount: 1},
 		},
 		{
-			name:        "SizeBasedListCostEstimate=false,SetKeys=true",
-			setKeys:     true,
-			estimate:    false,
-			expectStats: storage.Stats{ObjectCount: 1},
-		},
-		{
-			name:        "SizeBasedListCostEstimate=true,SetKeys=false",
-			setKeys:     false,
-			estimate:    true,
-			expectStats: storage.Stats{ObjectCount: 1},
-		},
-		{
-			name:        "SizeBasedListCostEstimate=true,SetKeys=true",
-			setKeys:     true,
+			name:        "Estimate=true",
 			estimate:    true,
 			expectStats: storage.Stats{ObjectCount: 1, EstimatedAverageObjectSizeBytes: 3},
 		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.name, func(t *testing.T) {
-			featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.SizeBasedListCostEstimate, tc.estimate)
 			ctx, store, c := testSetup(t, withPrefix("/registry"), withResourcePrefix("pods"))
-			if tc.setKeys {
-				store.SetKeysFunc(store.getKeys)
+			if tc.estimate {
+				err := store.EnableResourceSizeEstimation(store.getKeys)
+				if err != nil {
+					t.Fatal(err)
+				}
 			}
 			_, err := c.KV.Put(ctx, "key", "a")
 			if err != nil {

staging/src/k8s.io/apiserver/pkg/storage/etcd3/watcher.go

@@ -69,30 +69,30 @@ func TestOnlySetFatalOnDecodeError(b bool) {
 }
 
 type watcher struct {
-	client              *clientv3.Client
-	codec               runtime.Codec
-	newFunc             func() runtime.Object
-	objectType          string
-	groupResource       schema.GroupResource
-	versioner           storage.Versioner
-	transformer         value.Transformer
-	getCurrentStorageRV func(context.Context) (uint64, error)
-	stats               *statsCache
+	client                   *clientv3.Client
+	codec                    runtime.Codec
+	newFunc                  func() runtime.Object
+	objectType               string
+	groupResource            schema.GroupResource
+	versioner                storage.Versioner
+	transformer              value.Transformer
+	getCurrentStorageRV      func(context.Context) (uint64, error)
+	getResourceSizeEstimator func() *resourceSizeEstimator
 }
 
 // watchChan implements watch.Interface.
 type watchChan struct {
-	watcher           *watcher
-	key               string
-	initialRev        int64
-	recursive         bool
-	progressNotify    bool
-	internalPred      storage.SelectionPredicate
-	ctx               context.Context
-	cancel            context.CancelFunc
-	incomingEventChan chan *event
-	resultChan        chan watch.Event
-	stats             *statsCache
+	watcher                  *watcher
+	key                      string
+	initialRev               int64
+	recursive                bool
+	progressNotify           bool
+	internalPred             storage.SelectionPredicate
+	ctx                      context.Context
+	cancel                   context.CancelFunc
+	incomingEventChan        chan *event
+	resultChan               chan watch.Event
+	getResourceSizeEstimator func() *resourceSizeEstimator
 }
 
 // Watch watches on a key and returns a watch.Interface that transfers relevant notifications.
@@ -128,15 +128,15 @@ func (w *watcher) Watch(ctx context.Context, key string, rev int64, opts storage
 
 func (w *watcher) createWatchChan(ctx context.Context, key string, rev int64, recursive, progressNotify bool, pred storage.SelectionPredicate) *watchChan {
 	wc := &watchChan{
-		watcher:           w,
-		key:               key,
-		initialRev:        rev,
-		recursive:         recursive,
-		progressNotify:    progressNotify,
-		internalPred:      pred,
-		incomingEventChan: make(chan *event, incomingBufSize),
-		resultChan:        make(chan watch.Event, outgoingBufSize),
-		stats:             w.stats,
+		watcher:                  w,
+		key:                      key,
+		initialRev:               rev,
+		recursive:                recursive,
+		progressNotify:           progressNotify,
+		internalPred:             pred,
+		incomingEventChan:        make(chan *event, incomingBufSize),
+		resultChan:               make(chan watch.Event, outgoingBufSize),
+		getResourceSizeEstimator: w.getResourceSizeEstimator,
 	}
 	if pred.Empty() {
 		// The filter doesn't filter out any object.
@@ -385,6 +385,7 @@ func (wc *watchChan) startWatching(watchClosedCh chan struct{}, initialEventsEnd
 		opts = append(opts, clientv3.WithProgressNotify())
 	}
 	wch := wc.watcher.client.Watch(wc.ctx, wc.key, opts...)
+	estimator := wc.getResourceSizeEstimator()
 	for wres := range wch {
 		if wres.Err() != nil {
 			err := wres.Err()
@@ -405,12 +406,12 @@ func (wc *watchChan) startWatching(watchClosedCh chan struct{}, initialEventsEnd
 		}
 
 		for _, e := range wres.Events {
-			if wc.stats != nil {
+			if estimator != nil {
 				switch e.Type {
 				case clientv3.EventTypePut:
-					wc.stats.UpdateKey(e.Kv)
+					estimator.UpdateKey(e.Kv)
 				case clientv3.EventTypeDelete:
-					wc.stats.DeleteKey(e.Kv)
+					estimator.DeleteKey(e.Kv)
 				}
 			}
 			metrics.RecordEtcdEvent(wc.watcher.groupResource)

staging/src/k8s.io/apiserver/pkg/storage/interfaces.go

@@ -268,9 +268,8 @@ type Interface interface {
 	// This method issues an empty list request and reads only the ResourceVersion from the object metadata
 	GetCurrentResourceVersion(ctx context.Context) (uint64, error)
 
-	// SetKeysFunc allows to override the function used to get keys from storage.
-	// This allows to replace default function that fetches keys from storage with one using cache.
-	SetKeysFunc(KeysFunc)
+	// EnableResourceSizeEstimation enables estimating resource size by providing function get keys from storage.
+	EnableResourceSizeEstimation(KeysFunc) error
 
 	// CompactRevision returns latest observed revision that was compacted.
 	// Without ListFromCacheSnapshot enabled only locally executed compaction will be observed.