Quota

quota.go

package quota

import (
  "errors"
	"sync"
	"time"
)

var QuotaExceededErr = errors.New("quota exceeded")

type Quota interface {
	// Inc increments the quota counter by 1.
	// It returns the remaining quota increments, or QuotaExceededErr.
	Inc() (int, error)
	// IncsLeft returns the remaining increments.
	IncsLeft() int
	// Reset resets the quota.
	Reset()
}

type quota struct {
	list        *inc
	numIncs     int
	perDuration time.Duration
	mtx         sync.Mutex
}

func NewQuota(numIncs int, perDuration time.Duration) Quota {
	if perDuration <= 0 {
		panic("invalid quota duration")
	}
	return &quota{
		numIncs:     numIncs,
		perDuration: perDuration,
	}
}

func (q *quota) Inc() (int, error) {
	q.mtx.Lock()
	defer q.mtx.Unlock()

	now := time.Now()
	li := &inc{
		ts: now.UnixNano(),
	}
	if q.list == nil {
		q.list = li
	} else {
		root := q.list
		li.next = root
		q.list = li
	}

	n := 0
	limit := now.Add(-q.perDuration).UnixNano()

	node := q.list
	for node != nil {
		if node.ts >= limit {
			n++
		} else {
			node.next = nil
		}
		if n > q.numIncs {
			return 0, QuotaExceededErr
		}
		node = node.next
	}

	return q.IncsLeft(), nil
}

func (q *quota) IncsLeft() int {
	n := q.numIncs
	if q.list != nil {
		n -= q.list.len()
	}
	if n < 0 {
		return 0
	}
	return n
}

func (q *quota) Reset() {
	q.list = nil
}

type inc struct {
	ts   int64
	next *inc
}

// len returns the list length starting from the current node
func (i *inc) len() int {
	if i.next != nil {
		return 1 + i.next.len()
	}
	return 1
}