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

go.mod

@@ -57,6 +57,7 @@ require (
 	golang.org/x/sys v0.28.0
 	golang.org/x/term v0.27.0
 	golang.org/x/text v0.21.0
+	golang.org/x/time v0.10.0
 	gopkg.in/guregu/null.v4 v4.0.0
 	gopkg.in/yaml.v2 v2.4.0
 )

go.sum

@@ -954,6 +954,8 @@ golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
 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/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4=
+golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 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=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=

graphql/schema/types/stash-box.graphql

@@ -2,12 +2,15 @@ type StashBox {
   endpoint: String!
   api_key: String!
   name: String!
+  max_requests_per_minute: Int!
 }
 
 input StashBoxInput {
   endpoint: String!
   api_key: String!
   name: String!
+  # defaults to 240
+  max_requests_per_minute: Int
 }
 
 type StashID {

internal/api/stash_box.go

@@ -11,7 +11,7 @@ import (
 )
 
 func (r *Resolver) newStashBoxClient(box models.StashBox) *stashbox.Client {
-	return stashbox.NewClient(box, manager.GetInstance().Config.GetScraperExcludeTagPatterns())
+	return stashbox.NewClient(box, stashbox.ExcludeTagPatterns(manager.GetInstance().Config.GetScraperExcludeTagPatterns()))
 }
 
 func resolveStashBoxFn(indexField, endpointField string) func(index *int, endpoint *string) (*models.StashBox, error) {

internal/manager/config/config.go

@@ -1105,9 +1105,10 @@ func stashBoxValidate(str string) bool {
 }
 
 type StashBoxInput struct {
-	Endpoint string `json:"endpoint"`
+	Endpoint             string `json:"endpoint"`
-	APIKey   string `json:"api_key"`
+	APIKey               string `json:"api_key"`
-	Name     string `json:"name"`
+	Name                 string `json:"name"`
+	MaxRequestsPerMinute int    `json:"max_requests_per_minute"`
 }
 
 func (i *Config) ValidateStashBoxes(boxes []*StashBoxInput) error {

internal/manager/task_identify.go

@@ -180,7 +180,7 @@ func (j *IdentifyJob) getSources() ([]identify.ScraperSource, error) {
 			src = identify.ScraperSource{
 				Name: "stash-box: " + stashBox.Endpoint,
 				Scraper: stashboxSource{
-					Client:                 stashbox.NewClient(*stashBox, instance.Config.GetScraperExcludeTagPatterns()),
+					Client:                 stashbox.NewClient(*stashBox, stashbox.ExcludeTagPatterns(instance.Config.GetScraperExcludeTagPatterns())),
 					endpoint:               stashBox.Endpoint,
 					txnManager:             instance.Repository.TxnManager,
 					sceneFingerprintGetter: instance.SceneService,

internal/manager/task_stash_box_tag.go

@@ -96,7 +96,7 @@ func (t *StashBoxBatchTagTask) findStashBoxPerformer(ctx context.Context) (*mode
 
 	r := instance.Repository
 
-	client := stashbox.NewClient(*t.box, instance.Config.GetScraperExcludeTagPatterns())
+	client := stashbox.NewClient(*t.box, stashbox.ExcludeTagPatterns(instance.Config.GetScraperExcludeTagPatterns()))
 
 	if t.refresh {
 		var remoteID string
@@ -295,7 +295,7 @@ func (t *StashBoxBatchTagTask) findStashBoxStudio(ctx context.Context) (*models.
 
 	r := instance.Repository
 
-	client := stashbox.NewClient(*t.box, instance.Config.GetScraperExcludeTagPatterns())
+	client := stashbox.NewClient(*t.box, stashbox.ExcludeTagPatterns(instance.Config.GetScraperExcludeTagPatterns()))
 
 	if t.refresh {
 		var remoteID string

pkg/models/stash_box.go

@@ -7,7 +7,8 @@ type StashBoxFingerprint struct {
 }
 
 type StashBox struct {
-	Endpoint string `json:"endpoint"`
+	Endpoint             string `json:"endpoint"`
-	APIKey   string `json:"api_key"`
+	APIKey               string `json:"api_key"`
-	Name     string `json:"name"`
+	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.
+type ClientOption func(*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 {
+func ExcludeTagPatterns(patterns []string) ClientOption {
-		req.Header.Set("ApiKey", box.APIKey)
+	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{
+	ret.client = client
-		client:       client,
+
-		box:          box,
+	return ret
-		excludeTagRE: scraper.CompileExclusionRegexps(excludeTagPatterns),
-	}
 }
 
 func (c Client) getHTTPClient() *http.Client {

ui/v2.5/graphql/data/config.graphql

@@ -49,6 +49,7 @@ fragment ConfigGeneralData on ConfigGeneralResult {
     name
     endpoint
     api_key
+    max_requests_per_minute
   }
   pythonPath
   transcodeInputArgs

ui/v2.5/src/components/Settings/StashBoxConfiguration.tsx

@@ -10,6 +10,8 @@ export interface IStashBoxModal {
   close: (v?: GQL.StashBoxInput) => void;
 }
 
+const defaultMaxRequestsPerMinute = 240;
+
 export const StashBoxModal: React.FC<IStashBoxModal> = ({ value, close }) => {
   const intl = useIntl();
   const endpoint = useRef<HTMLInputElement | null>(null);
@@ -114,6 +116,38 @@ export const StashBoxModal: React.FC<IStashBoxModal> = ({ value, close }) => {
               </b>
             )}
           </Form.Group>
+
+          <Form.Group id="stashbox-max-requests-per-minute">
+            <h6>
+              {intl.formatMessage({
+                id: "config.stashbox.max_requests_per_minute",
+              })}
+            </h6>
+            <Form.Control
+              placeholder={intl.formatMessage({
+                id: "config.stashbox.max_requests_per_minute",
+              })}
+              className="text-input"
+              value={v?.max_requests_per_minute ?? defaultMaxRequestsPerMinute}
+              isValid={
+                (v?.max_requests_per_minute ?? defaultMaxRequestsPerMinute) >= 0
+              }
+              type="number"
+              onChange={(e: React.ChangeEvent<HTMLInputElement>) =>
+                setValue({
+                  ...v!,
+                  max_requests_per_minute: parseInt(e.currentTarget.value),
+                })
+              }
+              ref={apiKey}
+            />
+            <div className="sub-heading">
+              <FormattedMessage
+                id="config.stashbox.max_requests_per_minute_description"
+                values={{ defaultValue: defaultMaxRequestsPerMinute }}
+              />
+            </div>
+          </Form.Group>
         </>
       )}
       close={close}

ui/v2.5/src/locales/en-GB.json

@@ -448,6 +448,8 @@
       "description": "Stash-box facilitates automated tagging of scenes and performers based on fingerprints and filenames.\nEndpoint and API key can be found on your account page on the stash-box instance. Names are required when more than one instance is added.",
       "endpoint": "Endpoint",
       "graphql_endpoint": "GraphQL endpoint",
+      "max_requests_per_minute": "Max requests per minute",
+      "max_requests_per_minute_description": "Uses default value of {defaultValue} if set to 0",
       "name": "Name",
       "title": "Stash-box Endpoints"
     },