From f519645dd35430a7c5be6886b6e2f61307fc1ddb Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Sun, 24 Jan 2021 14:58:13 +0100 Subject: [PATCH 1/2] Add global rate limit to HTTP calls Fixes #156 --- common/client.go | 8 ++++++++ common/http.go | 3 +++ go.mod | 3 ++- go.sum | 3 ++- provider/provider.go | 9 +++++++++ 5 files changed, 24 insertions(+), 2 deletions(-) diff --git a/common/client.go b/common/client.go index 4312ff0d80..d639be448a 100644 --- a/common/client.go +++ b/common/client.go @@ -12,6 +12,8 @@ import ( "sync" "time" + "golang.org/x/time/rate" + "github.com/hashicorp/go-retryablehttp" "github.com/mitchellh/go-homedir" "gopkg.in/ini.v1" @@ -30,6 +32,8 @@ type DatabricksClient struct { TimeoutSeconds int DebugTruncateBytes int DebugHeaders bool + RateLimit int + rateLimiter *rate.Limiter httpClient *retryablehttp.Client authMutex sync.Mutex authVisitor func(r *http.Request) error @@ -177,6 +181,10 @@ func (c *DatabricksClient) configureHTTPCLient() { if c.TimeoutSeconds == 0 { c.TimeoutSeconds = 60 } + if c.RateLimit == 0 { + c.RateLimit = 1200 + } + c.rateLimiter = rate.NewLimiter(rate.Every(1*time.Minute), c.RateLimit) // Set up a retryable HTTP Client to handle cases where the service returns // a transient error on initial creation retryDelayDuration := 10 * time.Second diff --git a/common/http.go b/common/http.go index c18c7c32b4..b302dbbc23 100644 --- a/common/http.go +++ b/common/http.go @@ -413,6 +413,9 @@ func (c *DatabricksClient) genericQuery(ctx context.Context, method, requestURL if c.httpClient == nil { return nil, fmt.Errorf("DatabricksClient is not configured") } + if err = c.rateLimiter.Wait(ctx); err != nil { + return nil, err + } requestBody, err := makeRequestBody(method, &requestURL, data, true) if err != nil { return nil, err diff --git a/go.mod b/go.mod index 41c53be65b..6b74187fe6 100644 --- a/go.mod +++ b/go.mod @@ -17,5 +17,6 @@ require ( github.com/pkg/errors v0.9.1 github.com/smartystreets/goconvey v1.6.4 // indirect github.com/stretchr/testify v1.6.1 + golang.org/x/time v0.0.0-20191024005414-555d28b269f0 gopkg.in/ini.v1 v1.62.0 -) \ No newline at end of file +) diff --git a/go.sum b/go.sum index dde1b01a37..1fdf0f3e0a 100644 --- a/go.sum +++ b/go.sum @@ -468,6 +468,7 @@ golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= +golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= @@ -616,4 +617,4 @@ honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k= rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8= rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0= -rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA= \ No newline at end of file +rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA= diff --git a/provider/provider.go b/provider/provider.go index 3066614f3c..6c5547dadd 100644 --- a/provider/provider.go +++ b/provider/provider.go @@ -212,6 +212,12 @@ func DatabricksProvider() *schema.Provider { Description: "Debug HTTP headers of requests made by the provider. Default is false. Visible only when TF_LOG=DEBUG is set", DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_DEBUG_HEADERS", false), }, + "rate_limit": { + Optional: true, + Type: schema.TypeInt, + Description: "Maximum number of requests per minute made to Databricks REST API by Terraform.", + DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_RATE_LIMIT", 1200), + }, }, ConfigureContextFunc: func(c context.Context, d *schema.ResourceData) (interface{}, diag.Diagnostics) { pc := common.DatabricksClient{} @@ -286,6 +292,9 @@ func DatabricksProvider() *schema.Provider { if v, ok := d.GetOk("debug_truncate_bytes"); ok { pc.DebugTruncateBytes = v.(int) } + if v, ok := d.GetOk("rate_limit"); ok { + pc.RateLimit = v.(int) + } if v, ok := d.GetOk("debug_headers"); ok { pc.DebugHeaders = v.(bool) } From 089cb08f8e4531b08b5d545804eb25e5b97a54b4 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Sun, 24 Jan 2021 18:29:01 +0100 Subject: [PATCH 2/2] Extracted constants --- common/client.go | 13 ++++++++++--- provider/provider.go | 4 ++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/common/client.go b/common/client.go index d639be448a..3db30f0a50 100644 --- a/common/client.go +++ b/common/client.go @@ -19,6 +19,13 @@ import ( "gopkg.in/ini.v1" ) +// Default settings +const ( + DefaultRateLimit = 1200 + DebugTruncateBytes = 96 + DefaultTimeoutSeconds = 60 +) + // DatabricksClient is the client struct that contains clients for all the services available on Databricks type DatabricksClient struct { Host string @@ -45,7 +52,7 @@ func (c *DatabricksClient) Configure() error { c.configureHTTPCLient() c.AzureAuth.databricksClient = c if c.DebugTruncateBytes == 0 { - c.DebugTruncateBytes = 96 + c.DebugTruncateBytes = DebugTruncateBytes } return nil } @@ -179,10 +186,10 @@ func (c *DatabricksClient) encodeBasicAuth(username, password string) string { func (c *DatabricksClient) configureHTTPCLient() { if c.TimeoutSeconds == 0 { - c.TimeoutSeconds = 60 + c.TimeoutSeconds = DefaultTimeoutSeconds } if c.RateLimit == 0 { - c.RateLimit = 1200 + c.RateLimit = DefaultRateLimit } c.rateLimiter = rate.NewLimiter(rate.Every(1*time.Minute), c.RateLimit) // Set up a retryable HTTP Client to handle cases where the service returns diff --git a/provider/provider.go b/provider/provider.go index 6c5547dadd..db7acfda65 100644 --- a/provider/provider.go +++ b/provider/provider.go @@ -204,7 +204,7 @@ func DatabricksProvider() *schema.Provider { Optional: true, Type: schema.TypeInt, Description: "Truncate JSON fields in JSON above this limit. Default is 96. Visible only when TF_LOG=DEBUG is set", - DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_DEBUG_TRUNCATE_BYTES", 96), + DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_DEBUG_TRUNCATE_BYTES", common.DebugTruncateBytes), }, "debug_headers": { Optional: true, @@ -216,7 +216,7 @@ func DatabricksProvider() *schema.Provider { Optional: true, Type: schema.TypeInt, Description: "Maximum number of requests per minute made to Databricks REST API by Terraform.", - DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_RATE_LIMIT", 1200), + DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_RATE_LIMIT", common.DefaultRateLimit), }, }, ConfigureContextFunc: func(c context.Context, d *schema.ResourceData) (interface{}, diag.Diagnostics) {