Merge pull request #26140 from zmerlynn/rate-limit-everything

GCE provider: Rate limit all API calls

pkg/cloudprovider/providers/gce/gce.go

@@ -85,7 +85,6 @@ type GCECloud struct {
 	networkURL        string
 	nodeTags          []string // List of tags to use on firewall rules for load balancers
 	useMetadataServer bool
-	operationPollRateLimiter flowcontrol.RateLimiter
 }
 
 type Config struct {
@@ -113,6 +112,16 @@ func (g *GCECloud) GetComputeService() *compute.Service {
 	return g.service
 }
 
+type rateLimitedRoundTripper struct {
+	rt      http.RoundTripper
+	limiter flowcontrol.RateLimiter
+}
+
+func (rl *rateLimitedRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	rl.limiter.Accept()
+	return rl.rt.RoundTrip(req)
+}
+
 func getProjectAndZone() (string, string, error) {
 	result, err := metadata.Get("instance/zone")
 	if err != nil {
@@ -283,6 +292,11 @@ func CreateGCECloud(projectID, region, zone string, managedZones []string, netwo
 	}
 
 	client := oauth2.NewClient(oauth2.NoContext, tokenSource)
+	// Override the transport to make it rate-limited.
+	client.Transport = &rateLimitedRoundTripper{
+		rt:      client.Transport,
+		limiter: flowcontrol.NewTokenBucketRateLimiter(10, 100), // 10 qps, 100 bucket size.
+	}
 	svc, err := compute.New(client)
 	if err != nil {
 		return nil, err
@@ -311,8 +325,6 @@ func CreateGCECloud(projectID, region, zone string, managedZones []string, netwo
 		glog.Infof("managing multiple zones: %v", managedZones)
 	}
 
-	operationPollRateLimiter := flowcontrol.NewTokenBucketRateLimiter(10, 100) // 10 qps, 100 bucket size.
-
 	return &GCECloud{
 		service:           svc,
 		containerService:  containerSvc,
@@ -323,7 +335,6 @@ func CreateGCECloud(projectID, region, zone string, managedZones []string, netwo
 		networkURL:        networkURL,
 		nodeTags:          nodeTags,
 		useMetadataServer: useMetadataServer,
-		operationPollRateLimiter: operationPollRateLimiter,
 	}, nil
 }
 
@@ -404,7 +415,6 @@ func (gce *GCECloud) waitForOp(op *compute.Operation, getOperation func(operatio
 	opName := op.Name
 	return wait.Poll(operationPollInterval, operationPollTimeoutDuration, func() (bool, error) {
 		start := time.Now()
-		gce.operationPollRateLimiter.Accept()
 		duration := time.Now().Sub(start)
 		if duration > 5*time.Second {
 			glog.Infof("pollOperation: waited %v for %v", duration, opName)

pkg/cloudprovider/providers/gce/gce_test.go

@@ -17,11 +17,15 @@ limitations under the License.
 package gce
 
 import (
+	"net/http"
+	"net/http/httptest"
 	"reflect"
 	"testing"
 
 	compute "google.golang.org/api/compute/v1"
+	"k8s.io/kubernetes/pkg/util/flowcontrol"
 	"k8s.io/kubernetes/pkg/util/rand"
+	utiltesting "k8s.io/kubernetes/pkg/util/testing"
 )
 
 func TestGetRegion(t *testing.T) {
@@ -260,3 +264,31 @@ func TestComputeUpdate(t *testing.T) {
 		// }
 	}
 }
+
+func TestRateLimitedRoundTripper(t *testing.T) {
+	handler := utiltesting.FakeHandler{StatusCode: 200}
+	server := httptest.NewServer(&handler)
+	defer server.Close()
+
+	method := "GET"
+	path := "/foo/bar"
+
+	req, err := http.NewRequest(method, server.URL+path, nil)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	// TODO(zmerlynn): Validate the rate limiter is actually getting called.
+	client := http.Client{
+		Transport: &rateLimitedRoundTripper{
+			rt:      http.DefaultTransport,
+			limiter: flowcontrol.NewFakeAlwaysRateLimiter(),
+		},
+	}
+	_, err = client.Do(req)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	handler.ValidateRequest(t, path, method, nil)
+}