From 7cca3d65360b8bbe8d3dbf2fc140d59c22fff0e4 Mon Sep 17 00:00:00 2001
From: Peter Schuurman <psch@google.com>
Date: Fri, 20 May 2022 17:56:46 -0700
Subject: [PATCH] Add implicit ListVolumesResponse#Entry pagination limit

---
 pkg/gce-pd-csi-driver/controller.go      |  9 ++-
 pkg/gce-pd-csi-driver/controller_test.go | 76 +++++++++++++++++++++++-
 2 files changed, 81 insertions(+), 4 deletions(-)

diff --git a/pkg/gce-pd-csi-driver/controller.go b/pkg/gce-pd-csi-driver/controller.go
index c65a9cec2..92941a5d5 100644
--- a/pkg/gce-pd-csi-driver/controller.go
+++ b/pkg/gce-pd-csi-driver/controller.go
@@ -81,6 +81,13 @@ const (
 
 	replicationTypeNone       = "none"
 	replicationTypeRegionalPD = "regional-pd"
+
+	// The maximum number of entries that we can include in the
+	// ListVolumesResposne
+	// In reality, the limit here is 4MB (based on gRPC client response limits),
+	// but 500 is a good proxy (gives ~8KB of data per ListVolumesResponse#Entry)
+	// See https://github.com/grpc/grpc/blob/master/include/grpc/impl/codegen/grpc_types.h#L503)
+	maxListVolumesResponseEntries = 500
 )
 
 func isDiskReady(disk *gce.CloudDisk) (bool, error) {
@@ -708,7 +715,7 @@ func (gceCS *GCEControllerServer) ListVolumes(ctx context.Context, req *csi.List
 
 	var maxEntries int = int(req.MaxEntries)
 	if maxEntries == 0 {
-		maxEntries = len(gceCS.disks)
+		maxEntries = maxListVolumesResponseEntries
 	}
 
 	entries := []*csi.ListVolumesResponse_Entry{}
diff --git a/pkg/gce-pd-csi-driver/controller_test.go b/pkg/gce-pd-csi-driver/controller_test.go
index 649e32707..d92d390b8 100644
--- a/pkg/gce-pd-csi-driver/controller_test.go
+++ b/pkg/gce-pd-csi-driver/controller_test.go
@@ -15,6 +15,7 @@ limitations under the License.
 package gceGCEDriver
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"reflect"
@@ -25,8 +26,6 @@ import (
 	"github.com/GoogleCloudPlatform/k8s-cloud-provider/pkg/cloud/meta"
 	"github.com/golang/protobuf/ptypes"
 
-	"context"
-
 	compute "google.golang.org/api/compute/v1"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -840,8 +839,79 @@ func TestCreateVolumeArguments(t *testing.T) {
 	}
 }
 
+func TestListVolumePagination(t *testing.T) {
+	testCases := []struct {
+		name            string
+		diskCount       int
+		maxEntries      int32
+		expectedEntries []int
+	}{
+		{
+			name:            "no pagination (implicit)",
+			diskCount:       325,
+			expectedEntries: []int{325},
+		},
+		{
+			name:            "no pagination (explicit)",
+			diskCount:       2500,
+			maxEntries:      2500,
+			expectedEntries: []int{2500},
+		},
+		{
+			name:            "pagination (implicit)",
+			diskCount:       1327,
+			expectedEntries: []int{500, 500, 327},
+		},
+		{
+			name:            "pagination (explicit)",
+			diskCount:       723,
+			maxEntries:      200,
+			expectedEntries: []int{200, 200, 200, 123},
+		},
+		{
+			name:            "pagination (explicit)",
+			diskCount:       3253,
+			maxEntries:      1000,
+			expectedEntries: []int{1000, 1000, 1000, 253},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Setup new driver each time so no interference
+			var d []*gce.CloudDisk
+			for i := 0; i < tc.diskCount; i++ {
+				// Create diskCount dummy disks
+				d = append(d, gce.CloudDiskFromV1(&compute.Disk{Name: fmt.Sprintf("%v", i)}))
+			}
+			gceDriver := initGCEDriver(t, d)
+			tok := ""
+			for i, expectedEntry := range tc.expectedEntries {
+				lvr := &csi.ListVolumesRequest{
+					MaxEntries:    tc.maxEntries,
+					StartingToken: tok,
+				}
+				resp, err := gceDriver.cs.ListVolumes(context.TODO(), lvr)
+				if err != nil {
+					t.Fatalf("Got error %v", err)
+					return
+				}
+
+				if len(resp.Entries) != expectedEntry {
+					t.Fatalf("Got %v entries, expected %v on call # %d", len(resp.Entries), expectedEntry, i+1)
+				}
+
+				tok = resp.NextToken
+			}
+			if len(tok) != 0 {
+				t.Fatalf("Expected no more entries, but got NextToken in response: %s", tok)
+			}
+		})
+	}
+}
+
 func TestListVolumeArgs(t *testing.T) {
-	diskCount := 600
+	diskCount := 500
 	testCases := []struct {
 		name            string
 		maxEntries      int32