batching: added batcher

Author: brianshannan-wf

batcher/batcher.go

@@ -0,0 +1,143 @@
+package batcher
+
+import (
+	"errors"
+	"sync"
+	"time"
+)
+
+// Batcher provides an API for accumulating items into a batch for processing.
+type Batcher interface {
+	// Put adds items to the batcher.
+	Put(interface{}) error
+
+	// Get retrieves a batch from the batcher. This call will block until
+	// one of the conditions for a "complete" batch is reached.
+	Get() ([]interface{}, error)
+
+	// Dispose will dispose of the batcher. Any calls to Put or Get
+	// will return errors.
+	Dispose()
+}
+
+// ErrDisposed is the error returned for a disposed Batcher
+var ErrDisposed = errors.New("batcher: disposed")
+
+// CalculateBytes evaluates the number of bytes in an item added to a Batcher.
+type CalculateBytes func(interface{}) uint
+
+type basicBatcher struct {
+	maxTime        time.Duration
+	maxItems       uint
+	maxBytes       uint
+	queueLen       uint
+	calculateBytes CalculateBytes
+	disposed       bool
+	items          []interface{}
+	lock           sync.RWMutex
+	batchChan      chan []interface{}
+	availableBytes uint
+}
+
+// New creates a new Batcher using the provided arguments.
+// Batch readiness can be determined in three ways:
+//   - Maximum number of bytes per batch
+//   - Maximum number of items per batch
+//   - Maximum amount of time waiting for a batch
+// Values of zero for one of these fields indicate they should not be
+// taken into account when evaluating the readiness of a batch.
+func New(maxTime time.Duration, maxItems, maxBytes, queueLen uint, calculate CalculateBytes) (Batcher, error) {
+	if maxBytes > 0 && calculate == nil {
+		return nil, errors.New("batcher: must provide CalculateBytes function")
+	}
+
+	return &basicBatcher{
+		maxTime:        maxTime,
+		maxItems:       maxItems,
+		maxBytes:       maxBytes,
+		queueLen:       queueLen,
+		calculateBytes: calculate,
+		items:          make([]interface{}, 0, maxItems),
+		batchChan:      make(chan []interface{}, queueLen),
+	}, nil
+}
+
+// Put adds items to the batcher.
+func (b *basicBatcher) Put(item interface{}) error {
+	// TODO make this non-blocking and drop if batchChan is full?
+	// TODO Should there be a length check here and return an error if that happens without trying to add?
+
+	b.lock.Lock()
+	if b.disposed {
+		b.lock.Unlock()
+		return ErrDisposed
+	}
+
+	b.items = append(b.items, item)
+	b.availableBytes += b.calculateBytes(item)
+	if b.ready() {
+		b.batchChan <- b.items
+		b.items = make([]interface{}, 0, b.maxItems)
+		b.availableBytes = 0
+	}
+
+	b.lock.Unlock()
+	return nil
+}
+
+// Get retrieves a batch from the batcher. This call will block until
+// one of the conditions for a "complete" batch is reached.
+func (b *basicBatcher) Get() ([]interface{}, error) {
+	b.lock.RLock()
+	if b.disposed {
+		b.lock.RUnlock()
+		return nil, ErrDisposed
+	}
+
+	var timeout <-chan time.Time
+	if b.maxTime > 0 {
+		timeout = time.After(b.maxTime)
+	}
+	b.lock.RUnlock()
+
+	select {
+	case items := <-b.batchChan:
+		b.lock.RLock()
+		if b.disposed {
+			return nil, ErrDisposed
+		}
+		b.lock.RUnlock()
+		return items, nil
+	case <-timeout:
+		b.lock.Lock()
+		if b.disposed {
+			b.lock.Unlock()
+			return nil, ErrDisposed
+		}
+		items := b.items
+		b.items = make([]interface{}, 0, b.maxItems)
+		b.availableBytes = 0
+		b.lock.Unlock()
+		return items, nil
+	}
+}
+
+// Dispose will dispose of the batcher. Any calls to Put or Get
+// will return errors.
+func (b *basicBatcher) Dispose() {
+	b.lock.Lock()
+	b.disposed = true
+	b.items = nil
+	b.batchChan = nil
+	b.lock.Unlock()
+}
+
+func (b *basicBatcher) ready() bool {
+	if b.maxItems != 0 && uint(len(b.items)) >= b.maxItems {
+		return true
+	}
+	if b.maxBytes != 0 && b.availableBytes >= b.maxBytes {
+		return true
+	}
+	return false
+}

batcher/batcher_test.go

@@ -0,0 +1,121 @@
+package batcher
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNoCalculateBytes(t *testing.T) {
+	_, err := New(0, 0, 100, 5, nil)
+	assert.Error(t, err)
+}
+
+func TestMaxItems(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
+		return uint(len(str.(string)))
+	})
+	assert.Nil(err)
+
+	go func() {
+		for i := 0; i < 1000; i++ {
+			assert.Nil(b.Put("foo bar baz"))
+		}
+	}()
+
+	batch, err := b.Get()
+	assert.Len(batch, 100)
+	assert.Nil(err)
+}
+
+func TestMaxBytes(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(0, 10000, 100, 10, func(str interface{}) uint {
+		return uint(len(str.(string)))
+	})
+	assert.Nil(err)
+
+	go func() {
+		for i := 0; i < 1000; i++ {
+			b.Put("a")
+		}
+	}()
+
+	batch, err := b.Get()
+	assert.Len(batch, 100)
+	assert.Nil(err)
+}
+
+func TestMaxTime(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(time.Millisecond*200, 100000, 100000, 10,
+		func(str interface{}) uint {
+			return uint(len(str.(string)))
+		},
+	)
+	assert.Nil(err)
+
+	go func() {
+		for i := 0; i < 10000; i++ {
+			b.Put("a")
+			time.Sleep(time.Millisecond)
+		}
+	}()
+
+	before := time.Now()
+	batch, err := b.Get()
+	assert.InDelta(200, time.Since(before).Seconds()*1000, 2)
+	assert.True(len(batch) > 0)
+	assert.Nil(err)
+}
+
+func TestMultiConsumer(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
+		return uint(len(str.(string)))
+	})
+	assert.Nil(err)
+
+	var wg sync.WaitGroup
+	wg.Add(5)
+	for i := 0; i < 5; i++ {
+		go func() {
+			batch, err := b.Get()
+			assert.Len(batch, 100)
+			assert.Nil(err)
+			wg.Done()
+		}()
+	}
+
+	go func() {
+		for i := 0; i < 500; i++ {
+			b.Put("a")
+		}
+	}()
+
+	wg.Wait()
+}
+
+func TestDispose(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(0, 100000, 100000, 10, func(str interface{}) uint {
+		return uint(len(str.(string)))
+	})
+	assert.Nil(err)
+	b.Put("a")
+	wait := make(chan bool)
+	go func() {
+		_, err := b.Get()
+		assert.Equal(ErrDisposed, err)
+		wait <- true
+	}()
+
+	b.Dispose()
+
+	assert.Equal(ErrDisposed, b.Put("a"))
+
+	<-wait
+}

mock/batcher.go

@@ -0,0 +1,27 @@
+package mock
+
+import (
+	"github.com/stretchr/testify/mock"
+)
+
+type MockBatcher struct {
+	mock.Mock
+	PutChan chan bool
+}
+
+func (m *MockBatcher) Put(items interface{}) error {
+	args := m.Called(items)
+	if m.PutChan != nil {
+		m.PutChan <- true
+	}
+	return args.Error(0)
+}
+
+func (m *MockBatcher) Get() ([]interface{}, error) {
+	args := m.Called()
+	return args.Get(0).([]interface{}), args.Error(1)
+}
+
+func (m *MockBatcher) Dispose() {
+	m.Called()
+}