Initial implementation

.github/workflows/main.yml

@@ -0,0 +1,47 @@
+name: Go package
+on: [push]
+jobs:
+  build:
+    name: 'Go Build (1.21)'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
+        with:
+          go-version: '1.21'
+      - name: Install dependencies
+        run: go get .
+      - name: Build
+        run: go build ./...
+  static:
+    name: 'Go Static (1.21)'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
+        with:
+          go-version: '1.21'
+      - run: 'go install honnef.co/go/tools/cmd/staticcheck@latest'
+      - run: 'go vet ./...'
+      - run: 'staticcheck ./...'
+  test:
+    name: 'Go Test (1.21)'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
+        with:
+          go-version: '1.21'
+      - name: Install dependencies
+        run: go get .
+      - name: Test
+        run: go test -cover -vet all -coverprofile cover.out .
+      - name: Coverage Check
+        run: |
+          go tool cover -func ./cover.out
+          val=$(go tool cover -func cover.out | fgrep total | awk '{print $3}')
+          if [[ "100.0%" != $val ]] 
+          then
+            echo 'Test coverage is less than 100.0%'
+            exit 1
+          fi

.gitignore

@@ -19,3 +19,6 @@
 
 # Go workspace file
 go.work
+
+# Custom
+/.idea

README.md

@@ -1,2 +1,13 @@
 # go-leaky-bucket
-Leaky bucket meter implementation in Go
+[Leaky bucket meter](https://en.wikipedia.org/wiki/Leaky_bucket#As_a_meter) implementation in Go.
+
+This implementation atomically drains when its value is being mutated rather than using a timer or continual
+drain. Before mutation, the bucket will drain the supplied number of units as many times as necessary to match
+the precision of the supplied interval.
+
+For example, if a bucket is created which drains 5 units every 2 minutes, then after 2.5 minutes only 5 units will
+be drained. However, the 30 seconds of "unused" drain time will be accounted for to ensure future drains are kept
+accurate. If another 1.5 minutes were to pass, the bucket will drain by another 5 units because the unused time
+was recorded.
+
+See [`./examples`](./examples) for usage and inspiration.

bucket.go

@@ -0,0 +1,111 @@
+package leaky
+
+import (
+	"encoding/gob"
+	"errors"
+	"sync"
+	"time"
+)
+
+var ErrBucketFull = errors.New("leaky: bucket full or would overflow")
+
+func init() {
+	gob.Register(&Bucket{})
+}
+
+type Bucket struct {
+	DrainBy       int64
+	DrainInterval time.Duration
+	Capacity      int64
+
+	value     int64
+	lastDrain time.Time
+	lock      sync.Mutex
+}
+
+func NewBucket(drainBy int64, drainEvery time.Duration, capacity int64) (*Bucket, error) {
+	if drainBy <= 0 || drainEvery <= 0 {
+		return nil, errors.New("leaky: bucket never drains")
+	}
+	if capacity <= 0 {
+		return nil, errors.New("leaky: bucket can never fill")
+	}
+	return &Bucket{
+		DrainBy:       drainBy,
+		DrainInterval: drainEvery,
+		Capacity:      capacity,
+		value:         0,
+		lastDrain:     time.Now(),
+		lock:          sync.Mutex{},
+	}, nil
+}
+
+func (b *Bucket) drain() {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	if b.lastDrain.IsZero() {
+		b.lastDrain = time.Now() // assume we've never drained
+	}
+
+	if b.value <= 0 {
+		b.value = 0
+		b.lastDrain = time.Now()
+		return // nothing to drain, so don't bother
+	}
+
+	since := time.Since(b.lastDrain)
+	drainTime := since.Truncate(b.DrainInterval)
+	leaks := int64(drainTime.Abs() / b.DrainInterval.Abs())
+	b.value -= b.DrainBy * leaks
+	if b.value < 0 {
+		b.value = 0
+	}
+	b.lastDrain = time.Now().Add((since - drainTime) * -1)
+}
+
+func (b *Bucket) Peek() int64 {
+	return b.value
+}
+
+func (b *Bucket) Value() int64 {
+	b.drain()
+	return b.value
+}
+
+func (b *Bucket) Remaining() int64 {
+	b.drain()
+	return b.Capacity - b.value
+}
+
+func (b *Bucket) Add(amount int64) error {
+	b.drain()
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	newValue := b.value + amount
+	if newValue > b.Capacity {
+		return ErrBucketFull
+	}
+	b.value = newValue
+	return nil
+}
+
+func (b *Bucket) Set(value int64, resetDrain bool) error {
+	if value < 0 {
+		return errors.New("leaky: bucket value cannot be negative")
+	}
+	if value > b.Capacity {
+		return errors.New("leaky: bucket value cannot exceed capacity")
+	}
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	b.value = value
+	if resetDrain {
+		b.lastDrain = time.Now()
+	}
+	return nil
+}