Check batch channel length when draining channel

Author: Steven.Osborne

batcher/batcher.go

@@ -167,7 +167,7 @@ func (b *basicBatcher) Dispose() {
 	b.items = nil
 
 	// Drain the batch channel and all routines waiting to put on the channel
-	for atomic.LoadInt32(&b.waiting) > 0 {
+	for len(b.batchChan) > 0 || atomic.LoadInt32(&b.waiting) > 0 {
 		<-b.batchChan
 	}
 	close(b.batchChan)

batcher/batcher_test.go

@@ -137,7 +137,7 @@ func TestMultiConsumer(t *testing.T) {
 
 func TestDispose(t *testing.T) {
 	assert := assert.New(t)
-	b, err := New(1, 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)
@@ -157,12 +157,18 @@ func TestDispose(t *testing.T) {
 	assert.Nil(err)
 
 	b.Put("d")
+	b.Put("e")
+	b.Put("f")
+	b.Put("g")
+	b.Put("h")
+	b.Put("i")
+
 	b.Dispose()
 
 	_, err = b.Get()
 	assert.Equal(ErrDisposed, err)
 
-	assert.Equal(ErrDisposed, b.Put("e"))
+	assert.Equal(ErrDisposed, b.Put("j"))
 	assert.Equal(ErrDisposed, b.Flush())
 
 }