Support checkLimit for multiple pages

This change implements paging when listing volumes for account limit
checking, thereby supporting limits beyond the size of a single page.

We also refactor checkLimits to not return gRPC codes but leave that to
the calling function. This also allows us to distinguish between true
errors and limit violations, and better handle each case.

driver/controller.go

@@ -176,8 +176,14 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 	}
 
 	ll.Info("checking volume limit")
-	if err := d.checkLimit(ctx); err != nil {
-		return nil, err
+	details, err := d.checkLimit(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to check volume limit: %s", err)
+	}
+	if details != nil {
+		return nil, status.Errorf(codes.ResourceExhausted,
+			"volume limit (%d) has been reached. Current number of volumes: %d. Please contact support.",
+			details.limit, details.numVolumes)
 	}
 
 	ll.WithField("volume_req", volumeReq).Info("creating volume")
@@ -982,41 +988,64 @@ func (d *Driver) waitAction(ctx context.Context, volumeId string, actionId int)
 	return err
 }
 
-// checkLimit checks whether the user hit their volume limit to ensure.
-func (d *Driver) checkLimit(ctx context.Context) error {
+type limitDetails struct {
+	limit      int
+	numVolumes int
+}
+
+// checkLimit checks whether the user hit their account volume limit.
+func (d *Driver) checkLimit(ctx context.Context) (*limitDetails, error) {
 	// only one provisioner runs, we can make sure to prevent burst creation
 	d.readyMu.Lock()
 	defer d.readyMu.Unlock()
 
 	account, _, err := d.account.Get(ctx)
 	if err != nil {
-		return status.Errorf(codes.Internal,
-			"couldn't get account information to check volume limit: %s", err.Error())
+		return nil, fmt.Errorf("failed to get account information: %s", err)
 	}
 
 	// administrative accounts might have zero length limits, make sure to not check them
 	if account.VolumeLimit == 0 {
-		return nil //  hail to the king!
+		return nil, nil //  hail to the king!
 	}
 
-	// NOTE(arslan): the API returns the limit for *all* regions, so passing
-	// the region down as a parameter doesn't change the response.
-	// Nevertheless, this is something we should be aware of.
-	volumes, _, err := d.storage.ListVolumes(ctx, &godo.ListVolumeParams{
-		Region: d.region,
-	})
-	if err != nil {
-		return status.Errorf(codes.Internal,
-			"couldn't get fetch volume list to check volume limit: %s", err.Error())
+	opt := &godo.ListOptions{
+		Page:    1,
+		PerPage: 50,
 	}
+	var numVolumes int
+	for {
+		// The API returns the limit for *all* regions, so passing the region
+		// down as a parameter doesn't change the response. Nevertheless, this
+		// is something we should be aware of.
+		volumes, resp, err := d.storage.ListVolumes(ctx, &godo.ListVolumeParams{
+			Region:      d.region,
+			ListOptions: opt,
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to list volumes at page %d: %s", opt.Page, err)
+		}
 
-	if account.VolumeLimit <= len(volumes) {
-		return status.Errorf(codes.ResourceExhausted,
-			"volume limit (%d) has been reached. Current number of volumes: %d. Please contact support.",
-			account.VolumeLimit, len(volumes))
+		numVolumes += len(volumes)
+		if account.VolumeLimit <= numVolumes {
+			return &limitDetails{
+				limit:      account.VolumeLimit,
+				numVolumes: numVolumes,
+			}, nil
+		}
+
+		if resp.Links == nil || resp.Links.IsLastPage() {
+			break
+		}
+
+		page, err := resp.Links.CurrentPage()
+		if err != nil {
+			return nil, fmt.Errorf("failed to get current page: %s", err)
+		}
+		opt.Page = page + 1
 	}
 
-	return nil
+	return nil, nil
 }
 
 // toCSISnapshot converts a DO Snapshot struct into a csi.Snapshot struct

driver/controller_test.go

@@ -4,11 +4,13 @@ import (
 	"context"
 	"errors"
 	"net/http"
+	"strconv"
 	"testing"
 	"time"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/digitalocean/godo"
+	"github.com/google/go-cmp/cmp"
 	"github.com/magiconair/properties/assert"
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc/codes"
@@ -255,6 +257,65 @@ func TestControllerExpandVolume(t *testing.T) {
 	}
 }
 
+func TestCheckLimit(t *testing.T) {
+	tests := []struct {
+		name        string
+		limit       int
+		numVolumes  int
+		wantErr     error
+		wantDetails *limitDetails
+	}{
+		{
+			name:       "limit insufficient",
+			limit:      25,
+			numVolumes: 30,
+			wantDetails: &limitDetails{
+				limit:      25,
+				numVolumes: 30,
+			},
+		},
+		{
+			name:        "limit sufficient",
+			limit:       100,
+			numVolumes:  25,
+			wantDetails: nil,
+		},
+		{
+			name:        "administrative account",
+			limit:       0,
+			numVolumes:  1000,
+			wantDetails: nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			storage := &fakeStorageDriver{
+				volumes: map[string]*godo.Volume{},
+			}
+			for i := 0; i < test.numVolumes; i++ {
+				storage.volumes[strconv.Itoa(i)] = &godo.Volume{}
+			}
+
+			d := Driver{
+				account: &fakeAccountDriver{
+					volumeLimit: test.limit,
+				},
+				storage: storage,
+			}
+
+			gotDetails, err := d.checkLimit(context.Background())
+			if err != nil {
+				t.Fatalf("got error: %s", err)
+			}
+
+			if diff := cmp.Diff(gotDetails, test.wantDetails, cmp.AllowUnexported(limitDetails{})); diff != "" {
+				t.Errorf("details mismatch (-got +want):\n%s", diff)
+			}
+		})
+	}
+}
+
 type fakeStorageAction struct {
 	*fakeStorageActionsDriver
 	storageGetValsFunc func(invocation int) (*godo.Action, *godo.Response, error)

driver/driver.go

@@ -205,8 +205,15 @@ func (d *Driver) Run() error {
 	// something is wrong in the logs. Only check if the driver is running with
 	// a token (i.e: controller)
 	if d.isController {
-		if err := d.checkLimit(context.Background()); err != nil {
-			d.log.WithError(err).Warn("CSI plugin will not function correctly, please resolve volume limit")
+		details, err := d.checkLimit(context.Background())
+		if err != nil {
+			return fmt.Errorf("failed to check volumes limits on startup: %s", err)
+		}
+		if details != nil {
+			d.log.WithFields(logrus.Fields{
+				"limit":       details.limit,
+				"num_volumes": details.numVolumes,
+			}).Warn("CSI plugin will not function correctly, please resolve volume limit")
 		}
 	}
 

driver/driver_test.go

@@ -94,10 +94,14 @@ func TestDriverSuite(t *testing.T) {
 	sanity.Test(t, cfg)
 }
 
-type fakeAccountDriver struct{}
+type fakeAccountDriver struct {
+	volumeLimit int
+}
 
 func (f *fakeAccountDriver) Get(context.Context) (*godo.Account, *godo.Response, error) {
-	return &godo.Account{}, godoResponse(), nil
+	return &godo.Account{
+		VolumeLimit: f.volumeLimit,
+	}, godoResponse(), nil
 }
 
 type fakeStorageDriver struct {
@@ -114,7 +118,11 @@ func (f *fakeStorageDriver) ListVolumes(ctx context.Context, param *godo.ListVol
 
 	if param != nil && param.ListOptions != nil && param.ListOptions.PerPage != 0 {
 		perPage := param.ListOptions.PerPage
-		vols := volumes[:perPage]
+		chunkSize := perPage
+		if len(volumes) < perPage {
+			chunkSize = len(volumes)
+		}
+		vols := volumes[:chunkSize]
 		for _, vol := range vols {
 			delete(f.volumes, vol.ID)
 		}

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/digitalocean/godo v1.13.0
 	github.com/golang/groupcache v0.0.0-20181024230925-c65c006176ff // indirect
 	github.com/golang/protobuf v1.3.1
+	github.com/google/go-cmp v0.3.0
 	github.com/google/go-querystring v0.0.0-20170111101155-53e6ce116135 // indirect
 	github.com/googleapis/gnostic v0.2.0 // indirect
 	github.com/imdario/mergo v0.3.6 // indirect