From efef32652af0af08a0b9c9bc547a4dce4a95f9f5 Mon Sep 17 00:00:00 2001
From: Lukasz Szaszkiewicz <lukasz.szaszkiewicz@gmail.com>
Date: Mon, 17 Jun 2024 16:47:19 +0200
Subject: [PATCH] apiserver/storage: factor PrepareContinueToken to a new
 function (#125548)

* apiserver/storage/continue: intro PrepareContinueToken

PrepareContinueToken prepares optional
parameters for retrieving additional results for a paginated request.

This function sets up parameters that a client can use to fetch the remaining results
from the server if they are available.

* apiserver/storage/etcd3: refactor to use PrepareContinueToken
---
 .../k8s.io/apiserver/pkg/storage/continue.go  | 27 +++++++++++++++++++
 .../apiserver/pkg/storage/etcd3/store.go      | 24 +++--------------
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/staging/src/k8s.io/apiserver/pkg/storage/continue.go b/staging/src/k8s.io/apiserver/pkg/storage/continue.go
index fcd95af9b3d..1d7976690a8 100644
--- a/staging/src/k8s.io/apiserver/pkg/storage/continue.go
+++ b/staging/src/k8s.io/apiserver/pkg/storage/continue.go
@@ -91,3 +91,30 @@ func EncodeContinue(key, keyPrefix string, resourceVersion int64) (string, error
 	}
 	return base64.RawURLEncoding.EncodeToString(out), nil
 }
+
+// PrepareContinueToken prepares optional
+// parameters for retrieving additional results for a paginated request.
+//
+// This function sets up parameters that a client can use to fetch the remaining results
+// from the server if they are available.
+func PrepareContinueToken(keyLastItem, keyPrefix string, resourceVersion int64, itemsCount int64, hasMoreItems bool, opts ListOptions) (string, *int64, error) {
+	var remainingItemCount *int64
+	var continueValue string
+	var err error
+
+	if hasMoreItems {
+		// Instruct the client to begin querying from immediately after the last key.
+		continueValue, err = EncodeContinue(keyLastItem+"\x00", keyPrefix, resourceVersion)
+		if err != nil {
+			return "", remainingItemCount, err
+		}
+		// Etcd response counts in objects that do not match the pred.
+		// Instead of returning inaccurate count for non-empty selectors, we return nil.
+		// We only set remainingItemCount if the predicate is empty.
+		if opts.Predicate.Empty() {
+			remainingItems := itemsCount - opts.Predicate.Limit
+			remainingItemCount = &remainingItems
+		}
+	}
+	return continueValue, remainingItemCount, err
+}
diff --git a/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store.go b/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store.go
index 6b1d2d2cab4..efd91dc6995 100644
--- a/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store.go
+++ b/staging/src/k8s.io/apiserver/pkg/storage/etcd3/store.go
@@ -811,27 +811,11 @@ func (s *store) GetList(ctx context.Context, key string, opts storage.ListOption
 		v.Set(reflect.MakeSlice(v.Type(), 0, 0))
 	}
 
-	// instruct the client to begin querying from immediately after the last key we returned
-	// we never return a key that the client wouldn't be allowed to see
-	if hasMore {
-		// we want to start immediately after the last key
-		next, err := storage.EncodeContinue(string(lastKey)+"\x00", keyPrefix, withRev)
-		if err != nil {
-			return err
-		}
-		var remainingItemCount *int64
-		// getResp.Count counts in objects that do not match the pred.
-		// Instead of returning inaccurate count for non-empty selectors, we return nil.
-		// Only set remainingItemCount if the predicate is empty.
-		if opts.Predicate.Empty() {
-			c := int64(getResp.Count - opts.Predicate.Limit)
-			remainingItemCount = &c
-		}
-		return s.versioner.UpdateList(listObj, uint64(withRev), next, remainingItemCount)
+	continueValue, remainingItemCount, err := storage.PrepareContinueToken(string(lastKey), keyPrefix, withRev, getResp.Count, hasMore, opts)
+	if err != nil {
+		return err
 	}
-
-	// no continuation
-	return s.versioner.UpdateList(listObj, uint64(withRev), "", nil)
+	return s.versioner.UpdateList(listObj, uint64(withRev), continueValue, remainingItemCount)
 }
 
 // growSlice takes a slice value and grows its capacity up