minstrel/internal/acquire/pace.go

package acquire

import (
	"context"
	"sync"
	"time"

	"git.fabledsword.com/bvandeusen/minstrel/internal/indexer"
)

// SearchPace is the one Pacer every search shares. Process-wide because a
// Fanout is built fresh for each search; a Pacer per Fanout would start empty
// and pace nothing.
var SearchPace = &Pacer{}

// holdWithoutRetryAfter is how long a source that answered 429 without saying
// for how long is left alone: one schedule tick.
const holdWithoutRetryAfter = 10 * time.Minute

// Pacer spaces the searches sent to each source by its searches_per_hour: a
// spacing, not a burst bucket, because Prowlarr answers 429 to a few requests
// in quick succession. Every search spends, the schedule's and a
// person's alike; only the schedule waits.
type Pacer struct {
	// Now and Sleep are the clock, replaceable in tests.
	Now   func() time.Time
	Sleep func(context.Context, time.Duration) error

	mu   sync.Mutex
	next map[string]time.Time
}

func (p *Pacer) now() time.Time {
	if p.Now != nil {
		return p.Now()
	}
	return time.Now()
}

func (p *Pacer) sleep(ctx context.Context, d time.Duration) error {
	if p.Sleep != nil {
		return p.Sleep(ctx, d)
	}
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-t.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// gap is the spacing a source's number asks for; zero is unpaced.
func gap(perHour int) time.Duration {
	if perHour <= 0 {
		return 0
	}
	return time.Hour / time.Duration(perHour)
}

// readyAt is when every source may next be asked.
func (p *Pacer) readyAt(sources []Source) time.Time {
	p.mu.Lock()
	defer p.mu.Unlock()
	var at time.Time
	for _, s := range sources {
		if s.Search == nil {
			continue
		}
		if n := p.next[s.Code]; n.After(at) {
			at = n
		}
	}
	return at
}

// Wait holds until every source may be asked, and answers false, without
// waiting, when that is after the deadline or ctx ends first. All of them,
// not any: a search while one is held back would be narrowed, and the want
// recorded as looked for.
func (p *Pacer) Wait(ctx context.Context, sources []Source, deadline time.Time) bool {
	if p == nil {
		return true
	}
	at := p.readyAt(sources)
	if !deadline.IsZero() && at.After(deadline) {
		return false
	}
	if d := at.Sub(p.now()); d > 0 {
		return p.sleep(ctx, d) == nil
	}
	return ctx.Err() == nil
}

// spend takes one search's slot from a source.
func (p *Pacer) spend(s Source) {
	g := gap(s.PerHour)
	if p == nil || g == 0 {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.next == nil {
		p.next = map[string]time.Time{}
	}
	at := p.next[s.Code]
	if t := p.now(); at.Before(t) {
		at = t
	}
	p.next[s.Code] = at.Add(g)
}

// hold leaves a source alone after it answered 429, for as long as it asked,
// whatever its searches_per_hour says.
func (p *Pacer) hold(code string, rl *indexer.RateLimited) {
	if p == nil {
		return
	}
	d := rl.RetryAfter
	if d <= 0 {
		d = holdWithoutRetryAfter
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.next == nil {
		p.next = map[string]time.Time{}
	}
	if until := p.now().Add(d); until.After(p.next[code]) {
		p.next[code] = until
	}
}

// Ready holds until every source this service searches may be asked again,
// and answers false when that is past the deadline. The schedule calls it
// before each want; a person's search never does.
func (s *Service) Ready(ctx context.Context, deadline time.Time) bool {
	if s == nil {
		return ctx.Err() == nil
	}
	if f, ok := s.Search.(*Fanout); ok && f != nil {
		return f.Pace.Wait(ctx, f.Sources, deadline)
	}
	return ctx.Err() == nil
}