Merge pull request #109 from Workiva/batcher_deadlock

Avoid deadlock between Put and Get

Author: tylertreat-wf

batcher/batcher.go

@@ -19,9 +19,15 @@ package batcher
 import (
 	"errors"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
+const (
+	batcherActive   = uint32(0)
+	batcherDisposed = uint32(1)
+)
+
 // Batcher provides an API for accumulating items into a batch for processing.
 type Batcher interface {
 	// Put adds items to the batcher.
@@ -54,11 +60,13 @@ type basicBatcher struct {
 	maxItems       uint
 	maxBytes       uint
 	calculateBytes CalculateBytes
-	disposed       bool
+	disposed       uint32
 	items          []interface{}
 	lock           sync.RWMutex
 	batchChan      chan []interface{}
+	disposeChan    chan struct{}
 	availableBytes uint
+	waiting        int32
 }
 
 // New creates a new Batcher using the provided arguments.
@@ -80,17 +88,19 @@ func New(maxTime time.Duration, maxItems, maxBytes, queueLen uint, calculate Cal
 		calculateBytes: calculate,
 		items:          make([]interface{}, 0, maxItems),
 		batchChan:      make(chan []interface{}, queueLen),
+		disposeChan:    make(chan struct{}),
 	}, nil
 }
 
-// Put adds items to the batcher.
+// 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 {
-		b.lock.Unlock()
+	// Check to see if disposed before putting
+	if b.IsDisposed() {
 		return ErrDisposed
 	}
-
+	b.lock.Lock()
 	b.items = append(b.items, item)
 	if b.calculateBytes != nil {
 		b.availableBytes += b.calculateBytes(item)
@@ -104,7 +114,10 @@ func (b *basicBatcher) Put(item interface{}) error {
 }
 
 // Get retrieves a batch from the batcher. This call will block until
-// one of the conditions for a "complete" batch is reached.
+// 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.
@@ -114,18 +127,25 @@ func (b *basicBatcher) Get() ([]interface{}, error) {
 		timeout = time.After(b.maxTime)
 	}
 
+	// Check to see if disposed before blocking
+	if b.IsDisposed() {
+		return nil, ErrDisposed
+	}
+
 	select {
-	case items, ok := <-b.batchChan:
+	case items := <-b.batchChan:
+		return items, nil
+	case _, ok := <-b.disposeChan:
 		if !ok {
 			return nil, ErrDisposed
 		}
-		return items, nil
+		return nil, nil
 	case <-timeout:
-		b.lock.Lock()
-		if b.disposed {
-			b.lock.Unlock()
+		// Check to see if disposed before getting lock
+		if b.IsDisposed() {
 			return nil, ErrDisposed
 		}
+		b.lock.Lock()
 		items := b.items
 		b.items = make([]interface{}, 0, b.maxItems)
 		b.availableBytes = 0
@@ -136,11 +156,10 @@ func (b *basicBatcher) Get() ([]interface{}, error) {
 
 // Flush forcibly completes the batch currently being built
 func (b *basicBatcher) Flush() error {
-	b.lock.Lock()
-	if b.disposed {
-		b.lock.Unlock()
+	if b.IsDisposed() {
 		return ErrDisposed
 	}
+	b.lock.Lock()
 	b.flush()
 	b.lock.Unlock()
 	return nil
@@ -150,30 +169,44 @@ func (b *basicBatcher) Flush() error {
 // will return ErrDisposed, calls to Get will return an error iff
 // there are no more ready batches.
 func (b *basicBatcher) Dispose() {
-	b.lock.Lock()
-	if b.disposed {
-		b.lock.Unlock()
+	// Check to see if disposed before attempting to dispose
+	if atomic.CompareAndSwapUint32(&b.disposed, batcherActive, batcherDisposed) {
 		return
 	}
+	b.lock.Lock()
 	b.flush()
-	b.disposed = true
 	b.items = nil
+	close(b.disposeChan)
+
+	// Drain the batch channel and all routines waiting to put on the channel
+	for len(b.batchChan) > 0 || atomic.LoadInt32(&b.waiting) > 0 {
+		<-b.batchChan
+	}
 	close(b.batchChan)
 	b.lock.Unlock()
 }
 
 // IsDisposed will determine if the batcher is disposed
 func (b *basicBatcher) IsDisposed() bool {
-	b.lock.RLock()
-	disposed := b.disposed
-	b.lock.RUnlock()
-	return disposed
+	return atomic.LoadUint32(&b.disposed) == batcherDisposed
 }
 
 // flush adds the batch currently being built to the queue of completed batches.
 // flush is not threadsafe, so should be synchronized externally.
 func (b *basicBatcher) flush() {
-	b.batchChan <- b.items
+	// Note: This needs to be in a go-routine to avoid locking out gets when
+	// the batch channel is full.
+	cpItems := make([]interface{}, len(b.items))
+	for i, val := range b.items {
+		cpItems[i] = val
+	}
+	// 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

@@ -36,11 +36,9 @@ func TestMaxItems(t *testing.T) {
 	})
 	assert.Nil(err)
 
-	go func() {
-		for i := 0; i < 1000; i++ {
-			assert.Nil(b.Put("foo bar baz"))
-		}
-	}()
+	for i := 0; i < 1000; i++ {
+		assert.Nil(b.Put("foo bar baz"))
+	}
 
 	batch, err := b.Get()
 	assert.Len(batch, 100)
@@ -139,32 +137,46 @@ func TestMultiConsumer(t *testing.T) {
 
 func TestDispose(t *testing.T) {
 	assert := assert.New(t)
-	b, err := New(0, 2, 100000, 10, func(str interface{}) uint {
+	b, err := New(1, 2, 100000, 2, func(str interface{}) uint {
 		return uint(len(str.(string)))
 	})
 	assert.Nil(err)
 	b.Put("a")
 	b.Put("b")
 	b.Put("c")
-	wait := make(chan bool)
-	go func() {
-		batch1, err := b.Get()
-		assert.Equal([]interface{}{"a", "b"}, batch1)
-		assert.Nil(err)
-		batch2, err := b.Get()
-		assert.Equal([]interface{}{"c"}, batch2)
-		assert.Nil(err)
-		_, err = b.Get()
-		assert.Equal(ErrDisposed, err)
-		wait <- true
-	}()
+
+	possibleBatches := [][]interface{}{
+		[]interface{}{"a", "b"},
+		[]interface{}{"c"},
+	}
+
+	// Wait for items to get to the channel
+	for len(b.(*basicBatcher).batchChan) == 0 {
+		time.Sleep(1 * time.Millisecond)
+	}
+	batch1, err := b.Get()
+	assert.Contains(possibleBatches, batch1)
+	assert.Nil(err)
+
+	batch2, err := b.Get()
+	assert.Contains(possibleBatches, batch2)
+	assert.Nil(err)
+
+	b.Put("d")
+	b.Put("e")
+	b.Put("f")
+	b.Put("g")
+	b.Put("h")
+	b.Put("i")
 
 	b.Dispose()
 
-	assert.Equal(ErrDisposed, b.Put("d"))
+	_, err = b.Get()
+	assert.Equal(ErrDisposed, err)
+
+	assert.Equal(ErrDisposed, b.Put("j"))
 	assert.Equal(ErrDisposed, b.Flush())
 
-	<-wait
 }
 
 func TestIsDisposed(t *testing.T) {