new-api/controller/plugin_protocol_limiter.go

package controller

import (
	"errors"
	"fmt"
	"strings"
	"sync"
)

var (
	errPluginProtocolObservationLimitExceeded   = errors.New("plugin protocol observation limit exceeded")
	errInvalidPluginProtocolObservationIdentity = errors.New("invalid plugin protocol observation identity")
)

type pluginProtocolObservationLimits struct {
	global    int
	perPlugin int
	perUser   int
	perToken  int
}

var defaultPluginProtocolObservationLimits = pluginProtocolObservationLimits{
	global:    128,
	perPlugin: 32,
	perUser:   4,
	perToken:  2,
}

var pluginProtocolObservationAdmissions = newPluginProtocolObservationLimiter(
	defaultPluginProtocolObservationLimits,
)

type pluginProtocolObservationLimitError struct {
	scope string
	limit int
}

func (e *pluginProtocolObservationLimitError) Error() string {
	return fmt.Sprintf("%s: %s capacity is %d", errPluginProtocolObservationLimitExceeded, e.scope, e.limit)
}

func (e *pluginProtocolObservationLimitError) Unwrap() error {
	return errPluginProtocolObservationLimitExceeded
}

type pluginProtocolObservationLimiter struct {
	mu sync.Mutex

	limits pluginProtocolObservationLimits
	global int
	plugin map[string]int
	user   map[int]int
	token  map[int]int
}

func newPluginProtocolObservationLimiter(limits pluginProtocolObservationLimits) *pluginProtocolObservationLimiter {
	return &pluginProtocolObservationLimiter{
		limits: limits,
		plugin: make(map[string]int),
		user:   make(map[int]int),
		token:  make(map[int]int),
	}
}

func (l *pluginProtocolObservationLimiter) acquire(
	pluginKey string,
	userID int,
	tokenID int,
) (func(), error) {
	pluginKey = strings.TrimSpace(pluginKey)
	switch {
	case pluginKey == "":
		return nil, fmt.Errorf("%w: plugin key is required", errInvalidPluginProtocolObservationIdentity)
	case userID <= 0:
		return nil, fmt.Errorf("%w: user id must be positive", errInvalidPluginProtocolObservationIdentity)
	case tokenID <= 0:
		return nil, fmt.Errorf("%w: token id must be positive", errInvalidPluginProtocolObservationIdentity)
	}

	l.mu.Lock()
	if l.global >= l.limits.global {
		l.mu.Unlock()
		return nil, &pluginProtocolObservationLimitError{
			scope: "global",
			limit: l.limits.global,
		}
	}
	l.global++

	if l.plugin[pluginKey] >= l.limits.perPlugin {
		l.global--
		l.mu.Unlock()
		return nil, &pluginProtocolObservationLimitError{
			scope: "plugin",
			limit: l.limits.perPlugin,
		}
	}
	l.plugin[pluginKey]++

	if l.user[userID] >= l.limits.perUser {
		l.global--
		l.plugin[pluginKey]--
		if l.plugin[pluginKey] == 0 {
			delete(l.plugin, pluginKey)
		}
		l.mu.Unlock()
		return nil, &pluginProtocolObservationLimitError{
			scope: "user",
			limit: l.limits.perUser,
		}
	}
	l.user[userID]++

	if l.token[tokenID] >= l.limits.perToken {
		l.global--
		l.plugin[pluginKey]--
		if l.plugin[pluginKey] == 0 {
			delete(l.plugin, pluginKey)
		}
		l.user[userID]--
		if l.user[userID] == 0 {
			delete(l.user, userID)
		}
		l.mu.Unlock()
		return nil, &pluginProtocolObservationLimitError{
			scope: "token",
			limit: l.limits.perToken,
		}
	}
	l.token[tokenID]++
	l.mu.Unlock()

	var once sync.Once
	return func() {
		once.Do(func() {
			l.mu.Lock()
			defer l.mu.Unlock()

			l.global--

			l.plugin[pluginKey]--
			if l.plugin[pluginKey] == 0 {
				delete(l.plugin, pluginKey)
			}

			l.user[userID]--
			if l.user[userID] == 0 {
				delete(l.user, userID)
			}

			l.token[tokenID]--
			if l.token[tokenID] == 0 {
				delete(l.token, tokenID)
			}
		})
	}, nil
}