Add rate limit to stashbox connection (#5764)

* Add max requests per minute stashbox option
* Implement rate limiting
* Add requests per minute to stashbox config
* Add UI setting

pkg/models/stash_box.go

@@ -7,7 +7,8 @@ type StashBoxFingerprint struct {
 }
 
 type StashBox struct {
-	Endpoint string `json:"endpoint"`
-	APIKey   string `json:"api_key"`
-	Name     string `json:"name"`
+	Endpoint             string `json:"endpoint"`
+	APIKey               string `json:"api_key"`
+	Name                 string `json:"name"`
+	MaxRequestsPerMinute int    `json:"max_requests_per_minute" koanf:"max_requests_per_minute"`
 }

pkg/stashbox/client.go

@@ -10,33 +10,86 @@ import (
 	"github.com/stashapp/stash/pkg/models"
 	"github.com/stashapp/stash/pkg/scraper"
 	"github.com/stashapp/stash/pkg/stashbox/graphql"
+
+	"golang.org/x/time/rate"
 )
 
+// DefaultMaxRequestsPerMinute is the default maximum number of requests per minute.
+const DefaultMaxRequestsPerMinute = 240
+
 // Client represents the client interface to a stash-box server instance.
 type Client struct {
 	client *graphql.Client
 	box    models.StashBox
 
+	maxRequestsPerMinute int
+
 	// tag patterns to be excluded
 	excludeTagRE []*regexp.Regexp
 }
 
-// NewClient returns a new instance of a stash-box client.
-func NewClient(box models.StashBox, excludeTagPatterns []string) *Client {
-	authHeader := func(ctx context.Context, req *http.Request, gqlInfo *clientv2.GQLRequestInfo, res interface{}, next clientv2.RequestInterceptorFunc) error {
-		req.Header.Set("ApiKey", box.APIKey)
+type ClientOption func(*Client)
+
+func ExcludeTagPatterns(patterns []string) ClientOption {
+	return func(c *Client) {
+		c.excludeTagRE = scraper.CompileExclusionRegexps(patterns)
+	}
+}
+
+func MaxRequestsPerMinute(n int) ClientOption {
+	return func(c *Client) {
+		if n > 0 {
+			c.maxRequestsPerMinute = n
+		}
+	}
+}
+
+func setApiKeyHeader(apiKey string) clientv2.RequestInterceptor {
+	return func(ctx context.Context, req *http.Request, gqlInfo *clientv2.GQLRequestInfo, res interface{}, next clientv2.RequestInterceptorFunc) error {
+		req.Header.Set("ApiKey", apiKey)
 		return next(ctx, req, gqlInfo, res)
 	}
+}
+
+func rateLimit(n int) clientv2.RequestInterceptor {
+	perSec := float64(n) / 60
+	limiter := rate.NewLimiter(rate.Limit(perSec), 1)
+
+	return func(ctx context.Context, req *http.Request, gqlInfo *clientv2.GQLRequestInfo, res interface{}, next clientv2.RequestInterceptorFunc) error {
+		if err := limiter.Wait(ctx); err != nil {
+			// should only happen if the context is canceled
+			return err
+		}
+
+		return next(ctx, req, gqlInfo, res)
+	}
+}
+
+// NewClient returns a new instance of a stash-box client.
+func NewClient(box models.StashBox, options ...ClientOption) *Client {
+	ret := &Client{
+		box:                  box,
+		maxRequestsPerMinute: DefaultMaxRequestsPerMinute,
+	}
+
+	if box.MaxRequestsPerMinute > 0 {
+		ret.maxRequestsPerMinute = box.MaxRequestsPerMinute
+	}
+
+	for _, option := range options {
+		option(ret)
+	}
+
+	authHeader := setApiKeyHeader(box.APIKey)
+	limitRequests := rateLimit(ret.maxRequestsPerMinute)
 
 	client := &graphql.Client{
-		Client: clientv2.NewClient(http.DefaultClient, box.Endpoint, nil, authHeader),
+		Client: clientv2.NewClient(http.DefaultClient, box.Endpoint, nil, authHeader, limitRequests),
 	}
 
-	return &Client{
-		client:       client,
-		box:          box,
-		excludeTagRE: scraper.CompileExclusionRegexps(excludeTagPatterns),
-	}
+	ret.client = client
+
+	return ret
 }
 
 func (c Client) getHTTPClient() *http.Client {