package phb
import (
"errors"
"fmt"
"maps"
"slices"
"sync"
"time"
)
type ErrBackoffOverLimitTyp struct{}
func (*ErrBackoffOverLimitTyp) Error() string {
return "backoff over limit"
}
var (
ErrZeroWindow = errors.New("window must be at least 1 second")
ErrDuplicateWindow = errors.New("only one window of same length allowed in policy")
ErrAnachronisticIncrement = errors.New("an increment occurred that included writing to an old time slot. skipped")
ErrZeroSplit = errors.New("must have at least 1 split")
ErrNoSuchCounter = errors.New("no such counter")
ErrAtCapacity = errors.New("phb instance at capacity")
ErrRegisterNoPolicy = errors.New("must pass a policy (or set a default on registry) to register a counter")
ErrNoCounterOpen = errors.New("could not find a open slot in counter")
ErrAlreadyRegistered = errors.New("noun, verb, subject already registered")
ErrBackoffOverLimit = &ErrBackoffOverLimitTyp{}
ErrImplicitRegister = errors.New("error in implicit register from an increment of an event not registered")
)
type BucketPolicy struct {
// in seconds. If window is 60, it means after a full 60 seconds, count is zero
Window uint `json:"window" yaml:"splits"`
// number of splits within the window. If window is 60 and splits is 10,
// then there is partial refresh every 10 seconds
Splits uint `json:"splits" yaml:"splits"`
// max in a given window
Max uint `json:"max" yaml:"max"`
}
func (b BucketPolicy) Validate() error {
if b.Window == 0 {
return ErrZeroWindow
}
if b.Splits == 0 {
return ErrZeroSplit
}
return nil
}
type CounterPolicy struct {
DetailedEnforceError bool `json:"detailed_enforce_error" yaml:"detailed_enforce_error"`
SkipEnforceOnIncrement bool `json:"skip_enforce_on_increment" yaml:"skip_enforce_on_increment"`
Buckets []BucketPolicy `json:"time_ranges" yaml:"time_ranges"`
}
func (c CounterPolicy) Validate() error {
errs := []error{}
seenWindows := map[uint]struct{}{}
for _, v := range c.Buckets {
if _, isDupe := seenWindows[v.Window]; isDupe {
errs = append(errs, ErrDuplicateWindow)
}
if err := v.Validate(); err != nil {
errs = append(errs, err)
}
seenWindows[v.Window] = struct{}{}
}
return errors.Join(errs...)
}
type Bucket struct {
max uint
window uint
splitCounts []uint
splitIndices []uint
}
func (b Bucket) getTotal(start, now time.Time) uint {
splits := len(b.splitCounts)
writeIndex := toWriteIndex(start, now, b.window, uint(splits))
var lastRelevantIndex uint = 0
if (int(writeIndex) - splits) > 0 {
lastRelevantIndex = writeIndex - uint(splits)
}
var total uint = 0
for i := range b.splitCounts {
if b.splitIndices[i] <= lastRelevantIndex {
total += b.splitCounts[i]
}
}
return total
}
type Counter struct {
skipEnforceOnIncrement bool
detailedEnforceError bool
start time.Time
mu sync.RWMutex
buckets map[uint]*Bucket
}
func toWriteIndex(
start, now time.Time,
window, splits uint,
) uint {
diff := now.Sub(start)
nextEvery := time.Second * time.Duration((window*1000)/(splits*1000))
return uint((diff / nextEvery))
}
func writeIndexToSlot(index, splits uint) uint { return index % splits }
type Registry struct {
keyMap map[string]int
mu sync.RWMutex
cap int
defaultPolicy *CounterPolicy
}
const sep = "__PHB__"
func toKey(n, v, s string) string {
return n + sep + v + sep + s
}
func (r *Registry) SetDefaultPolicy(policy *CounterPolicy) error {
if policy == nil {
r.mu.Lock()
defer r.mu.Unlock()
r.defaultPolicy = policy
return nil
}
if err := policy.Validate(); err != nil {
return err
}
r.mu.Lock()
defer r.mu.Unlock()
r.defaultPolicy = policy
return nil
}
func Register(r *Registry, cs []Counter, n, v, s string, policy *CounterPolicy) error {
_, err := register(r, cs, n, v, s, policy, time.Now())
return err
}
func register(r *Registry, cs []Counter, n, v, s string, policy *CounterPolicy, now time.Time) (int, error) {
r.mu.Lock()
defer r.mu.Unlock()
key := toKey(n, v, s)
if counterI, alreadyExists := r.keyMap[key]; alreadyExists {
return counterI, ErrAlreadyRegistered
}
if policy == nil && r.defaultPolicy == nil {
return 0, ErrRegisterNoPolicy
}
if len(r.keyMap) > r.cap {
return 0, ErrAtCapacity
}
effectivePolicy := policy
if effectivePolicy == nil {
effectivePolicy = r.defaultPolicy
}
if err := effectivePolicy.Validate(); err != nil {
return -1, err
}
allTakenSlots := slices.Collect(maps.Values(r.keyMap))
slices.Sort(allTakenSlots)
taken := map[int]struct{}{}
for _, v := range r.keyMap {
taken[v] = struct{}{}
}
i := 0
for i < r.cap {
if _, alreadyTaken := taken[i]; !alreadyTaken {
cs[i].mu.Lock()
defer cs[i].mu.Unlock()
cs[i].start = now
cs[i].buckets = map[uint]*Bucket{}
cs[i].detailedEnforceError = effectivePolicy.DetailedEnforceError
cs[i].skipEnforceOnIncrement = effectivePolicy.SkipEnforceOnIncrement
for _, bucketPolicy := range effectivePolicy.Buckets {
if err := bucketPolicy.Validate(); err != nil {
return -1, err
}
cs[i].buckets[bucketPolicy.Window] = &Bucket{
max: bucketPolicy.Max,
window: bucketPolicy.Window,
splitCounts: make([]uint, bucketPolicy.Splits),
splitIndices: make([]uint, bucketPolicy.Splits),
}
}
r.keyMap[key] = i
return i, nil
}
i += 1
}
return -1, ErrNoCounterOpen
}
func Unregister(r *Registry, cs []Counter, n, v, s string) {
r.mu.Lock()
defer r.mu.Unlock()
key := toKey(n, v, s)
delete(r.keyMap, key)
}
func IncrementMoment(r *Registry, cs []Counter, n, v, s string, moment time.Time, amount uint) error {
key := toKey(n, v, s)
r.mu.RLock()
counterI, exists := r.keyMap[key]
r.mu.RUnlock()
if !exists {
var err error
counterI, err = register(r, cs, n, v, s, nil, moment)
if err != nil && err != ErrAlreadyRegistered {
// TOCTOU issue: if two goroutines call increment at same time
// and neither sees it registered, both call register.
// we rely on register returning ErrAlreadyRegistered with the registry handle.
// This way, we avoid having to write lock the registry on every request
return fmt.Errorf("%w: %w", ErrImplicitRegister, err)
}
}
cs[counterI].mu.Lock()
defer cs[counterI].mu.Unlock()
errsIncrement := []error{}
for window, bucket := range cs[counterI].buckets {
splits := uint(len(bucket.splitCounts))
writeIndex := toWriteIndex(cs[counterI].start, moment, window, splits)
writeSlot := writeIndexToSlot(writeIndex, splits)
if existingWriteIndex := bucket.splitIndices[writeSlot]; existingWriteIndex == writeIndex {
bucket.splitCounts[writeSlot] += amount
} else {
if writeIndex < existingWriteIndex {
err := fmt.Errorf(
"%w: bucket %d: attempt_index: %d, existing_index: %d, moment: %s, start: %s",
ErrAnachronisticIncrement,
window,
writeIndex,
existingWriteIndex,
moment.Format(time.RFC3339),
cs[counterI].start.Format(time.RFC3339),
)
errsIncrement = append(errsIncrement, err)
}
bucket.splitIndices[writeSlot] = writeIndex
bucket.splitCounts[writeSlot] = amount
}
}
if cs[counterI].skipEnforceOnIncrement {
return nil
}
return errors.Join(append(errsIncrement, enforceNoLock(cs, n, v, s, counterI, moment))...)
}
// example: r, cs, "bob", "GET", "/api/users", 1
func Increment(r *Registry, cs []Counter, n, v, s string, amount uint) error {
return IncrementMoment(r, cs, n, v, s, time.Now().UTC(), amount)
}
func Enforce(r *Registry, cs []Counter, n, v, s string, moment time.Time) error {
key := toKey(n, v, s)
r.mu.RLock()
counterI, exists := r.keyMap[key]
r.mu.RUnlock()
if !exists {
return fmt.Errorf("%w: \"%s %s %s\"", ErrNoSuchCounter, n, v, s)
}
cs[counterI].mu.RLock()
defer cs[counterI].mu.RUnlock()
return enforceNoLock(cs, n, v, s, counterI, moment)
}
func enforceNoLock(cs []Counter, n, v, s string, counterI int, moment time.Time) error {
maxErrs := []error{}
for window, bucket := range cs[counterI].buckets {
splits := uint(len(bucket.splitCounts))
bucketTotal := bucket.getTotal(cs[counterI].start, moment)
if bucketTotal > bucket.max {
var err error = ErrBackoffOverLimit
if cs[counterI].detailedEnforceError {
err = fmt.Errorf("%w: policy \"%s %s %s\" bucket \"%d seconds split every %d seconds\" at: %d, but allowed: %d", err, n, v, s, window, splits, bucketTotal, bucket.max)
}
maxErrs = append(maxErrs, err)
}
}
return errors.Join(maxErrs...)
}
func GetCounts(r *Registry, cs []Counter, n, v, s string, moment time.Time) (map[uint]uint, bool) {
key := toKey(n, v, s)
r.mu.RLock()
counterI, exists := r.keyMap[key]
r.mu.RUnlock()
if !exists {
return nil, false
}
cs[counterI].mu.RLock()
defer cs[counterI].mu.RUnlock()
counts := map[uint]uint{}
for window, bucket := range cs[counterI].buckets {
counts[window] = bucket.getTotal(cs[counterI].start, moment)
}
return counts, true
}
func New(cap int) (*Registry, []Counter) {
newRegistry := Registry{keyMap: map[string]int{}, cap: cap}
return &newRegistry, make([]Counter, cap)
}