Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add global rate limit to HTTP calls #465

Merged
merged 2 commits into from
Jan 24, 2021
Merged

Add global rate limit to HTTP calls #465

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f519645
Add global rate limit to HTTP calls
nfx Jan 24, 2021
089cb08
Extracted constants
nfx Jan 24, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions common/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -177,6 +181,10 @@ func (c *DatabricksClient) configureHTTPCLient() {
if c.TimeoutSeconds == 0 {
c.TimeoutSeconds = 60
}
if c.RateLimit == 0 {
c.RateLimit = 1200
nfx marked this conversation as resolved.
Show resolved Hide resolved
}
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
Expand Down
3 changes: 3 additions & 0 deletions common/http.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
)
3 changes: 2 additions & 1 deletion go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=
9 changes: 9 additions & 0 deletions provider/provider.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down Expand Up @@ -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)
}
Expand Down