quintodrome/core/stream/limiter.go

package stream

import (
	"context"
	"errors"
	"io"
	"sync"
	"sync/atomic"
)

// ErrTooManyTranscodes is returned by TranscodeLimiter.Acquire when the
// configured concurrency cap has been reached. Callers should translate this
// into an HTTP 429 response so well-behaved clients back off and retry.
var ErrTooManyTranscodes = errors.New("too many concurrent transcodes")

// RetryAfterSeconds is the value returned in the HTTP Retry-After header when
// a request is rejected with ErrTooManyTranscodes. Most transcodes finish well
// within this window, so retrying after this delay typically succeeds.
const RetryAfterSeconds = 5

// TranscodeLimiter gates the number of concurrent ffmpeg transcodes. It enforces
// both a global cap (to protect the host from process exhaustion) and an optional
// per-user cap (to keep one client from starving the others). Acquire never
// blocks: it either reserves a slot or returns ErrTooManyTranscodes immediately.
type TranscodeLimiter interface {
	// Acquire reserves a slot for the given user. On success it returns a release
	// function that must be called exactly once when the transcode is done.
	// Calling release more than once is safe and idempotent.
	Acquire(ctx context.Context, user string) (release func(), err error)

	// Enabled reports whether the limiter actually enforces any cap. Callers
	// can use it to decide whether to bind ffmpeg's lifetime to the request
	// context so disconnects free slots quickly, rather than letting the
	// process drain to completion in the background.
	Enabled() bool
}

// NewTranscodeLimiter returns a limiter enforcing the given caps. Each cap is
// independent: a value of zero or less disables that cap. When both caps are
// disabled the limiter is a no-op.
func NewTranscodeLimiter(maxConcurrent, maxPerUser int) TranscodeLimiter {
	if maxConcurrent <= 0 && maxPerUser <= 0 {
		return noopLimiter{}
	}
	l := &transcodeLimiter{maxPerUser: maxPerUser}
	if maxConcurrent > 0 {
		l.global = make(chan struct{}, maxConcurrent)
	}
	if maxPerUser > 0 {
		l.perUser = make(map[string]int)
	}
	return l
}

// releasingReadCloser wraps an io.ReadCloser so that closing it also releases
// the limiter slot exactly once. release must be the function returned by
// TranscodeLimiter.Acquire; its own idempotency makes double-Close safe too.
type releasingReadCloser struct {
	io.ReadCloser
	release func()
}

func (r *releasingReadCloser) Close() error {
	err := r.ReadCloser.Close()
	r.release()
	return err
}

type noopLimiter struct{}

func (noopLimiter) Acquire(context.Context, string) (func(), error) {
	return func() {}, nil
}

func (noopLimiter) Enabled() bool { return false }

type transcodeLimiter struct {
	maxPerUser int
	global     chan struct{}

	mu      sync.Mutex
	perUser map[string]int
}

func (*transcodeLimiter) Enabled() bool { return true }

func (l *transcodeLimiter) Acquire(_ context.Context, user string) (func(), error) {
	// Reserve a per-user slot first so a noisy user can't burn through
	// global slots only to be rejected later. An empty user key means
	// "anonymous" (e.g. public share viewers); we skip the per-user cap
	// entirely so unrelated anonymous clients do not share a bucket.
	perUserActive := l.maxPerUser > 0 && user != ""
	if perUserActive {
		l.mu.Lock()
		if l.perUser[user] >= l.maxPerUser {
			l.mu.Unlock()
			return nil, ErrTooManyTranscodes
		}
		l.perUser[user]++
		l.mu.Unlock()
	}

	if l.global != nil {
		select {
		case l.global <- struct{}{}:
		default:
			if perUserActive {
				l.releasePerUser(user)
			}
			return nil, ErrTooManyTranscodes
		}
	}

	var released atomic.Bool
	return func() {
		if !released.CompareAndSwap(false, true) {
			return
		}
		if l.global != nil {
			<-l.global
		}
		if perUserActive {
			l.releasePerUser(user)
		}
	}, nil
}

func (l *transcodeLimiter) releasePerUser(user string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.perUser[user]--
	if l.perUser[user] <= 0 {
		delete(l.perUser, user)
	}
}