feat: move bucket into internal package (#2)

* feat: move bucket into internal package

- this isn't part of the public api, so enforce it by using
go idioms making it easier to test as well.
- rename `maxTokens` -> `max`

* fix: update changelog & readme

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Readme updates.
+- `internal` package - has no change to library functionality.
+- Rename `maxTokens` to `max` when constructing new rate limiter.
+
 ## [0.1.0] - 2023-03-30
 
 ### Added

README.md

@@ -1,4 +1,4 @@
-## ratelimit [![Go Reference](https://pkg.go.dev/badge/github.com/vivangkumar/ratelimit.svg)](https://pkg.go.dev/github.com/vivangkumar/ratelimit) ![CI](https://github.com/vivangkumar/ratelimit/actions/workflows/ci.yaml/badge.svg?branch=main) 
+## ratelimit [![Go Reference](https://pkg.go.dev/badge/github.com/vivangkumar/ratelimit.svg)](https://pkg.go.dev/github.com/vivangkumar/ratelimit) ![CI](https://github.com/vivangkumar/ratelimit/actions/workflows/ci.yaml/badge.svg?branch=main)
 
 Ratelimit implements a [token bucket](https://en.wikipedia.org/wiki/Token_bucket) based rate limiter.
 
@@ -74,4 +74,4 @@ make build
 
 ## Changelog
 
-Please add changes between release to the changelog [Changelog](CHANGELOG.md).
+Please add changes between release to the [Changelog](CHANGELOG.md).

bucket.go → internal/bucket/bucket.go

@@ -1,9 +1,9 @@
-package ratelimiter
+package bucket
 
-// bucket represents a token bucket.
+// Bucket represents a token bucket.
 //
 // It is not safe for concurrent use.
-type bucket struct {
+type Bucket struct {
 	// size is the max tokens the bucket can hold.
 	size uint64
 
@@ -12,28 +12,28 @@ type bucket struct {
 	available uint64
 }
 
-// newBucket returns a token bucket configured with
+// New returns a token bucket configured with
 // size number of max tokens.
-func newBucket(size uint64) *bucket {
-	return &bucket{
+func New(size uint64) *Bucket {
+	return &Bucket{
 		size:      size,
 		available: size,
 	}
 }
 
-// take uses up a single token from the bucket.
+// Take uses up a single token from the bucket.
 //
 // It returns true if a token could be acquired.
 // Otherwise, it returns false.
-func (b *bucket) take() bool {
-	return b.takeN(1)
+func (b *Bucket) Take() bool {
+	return b.TakeN(1)
 }
 
-// takeN acquires n tokens from the bucket, if available.
+// TakeN acquires n tokens from the bucket, if available.
 //
 // If n tokens are not available, no tokens are removed
 // from the bucket.
-func (b *bucket) takeN(n uint64) bool {
+func (b *Bucket) TakeN(n uint64) bool {
 	if b.available >= n {
 		b.available -= n
 		return true
@@ -42,14 +42,24 @@ func (b *bucket) takeN(n uint64) bool {
 	return false
 }
 
-// refill refills the bucket with n tokens.
+// Refill refills the bucket with n tokens.
 //
 // If the quantity to refill exceeds the size of the bucket,
 // the bucket is refilled upto its size.
-func (b *bucket) refill(n uint64) {
+func (b *Bucket) Refill(n uint64) {
 	t := b.available + n
 	if t > b.size {
 		t = b.size
 	}
 	b.available = t
 }
+
+// Available returns the tokens currently available in the bucket.
+func (b *Bucket) Available() uint64 {
+	return b.available
+}
+
+// Size returns the max tokens a bucket can have.
+func (b *Bucket) Size() uint64 {
+	return b.size
+}

internal/bucket/bucket_test.go

@@ -0,0 +1,37 @@
+package bucket_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/vivangkumar/ratelimit/internal/bucket"
+)
+
+func TestBucket_Take(t *testing.T) {
+	b := bucket.New(1)
+
+	assert.True(t, b.Take())
+	assert.EqualValues(t, b.Available(), 0)
+	assert.False(t, b.Take())
+}
+
+func TestBucket_TakeN(t *testing.T) {
+	b := bucket.New(10)
+
+	assert.True(t, b.TakeN(10))
+	assert.EqualValues(t, b.Available(), 0)
+	assert.False(t, b.TakeN(10))
+}
+
+func TestBucket_Refill(t *testing.T) {
+	b := bucket.New(10)
+	assert.True(t, b.TakeN(10))
+
+	b.Refill(1)
+	assert.EqualValues(t, b.Available(), 1)
+
+	// Fill to max.
+	b.Refill(20)
+	assert.EqualValues(t, b.Available(), 10)
+}

ratelimiter.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"sync"
 	"time"
+
+	"github.com/vivangkumar/ratelimit/internal/bucket"
 )
 
 // NowFunc helps with mocking time.
@@ -21,7 +23,7 @@ type RateLimiter struct {
 	m sync.Mutex
 	// bucket is the underlying storage structure
 	// for the rate limiter.
-	bucket *bucket
+	bucket *bucket.Bucket
 
 	// lastRefillAt keeps track of the time when the last
 	// refresh of tokens took place.
@@ -37,7 +39,7 @@ type RateLimiter struct {
 	now         NowFunc
 }
 
-// New constructs a rate limiter that accepts the max tokens that
+// New constructs a rate limiter that accepts the max tokens (size) that
 // the limiter holds, along with the limit per duration.
 //
 // For example: if the bucket is configured with a max of 100 tokens
@@ -46,7 +48,7 @@ type RateLimiter struct {
 // This refills the bucket with 1 token every 6s, while giving us a
 // max "burst" of 100 tokens.
 func New(
-	maxTokens uint64,
+	max uint64,
 	limit uint64,
 	per time.Duration,
 	opts ...Opt,
@@ -56,7 +58,7 @@ func New(
 	}
 
 	r := &RateLimiter{
-		bucket:      newBucket(maxTokens),
+		bucket:      bucket.New(max),
 		refillEvery: time.Duration(float64(per) / float64(limit)),
 		now:         time.Now,
 	}
@@ -83,7 +85,7 @@ func (r *RateLimiter) refill() {
 	t := (now.UnixMilli() - r.lastRefillAt) / r.refillEvery.Milliseconds()
 	if t > 0 {
 		r.lastRefillAt = now.UnixMilli()
-		r.bucket.refill(uint64(t))
+		r.bucket.Refill(uint64(t))
 	}
 }
 
@@ -105,7 +107,7 @@ func (r *RateLimiter) AddN(n uint64) bool {
 	r.refill()
 
 	r.m.Lock()
-	ok := r.bucket.takeN(n)
+	ok := r.bucket.TakeN(n)
 	r.m.Unlock()
 
 	return ok
@@ -130,7 +132,7 @@ func (r *RateLimiter) Wait(ctx context.Context) error {
 //
 // This method also consumes n tokens, if successful.
 func (r *RateLimiter) WaitN(ctx context.Context, n uint64) error {
-	if n > r.bucket.size {
+	if n > r.bucket.Size() {
 		return fmt.Errorf("tokens requested exceeds max tokens")
 	}