From 455fcac51c2b1b18e928dac48bb6f5c8512b5010 Mon Sep 17 00:00:00 2001
From: Rohit Ramkumar <rramkumar@google.com>
Date: Mon, 26 Mar 2018 15:45:34 -0700
Subject: [PATCH] Add support for setting a custom rate limiter in gce cloud
 provider

---
 pkg/cloudprovider/providers/gce/gce.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/pkg/cloudprovider/providers/gce/gce.go b/pkg/cloudprovider/providers/gce/gce.go
index 80ba27f8345..a3cfc029886 100644
--- a/pkg/cloudprovider/providers/gce/gce.go
+++ b/pkg/cloudprovider/providers/gce/gce.go
@@ -152,6 +152,9 @@ type GCECloud struct {
 
 	// New code generated interface to the GCE compute library.
 	c cloud.Cloud
+
+	// Keep a reference of this around so we can inject a new cloud.RateLimiter implementation.
+	s *cloud.Service
 }
 
 // TODO: replace gcfg with json
@@ -508,17 +511,27 @@ func CreateGCECloud(config *CloudConfig) (*GCECloud, error) {
 	}
 
 	gce.manager = &gceServiceManager{gce}
-	gce.c = cloud.NewGCE(&cloud.Service{
+	gce.s = &cloud.Service{
 		GA:            service,
 		Alpha:         serviceAlpha,
 		Beta:          serviceBeta,
 		ProjectRouter: &gceProjectRouter{gce},
 		RateLimiter:   &gceRateLimiter{gce},
-	})
+	}
+	gce.c = cloud.NewGCE(gce.s)
 
 	return gce, nil
 }
 
+// SetRateLimiter adds a custom cloud.RateLimiter implementation.
+// WARNING: Calling this could have unexpected behavior if you have in-flight
+// requests. It is best to use this immediately after creating a GCECloud.
+func (g *GCECloud) SetRateLimiter(rl cloud.RateLimiter) {
+	if rl != nil {
+		g.s.RateLimiter = rl
+	}
+}
+
 // determineSubnetURL queries for all subnetworks in a region for a given network and returns
 // the URL of the subnetwork which exists in the auto-subnet range.
 func determineSubnetURL(service *compute.Service, networkProjectID, networkName, region string) (string, error) {