diff --git a/pkg/registry/generic/registry/dryrun.go b/pkg/registry/generic/registry/dryrun.go index ecf545274..dd6a2c78e 100644 --- a/pkg/registry/generic/registry/dryrun.go +++ b/pkg/registry/generic/registry/dryrun.go @@ -107,8 +107,8 @@ func (s *DryRunnableStorage) GuaranteedUpdate( return s.Storage.GuaranteedUpdate(ctx, key, destination, ignoreNotFound, preconditions, tryUpdate, cachedExistingObject) } -func (s *DryRunnableStorage) Count(key string) (int64, error) { - return s.Storage.Count(key) +func (s *DryRunnableStorage) Count(ctx context.Context, key string) (int64, error) { + return s.Storage.Count(ctx, key) } func (s *DryRunnableStorage) copyInto(in, out runtime.Object) error { diff --git a/pkg/registry/generic/registry/store.go b/pkg/registry/generic/registry/store.go index beb1e19ba..e516cb7d8 100644 --- a/pkg/registry/generic/registry/store.go +++ b/pkg/registry/generic/registry/store.go @@ -1658,12 +1658,13 @@ func (e *Store) CompleteWithOptions(options *generic.StoreOptions) error { // startObservingCount starts monitoring given prefix and periodically updating metrics. It returns a function to stop collection. func (e *Store) startObservingCount(period time.Duration, objectCountTracker flowcontrolrequest.StorageObjectCountTracker) func() { - prefix := e.KeyRootFunc(genericapirequest.NewContext()) + ctx := genericapirequest.NewContext() + prefix := e.KeyRootFunc(ctx) resourceName := e.DefaultQualifiedResource.String() klog.V(2).InfoS("Monitoring resource count at path", "resource", resourceName, "path", "/"+prefix) stopCh := make(chan struct{}) go wait.JitterUntil(func() { - count, err := e.Storage.Count(prefix) + count, err := e.Storage.Count(ctx, prefix) if err != nil { klog.V(5).InfoS("Failed to update storage count metric", "err", err) count = -1 diff --git a/pkg/storage/cacher/cacher_whitebox_test.go b/pkg/storage/cacher/cacher_whitebox_test.go index c41f9df24..ada250236 100644 --- a/pkg/storage/cacher/cacher_whitebox_test.go +++ b/pkg/storage/cacher/cacher_whitebox_test.go @@ -189,7 +189,7 @@ func (d *dummyStorage) GetList(ctx context.Context, resPrefix string, opts stora func (d *dummyStorage) GuaranteedUpdate(_ context.Context, _ string, _ runtime.Object, _ bool, _ *storage.Preconditions, _ storage.UpdateFunc, _ runtime.Object) error { return fmt.Errorf("unimplemented") } -func (d *dummyStorage) Count(_ string) (int64, error) { +func (d *dummyStorage) Count(_ context.Context, _ string) (int64, error) { return 0, fmt.Errorf("unimplemented") } func (d *dummyStorage) ReadinessCheck() error { diff --git a/pkg/storage/cacher/delegator.go b/pkg/storage/cacher/delegator.go index cb7eb3399..176a6d824 100644 --- a/pkg/storage/cacher/delegator.go +++ b/pkg/storage/cacher/delegator.go @@ -265,8 +265,8 @@ func (c *CacheDelegator) GuaranteedUpdate(ctx context.Context, key string, desti return c.storage.GuaranteedUpdate(ctx, key, destination, ignoreNotFound, preconditions, tryUpdate, nil) } -func (c *CacheDelegator) Count(pathPrefix string) (int64, error) { - return c.storage.Count(pathPrefix) +func (c *CacheDelegator) Count(ctx context.Context, pathPrefix string) (int64, error) { + return c.storage.Count(ctx, pathPrefix) } func (c *CacheDelegator) ReadinessCheck() error { diff --git a/pkg/storage/etcd3/store.go b/pkg/storage/etcd3/store.go index b2b42d932..56610b208 100644 --- a/pkg/storage/etcd3/store.go +++ b/pkg/storage/etcd3/store.go @@ -617,7 +617,7 @@ func getNewItemFunc(listObj runtime.Object, v reflect.Value) func() runtime.Obje } } -func (s *store) Count(key string) (int64, error) { +func (s *store) Count(ctx context.Context, key string) (int64, error) { preparedKey, err := s.prepareKey(key) if err != nil { return 0, err @@ -631,7 +631,7 @@ func (s *store) Count(key string) (int64, error) { } startTime := time.Now() - count, err := s.client.Kubernetes.Count(context.Background(), preparedKey, kubernetes.CountOptions{}) + count, err := s.client.Kubernetes.Count(ctx, preparedKey, kubernetes.CountOptions{}) metrics.RecordEtcdRequest("listWithCount", s.groupResource, err, startTime) if err != nil { return 0, err diff --git a/pkg/storage/etcd3/store_test.go b/pkg/storage/etcd3/store_test.go index d5b630489..dff4930a7 100644 --- a/pkg/storage/etcd3/store_test.go +++ b/pkg/storage/etcd3/store_test.go @@ -692,7 +692,7 @@ func TestInvalidKeys(t *testing.T) { expectInvalidKey("Get", store.Get(ctx, invalidKey, storage.GetOptions{}, nil)) expectInvalidKey("GetList", store.GetList(ctx, invalidKey, storage.ListOptions{}, nil)) expectInvalidKey("GuaranteedUpdate", store.GuaranteedUpdate(ctx, invalidKey, nil, true, nil, nil, nil)) - _, countErr := store.Count(invalidKey) + _, countErr := store.Count(t.Context(), invalidKey) expectInvalidKey("Count", countErr) } diff --git a/pkg/storage/interfaces.go b/pkg/storage/interfaces.go index 0befb3fc6..d1f7ce31c 100644 --- a/pkg/storage/interfaces.go +++ b/pkg/storage/interfaces.go @@ -244,7 +244,7 @@ type Interface interface { preconditions *Preconditions, tryUpdate UpdateFunc, cachedExistingObject runtime.Object) error // Count returns number of different entries under the key (generally being path prefix). - Count(key string) (int64, error) + Count(ctx context.Context, key string) (int64, error) // ReadinessCheck checks if the storage is ready for accepting requests. ReadinessCheck() error diff --git a/pkg/storage/testing/store_tests.go b/pkg/storage/testing/store_tests.go index acbafd9bf..3eccde4a5 100644 --- a/pkg/storage/testing/store_tests.go +++ b/pkg/storage/testing/store_tests.go @@ -3130,7 +3130,7 @@ func RunTestCount(ctx context.Context, t *testing.T, store storage.Interface) { } } - resourceACountGot, err := store.Count(resourceA) + resourceACountGot, err := store.Count(t.Context(), resourceA) if err != nil { t.Fatalf("store.Count failed: %v", err) }