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) {