feat: separate rate limiting from timeouts

Author: sentriz

clientutil/clientutil.go

@@ -3,14 +3,12 @@
 package clientutil
 
 import (
-	"context"
 	"log/slog"
 	"net/http"
 	"sync"
 	"time"
 
 	"github.com/gregjones/httpcache"
-	"golang.org/x/time/rate"
 )
 
 type Middleware func(http.RoundTripper) http.RoundTripper
@@ -36,34 +34,6 @@ func WithCache() Middleware {
 	}
 }
 
-func WithTimeout(timeout time.Duration) Middleware {
-	if timeout == 0 {
-		return Passthrough
-	}
-	return func(next http.RoundTripper) http.RoundTripper {
-		return RoundTripFunc(func(r *http.Request) (*http.Response, error) {
-			ctx, cancel := context.WithTimeout(r.Context(), timeout)
-			defer cancel()
-			return next.RoundTrip(r.WithContext(ctx))
-		})
-	}
-}
-
-func WithRateLimit(interval time.Duration) Middleware {
-	if interval == 0 {
-		return Passthrough
-	}
-	return func(next http.RoundTripper) http.RoundTripper {
-		limiter := rate.NewLimiter(rate.Every(interval), 1)
-		return RoundTripFunc(func(r *http.Request) (*http.Response, error) {
-			if err := limiter.Wait(r.Context()); err != nil {
-				return nil, err
-			}
-			return next.RoundTrip(r)
-		})
-	}
-}
-
 func WithLogging(logger *slog.Logger) Middleware {
 	return func(next http.RoundTripper) http.RoundTripper {
 		return RoundTripFunc(func(r *http.Request) (*http.Response, error) {

cmd/internal/wrtagflag/wrtagflag.go

@@ -22,6 +22,7 @@ import (
 	"go.senan.xyz/wrtag/notifications"
 	"go.senan.xyz/wrtag/pathformat"
 	"go.senan.xyz/wrtag/researchlink"
+	"golang.org/x/time/rate"
 
 	_ "go.senan.xyz/wrtag/addon/lyrics"
 	_ "go.senan.xyz/wrtag/addon/musicdesc"
@@ -83,10 +84,20 @@ func Config() *wrtag.Config {
 	flag.Var(&tagConfigParser{&cfg.TagConfig}, "tag-config", "Specify tag keep and drop rules when writing new tag revisions (see [Tagging](#tagging)) (stackable)")
 
 	flag.StringVar(&cfg.MusicBrainzClient.BaseURL, "mb-base-url", `https://musicbrainz.org/ws/2/`, "MusicBrainz base URL")
-	flag.DurationVar(&cfg.MusicBrainzClient.RateLimit, "mb-rate-limit", 1*time.Second, "MusicBrainz rate limit duration")
+
+	cfg.MusicBrainzClient.Limiter = rate.NewLimiter(rate.Every(1*time.Second), 1)
+	flag.Var(&rateLimitParser{cfg.MusicBrainzClient.Limiter}, "mb-rate-limit", "MusicBrainz rate limit duration")
+
+	cfg.MusicBrainzClient.HTTPClient = &http.Client{Timeout: 30 * time.Second}
 
 	flag.StringVar(&cfg.CoverArtArchiveClient.BaseURL, "caa-base-url", `https://coverartarchive.org/`, "CoverArtArchive base URL")
-	flag.DurationVar(&cfg.CoverArtArchiveClient.RateLimit, "caa-rate-limit", 0, "CoverArtArchive rate limit duration")
+
+	cfg.CoverArtArchiveClient.Limiter = rate.NewLimiter(rate.Inf, 0)
+	flag.Var(&rateLimitParser{cfg.CoverArtArchiveClient.Limiter}, "caa-rate-limit", "CoverArtArchive rate limit duration")
+
+	cfg.CoverArtArchiveClient.HTTPClient = clientutil.Wrap(&http.Client{Timeout: 30 * time.Second},
+		clientutil.WithCache(),
+	)
 
 	flag.BoolVar(&cfg.UpgradeCover, "cover-upgrade", false, "Fetch new cover art even if it exists locally")
 
@@ -311,3 +322,26 @@ func (a addonsParser) String() string {
 	}
 	return strings.Join(parts, ", ")
 }
+
+type rateLimitParser struct {
+	l *rate.Limiter
+}
+
+func (rl rateLimitParser) Set(value string) error {
+	dur, err := time.ParseDuration(value)
+	if err != nil {
+		return err
+	}
+	*rl.l = *rate.NewLimiter(rate.Every(dur), 1)
+	return nil
+}
+func (rl *rateLimitParser) String() string {
+	if rl.l == nil {
+		return ""
+	}
+	dur := time.Duration((1.0 / float64(rl.l.Limit())) * float64(time.Second))
+	if dur == 0 {
+		return ""
+	}
+	return dur.String()
+}

lyrics/lyrics.go

@@ -6,10 +6,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/http"
 	"strings"
 	"time"
 
 	"golang.org/x/net/html"
+	"golang.org/x/time/rate"
 )
 
 type Source interface {
@@ -19,11 +21,11 @@ type Source interface {
 func NewSource(name string) (Source, error) {
 	switch name {
 	case "genius":
-		return &Genius{RateLimit: 500 * time.Millisecond}, nil
+		return &Genius{&http.Client{}, rate.NewLimiter(rate.Every(500*time.Millisecond), 1)}, nil
 	case "musixmatch":
-		return &Musixmatch{RateLimit: 500 * time.Millisecond}, nil
+		return &Musixmatch{&http.Client{}, rate.NewLimiter(rate.Every(500*time.Millisecond), 1)}, nil
 	case "lrclib":
-		return &LRCLib{RateLimit: 100 * time.Millisecond}, nil
+		return &LRCLib{&http.Client{}, rate.NewLimiter(rate.Every(100*time.Millisecond), 1)}, nil
 	default:
 		return nil, errors.New("unknown source")
 	}

lyrics/lyrics_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.senan.xyz/wrtag/lyrics"
+	"golang.org/x/time/rate"
 )
 
 //go:embed testdata
@@ -19,6 +20,7 @@ func TestMusixmatch(t *testing.T) {
 
 	var src lyrics.Musixmatch
 	src.HTTPClient = fsClient(responses, "testdata/musixmatch")
+	src.Limiter = rate.NewLimiter(rate.Inf, 0)
 
 	resp, err := src.Search(t.Context(), "The Fall", "Wings", 0)
 	require.NoError(t, err)
@@ -41,6 +43,7 @@ func TestGenius(t *testing.T) {
 
 	var src lyrics.Genius
 	src.HTTPClient = fsClient(responses, "testdata/genius")
+	src.Limiter = rate.NewLimiter(rate.Inf, 0)
 
 	resp, err := src.Search(t.Context(), "the fall", "totally wired", 0)
 	require.NoError(t, err)
@@ -64,6 +67,7 @@ func TestGeniusLineBreak(t *testing.T) {
 
 	var src lyrics.Genius
 	src.HTTPClient = fsClient(responses, "testdata/genius")
+	src.Limiter = rate.NewLimiter(rate.Inf, 0)
 
 	resp, err := src.Search(t.Context(), "pink floyd", "breathe in the air", 0)
 	require.NoError(t, err)

lyrics/source_genius.go

@@ -6,12 +6,11 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/andybalholm/cascadia"
-	"go.senan.xyz/wrtag/clientutil"
 	"golang.org/x/net/html"
+	"golang.org/x/time/rate"
 )
 
 var geniusBaseURL = `https://genius.com`
@@ -32,18 +31,14 @@ var geniusEsc = strings.NewReplacer(
 )
 
 type Genius struct {
-	RateLimit time.Duration
-
-	initOnce   sync.Once
 	HTTPClient *http.Client
+	Limiter    *rate.Limiter
 }
 
 func (g *Genius) Search(ctx context.Context, artist, song string, duration time.Duration) (string, error) {
-	g.initOnce.Do(func() {
-		g.HTTPClient = clientutil.Wrap(g.HTTPClient, clientutil.Chain(
-			clientutil.WithRateLimit(g.RateLimit),
-		))
-	})
+	if err := g.Limiter.Wait(ctx); err != nil {
+		return "", err
+	}
 
 	// use genius case rules to miminise redirects
 	page := fmt.Sprintf("%s-%s-lyrics", artist, song)

lyrics/source_lrclib.go

@@ -6,27 +6,22 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
-	"sync"
 	"time"
 
-	"go.senan.xyz/wrtag/clientutil"
+	"golang.org/x/time/rate"
 )
 
 var lrclibBaseURL = `https://lrclib.net/api/get`
 
 type LRCLib struct {
-	RateLimit time.Duration
-
-	initOnce   sync.Once
 	HTTPClient *http.Client
+	Limiter    *rate.Limiter
 }
 
 func (l *LRCLib) Search(ctx context.Context, artist, song string, duration time.Duration) (string, error) {
-	l.initOnce.Do(func() {
-		l.HTTPClient = clientutil.Wrap(l.HTTPClient, clientutil.Chain(
-			clientutil.WithRateLimit(l.RateLimit),
-		))
-	})
+	if err := l.Limiter.Wait(ctx); err != nil {
+		return "", err
+	}
 
 	u, _ := url.Parse(lrclibBaseURL)
 	q := u.Query()

lyrics/source_musixmatch.go

@@ -6,12 +6,11 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/andybalholm/cascadia"
-	"go.senan.xyz/wrtag/clientutil"
 	"golang.org/x/net/html"
+	"golang.org/x/time/rate"
 )
 
 var musixmatchBaseURL = `https://www.musixmatch.com/lyrics`
@@ -33,18 +32,14 @@ var musixmatchEsc = strings.NewReplacer(
 )
 
 type Musixmatch struct {
-	RateLimit time.Duration
-
-	initOnce   sync.Once
 	HTTPClient *http.Client
+	Limiter    *rate.Limiter
 }
 
 func (mm *Musixmatch) Search(ctx context.Context, artist, song string, duration time.Duration) (string, error) {
-	mm.initOnce.Do(func() {
-		mm.HTTPClient = clientutil.Wrap(mm.HTTPClient, clientutil.Chain(
-			clientutil.WithRateLimit(mm.RateLimit),
-		))
-	})
+	if err := mm.Limiter.Wait(ctx); err != nil {
+		return "", err
+	}
 
 	url, _ := url.Parse(musixmatchBaseURL)
 	url = url.JoinPath(musixmatchEsc.Replace(artist))

musicbrainz/coverartarchive.go

@@ -6,18 +6,14 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"sync"
-	"time"
 
-	"go.senan.xyz/wrtag/clientutil"
+	"golang.org/x/time/rate"
 )
 
 type CAAClient struct {
-	BaseURL   string
-	RateLimit time.Duration
-
-	initOnce   sync.Once
+	BaseURL    string
 	HTTPClient *http.Client
+	Limiter    *rate.Limiter
 }
 
 func (c *CAAClient) GetCoverURL(ctx context.Context, release *Release) (string, error) {
@@ -74,13 +70,9 @@ type caaResponse struct {
 }
 
 func (c *CAAClient) request(ctx context.Context, r *http.Request, dest any) error {
-	c.initOnce.Do(func() {
-		c.HTTPClient = clientutil.Wrap(c.HTTPClient, clientutil.Chain(
-			clientutil.WithCache(),
-			clientutil.WithRateLimit(c.RateLimit),
-			clientutil.WithTimeout(30*time.Second),
-		))
-	})
+	if err := c.Limiter.Wait(ctx); err != nil {
+		return err
+	}
 
 	r = r.WithContext(ctx)
 	resp, err := c.HTTPClient.Do(r)

musicbrainz/musicbrainz.go

@@ -15,22 +15,19 @@ import (
 	"slices"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 	"unicode"
 
 	"github.com/araddon/dateparse"
-	"go.senan.xyz/wrtag/clientutil"
+	"golang.org/x/time/rate"
 )
 
 var ErrNoResults = errors.New("no results")
 
 type MBClient struct {
-	BaseURL   string
-	RateLimit time.Duration
-
-	initOnce   sync.Once
+	BaseURL    string
 	HTTPClient *http.Client
+	Limiter    *rate.Limiter
 }
 
 func (c *MBClient) GetRelease(ctx context.Context, mbid string) (*Release, error) {
@@ -146,12 +143,9 @@ func (c *MBClient) SearchRelease(ctx context.Context, q ReleaseQuery) (*Release,
 }
 
 func (c *MBClient) request(ctx context.Context, r *http.Request, dest any) error {
-	c.initOnce.Do(func() {
-		c.HTTPClient = clientutil.Wrap(c.HTTPClient, clientutil.Chain(
-			clientutil.WithRateLimit(c.RateLimit),
-			clientutil.WithTimeout(30*time.Second),
-		))
-	})
+	if err := c.Limiter.Wait(ctx); err != nil {
+		return err
+	}
 
 	r = r.WithContext(ctx)
 	resp, err := c.HTTPClient.Do(r)