Drain the batch channel more intelligently. Add order comments

Author: Steven.Osborne

batcher/batcher.go

@@ -19,6 +19,7 @@ package batcher
 import (
 	"errors"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -59,6 +60,7 @@ type basicBatcher struct {
 	lock           sync.RWMutex
 	batchChan      chan []interface{}
 	availableBytes uint
+	waiting        int32
 }
 
 // New creates a new Batcher using the provided arguments.
@@ -85,6 +87,7 @@ func New(maxTime time.Duration, maxItems, maxBytes, queueLen uint, calculate Cal
 
 // Put adds items to the batcher. If Put is continually called without calls to
 // Get, an unbounded number of go-routines will be generated.
+// Note: there is no order guarantee for items entering/leaving the batcher.
 func (b *basicBatcher) Put(item interface{}) error {
 	b.lock.Lock()
 	if b.disposed {
@@ -108,6 +111,7 @@ func (b *basicBatcher) Put(item interface{}) error {
 // one of the conditions for a "complete" batch is reached. If Put is
 // continually called without calls to Get, an unbounded number of
 // go-routines will be generated.
+// Note: there is no order guarantee for items entering/leaving the batcher.
 func (b *basicBatcher) Get() ([]interface{}, error) {
 	// Don't check disposed yet so any items remaining in the queue
 	// will be returned properly.
@@ -163,15 +167,9 @@ func (b *basicBatcher) Dispose() {
 	b.items = nil
 
 	// Drain the batch channel and all routines waiting to put on the channel
-DrainLoop:
-	for {
-		select {
-		case <-b.batchChan:
-		case <-time.After(5 * time.Millisecond):
-			break DrainLoop
-		}
+	for atomic.LoadInt32(&b.waiting) > 0 {
+		<-b.batchChan
 	}
-
 	close(b.batchChan)
 	b.lock.Unlock()
 }
@@ -193,7 +191,13 @@ func (b *basicBatcher) flush() {
 	for i, val := range b.items {
 		cpItems[i] = val
 	}
-	go func() { b.batchChan <- cpItems }()
+	// Signal one more waiter for the batch channel
+	atomic.AddInt32(&b.waiting, 1)
+	// Don't block on the channel put
+	go func() {
+		b.batchChan <- cpItems
+		atomic.AddInt32(&b.waiting, -1)
+	}()
 	b.items = make([]interface{}, 0, b.maxItems)
 	b.availableBytes = 0
 }

batcher/batcher_test.go

@@ -145,19 +145,24 @@ func TestDispose(t *testing.T) {
 	b.Put("b")
 	b.Put("c")
 
+	possibleBatches := [][]interface{}{
+		[]interface{}{"a", "b"},
+		[]interface{}{"c"},
+	}
 	batch1, err := b.Get()
-	assert.Equal([]interface{}{"a", "b"}, batch1)
+	assert.Contains(possibleBatches, batch1)
 	assert.Nil(err)
 	batch2, err := b.Get()
-	assert.Equal([]interface{}{"c"}, batch2)
+	assert.Contains(possibleBatches, batch2)
 	assert.Nil(err)
 
+	b.Put("d")
 	b.Dispose()
 
 	_, err = b.Get()
 	assert.Equal(ErrDisposed, err)
 
-	assert.Equal(ErrDisposed, b.Put("d"))
+	assert.Equal(ErrDisposed, b.Put("e"))
 	assert.Equal(ErrDisposed, b.Flush())
 
 }