Limits

util/limit.go

@@ -0,0 +1,65 @@
+package util
+
+import (
+	"errors"
+	"sync"
+)
+
+// ErrLimitReached is the error returned by the Limiter and LimitWriter when the predefined limit has been reached
+var ErrLimitReached = errors.New("limit reached")
+
+// Limiter is a helper that allows adding values up to a well-defined limit. Once the limit is reached
+// ErrLimitReached will be returned. Limiter may be used by multiple goroutines.
+type Limiter struct {
+	value int64
+	limit int64
+	mu    sync.Mutex
+}
+
+// NewLimiter creates a new Limiter
+func NewLimiter(limit int64) *Limiter {
+	return &Limiter{
+		limit: limit,
+	}
+}
+
+// Add adds n to the limiters internal value, but only if the limit has not been reached. If the limit would be
+// exceeded after adding n, ErrLimitReached is returned.
+func (l *Limiter) Add(n int64) error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.limit == 0 {
+		l.value += n
+		return nil
+	} else if l.value+n <= l.limit {
+		l.value += n
+		return nil
+	} else {
+		return ErrLimitReached
+	}
+}
+
+// Sub subtracts a value from the limiters internal value
+func (l *Limiter) Sub(n int64) {
+	l.Add(-n)
+}
+
+// Set sets the value of the limiter to n. This function ignores the limit. It is meant to set the value
+// based on reality.
+func (l *Limiter) Set(n int64) {
+	l.mu.Lock()
+	l.value = n
+	l.mu.Unlock()
+}
+
+// Value returns the internal value of the limiter
+func (l *Limiter) Value() int64 {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.value
+}
+
+// Limit returns the defined limit
+func (l *Limiter) Limit() int64 {
+	return l.limit
+}