From 8472e1bc13b0fa8aa3a67e5fcf2f13bfd7974cd3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Tyczy=C5=84ski?= <wojtekt@google.com>
Date: Wed, 26 Oct 2022 11:30:32 +0200
Subject: [PATCH] Refactor GetReads

---
 .../apiserver/pkg/storage/etcd3/store_test.go  | 18 +++++++-----------
 .../apiserver/pkg/storage/testing/utils.go     |  8 ++------
 2 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store_test.go b/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store_test.go
index 8160a94e705..45e226d0210 100644
--- a/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store_test.go
@@ -273,13 +273,12 @@ func TestListContinuation(t *testing.T) {
 		t.Fatalf("No continuation token set")
 	}
 	storagetesting.ExpectNoDiff(t, "incorrect first page", []example.Pod{*preset[0].storedObj}, out.Items)
-	if reads := transformer.GetReads(); reads != 1 {
+	if reads := transformer.GetReadsAndReset(); reads != 1 {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	if recorder.reads != 1 {
 		t.Errorf("unexpected reads: %d", recorder.reads)
 	}
-	transformer.ResetReads()
 	recorder.resetReads()
 
 	continueFromSecondItem := out.Continue
@@ -300,13 +299,12 @@ func TestListContinuation(t *testing.T) {
 	key, rv, err := storage.DecodeContinue(continueFromSecondItem, "/")
 	t.Logf("continue token was %d %s %v", rv, key, err)
 	storagetesting.ExpectNoDiff(t, "incorrect second page", []example.Pod{*preset[1].storedObj, *preset[2].storedObj}, out.Items)
-	if reads := transformer.GetReads(); reads != 2 {
+	if reads := transformer.GetReadsAndReset(); reads != 2 {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	if recorder.reads != 1 {
 		t.Errorf("unexpected reads: %d", recorder.reads)
 	}
-	transformer.ResetReads()
 	recorder.resetReads()
 
 	// limit, should get two more pages
@@ -323,13 +321,12 @@ func TestListContinuation(t *testing.T) {
 		t.Fatalf("No continuation token set")
 	}
 	storagetesting.ExpectNoDiff(t, "incorrect second page", []example.Pod{*preset[1].storedObj}, out.Items)
-	if reads := transformer.GetReads(); reads != 1 {
+	if reads := transformer.GetReadsAndReset(); reads != 1 {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	if recorder.reads != 1 {
 		t.Errorf("unexpected reads: %d", recorder.reads)
 	}
-	transformer.ResetReads()
 	recorder.resetReads()
 
 	continueFromThirdItem := out.Continue
@@ -347,7 +344,7 @@ func TestListContinuation(t *testing.T) {
 		t.Fatalf("Unexpected continuation token set")
 	}
 	storagetesting.ExpectNoDiff(t, "incorrect third page", []example.Pod{*preset[2].storedObj}, out.Items)
-	if reads := transformer.GetReads(); reads != 1 {
+	if reads := transformer.GetReadsAndReset(); reads != 1 {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	if recorder.reads != 1 {
@@ -396,7 +393,7 @@ func TestListPaginationRareObject(t *testing.T) {
 	if len(out.Items) != 1 || !reflect.DeepEqual(&out.Items[0], pods[999]) {
 		t.Fatalf("Unexpected first page: %#v", out.Items)
 	}
-	if reads := transformer.GetReads(); reads != uint64(podCount) {
+	if reads := transformer.GetReadsAndReset(); reads != uint64(podCount) {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	// We expect that kube-apiserver will be increasing page sizes
@@ -491,13 +488,12 @@ func TestListContinuationWithFilter(t *testing.T) {
 		t.Errorf("No continuation token set")
 	}
 	storagetesting.ExpectNoDiff(t, "incorrect first page", []example.Pod{*preset[0].storedObj, *preset[2].storedObj}, out.Items)
-	if reads := transformer.GetReads(); reads != 3 {
+	if reads := transformer.GetReadsAndReset(); reads != 3 {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	if recorder.reads != 2 {
 		t.Errorf("unexpected reads: %d", recorder.reads)
 	}
-	transformer.ResetReads()
 	recorder.resetReads()
 
 	// the rest of the test does not make sense if the previous call failed
@@ -523,7 +519,7 @@ func TestListContinuationWithFilter(t *testing.T) {
 		t.Errorf("Unexpected continuation token set")
 	}
 	storagetesting.ExpectNoDiff(t, "incorrect second page", []example.Pod{*preset[3].storedObj}, out.Items)
-	if reads := transformer.GetReads(); reads != 1 {
+	if reads := transformer.GetReadsAndReset(); reads != 1 {
 		t.Errorf("unexpected reads: %d", reads)
 	}
 	if recorder.reads != 1 {
diff --git a/staging/src/k8s.io/apiserver/pkg/storage/testing/utils.go b/staging/src/k8s.io/apiserver/pkg/storage/testing/utils.go
index 42db0cbffe5..c6bca507b8f 100644
--- a/staging/src/k8s.io/apiserver/pkg/storage/testing/utils.go
+++ b/staging/src/k8s.io/apiserver/pkg/storage/testing/utils.go
@@ -232,10 +232,6 @@ func (p *PrefixTransformer) TransformToStorage(ctx context.Context, data []byte,
 	return data, p.err
 }
 
-func (p *PrefixTransformer) GetReads() uint64 {
-	return atomic.LoadUint64(&p.reads)
-}
-
-func (p *PrefixTransformer) ResetReads() {
-	atomic.StoreUint64(&p.reads, 0)
+func (p *PrefixTransformer) GetReadsAndReset() uint64 {
+	return atomic.SwapUint64(&p.reads, 0)
 }